refactor: 21915 Introduce StateLifecycleManager

hedera-node/hedera-app/src/main/java/com/hedera/node/app/HederaVirtualMapState.java

@@ -43,7 +43,7 @@ public class HederaVirtualMapState extends VirtualMapState<HederaVirtualMapState
      */
     public HederaVirtualMapState(
             @NonNull final Configuration configuration, @NonNull final Metrics metrics, @NonNull final Time time) {
-        super(configuration, metrics, time);
+        super(configuration, metrics);
     }
 
     /**
@@ -55,7 +55,7 @@ public HederaVirtualMapState(
      */
     public HederaVirtualMapState(
             @NonNull final VirtualMap virtualMap, @NonNull final Metrics metrics, @NonNull final Time time) {
-        super(virtualMap, metrics, time);
+        super(virtualMap, metrics);
     }
 
     protected HederaVirtualMapState(@NonNull final HederaVirtualMapState from) {
@@ -70,28 +70,6 @@ protected HederaVirtualMapState copyingConstructor() {
         return new HederaVirtualMapState(this);
     }
 
-    /**
-     * Creates a new instance of {@link HederaVirtualMapState} with the specified {@link VirtualMap}.
-     *
-     * @param virtualMap the virtual map whose metrics must already be registered
-     * @param metrics    the platform metric instance to use when creating the new instance of state
-     * @param time       the time instance to use when creating the new instance of state
-     * @return a new instance of {@link HederaVirtualMapState}
-     */
-    @Override
-    protected HederaVirtualMapState newInstance(
-            @NonNull final VirtualMap virtualMap, @NonNull final Metrics metrics, @NonNull final Time time) {
-        return new HederaVirtualMapState(virtualMap, metrics, time);
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    public long getRound() {
-        return DEFAULT_PLATFORM_STATE_FACADE.roundOf(this);
-    }
-
     @Override
     public String getInfoJson() {
         final JSONObject rootJson = new JSONObject();
@@ -152,4 +130,9 @@ public String getInfoJson() {
 
         return rootJson.toString();
     }
+
+    @Override
+    public String toString() {
+        return "HederaVirtualMapState[round=%d]".formatted(DEFAULT_PLATFORM_STATE_FACADE.roundOf(this));
+    }
 }

hedera-node/hedera-app/src/test/java/com/hedera/node/app/state/merkle/SerializationTest.java

@@ -11,18 +11,22 @@
 import com.hedera.node.app.spi.fixtures.TestSchema;
 import com.hedera.node.app.spi.migrate.StartupNetworks;
 import com.hedera.node.config.data.HederaConfig;
+import com.swirlds.base.test.fixtures.time.FakeTime;
 import com.swirlds.common.io.utility.LegacyTemporaryFileBuilder;
 import com.swirlds.common.merkle.utility.MerkleTreeSnapshotReader;
+import com.swirlds.common.metrics.noop.NoOpMetrics;
 import com.swirlds.config.api.Configuration;
 import com.swirlds.config.extensions.sources.SimpleConfigSource;
 import com.swirlds.config.extensions.test.fixtures.TestConfigBuilder;
 import com.swirlds.platform.state.signed.SignedState;
 import com.swirlds.platform.test.fixtures.state.RandomSignedStateGenerator;
 import com.swirlds.state.MerkleNodeState;
 import com.swirlds.state.State;
+import com.swirlds.state.StateLifecycleManager;
 import com.swirlds.state.lifecycle.MigrationContext;
 import com.swirlds.state.lifecycle.Schema;
 import com.swirlds.state.lifecycle.StateDefinition;
+import com.swirlds.state.merkle.StateLifecycleManagerImpl;
 import com.swirlds.state.merkle.disk.OnDiskReadableKVState;
 import com.swirlds.state.merkle.disk.OnDiskWritableKVState;
 import com.swirlds.state.spi.ReadableKVState;
@@ -137,7 +141,9 @@ private void forceFlush(ReadableKVState<?, ?> state) {
 
                 if (vm.size() > 1) {
                     vm.enableFlush();
-                    vm.release();
+                    if (vm.getReservationCount() > 0) {
+                        vm.release();
+                    }
                     vm.waitUntilFlushed();
                 }
             } catch (IllegalAccessException | NoSuchFieldException | InterruptedException e) {
@@ -157,11 +163,12 @@ private void forceFlush(ReadableKVState<?, ?> state) {
     @ParameterizedTest
     @ValueSource(booleans = {true, false})
     void simpleReadAndWrite(boolean forceFlush) throws IOException, ConstructableRegistryException {
-        final var schemaV1 = createV1Schema();
-        final var originalTree = createMerkleHederaState(schemaV1);
+        final Schema schemaV1 = createV1Schema();
+        final StateLifecycleManager stateLifecycleManager = createStateLifecycleManager(schemaV1);
+        final MerkleNodeState originalTree = stateLifecycleManager.getMutableState();
 
         // When we serialize it to bytes and deserialize it back into a tree
-        MerkleNodeState copy = originalTree.copy(); // make a copy to make VM flushable
+        MerkleNodeState copy = stateLifecycleManager.copyMutableState(); // make a copy to make VM flushable
         final byte[] serializedBytes;
         if (forceFlush) {
             // Force flush the VMs to disk to test serialization and deserialization
@@ -188,21 +195,22 @@ void simpleReadAndWrite(boolean forceFlush) throws IOException, ConstructableReg
 
     @Test
     void snapshot() throws IOException {
-        final var schemaV1 = createV1Schema();
-        final var originalTree = createMerkleHederaState(schemaV1);
-        final var tempDir = LegacyTemporaryFileBuilder.buildTemporaryDirectory(config);
+        final Schema<SemanticVersion> schemaV1 = createV1Schema();
+        final StateLifecycleManager stateLifecycleManager = createStateLifecycleManager(schemaV1);
+        final Path tempDir = LegacyTemporaryFileBuilder.buildTemporaryDirectory(config);
+        final MerkleNodeState originalTree = stateLifecycleManager.getLatestImmutableState();
 
         // prepare the tree and create a snapshot
-        originalTree.copy().release();
+        stateLifecycleManager.getMutableState().release();
         originalTree.computeHash();
-        originalTree.createSnapshot(tempDir);
+        stateLifecycleManager.createSnapshot(originalTree, tempDir);
+        originalTree.release();
 
         final MerkleNodeState state =
-                originalTree.loadSnapshot(tempDir.resolve(MerkleTreeSnapshotReader.SIGNED_STATE_FILE_NAME));
+                stateLifecycleManager.loadSnapshot(tempDir.resolve(MerkleTreeSnapshotReader.SIGNED_STATE_FILE_NAME));
         initServices(schemaV1, state);
         assertTree(state);
 
-        originalTree.release();
         state.release();
     }
 
@@ -212,14 +220,17 @@ void snapshot() throws IOException {
      * After it gets saved to disk again, and then loaded back in, it results in ClassCastException due to incorrect classId.
      */
     @Test
-    void dualReadAndWrite() throws IOException, ConstructableRegistryException {
-        final var schemaV1 = createV1Schema();
-        final var originalTree = createMerkleHederaState(schemaV1);
+    void dualReadAndWrite() throws IOException {
+        final Schema<SemanticVersion> schemaV1 = createV1Schema();
+        final StateLifecycleManager stateLifecycleManager = createStateLifecycleManager(schemaV1);
+        final MerkleNodeState originalTree = stateLifecycleManager.getMutableState();
 
-        MerkleNodeState copy = originalTree.copy(); // make a copy to make VM flushable
+        MerkleNodeState copy = stateLifecycleManager.copyMutableState(); // make a copy to make VM flushable
 
+        stateLifecycleManager
+                .copyMutableState()
+                .release(); // make a fast copy because we can only write to disk an immutable copy
         forceFlush(originalTree.getReadableStates(FIRST_SERVICE).get(FRUIT_STATE_ID));
-        copy.copy().release(); // make a fast copy because we can only write to disk an immutable copy
         copy.getRoot().getHash();
         final byte[] serializedBytes = writeTree(copy.getRoot(), dir);
 
@@ -275,11 +286,11 @@ private void initServices(Schema<SemanticVersion> schemaV1, MerkleNodeState load
         loadedTree.getRoot().migrate(MINIMUM_SUPPORTED_VERSION);
     }
 
-    private MerkleNodeState createMerkleHederaState(Schema schemaV1) {
+    private StateLifecycleManager createStateLifecycleManager(Schema schemaV1) {
         final SignedState randomState =
                 new RandomSignedStateGenerator().setRound(1).build();
 
-        final var originalTree = randomState.getState();
+        final MerkleNodeState originalTree = randomState.getState();
         // the state is not hashed yet
         final var originalTreeCopy = originalTree.copy();
         originalTree.release();
@@ -295,7 +306,12 @@ private MerkleNodeState createMerkleHederaState(Schema schemaV1) {
                 migrationStateChanges,
                 startupNetworks,
                 TEST_PLATFORM_STATE_FACADE);
-        return originalTreeCopy;
+
+        final StateLifecycleManager stateLifecycleManager =
+                new StateLifecycleManagerImpl(new NoOpMetrics(), new FakeTime(), TestVirtualMapState::new);
+
+        stateLifecycleManager.initState(originalTreeCopy, true);
+        return stateLifecycleManager;
     }
 
     private static void populateVmCache(State loadedTree) {

hedera-node/hedera-app/src/testFixtures/java/com/hedera/node/app/fixtures/state/FakeState.java

@@ -33,7 +33,6 @@
 import com.swirlds.state.test.fixtures.MapWritableStates;
 import com.swirlds.state.test.fixtures.merkle.TestVirtualMapState;
 import edu.umd.cs.findbugs.annotations.NonNull;
-import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -303,9 +302,4 @@ public long queueElementPath(final int stateId, @NonNull final Bytes expectedVal
     public void initializeState(@NonNull final StateMetadata<?, ?> md) {
         // do nothing
     }
-
-    @Override
-    public MerkleNodeState loadSnapshot(@NonNull final Path targetPath) {
-        throw new UnsupportedOperationException();
-    }
 }

hedera-state-validator/src/main/java/com/hedera/statevalidation/blockstream/BlockStreamRecoveryWorkflow.java

@@ -8,6 +8,7 @@
 import com.hedera.hapi.block.stream.Block;
 import com.hedera.hapi.block.stream.BlockItem;
 import com.hedera.hapi.block.stream.output.StateChanges;
+import com.hedera.node.app.HederaVirtualMapState;
 import com.hedera.node.app.hapi.utils.blocks.BlockStreamAccess;
 import com.hedera.node.app.hapi.utils.blocks.BlockStreamUtils;
 import com.hedera.statevalidation.util.PlatformContextHelper;
@@ -18,6 +19,8 @@
 import com.swirlds.platform.state.snapshot.DeserializedSignedState;
 import com.swirlds.platform.state.snapshot.SignedStateFileWriter;
 import com.swirlds.state.MerkleNodeState;
+import com.swirlds.state.StateLifecycleManager;
+import com.swirlds.state.merkle.StateLifecycleManagerImpl;
 import com.swirlds.state.spi.CommittableWritableStates;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import java.io.IOException;
@@ -149,9 +152,18 @@ public void applyBlocks(
                 false,
                 DEFAULT_PLATFORM_STATE_FACADE);
 
+        final StateLifecycleManager stateLifecycleManager = new StateLifecycleManagerImpl(
+                platformContext.getMetrics(),
+                platformContext.getTime(),
+                vm -> new HederaVirtualMapState(vm, platformContext.getMetrics(), platformContext.getTime()));
         try {
             SignedStateFileWriter.writeSignedStateFilesToDirectory(
-                    platformContext, selfId, outputPath, signedState, DEFAULT_PLATFORM_STATE_FACADE);
+                    platformContext,
+                    selfId,
+                    outputPath,
+                    signedState,
+                    DEFAULT_PLATFORM_STATE_FACADE,
+                    stateLifecycleManager);
         } catch (IOException e) {
             throw new RuntimeException(e);
         }

platform-sdk/consensus-otter-tests/src/testFixtures/java/org/hiero/otter/fixtures/app/OtterAppState.java

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 package org.hiero.otter.fixtures.app;
 
-import static com.swirlds.platform.state.service.PlatformStateFacade.DEFAULT_PLATFORM_STATE_FACADE;
 import static org.hiero.otter.fixtures.app.state.OtterStateInitializer.initOtterAppState;
 
 import com.hedera.hapi.node.base.SemanticVersion;
@@ -23,12 +22,12 @@ public class OtterAppState extends VirtualMapState<OtterAppState> implements Mer
 
     public OtterAppState(
             @NonNull final Configuration configuration, @NonNull final Metrics metrics, @NonNull final Time time) {
-        super(configuration, metrics, time);
+        super(configuration, metrics);
     }
 
     public OtterAppState(
             @NonNull final VirtualMap virtualMap, @NonNull final Metrics metrics, @NonNull final Time time) {
-        super(virtualMap, metrics, time);
+        super(virtualMap, metrics);
     }
 
     /**
@@ -82,20 +81,6 @@ protected OtterAppState copyingConstructor() {
         return new OtterAppState(this);
     }
 
-    @Override
-    protected OtterAppState newInstance(
-            @NonNull final VirtualMap virtualMap, @NonNull final Metrics metrics, @NonNull final Time time) {
-        return new OtterAppState(virtualMap, metrics, time);
-    }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    protected long getRound() {
-        return DEFAULT_PLATFORM_STATE_FACADE.roundOf(this);
-    }
-
     /**
      * Commit the state of all services.
      */

platform-sdk/platform-apps/demos/StatsDemo/src/main/java/com/swirlds/demo/stats/StatsDemoState.java

@@ -66,7 +66,7 @@ private StatsDemoState(final StatsDemoState sourceState) {
      * {@inheritDoc}
      */
     @Override
-    protected long getRound() {
+    public long getRound() {
         return DEFAULT_PLATFORM_STATE_FACADE.roundOf(this);
     }
 

platform-sdk/platform-apps/tests/ConsistencyTestingTool/src/main/java/com/swirlds/demo/consistency/ConsistencyTestingToolState.java

@@ -6,7 +6,6 @@
 import static com.swirlds.demo.consistency.V0680ConsistencyTestingToolSchema.STATE_LONG_STATE_ID;
 import static com.swirlds.logging.legacy.LogMarker.EXCEPTION;
 import static com.swirlds.logging.legacy.LogMarker.STARTUP;
-import static com.swirlds.platform.state.service.PlatformStateFacade.DEFAULT_PLATFORM_STATE_FACADE;
 import static org.hiero.base.utility.ByteUtils.byteArrayToLong;
 import static org.hiero.base.utility.NonCryptographicHashing.hash64;
 
@@ -84,7 +83,7 @@ public class ConsistencyTestingToolState extends VirtualMapState<ConsistencyTest
 
     public ConsistencyTestingToolState(
             @NonNull final Configuration configuration, @NonNull final Metrics metrics, @NonNull final Time time) {
-        super(configuration, metrics, time);
+        super(configuration, metrics);
         transactionHandlingHistory = new TransactionHandlingHistory();
         transactionsAwaitingPostHandle = ConcurrentHashMap.newKeySet();
         logger.info(STARTUP.getMarker(), "New State Constructed.");
@@ -95,7 +94,7 @@ public ConsistencyTestingToolState(
      */
     public ConsistencyTestingToolState(
             @NonNull final VirtualMap virtualMap, @NonNull final Metrics metrics, @NonNull final Time time) {
-        super(virtualMap, metrics, time);
+        super(virtualMap, metrics);
         transactionHandlingHistory = new TransactionHandlingHistory();
         transactionsAwaitingPostHandle = ConcurrentHashMap.newKeySet();
         logger.info(STARTUP.getMarker(), "New State Constructed.");
@@ -119,12 +118,6 @@ protected ConsistencyTestingToolState copyingConstructor() {
         return new ConsistencyTestingToolState(this);
     }
 
-    @Override
-    protected ConsistencyTestingToolState newInstance(
-            @NonNull final VirtualMap virtualMap, @NonNull final Metrics metrics, @NonNull final Time time) {
-        return new ConsistencyTestingToolState(virtualMap, metrics, time);
-    }
-
     /**
      * Initialize the state
      */
@@ -156,14 +149,6 @@ void initState(Path logFilePath) {
         transactionHandlingHistory.init(logFilePath);
     }
 
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    protected long getRound() {
-        return DEFAULT_PLATFORM_STATE_FACADE.roundOf(this);
-    }
-
     /**
      * @return the number of rounds handled
      */

platform-sdk/platform-apps/tests/ISSTestingTool/src/main/java/com/swirlds/demo/iss/ISSTestingToolState.java

@@ -16,7 +16,6 @@
 import static com.swirlds.demo.iss.V0680ISSTestingToolSchema.PLANNED_ISS_LIST_STATE_ID;
 import static com.swirlds.demo.iss.V0680ISSTestingToolSchema.PLANNED_LOG_ERROR_LIST_STATE_ID;
 import static com.swirlds.demo.iss.V0680ISSTestingToolSchema.RUNNING_SUM_STATE_ID;
-import static com.swirlds.platform.state.service.PlatformStateFacade.DEFAULT_PLATFORM_STATE_FACADE;
 
 import com.hedera.hapi.node.state.primitives.ProtoLong;
 import com.hedera.hapi.node.state.primitives.ProtoString;
@@ -74,12 +73,12 @@ public class ISSTestingToolState extends VirtualMapState<ISSTestingToolState> im
 
     public ISSTestingToolState(
             @NonNull final Configuration configuration, @NonNull final Metrics metrics, @NonNull final Time time) {
-        super(configuration, metrics, time);
+        super(configuration, metrics);
     }
 
     public ISSTestingToolState(
             @NonNull final VirtualMap virtualMap, @NonNull final Metrics metrics, @NonNull final Time time) {
-        super(virtualMap, metrics, time);
+        super(virtualMap, metrics);
     }
 
     /**
@@ -98,12 +97,6 @@ protected ISSTestingToolState copyingConstructor() {
         return new ISSTestingToolState(this);
     }
 
-    @Override
-    protected ISSTestingToolState newInstance(
-            @NonNull final VirtualMap virtualMap, @NonNull final Metrics metrics, @NonNull final Time time) {
-        return new ISSTestingToolState(virtualMap, metrics, time);
-    }
-
     public void initState(InitTrigger trigger, Platform platform) {
         throwIfImmutable();
 
@@ -193,12 +186,4 @@ List<PlannedIss> getPlannedIssList() {
     List<PlannedLogError> getPlannedLogErrorList() {
         return plannedLogErrorList;
     }
-
-    /**
-     * {@inheritDoc}
-     */
-    @Override
-    protected long getRound() {
-        return DEFAULT_PLATFORM_STATE_FACADE.roundOf(this);
-    }
 }