fix: 16748 Fixed serialization for AddressBookTestingToolState (#16799)…

… (#16848) Signed-off-by: Ivan Malygin <[email protected]>
hashgraph · Dec 2, 2024 · 34633b6 · 34633b6
1 parent cd75b68
commit 34633b6
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 6 deletions.
diff --git a/...okTestingTool/src/main/java/com/swirlds/demo/addressbook/AddressBookTestingToolState.java b/...okTestingTool/src/main/java/com/swirlds/demo/addressbook/AddressBookTestingToolState.java
@@ -58,6 +58,7 @@
 import com.swirlds.platform.system.address.AddressBookUtils;
 import com.swirlds.platform.system.events.ConsensusEvent;
 import com.swirlds.platform.system.transaction.ConsensusTransaction;
+import com.swirlds.state.merkle.singleton.StringLeaf;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
 import java.io.File;
@@ -98,6 +99,9 @@ private static class ClassVersion {
 
     private static final long CLASS_ID = 0xf052378c7364ef47L;
 
+    private static final int RUNNING_SUM_INDEX = 1;
+    private static final int ROUND_HANDLED_INDEX = 2;
+
     private NodeId selfId;
 
     /** false until the test scenario has been validated, true afterwards. */
@@ -185,6 +189,17 @@ public void init(
         throwIfImmutable();
 
         this.selfId = platform.getSelfId();
+
+        final StringLeaf runningSumLeaf = getChild(RUNNING_SUM_INDEX);
+        if (runningSumLeaf != null && runningSumLeaf.getLabel() != null) {
+            this.runningSum = Long.parseLong(runningSumLeaf.getLabel());
+            logger.info(STARTUP.getMarker(), "State initialized with state long {}.", runningSum);
+        }
+        final StringLeaf roundsHandledLeaf = getChild(ROUND_HANDLED_INDEX);
+        if (roundsHandledLeaf != null && roundsHandledLeaf.getLabel() != null) {
+            this.roundsHandled = Long.parseLong(roundsHandledLeaf.getLabel());
+            logger.info(STARTUP.getMarker(), "State initialized with {} rounds handled.", roundsHandled);
+        }
     }
 
     /**
@@ -206,6 +221,7 @@ public void handleConsensusRound(@NonNull final Round round, @NonNull final Plat
         }
 
         roundsHandled++;
+        setChild(ROUND_HANDLED_INDEX, new StringLeaf(Long.toString(roundsHandled)));
 
         final Iterator<ConsensusEvent> eventIterator = round.iterator();
 
@@ -236,6 +252,7 @@ private void handleTransaction(@NonNull final ConsensusTransaction transaction)
         final int delta =
                 ByteUtils.byteArrayToInt(transaction.getApplicationTransaction().toByteArray(), 0);
         runningSum += delta;
+        setChild(RUNNING_SUM_INDEX, new StringLeaf(Long.toString(runningSum)));
     }
 
     /**

diff --git a/...latform-core/src/main/java/com/swirlds/platform/state/snapshot/SignedStateFileReader.java b/...latform-core/src/main/java/com/swirlds/platform/state/snapshot/SignedStateFileReader.java
@@ -95,7 +95,7 @@ public static List<SavedStateInfo> getSavedStateFiles(
                     (final MerkleDataInputStream in) -> {
                         readAndCheckSigSetFileVersion(in);
                         final SigSet sigSet = in.readSerializable();
-                        return new MerkleTreeSnapshotReader.StateFileData(data.state(), data.hash(), sigSet);
+                        return new MerkleTreeSnapshotReader.StateFileData(data.stateRoot(), data.hash(), sigSet);
                     });
         } else {
             normalizedData = data;
@@ -104,7 +104,7 @@ public static List<SavedStateInfo> getSavedStateFiles(
         final SignedState newSignedState = new SignedState(
                 configuration,
                 CryptoStatic::verifySignature,
-                (MerkleRoot) normalizedData.state(),
+                (MerkleRoot) normalizedData.stateRoot(),
                 "SignedStateFileReader.readStateFile()",
                 false,
                 false,

diff --git a/platform-sdk/swirlds-state-impl/src/main/java/com/swirlds/state/merkle/MerkleStateRoot.java b/platform-sdk/swirlds-state-impl/src/main/java/com/swirlds/state/merkle/MerkleStateRoot.java
@@ -908,6 +908,7 @@ public void createSnapshot(@NonNull final Path targetPath) {
      */
     @Override
     public MerkleStateRoot<?> loadSnapshot(@NonNull Path targetPath) throws IOException {
-        return MerkleTreeSnapshotReader.readStateFileData(targetPath).state();
+        return (MerkleStateRoot<?>)
+                MerkleTreeSnapshotReader.readStateFileData(targetPath).stateRoot();
     }
 }
diff --git a/...k/swirlds-state-impl/src/main/java/com/swirlds/state/merkle/MerkleTreeSnapshotReader.java b/...k/swirlds-state-impl/src/main/java/com/swirlds/state/merkle/MerkleTreeSnapshotReader.java
@@ -20,6 +20,7 @@
 
 import com.swirlds.common.crypto.Hash;
 import com.swirlds.common.io.streams.MerkleDataInputStream;
+import com.swirlds.common.merkle.impl.PartialNaryMerkleInternal;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
 import java.io.BufferedInputStream;
@@ -61,11 +62,12 @@ public class MerkleTreeSnapshotReader {
 
     /**
      * This is a helper class to hold the data read from a state file.
-     * @param state the Merkle tree state
+     * @param stateRoot the root of Merkle tree state
      * @param hash the hash of the state
      * @param sigSet the signature set
      */
-    public record StateFileData(@NonNull MerkleStateRoot<?> state, @NonNull Hash hash, @Nullable SigSet sigSet) {}
+    public record StateFileData(
+            @NonNull PartialNaryMerkleInternal stateRoot, @NonNull Hash hash, @Nullable SigSet sigSet) {}
 
     /**
      * Reads a state file from disk
@@ -100,7 +102,7 @@ private static StateFileData readStateFileDataV1(
             @NonNull final Path stateFile, @NonNull final MerkleDataInputStream in, @NonNull final Path directory)
             throws IOException {
         try {
-            final MerkleStateRoot<?> state = in.readMerkleTree(directory, MAX_MERKLE_NODES_IN_STATE);
+            final PartialNaryMerkleInternal state = in.readMerkleTree(directory, MAX_MERKLE_NODES_IN_STATE);
             final Hash hash = in.readSerializable();
             final SigSet sigSet = in.readSerializable();
             return new StateFileData(state, hash, sigSet);