IGNITE-24364 Store zone-wide tx state storage per zone (#5159)

Author: rpuch

modules/distribution-zones/src/integrationTest/java/org/apache/ignite/internal/rebalance/ItRebalanceDistributedTest.java

@@ -1527,7 +1527,8 @@ private class Node {
                             clockService,
                             placementDriver,
                             schemaSyncService,
-                            systemDistributedConfiguration
+                            systemDistributedConfiguration,
+                            sharedTxStateStorage
                     ),
                     minTimeCollectorService,
                     systemDistributedConfiguration

modules/partition-replicator/build.gradle

@@ -41,6 +41,12 @@ dependencies {
     implementation project(':ignite-transactions')
     implementation project(':ignite-storage-api')
     implementation project(':ignite-low-watermark')
+    implementation project(':ignite-workers')
+
+    testImplementation testFixtures(project(':ignite-core'))
+    testImplementation libs.hamcrest.core
+    testImplementation libs.mockito.core
+    testImplementation libs.mockito.junit
 
     integrationTestImplementation testFixtures(project(':ignite-runner'))
     integrationTestImplementation testFixtures(project(':ignite-core'))

modules/partition-replicator/src/integrationTest/java/org/apache/ignite/internal/partition/replicator/fixtures/Node.java

@@ -549,7 +549,7 @@ public CompletableFuture<Boolean> invoke(
                 raftManager,
                 partitionRaftConfigurer,
                 view -> new LocalLogStorageFactory(),
-                ForkJoinPool.commonPool(),
+                threadPoolsManager.tableIoExecutor(),
                 t -> converter.apply(t),
                 replicaGrpId -> metaStorageManager.get(pendingPartAssignmentsKey((ZonePartitionId) replicaGrpId))
                         .thenApply(Entry::value)
@@ -590,6 +590,13 @@ public CompletableFuture<Boolean> invoke(
                 systemDistributedConfiguration
         );
 
+        sharedTxStateStorage = new TxStateRocksDbSharedStorage(
+                storagePath.resolve("tx-state"),
+                threadPoolsManager.commonScheduler(),
+                threadPoolsManager.tableIoExecutor(),
+                partitionsLogStorageFactory
+        );
+
         partitionReplicaLifecycleManager = new PartitionReplicaLifecycleManager(
                 catalogManager,
                 replicaManager,
@@ -603,21 +610,15 @@ public CompletableFuture<Boolean> invoke(
                 clockService,
                 placementDriver,
                 schemaSyncService,
-                systemDistributedConfiguration
+                systemDistributedConfiguration,
+                sharedTxStateStorage
         );
 
         StorageUpdateConfiguration storageUpdateConfiguration = clusterConfigRegistry
                 .getConfiguration(StorageUpdateExtensionConfiguration.KEY).storageUpdate();
 
         MinimumRequiredTimeCollectorService minTimeCollectorService = new MinimumRequiredTimeCollectorServiceImpl();
 
-        sharedTxStateStorage = new TxStateRocksDbSharedStorage(
-                storagePath.resolve("tx-state"),
-                threadPoolsManager.commonScheduler(),
-                threadPoolsManager.tableIoExecutor(),
-                partitionsLogStorageFactory
-        );
-
         tableManager = new TableManager(
                 name,
                 registry,
@@ -729,8 +730,8 @@ public void start() {
                 txManager,
                 dataStorageMgr,
                 schemaManager,
-                partitionReplicaLifecycleManager,
                 sharedTxStateStorage,
+                partitionReplicaLifecycleManager,
                 tableManager,
                 indexManager
         )).thenComposeAsync(componentFuts -> {

modules/partition-replicator/src/main/java/org/apache/ignite/internal/partition/replicator/PartitionReplicaLifecycleManager.java

@@ -40,6 +40,7 @@
 import static org.apache.ignite.internal.distributionzones.rebalance.ZoneRebalanceUtil.zoneAssignmentsGetLocally;
 import static org.apache.ignite.internal.distributionzones.rebalance.ZoneRebalanceUtil.zonePartitionAssignmentsGetLocally;
 import static org.apache.ignite.internal.hlc.HybridTimestamp.LOGICAL_TIME_BITS_SIZE;
+import static org.apache.ignite.internal.hlc.HybridTimestamp.nullableHybridTimestamp;
 import static org.apache.ignite.internal.lang.IgniteSystemProperties.getBoolean;
 import static org.apache.ignite.internal.metastorage.dsl.Conditions.notExists;
 import static org.apache.ignite.internal.metastorage.dsl.Operations.put;
@@ -128,8 +129,11 @@
 import org.apache.ignite.internal.replicator.ZonePartitionId;
 import org.apache.ignite.internal.replicator.listener.ReplicaListener;
 import org.apache.ignite.internal.schema.SchemaSyncService;
+import org.apache.ignite.internal.tx.storage.state.TxStatePartitionStorage;
+import org.apache.ignite.internal.tx.storage.state.rocksdb.TxStateRocksDbSharedStorage;
 import org.apache.ignite.internal.util.Cursor;
 import org.apache.ignite.internal.util.IgniteSpinBusyLock;
+import org.apache.ignite.internal.util.IgniteUtils;
 import org.apache.ignite.network.ClusterNode;
 import org.jetbrains.annotations.Nullable;
 
@@ -213,6 +217,8 @@ public class PartitionReplicaLifecycleManager extends
 
     private final ConcurrentMap<ZonePartitionId, ZonePartitionRaftListener> zonePartitionRaftListeners = new ConcurrentHashMap<>();
 
+    private final ZoneResourcesManager zoneResourcesManager;
+
     /**
      * The constructor.
      *
@@ -228,6 +234,7 @@ public class PartitionReplicaLifecycleManager extends
      * @param placementDriver Placement driver.
      * @param schemaSyncService Schema synchronization service.
      * @param systemDistributedConfiguration System distributed configuration.
+     * @param sharedTxStateStorage Shared tx state storage.
      */
     public PartitionReplicaLifecycleManager(
             CatalogManager catalogMgr,
@@ -242,7 +249,8 @@ public PartitionReplicaLifecycleManager(
             ClockService clockService,
             PlacementDriver placementDriver,
             SchemaSyncService schemaSyncService,
-            SystemDistributedConfiguration systemDistributedConfiguration
+            SystemDistributedConfiguration systemDistributedConfiguration,
+            TxStateRocksDbSharedStorage sharedTxStateStorage
     ) {
         this.catalogMgr = catalogMgr;
         this.replicaMgr = replicaMgr;
@@ -266,6 +274,8 @@ public PartitionReplicaLifecycleManager(
                 Integer::parseInt
         );
 
+        zoneResourcesManager = new ZoneResourcesManager(sharedTxStateStorage);
+
         pendingAssignmentsRebalanceListener = createPendingAssignmentsRebalanceListener();
         stableAssignmentsRebalanceListener = createStableAssignmentsRebalanceListener();
         assignmentsSwitchRebalanceListener = createAssignmentsSwitchRebalanceListener();
@@ -415,14 +425,17 @@ private CompletableFuture<Void> calculateZoneAssignmentsAndCreateReplicationNode
 
             return getOrCreateAssignments(zoneDescriptor, causalityToken, catalogVersion)
                     .thenCompose(assignments -> writeZoneAssignmentsToMetastore(zoneId, assignments))
-                    .thenCompose(assignments -> createZoneReplicationNodes(zoneId, assignments, causalityToken));
+                    .thenCompose(
+                            assignments -> createZoneReplicationNodes(zoneId, assignments, causalityToken, zoneDescriptor.partitions())
+                    );
         });
     }
 
     private CompletableFuture<Void> createZoneReplicationNodes(
             int zoneId,
             List<Assignments> assignments,
-            long revision
+            long revision,
+            int partitionCount
     ) {
         return inBusyLockAsync(busyLock, () -> {
             assert assignments != null : IgniteStringFormatter.format("Zone has empty assignments [id={}].", zoneId);
@@ -440,7 +453,8 @@ private CompletableFuture<Void> createZoneReplicationNodes(
                         zonePartitionId,
                         localMemberAssignment,
                         zoneAssignment,
-                        revision
+                        revision,
+                        partitionCount
                 );
             }
 
@@ -456,13 +470,15 @@ private CompletableFuture<Void> createZoneReplicationNodes(
      * @param localMemberAssignment Assignment of the local member, or null if local member is not part of the assignment.
      * @param stableAssignments Stable assignments.
      * @param revision Event's revision.
+     * @param partitionCount Number of partitions on the zone.
      * @return Future that completes when a replica is started.
      */
     private CompletableFuture<?> createZonePartitionReplicationNode(
             ZonePartitionId zonePartitionId,
             @Nullable Assignment localMemberAssignment,
             Assignments stableAssignments,
-            long revision
+            long revision,
+            int partitionCount
     ) {
         if (localMemberAssignment == null) {
             return nullCompletedFuture();
@@ -490,6 +506,12 @@ private CompletableFuture<?> createZonePartitionReplicationNode(
         );
 
         Supplier<CompletableFuture<Boolean>> startReplicaSupplier = () -> {
+            TxStatePartitionStorage txStatePartitionStorage = zoneResourcesManager.getOrCreatePartitionTxStateStorage(
+                    zonePartitionId.zoneId(),
+                    partitionCount,
+                    zonePartitionId.partitionId()
+            );
+
             try {
                 return replicaMgr.startReplica(
                                 zonePartitionId,
@@ -499,6 +521,9 @@ private CompletableFuture<?> createZonePartitionReplicationNode(
                                 stablePeersAndLearners,
                                 raftGroupListener,
                                 raftGroupEventsListener,
+                                // TODO: IGNITE-24371 - pass real isVolatile flag
+                                false,
+                                txStatePartitionStorage,
                                 busyLock
                         ).thenCompose(replica -> executeUnderZoneWriteLock(zonePartitionId.zoneId(), () -> {
                             replicationGroupIds.add(zonePartitionId);
@@ -998,11 +1023,19 @@ private CompletableFuture<Void> handleChangePendingAssignmentEvent(
         CompletableFuture<?> localServicesStartFuture;
 
         if (shouldStartLocalGroupNode) {
+            // We can safely access the Catalog at the timestamp because:
+            // 1. It is guaranteed that Catalog update bringing the Catalog version has been applied (as we are now handling
+            // a Metastorage event that was caused by that same Catalog version we need, so the Catalog version update is already
+            // handled), so no Schema sync is needed.
+            // 2. It is guaranteed that Catalog compactor cannot remove Catalog version corresponding to pending assignments timestamp.
+            CatalogZoneDescriptor zoneDescriptor = zoneDescriptorAt(replicaGrpId.zoneId(), pendingAssignments.timestamp());
+
             localServicesStartFuture = createZonePartitionReplicationNode(
                     replicaGrpId,
                     localMemberAssignment,
                     computedStableAssignments,
-                    revision
+                    revision,
+                    zoneDescriptor.partitions()
             );
         } else if (pendingAssignmentsAreForced && localMemberAssignment != null) {
             localServicesStartFuture = runAsync(() -> {
@@ -1041,6 +1074,16 @@ private CompletableFuture<Void> handleChangePendingAssignmentEvent(
                 }), ioExecutor);
     }
 
+    private CatalogZoneDescriptor zoneDescriptorAt(int zoneId, long timestamp) {
+        Catalog catalog = catalogMgr.activeCatalog(timestamp);
+        assert catalog != null : "Catalog is not available at " + nullableHybridTimestamp(timestamp);
+
+        CatalogZoneDescriptor zoneDescriptor = catalog.zone(zoneId);
+        assert zoneDescriptor != null : "Zone descriptor is not available at " + nullableHybridTimestamp(timestamp) + " for zone " + zoneId;
+
+        return zoneDescriptor;
+    }
+
     private CompletableFuture<Void> changePeersOnRebalance(
             ReplicaManager replicaMgr,
             ZonePartitionId replicaGrpId,
@@ -1191,6 +1234,12 @@ public CompletableFuture<Void> stopAsync(ComponentContext componentContext) {
             return nullCompletedFuture();
         }
 
+        try {
+            IgniteUtils.closeAllManually(zoneResourcesManager);
+        } catch (Exception e) {
+            return failedFuture(e);
+        }
+
         return nullCompletedFuture();
     }
 

modules/partition-replicator/src/main/java/org/apache/ignite/internal/partition/replicator/ZoneResourcesManager.java

@@ -0,0 +1,98 @@
+/*
+ * 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.partition.replicator;
+
+import static org.apache.ignite.internal.util.IgniteUtils.inBusyLock;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import org.apache.ignite.internal.close.ManuallyCloseable;
+import org.apache.ignite.internal.tx.storage.state.ThreadAssertingTxStateStorage;
+import org.apache.ignite.internal.tx.storage.state.TxStatePartitionStorage;
+import org.apache.ignite.internal.tx.storage.state.TxStateStorage;
+import org.apache.ignite.internal.tx.storage.state.rocksdb.TxStateRocksDbSharedStorage;
+import org.apache.ignite.internal.tx.storage.state.rocksdb.TxStateRocksDbStorage;
+import org.apache.ignite.internal.util.IgniteSpinBusyLock;
+import org.apache.ignite.internal.worker.ThreadAssertions;
+
+/**
+ * Manages resources of distribution zones; that is, allows creation of underlying storages and closes them on node stop.
+ */
+class ZoneResourcesManager implements ManuallyCloseable {
+    private final TxStateRocksDbSharedStorage sharedTxStateStorage;
+
+    /** Map from zone IDs to their resource holders. */
+    private final Map<Integer, ZoneResources> resourcesByZoneId = new ConcurrentHashMap<>();
+
+    private final IgniteSpinBusyLock busyLock = new IgniteSpinBusyLock();
+
+    ZoneResourcesManager(TxStateRocksDbSharedStorage sharedTxStateStorage) {
+        this.sharedTxStateStorage = sharedTxStateStorage;
+    }
+
+    /**
+     * Gets or creates a transaction state storage for a zone partition.
+     *
+     * @param zoneId ID of the zone.
+     * @param partitionCount Number of partitions in the zone.
+     * @param partitionId Partition ID.
+     */
+    TxStatePartitionStorage getOrCreatePartitionTxStateStorage(int zoneId, int partitionCount, int partitionId) {
+        return inBusyLock(busyLock, () -> {
+            ZoneResources zoneResources = resourcesByZoneId.computeIfAbsent(
+                    zoneId,
+                    id -> createZoneResources(id, partitionCount)
+            );
+
+            return zoneResources.txStateStorage.getOrCreatePartitionStorage(partitionId);
+        });
+    }
+
+    private ZoneResources createZoneResources(int zoneId, int partitionCount) {
+        return new ZoneResources(createTxStateStorage(zoneId, partitionCount));
+    }
+
+    private TxStateStorage createTxStateStorage(int zoneId, int partitionCount) {
+        TxStateStorage txStateStorage = new TxStateRocksDbStorage(zoneId, partitionCount, sharedTxStateStorage);
+
+        if (ThreadAssertions.enabled()) {
+            txStateStorage = new ThreadAssertingTxStateStorage(txStateStorage);
+        }
+
+        txStateStorage.start();
+
+        return txStateStorage;
+    }
+
+    @Override
+    public void close() {
+        busyLock.block();
+
+        for (ZoneResources zoneResources : resourcesByZoneId.values()) {
+            zoneResources.txStateStorage.close();
+        }
+    }
+
+    private static class ZoneResources {
+        private final TxStateStorage txStateStorage;
+
+        private ZoneResources(TxStateStorage txStateStorage) {
+            this.txStateStorage = txStateStorage;
+        }
+    }
+}