Perf: Massively improved the performance of the replayMilestones method

src/main/java/com/iota/iri/service/snapshot/Snapshot.java

@@ -71,4 +71,15 @@ public interface Snapshot extends SnapshotMetaData, SnapshotState {
      * @param snapshot the new snapshot details that shall overwrite the current ones
      */
     void update(Snapshot snapshot);
+
+    /**
+     * Creates a deep clone of the Snapshot which can be modified without affecting the values of the original
+     * one.<br />
+     * <br />
+     * Since the data structures inside the Snapshot are extremely big, this method is relatively expensive and should
+     * only be used when it is really necessary.<br />
+     *
+     * @return a deep clone of the snapshot
+     */
+    Snapshot clone();
 }

src/main/java/com/iota/iri/service/snapshot/SnapshotService.java

@@ -17,15 +17,15 @@
  */
 public interface SnapshotService {
     /**
-     * This method applies the balance changes that are introduced by future milestones to the current Snapshot.
-     *
+     * This method applies the balance changes that are introduced by future milestones to the current Snapshot.<br />
+     * <br />
      * It iterates over the milestone indexes starting from the current index to the target index and applies all found
      * milestone balances. If it can not find a milestone for a certain index it keeps track of these skipped
      * milestones, which allows us to revert the changes even if the missing milestone was received and processed in the
      * mean time. If the application of changes fails, we restore the state of the snapshot to the one it had before the
-     * application attempt so this method only modifies the Snapshot if it succeeds.
-     *
-     * Note: the changes done by this method can be reverted by using {@link #rollBackMilestones(Snapshot, int)}
+     * application attempt so this method only modifies the Snapshot if it succeeds.<br />
+     * <br />
+     * Note: the changes done by this method can be reverted by using {@link #rollBackMilestones(Snapshot, int)}<br />
      *
      * @param snapshot the Snapshot that shall get modified
      * @param targetMilestoneIndex the index of the milestone that should be applied

src/main/java/com/iota/iri/service/snapshot/impl/SnapshotImpl.java

@@ -9,6 +9,7 @@
 
 import java.util.HashSet;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
@@ -55,15 +56,17 @@ public SnapshotImpl(SnapshotState state, SnapshotMetaData metaData) {
     }
 
     /**
-     * Creates a deep clone of the passed in {@link Snapshot}.
+     * Creates a deep clone of the passed in snapshot.
      *
      * @param snapshot object that shall be cloned
      */
-    public SnapshotImpl(Snapshot snapshot) {
+    private SnapshotImpl(SnapshotImpl snapshot) {
         this(
-            new SnapshotStateImpl(((SnapshotImpl) snapshot).state),
-            new SnapshotMetaDataImpl(((SnapshotImpl) snapshot).metaData)
+            new SnapshotStateImpl(snapshot.state),
+            new SnapshotMetaDataImpl(snapshot.metaData)
         );
+
+        skippedMilestones.addAll(snapshot.skippedMilestones);
     }
 
     /**
@@ -129,8 +132,32 @@ public void update(Snapshot snapshot) {
         }
     }
 
-    //region [THREAD-SAFE METADATA METHODS] ////////////////////////////////////////////////////////////////////////////
+    @Override
+    public int hashCode() {
+        return Objects.hash(getClass(), state.hashCode(), metaData.hashCode(), skippedMilestones);
+    }
 
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == this) {
+            return true;
+        }
+
+        if (obj == null || !getClass().equals(obj.getClass())) {
+            return false;
+        }
+
+        return Objects.equals(state, ((SnapshotImpl) obj).state) &&
+               Objects.equals(metaData, ((SnapshotImpl) obj).metaData) &&
+               Objects.equals(skippedMilestones, ((SnapshotImpl) obj).skippedMilestones);
+    }
+
+    @Override
+    public SnapshotImpl clone() {
+        return new SnapshotImpl(this);
+    }
+
+    //region [THREAD-SAFE METADATA METHODS] ////////////////////////////////////////////////////////////////////////////
 
     /**
      * {@inheritDoc}

src/main/java/com/iota/iri/service/snapshot/impl/SnapshotMetaDataImpl.java

@@ -5,6 +5,7 @@
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * Implements the basic contract of the {@link SnapshotMetaData} interface.
@@ -250,4 +251,31 @@ public void update(SnapshotMetaData newMetaData) {
         setSolidEntryPoints(new HashMap<>(newMetaData.getSolidEntryPoints()));
         setSeenMilestones(new HashMap<>(newMetaData.getSeenMilestones()));
     }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(getClass(), initialHash, initialIndex, initialTimestamp, hash, index, timestamp,
+                solidEntryPoints, seenMilestones);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == this) {
+            return true;
+        }
+
+        if (obj == null || !getClass().equals(obj.getClass())) {
+            return false;
+        }
+
+        return Objects.equals(initialHash, ((SnapshotMetaDataImpl) obj).initialHash) &&
+               Objects.equals(initialIndex, ((SnapshotMetaDataImpl) obj).initialIndex) &&
+               Objects.equals(initialTimestamp, ((SnapshotMetaDataImpl) obj).initialTimestamp) &&
+               Objects.equals(hash, ((SnapshotMetaDataImpl) obj).hash) &&
+               Objects.equals(index, ((SnapshotMetaDataImpl) obj).index) &&
+               Objects.equals(timestamp, ((SnapshotMetaDataImpl) obj).timestamp) &&
+               Objects.equals(solidEntryPoints, ((SnapshotMetaDataImpl) obj).solidEntryPoints) &&
+               Objects.equals(seenMilestones, ((SnapshotMetaDataImpl) obj).seenMilestones);
+
+    }
 }

src/main/java/com/iota/iri/service/snapshot/impl/SnapshotProviderImpl.java

@@ -164,7 +164,7 @@ private void loadSnapshots() throws SnapshotException {
             initialSnapshot = loadBuiltInSnapshot();
         }
 
-        latestSnapshot = new SnapshotImpl(initialSnapshot);
+        latestSnapshot = initialSnapshot.clone();
     }
 
     /**
@@ -259,7 +259,7 @@ private Snapshot loadBuiltInSnapshot() throws SnapshotException {
             );
         }
 
-        return new SnapshotImpl(builtinSnapshot);
+        return builtinSnapshot.clone();
     }
 
     //endregion ////////////////////////////////////////////////////////////////////////////////////////////////////////

src/main/java/com/iota/iri/service/snapshot/impl/SnapshotServiceImpl.java

@@ -98,12 +98,17 @@ public SnapshotServiceImpl init(Tangle tangle, SnapshotProvider snapshotProvider
 
     /**
      * {@inheritDoc}
+     * <br />
+     * To increase the performance of this operation, we do not apply every single milestone separately but first
+     * accumulate all the necessary changes and then apply it to the snapshot in a single run. This allows us to
+     * modify its values without having to create a "copy" of the initial state to possibly roll back the changes if
+     * anything unexpected happens (creating a backup of the state requires a lot of memory).<br />
      */
     @Override
     public void replayMilestones(Snapshot snapshot, int targetMilestoneIndex) throws SnapshotException {
-        snapshot.lockWrite();
-
-        Snapshot snapshotBeforeChanges = new SnapshotImpl(snapshot);
+        Map<Hash, Long> balanceChanges = new HashMap<>();
+        Set<Integer> skippedMilestones = new HashSet<>();
+        MilestoneViewModel lastAppliedMilestone = null;
 
         try {
             for (int currentMilestoneIndex = snapshot.getIndex() + 1; currentMilestoneIndex <= targetMilestoneIndex;
@@ -113,30 +118,41 @@ public void replayMilestones(Snapshot snapshot, int targetMilestoneIndex) throws
                 if (currentMilestone != null) {
                     StateDiffViewModel stateDiffViewModel = StateDiffViewModel.load(tangle, currentMilestone.getHash());
                     if(!stateDiffViewModel.isEmpty()) {
-                        snapshot.applyStateDiff(new SnapshotStateDiffImpl(stateDiffViewModel.getDiff()));
+                        stateDiffViewModel.getDiff().forEach((address, change) -> {
+                            balanceChanges.compute(address, (k, balance) -> (balance == null ? 0 : balance) + change);
+                        });
                     }
 
-                    snapshot.setIndex(currentMilestone.index());
-                    snapshot.setHash(currentMilestone.getHash());
+                    lastAppliedMilestone = currentMilestone;
+                } else {
+                    skippedMilestones.add(currentMilestoneIndex);
+                }
+            }
+
+            if (lastAppliedMilestone != null) {
+                try {
+                    snapshot.lockWrite();
 
-                    TransactionViewModel currentMilestoneTransaction = TransactionViewModel.fromHash(tangle,
-                            currentMilestone.getHash());
+                    snapshot.applyStateDiff(new SnapshotStateDiffImpl(balanceChanges));
 
-                    if(currentMilestoneTransaction != null &&
-                            currentMilestoneTransaction.getType() != TransactionViewModel.PREFILLED_SLOT) {
+                    snapshot.setIndex(lastAppliedMilestone.index());
+                    snapshot.setHash(lastAppliedMilestone.getHash());
 
-                        snapshot.setTimestamp(currentMilestoneTransaction.getTimestamp());
+                    TransactionViewModel milestoneTransaction = TransactionViewModel.fromHash(tangle,
+                            lastAppliedMilestone.getHash());
+                    if(milestoneTransaction.getType() != TransactionViewModel.PREFILLED_SLOT) {
+                        snapshot.setTimestamp(milestoneTransaction.getTimestamp());
                     }
-                } else {
-                    snapshot.addSkippedMilestone(currentMilestoneIndex);
+
+                    for (int skippedMilestoneIndex : skippedMilestones) {
+                        snapshot.addSkippedMilestone(skippedMilestoneIndex);
+                    }
+                } finally {
+                    snapshot.unlockWrite();
                 }
             }
         } catch (Exception e) {
-            snapshot.update(snapshotBeforeChanges);
-
-            throw new SnapshotException("failed to replay the the state of the ledger", e);
-        } finally {
-            snapshot.unlockWrite();
+            throw new SnapshotException("failed to replay the state of the ledger", e);
         }
     }
 
@@ -151,7 +167,7 @@ public void rollBackMilestones(Snapshot snapshot, int targetMilestoneIndex) thro
 
         snapshot.lockWrite();
 
-        Snapshot snapshotBeforeChanges = new SnapshotImpl(snapshot);
+        Snapshot snapshotBeforeChanges = snapshot.clone();
 
         try {
             boolean rollbackSuccessful = true;
@@ -216,11 +232,11 @@ public Snapshot generateSnapshot(LatestMilestoneTracker latestMilestoneTracker,
                     targetMilestone.index());
 
             if (distanceFromInitialSnapshot <= distanceFromLatestSnapshot) {
-                snapshot = new SnapshotImpl(snapshotProvider.getInitialSnapshot());
+                snapshot = snapshotProvider.getInitialSnapshot().clone();
 
                 replayMilestones(snapshot, targetMilestone.index());
             } else {
-                snapshot = new SnapshotImpl(snapshotProvider.getLatestSnapshot());
+                snapshot = snapshotProvider.getLatestSnapshot().clone();
 
                 rollBackMilestones(snapshot, targetMilestone.index() + 1);
             }

src/main/java/com/iota/iri/service/snapshot/impl/SnapshotStateImpl.java

@@ -16,6 +16,7 @@
 import java.nio.file.Paths;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Objects;
 import java.util.stream.Collectors;
 
 /**
@@ -112,6 +113,10 @@ public void applyStateDiff(SnapshotStateDiff diff) throws SnapshotException {
             if (balances.computeIfPresent(addressHash, (hash, aLong) -> balance + aLong) == null) {
                 balances.putIfAbsent(addressHash, balance);
             }
+
+            if (balances.get(addressHash) == 0) {
+                balances.remove(addressHash);
+            }
         });
     }
 
@@ -130,6 +135,24 @@ public SnapshotState patchedState(SnapshotStateDiff snapshotStateDiff) {
         return new SnapshotStateImpl(patchedBalances);
     }
 
+    @Override
+    public int hashCode() {
+        return Objects.hash(getClass(), balances);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == this) {
+            return true;
+        }
+
+        if (obj == null || !getClass().equals(obj.getClass())) {
+            return false;
+        }
+
+        return Objects.equals(balances, ((SnapshotStateImpl) obj).balances);
+    }
+
     /**
      * Returns all addresses that have a negative balance.
      *

src/test/java/com/iota/iri/service/snapshot/impl/SnapshotMockUtils.java

@@ -54,7 +54,7 @@ public static SnapshotProvider mockSnapshotProvider(SnapshotProvider snapshotPro
                 new SnapshotMetaDataImpl(genesisHash, milestoneStartIndex, genesisTimestamp, solidEntryPoints,
                         new HashMap<>())
         );
-        Snapshot latestSnapshot = new SnapshotImpl(initialSnapshot);
+        Snapshot latestSnapshot = initialSnapshot.clone();
 
         Mockito.when(snapshotProvider.getInitialSnapshot()).thenReturn(initialSnapshot);
         Mockito.when(snapshotProvider.getLatestSnapshot()).thenReturn(latestSnapshot);