From e3434d1f9229e8141bdf8c84c70961bbecd1c721 Mon Sep 17 00:00:00 2001 From: Cody Littley <56973212+cody-littley@users.noreply.github.com> Date: Tue, 2 Jan 2024 11:39:21 -0600 Subject: [PATCH] fix: state leak (#10690) Signed-off-by: Cody Littley --- .../common/merkle/copy/MerkleInitialize.java | 7 ++- .../platform/state/SwirldStateManager.java | 11 ++-- .../test/merkle/MerkleSerializationTests.java | 55 ++++++++++++++++--- 3 files changed, 60 insertions(+), 13 deletions(-) diff --git a/platform-sdk/swirlds-common/src/main/java/com/swirlds/common/merkle/copy/MerkleInitialize.java b/platform-sdk/swirlds-common/src/main/java/com/swirlds/common/merkle/copy/MerkleInitialize.java index e0d84f23b79c..cf1d6726e2aa 100644 --- a/platform-sdk/swirlds-common/src/main/java/com/swirlds/common/merkle/copy/MerkleInitialize.java +++ b/platform-sdk/swirlds-common/src/main/java/com/swirlds/common/merkle/copy/MerkleInitialize.java @@ -82,7 +82,12 @@ public static MerkleNode initializeAndMigrateTreeAfterDeserialization( final int deserializationVersion = Objects.requireNonNull( deserializationVersions.get(root.getClassId()), "class not discovered during deserialization"); - return root.migrate(deserializationVersion); + final MerkleNode migratedRoot = root.migrate(deserializationVersion); + if (migratedRoot != root) { + root.release(); + } + + return migratedRoot; } /** diff --git a/platform-sdk/swirlds-platform-core/src/main/java/com/swirlds/platform/state/SwirldStateManager.java b/platform-sdk/swirlds-platform-core/src/main/java/com/swirlds/platform/state/SwirldStateManager.java index 6f8260f40440..4477012380e3 100644 --- a/platform-sdk/swirlds-platform-core/src/main/java/com/swirlds/platform/state/SwirldStateManager.java +++ b/platform-sdk/swirlds-platform-core/src/main/java/com/swirlds/platform/state/SwirldStateManager.java @@ -151,11 +151,14 @@ public void prehandleApplicationTransactions(final EventImpl event) { while (!immutableState.tryReserve()) { immutableState = latestImmutableState.get(); } - transactionHandler.preHandle(event, immutableState.getSwirldState()); - event.getBaseEvent().signalPrehandleCompletion(); - immutableState.release(); + try { + transactionHandler.preHandle(event, immutableState.getSwirldState()); + } finally { + event.getBaseEvent().signalPrehandleCompletion(); + immutableState.release(); - stats.preHandleTime(startTime, System.nanoTime()); + stats.preHandleTime(startTime, System.nanoTime()); + } } /** diff --git a/platform-sdk/swirlds-unit-tests/common/swirlds-common-test/src/test/java/com/swirlds/common/test/merkle/MerkleSerializationTests.java b/platform-sdk/swirlds-unit-tests/common/swirlds-common-test/src/test/java/com/swirlds/common/test/merkle/MerkleSerializationTests.java index 8ce0647079b3..78ca867a0764 100644 --- a/platform-sdk/swirlds-unit-tests/common/swirlds-common-test/src/test/java/com/swirlds/common/test/merkle/MerkleSerializationTests.java +++ b/platform-sdk/swirlds-unit-tests/common/swirlds-common-test/src/test/java/com/swirlds/common/test/merkle/MerkleSerializationTests.java @@ -18,11 +18,13 @@ import static com.swirlds.common.io.utility.FileUtils.deleteDirectory; import static com.swirlds.common.test.merkle.util.MerkleTestUtils.areTreesEqual; +import static com.swirlds.common.test.merkle.util.MerkleTestUtils.buildLessSimpleTree; import static com.swirlds.common.test.merkle.util.MerkleTestUtils.buildLessSimpleTreeExtended; import static com.swirlds.common.test.merkle.util.MerkleTestUtils.isFullyInitialized; import static com.swirlds.test.framework.ResourceLoader.getFile; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -61,6 +63,7 @@ import java.nio.file.Path; import java.util.LinkedList; import java.util.List; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Tag; @@ -219,22 +222,22 @@ void deserializeInvalidTrees() throws IOException { final List trees = new LinkedList<>(); // Too low of a version on a leaf - DummyMerkleInternal root = MerkleTestUtils.buildLessSimpleTree(); + DummyMerkleInternal root = buildLessSimpleTree(); ((DummyMerkleLeaf) root.asInternal().getChild(0)).setVersion(0); trees.add(root); // Too high of a version on a leaf - root = MerkleTestUtils.buildLessSimpleTree(); + root = buildLessSimpleTree(); ((DummyMerkleLeaf) root.asInternal().getChild(0)).setVersion(1234); trees.add(root); // Too low of a version on an internal node - root = MerkleTestUtils.buildLessSimpleTree(); + root = buildLessSimpleTree(); root.setVersion(0); trees.add(root); // Too high of a version on an internal node - root = MerkleTestUtils.buildLessSimpleTree(); + root = buildLessSimpleTree(); root.setVersion(1234); trees.add(root); @@ -245,9 +248,9 @@ void deserializeInvalidTrees() throws IOException { /** * This test asserts that the hash of a tree does not change due to future code changes. - * - * Although there is no contractual requirement not to change the hash between versions, - * we should at least be aware if a change occurs. + *

+ * Although there is no contractual requirement not to change the hash between versions, we should at least be aware + * if a change occurs. */ @Test @Tag(TestTypeTags.FUNCTIONAL) @@ -259,7 +262,7 @@ void testHashFromFile() throws IOException { new DataInputStream(ResourceLoader.loadFileAsStream("merkle/hashed-tree-merkle-v1.dat")); final Hash oldHash = new Hash(dataStream.readAllBytes(), DigestType.SHA_384); - final MerkleNode tree = MerkleTestUtils.buildLessSimpleTree(); + final MerkleNode tree = buildLessSimpleTree(); assertEquals( oldHash, @@ -397,6 +400,42 @@ void nodeMigrationTest() throws IOException { DummyMerkleInternal.resetMigrationMapper(); } + @Test + void migrateRootTest() throws IOException { + + // In this test, we will attempt to swap out originalRoot with newRoot + + final MerkleNode originalRoot = buildLessSimpleTreeExtended(); + final MerkleNode newRoot = buildLessSimpleTree(); + + final AtomicReference deserializedRootBeforeMigration = new AtomicReference<>(); + DummyMerkleInternal.setMigrationMapper((node, version) -> { + if (node.getDepth() == 0) { + assertNull(deserializedRootBeforeMigration.get(), "deserialized root should not have been set yet"); + deserializedRootBeforeMigration.set(node); + return newRoot; + } + return node; + }); + + final ByteArrayOutputStream byteOut = new ByteArrayOutputStream(); + final MerkleDataOutputStream out = new MerkleDataOutputStream(byteOut); + + out.writeMerkleTree(testDirectory, originalRoot); + out.flush(); + + final MerkleDataInputStream in = + new DebuggableMerkleDataInputStream(new ByteArrayInputStream(byteOut.toByteArray())); + + final MerkleInternal deserializedRoot = in.readMerkleTree(testDirectory, Integer.MAX_VALUE); + + assertTrue(areTreesEqual(newRoot, deserializedRoot), "deserialized tree should match new root"); + assertNotNull(deserializedRootBeforeMigration.get(), "deserialized root should have been set"); + assertTrue(deserializedRootBeforeMigration.get().isDestroyed(), "deserialized root should have been destroyed"); + + DummyMerkleInternal.resetMigrationMapper(); + } + /** * Disabled on Windows because it does not support changing a file or directory's readability and write-ability. *