IGNITE-24679. Implement mapping cache purging on zone primary replica expiration

modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/exec/mapping/MappingServiceImpl.java

@@ -117,18 +117,16 @@ public MappingServiceImpl(
     public CompletableFuture<Boolean> onPrimaryReplicaExpired(PrimaryReplicaEventParameters parameters) {
         assert parameters != null;
 
-        assert enabledColocation ? parameters.groupId() instanceof ZonePartitionId : parameters.groupId() instanceof TablePartitionId
-                : parameters.groupId();
+        int tableOrZoneId;
 
-        if (parameters.groupId() instanceof ZonePartitionId) {
-            // TODO: https://issues.apache.org/jira/browse/IGNITE-24679 - remove mappings from cache for zone partitions.
-            return CompletableFutures.falseCompletedFuture();
+        if (enabledColocation) {
+            tableOrZoneId = ((ZonePartitionId) parameters.groupId()).zoneId();
+        } else {
+            tableOrZoneId = ((TablePartitionId) parameters.groupId()).tableId();
         }
 
-        int tabId = ((TablePartitionId) parameters.groupId()).tableId();
-
         // TODO https://issues.apache.org/jira/browse/IGNITE-21201 Move complex computations to a different thread.
-        mappingsCache.removeIfValue(value -> value.tableIds.contains(tabId));
+        mappingsCache.removeIfValue(value -> value.tabelOrZoneIds.contains(tableOrZoneId));
 
         return CompletableFutures.falseCompletedFuture();
     }
@@ -146,36 +144,45 @@ public CompletableFuture<List<MappedFragment>> map(MultiStepPlan multiStepPlan,
         } else {
             mappedFragments = mappingsCache.compute(
                     new MappingsCacheKey(multiStepPlan.id(), mapOnBackups),
-                    (key, val) -> {
-                        if (val == null) {
-                            IntSet tableIds = new IntOpenHashSet();
-                            boolean topologyAware = false;
-
-                            for (Fragment fragment : template.fragments) {
-                                topologyAware = topologyAware || !fragment.systemViews().isEmpty();
-                                for (IgniteDataSource source : fragment.tables().values()) {
-                                    tableIds.add(source.id());
-                                }
-                            }
+                    (key, val) -> computeMappingCacheKey(val, template, mapOnBackups)
+            ).mappedFragments;
+        }
 
-                            long topVer = topologyAware ? logicalTopologyVerSupplier.get() : Long.MAX_VALUE;
+        return mappedFragments.thenApply(frags -> applyPartitionPruning(frags.fragments, parameters));
+    }
 
-                            assert nodeExclusionFilter == null;
+    private MappingsCacheValue computeMappingCacheKey(
+            MappingsCacheValue val,
+            FragmentsTemplate template,
+            boolean mapOnBackups
+    ) {
+        if (val == null) {
+            IntSet tableOrZoneIds = new IntOpenHashSet();
+            boolean topologyAware = false;
+
+            for (Fragment fragment : template.fragments) {
+                topologyAware = topologyAware || !fragment.systemViews().isEmpty();
+                for (IgniteTable source : fragment.tables().values()) {
+                    if (enabledColocation) {
+                        tableOrZoneIds.add(source.zoneId());
+                    } else {
+                        tableOrZoneIds.add(source.id());
+                    }
+                }
+            }
 
-                            return new MappingsCacheValue(topVer, tableIds, mapFragments(template, mapOnBackups, null));
-                        }
+            long topVer = topologyAware ? logicalTopologyVerSupplier.get() : Long.MAX_VALUE;
 
-                        long topologyVer = logicalTopologyVerSupplier.get();
+            return new MappingsCacheValue(topVer, tableOrZoneIds, mapFragments(template, mapOnBackups, null));
+        }
 
-                        if (val.topologyVersion < topologyVer) {
-                            return new MappingsCacheValue(topologyVer, val.tableIds, mapFragments(template, mapOnBackups, null));
-                        }
+        long topologyVer = logicalTopologyVerSupplier.get();
 
-                        return val;
-                    }).mappedFragments;
+        if (val.topologyVersion < topologyVer) {
+            return new MappingsCacheValue(topologyVer, val.tabelOrZoneIds, mapFragments(template, mapOnBackups, null));
         }
 
-        return mappedFragments.thenApply(frags -> applyPartitionPruning(frags.fragments, parameters));
+        return val;
     }
 
     CompletableFuture<DistributionHolder> composeDistributions(
@@ -396,12 +403,12 @@ private static class MappedFragments {
 
     private static class MappingsCacheValue {
         private final long topologyVersion;
-        private final IntSet tableIds;
+        private final IntSet tabelOrZoneIds;
         private final CompletableFuture<MappedFragments> mappedFragments;
 
-        MappingsCacheValue(long topologyVersion, IntSet tableIds, CompletableFuture<MappedFragments> mappedFragments) {
+        MappingsCacheValue(long topologyVersion, IntSet tabelOrZoneIds, CompletableFuture<MappedFragments> mappedFragments) {
             this.topologyVersion = topologyVersion;
-            this.tableIds = tableIds;
+            this.tabelOrZoneIds = tabelOrZoneIds;
             this.mappedFragments = mappedFragments;
         }
     }

modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/exec/mapping/MappingServiceImplTest.java

@@ -23,10 +23,10 @@
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.in;
 import static org.hamcrest.Matchers.not;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNotSame;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.junit.jupiter.api.Assumptions.assumeFalse;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anySet;
@@ -47,6 +47,7 @@
 import org.apache.ignite.internal.catalog.Catalog;
 import org.apache.ignite.internal.catalog.CatalogService;
 import org.apache.ignite.internal.catalog.descriptors.CatalogObjectDescriptor;
+import org.apache.ignite.internal.catalog.descriptors.CatalogZoneDescriptor;
 import org.apache.ignite.internal.hlc.ClockService;
 import org.apache.ignite.internal.hlc.HybridTimestamp;
 import org.apache.ignite.internal.hlc.TestClockService;
@@ -56,6 +57,7 @@
 import org.apache.ignite.internal.partitiondistribution.TokenizedAssignmentsImpl;
 import org.apache.ignite.internal.placementdriver.event.PrimaryReplicaEventParameters;
 import org.apache.ignite.internal.replicator.TablePartitionId;
+import org.apache.ignite.internal.replicator.ZonePartitionId;
 import org.apache.ignite.internal.sql.engine.framework.TestBuilders;
 import org.apache.ignite.internal.sql.engine.framework.TestCluster;
 import org.apache.ignite.internal.sql.engine.prepare.MultiStepPlan;
@@ -65,8 +67,10 @@
 import org.apache.ignite.internal.sql.engine.util.cache.CaffeineCacheFactory;
 import org.apache.ignite.internal.systemview.api.SystemViews;
 import org.apache.ignite.internal.testframework.BaseIgniteAbstractTest;
+import org.apache.ignite.internal.testframework.WithSystemProperty;
 import org.apache.ignite.internal.type.NativeTypes;
 import org.apache.ignite.internal.util.SubscriptionUtils;
+import org.jetbrains.annotations.Nullable;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
 
@@ -75,6 +79,9 @@
  */
 @SuppressWarnings("ThrowFromFinallyBlock")
 public class MappingServiceImplTest extends BaseIgniteAbstractTest {
+    private static final String ZONE_NAME_1 = "zone1";
+    private static final String ZONE_NAME_2 = "zone2";
+
     private static final MultiStepPlan PLAN;
     private static final MultiStepPlan PLAN_WITH_SYSTEM_VIEW;
     private static final TestCluster cluster;
@@ -88,13 +95,21 @@ public class MappingServiceImplTest extends BaseIgniteAbstractTest {
         // @formatter:off
         cluster = TestBuilders.cluster()
                 .nodes("N1")
+                .addZone()
+                        .name(ZONE_NAME_1)
+                        .end()
+                .addZone()
+                        .name(ZONE_NAME_2)
+                        .end()
                 .addTable()
                         .name("T1")
+                        .zoneName(ZONE_NAME_1)
                         .addKeyColumn("ID", NativeTypes.INT32)
                         .addColumn("VAL", NativeTypes.INT32)
                         .end()
                 .addTable()
                         .name("T2")
+                        .zoneName(ZONE_NAME_2)
                         .addKeyColumn("ID", NativeTypes.INT32)
                         .addColumn("VAL", NativeTypes.INT32)
                         .end()
@@ -236,9 +251,6 @@ public void testCacheInvalidationOnTopologyChange() {
 
     @Test
     public void testCacheInvalidationOnPrimaryExpiration() {
-        // TODO: https://issues.apache.org/jira/browse/IGNITE-24679 - remove this assumption.
-        assumeFalse(IgniteSystemProperties.enabledColocation());
-
         String localNodeName = "NODE";
         List<String> nodeNames = List.of(localNodeName, "NODE1");
 
@@ -286,6 +298,53 @@ public void testCacheInvalidationOnPrimaryExpiration() {
         verify(execProvider, times(2)).forTable(any(HybridTimestamp.class), any(IgniteTable.class), anyBoolean());
     }
 
+    @WithSystemProperty(key = IgniteSystemProperties.COLOCATION_FEATURE_FLAG, value = "true")
+    @Test
+    public void testCacheInvalidationOnPrimaryZoneExpiration() {
+        String localNodeName = "NODE";
+        List<String> nodeNames = List.of(localNodeName, "NODE1");
+
+        Function<String, PrimaryReplicaEventParameters> prepareEvtParams = (name) -> {
+            CatalogService catalogService = cluster.catalogManager();
+            Catalog catalog = catalogService.catalog(catalogService.latestCatalogVersion());
+
+            @Nullable CatalogZoneDescriptor zoneDescriptor = catalog.zone(name);
+
+            assertNotNull(zoneDescriptor);
+
+            return new PrimaryReplicaEventParameters(
+                    0, new ZonePartitionId(zoneDescriptor.id(), 0), new UUID(0, 0), "ignored", HybridTimestamp.MIN_VALUE);
+        };
+
+        // Initialize mapping service.
+        Supplier<Long> logicalTopologyVerSupplier = createStableTopologySupplier();
+        ExecutionDistributionProvider execProvider = Mockito.spy(new TestExecutionDistributionProvider(nodeNames));
+
+        MappingServiceImpl mappingService = Mockito.spy(new MappingServiceImpl(
+                localNodeName,
+                CLOCK_SERVICE,
+                CaffeineCacheFactory.INSTANCE,
+                100,
+                PARTITION_PRUNER,
+                logicalTopologyVerSupplier,
+                execProvider
+        ));
+
+        List<MappedFragment> mappedFragments = await(mappingService.map(PLAN, PARAMS));
+        verify(execProvider, times(1)).forTable(any(HybridTimestamp.class), any(IgniteTable.class), anyBoolean());
+
+        // Simulate expiration of the primary replica for non-mapped table - the cache entry should not be invalidated.
+        await(mappingService.onPrimaryReplicaExpired(prepareEvtParams.apply(ZONE_NAME_2)));
+        assertSame(mappedFragments, await(mappingService.map(PLAN, PARAMS)));
+
+        verify(mappingService, times(1)).composeDistributions(anySet(), anySet(), anyBoolean());
+
+        // Simulate expiration of the primary replica for mapped table - the cache entry should be invalidated.
+        await(mappingService.onPrimaryReplicaExpired(prepareEvtParams.apply(ZONE_NAME_1)));
+        assertNotSame(mappedFragments, await(mappingService.map(PLAN, PARAMS)));
+        verify(execProvider, times(2)).forTable(any(HybridTimestamp.class), any(IgniteTable.class), anyBoolean());
+    }
+
     private MappingServiceImpl createMappingServiceNoCache(String localNodeName, List<String> nodeNames) {
         return createMappingService(localNodeName, nodeNames, 0);
     }

modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/framework/ClusterZoneBuilder.java

@@ -0,0 +1,29 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.ignite.internal.sql.engine.framework;
+
+import org.apache.ignite.internal.catalog.CatalogCommand;
+import org.apache.ignite.internal.sql.engine.framework.TestBuilders.ClusterBuilder;
+import org.apache.ignite.internal.sql.engine.framework.TestBuilders.NestedBuilder;
+
+/**
+ * A builder interface for creating a test zone as a nested object within a test cluster.
+ */
+public interface ClusterZoneBuilder extends ZoneBuilderBase<ClusterZoneBuilder>, NestedBuilder<ClusterBuilder> {
+    CatalogCommand build();
+}

modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/framework/ClusterZoneBuilderImpl.java

@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.ignite.internal.sql.engine.framework;
+
+import static org.apache.ignite.internal.catalog.CatalogService.DEFAULT_STORAGE_PROFILE;
+
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.apache.ignite.internal.catalog.CatalogCommand;
+import org.apache.ignite.internal.catalog.commands.CreateZoneCommand;
+import org.apache.ignite.internal.catalog.commands.StorageProfileParams;
+import org.apache.ignite.internal.sql.engine.framework.TestBuilders.ClusterBuilder;
+import org.apache.ignite.internal.sql.engine.framework.TestBuilders.ClusterBuilderImpl;
+
+/**
+ * An implementation of the {@link ClusterZoneBuilder} interface.
+ */
+class ClusterZoneBuilderImpl implements ClusterZoneBuilder {
+    private final ClusterBuilderImpl parent;
+
+    private String name;
+
+    ClusterZoneBuilderImpl(ClusterBuilderImpl parent) {
+        this.parent = parent;
+    }
+
+    @Override
+    public ClusterBuilder end() {
+        parent.zoneBuilders().add(this);
+
+        return parent;
+    }
+
+    @Override
+    public ClusterZoneBuilder name(String name) {
+        this.name = name;
+
+        return this;
+    }
+
+    @Override
+    public CatalogCommand build() {
+        List<StorageProfileParams> storageProfileParams = Stream.of(DEFAULT_STORAGE_PROFILE)
+                .map(profileName -> StorageProfileParams.builder().storageProfile(profileName).build())
+                .collect(Collectors.toList());
+
+        return CreateZoneCommand.builder()
+                .zoneName(name)
+                .storageProfilesParams(storageProfileParams)
+                .build();
+    }
+}