feat: 24213 - Improve caching for hash chunk preload during reconnects (#24466)

Signed-off-by: Artur Kugal <artur.kugal@swirldslabs.com>

Author: akugal

platform-sdk/swirlds-common/src/main/java/com/swirlds/common/merkle/synchronization/TeachingSynchronizer.java

@@ -111,7 +111,7 @@ public void synchronize() throws InterruptedException {
 
         InterruptedException interruptException = null;
         try (view) {
-            view.startTeacherTasks(this, time, workGroup, in, out);
+            view.startTeacherTasks(time, workGroup, in, out);
             workGroup.waitForTermination();
         } catch (final InterruptedException e) { // NOSONAR: Exception is rethrown below after cleanup.
             interruptException = e;

platform-sdk/swirlds-common/src/main/java/com/swirlds/common/merkle/synchronization/views/TeacherTreeView.java

@@ -2,7 +2,6 @@
 package com.swirlds.common.merkle.synchronization.views;
 
 import com.swirlds.base.time.Time;
-import com.swirlds.common.merkle.synchronization.TeachingSynchronizer;
 import com.swirlds.common.merkle.synchronization.streams.AsyncInputStream;
 import com.swirlds.common.merkle.synchronization.streams.AsyncOutputStream;
 import com.swirlds.common.merkle.synchronization.utility.MerkleSynchronizationException;
@@ -23,17 +22,12 @@ public interface TeacherTreeView extends TeacherHandleQueue, TeacherResponseQueu
      * new custom tree views are encountered, they must be added to {@code subtrees}, although it isn't
      * currently supported by virtual tree views, as nested virtual maps are not supported.
      *
-     * @param teachingSynchronizer the teacher synchronizer
      * @param workGroup the work group to run teaching task(s) in
      * @param in the input stream to read data from learner
      * @param out the output stream to write data to learner
      */
     void startTeacherTasks(
-            final TeachingSynchronizer teachingSynchronizer,
-            final Time time,
-            final StandardWorkGroup workGroup,
-            final AsyncInputStream in,
-            final AsyncOutputStream out);
+            final Time time, final StandardWorkGroup workGroup, final AsyncInputStream in, final AsyncOutputStream out);
 
     /**
      * Write data for a merkle leaf to the stream.

platform-sdk/swirlds-virtualmap/src/main/java/com/swirlds/virtualmap/VirtualMap.java

@@ -34,6 +34,7 @@
 import com.swirlds.metrics.api.Metrics;
 import com.swirlds.virtualmap.config.VirtualMapConfig;
 import com.swirlds.virtualmap.config.VirtualMapReconnectMode;
+import com.swirlds.virtualmap.datasource.DataSourceHashChunkPreloader;
 import com.swirlds.virtualmap.datasource.VirtualDataSource;
 import com.swirlds.virtualmap.datasource.VirtualDataSourceBuilder;
 import com.swirlds.virtualmap.datasource.VirtualHashChunk;
@@ -51,7 +52,6 @@
 import com.swirlds.virtualmap.internal.reconnect.ConcurrentBlockingIterator;
 import com.swirlds.virtualmap.internal.reconnect.LearnerPullVirtualTreeView;
 import com.swirlds.virtualmap.internal.reconnect.LearnerPushVirtualTreeView;
-import com.swirlds.virtualmap.internal.reconnect.NodeTraversalOrder;
 import com.swirlds.virtualmap.internal.reconnect.ParallelSyncTraversalOrder;
 import com.swirlds.virtualmap.internal.reconnect.ReconnectHashLeafFlusher;
 import com.swirlds.virtualmap.internal.reconnect.ReconnectHashListener;
@@ -321,22 +321,13 @@ public static class ClassVersion {
      */
     private VirtualMapMetadata reconnectState;
 
-    private VirtualNodeCache reconnectCache;
-
-    /**
-     * The {@link RecordAccessor} for the state, cache, and data source needed during reconnect.
-     */
-    private RecordAccessor reconnectRecords;
-
     /**
      * During reconnect as a learner, this is the root node in the old learner merkle tree.
      */
     private VirtualMap originalMap;
 
     private ReconnectHashLeafFlusher reconnectFlusher;
 
-    private ReconnectNodeRemover nodeRemover;
-
     private final long fastCopyVersion;
 
     private VirtualMapStatistics statistics;
@@ -423,8 +414,6 @@ private VirtualMap(final VirtualMap source) {
         reconnectHashingFuture = null;
         reconnectHashingStarted = null;
         reconnectIterator = null;
-        reconnectCache = null;
-        reconnectRecords = null;
         pipeline = source.pipeline;
         flushCandidateThreshold.set(source.flushCandidateThreshold.get());
         statistics = source.statistics;
@@ -869,10 +858,6 @@ public void onShutdown(final boolean immediately) {
             hasher.shutdown();
         }
         cache.shutdown();
-        if (reconnectCache != null) {
-            reconnectCache.shutdown();
-            reconnectCache = null;
-        }
         closeDataSource();
     }
 
@@ -1214,6 +1199,19 @@ public void onHashChunkHashed(@NonNull VirtualHashChunk chunk) {
         statistics.recordHash(end - start);
     }
 
+    /**
+     * @return copy of underlying datasource with cache copy flushed into it, and running compaction
+     */
+    public VirtualDataSource detachAsDataSourceCopy() {
+        return pipeline.pausePipelineAndRun("detach", () -> {
+            final Path snapshotPath = dataSourceBuilder.snapshot(null, dataSource);
+            VirtualDataSource dataSourceCopy = dataSourceBuilder.build(getLabel(), snapshotPath, true, false);
+
+            flush(cache.snapshot(), metadata, dataSourceCopy);
+            return dataSourceCopy;
+        });
+    }
+
     /**
      * Prepares a read-only copy so that it may be used even when removed from the pipeline.
      * Can be called only on immutable hashed copy.
@@ -1291,31 +1289,7 @@ private void setupWithOriginalNode(@NonNull final VirtualMap originalMap) {
 
         // Start with empty state, it will be updated from the teacher during reconnect
         reconnectState = new VirtualMapMetadata();
-        reconnectRecords = originalMap.pipeline.pausePipelineAndRun("copy", () -> {
-            // shutdown background compaction on original data source as it is no longer needed to be running as all
-            // data
-            // in that data source is only there as a starting point for reconnect now. So compacting it further is not
-            // helpful and will just burn resources.
-            originalMap.dataSource.stopAndDisableBackgroundCompaction();
-
-            // Take a snapshot, and use the snapshot database as my data source
-            final Path snapshotPath = dataSourceBuilder.snapshot(null, originalMap.dataSource);
-            this.dataSource = dataSourceBuilder.build(originalMap.getLabel(), snapshotPath, true, false);
-
-            // The old map's cache is going to become immutable, but that's OK, because the old map
-            // will NEVER be updated again.
-            assert originalMap.isHashed() : "The system should have made sure this was hashed by this point!";
-            final VirtualNodeCache snapshotCache = originalMap.cache.snapshot();
-            flush(snapshotCache, originalMap.metadata, this.dataSource);
-
-            final int hashChunkHeight = dataSource.getHashChunkHeight();
-            reconnectCache = new VirtualNodeCache(
-                    virtualMapConfig,
-                    hashChunkHeight,
-                    dataSource::loadHashChunk,
-                    originalMap.cache.getFastCopyVersion());
-            return new RecordAccessor(reconnectState, hashChunkHeight, reconnectCache, dataSource);
-        });
+        this.dataSource = originalMap.detachAsDataSourceCopy();
 
         // Set up the VirtualHasher which we will use during reconnect.
         // Initial timeout is intentionally very long, timeout is reduced once we receive the first leaf in the tree.
@@ -1358,7 +1332,7 @@ public LearnerTreeView buildLearnerView(
         final VirtualMapMetadata originalState = originalMap.getMetadata();
         reconnectFlusher =
                 new ReconnectHashLeafFlusher(dataSource, virtualMapConfig.reconnectFlushInterval(), statistics);
-        nodeRemover = new ReconnectNodeRemover(
+        final ReconnectNodeRemover nodeRemover = new ReconnectNodeRemover(
                 originalMap.getRecords(),
                 originalState.getFirstLeafPath(),
                 originalState.getLastLeafPath(),
@@ -1367,42 +1341,36 @@ public LearnerTreeView buildLearnerView(
             case VirtualMapReconnectMode.PUSH ->
                 new LearnerPushVirtualTreeView(
                         this, originalMap.records, originalState, reconnectState, nodeRemover, mapStats);
-            case VirtualMapReconnectMode.PULL_TOP_TO_BOTTOM -> {
-                final NodeTraversalOrder topToBottom = new TopToBottomTraversalOrder();
-                yield new LearnerPullVirtualTreeView(
+            case VirtualMapReconnectMode.PULL_TOP_TO_BOTTOM ->
+                new LearnerPullVirtualTreeView(
                         reconnectConfig,
                         this,
                         originalMap.records,
                         originalState,
                         reconnectState,
                         nodeRemover,
-                        topToBottom,
+                        new TopToBottomTraversalOrder(),
                         mapStats);
-            }
-            case VirtualMapReconnectMode.PULL_TWO_PHASE_PESSIMISTIC -> {
-                final NodeTraversalOrder twoPhasePessimistic = new TwoPhasePessimisticTraversalOrder();
-                yield new LearnerPullVirtualTreeView(
+            case VirtualMapReconnectMode.PULL_TWO_PHASE_PESSIMISTIC ->
+                new LearnerPullVirtualTreeView(
                         reconnectConfig,
                         this,
                         originalMap.records,
                         originalState,
                         reconnectState,
                         nodeRemover,
-                        twoPhasePessimistic,
+                        new TwoPhasePessimisticTraversalOrder(),
                         mapStats);
-            }
-            case VirtualMapReconnectMode.PULL_PARALLEL_SYNC -> {
-                final NodeTraversalOrder parallelSync = new ParallelSyncTraversalOrder();
-                yield new LearnerPullVirtualTreeView(
+            case VirtualMapReconnectMode.PULL_PARALLEL_SYNC ->
+                new LearnerPullVirtualTreeView(
                         reconnectConfig,
                         this,
                         originalMap.records,
                         originalState,
                         reconnectState,
                         nodeRemover,
-                        parallelSync,
+                        new ParallelSyncTraversalOrder(),
                         mapStats);
-            }
             default ->
                 throw new UnsupportedOperationException("Unknown reconnect mode: " + virtualMapConfig.reconnectMode());
         };
@@ -1442,8 +1410,9 @@ public void handleReconnectLeaf(@NonNull final VirtualLeafBytes<?> leafRecord) {
 
     public void prepareReconnectHashing(final long firstLeafPath, final long lastLeafPath) {
         assert reconnectFlusher != null : "Cannot prepare reconnect hashing, since reconnect is not started";
+        final DataSourceHashChunkPreloader hashChunkPreloader = new DataSourceHashChunkPreloader(dataSource);
         // The hash listener will be responsible for flushing stuff to the reconnect data source
-        final ReconnectHashListener hashListener = new ReconnectHashListener(reconnectFlusher);
+        final ReconnectHashListener hashListener = new ReconnectHashListener(reconnectFlusher, hashChunkPreloader);
 
         // This background thread will be responsible for hashing the tree and sending the
         // data to the hash listener to flush.
@@ -1452,7 +1421,7 @@ public void prepareReconnectHashing(final long firstLeafPath, final long lastLea
                 .setThreadName("hasher")
                 .setRunnable(() -> reconnectHashingFuture.complete(hasher.hash(
                         dataSource.getHashChunkHeight(),
-                        reconnectCache::preloadHashChunk,
+                        hashChunkPreloader,
                         reconnectIterator,
                         firstLeafPath,
                         lastLeafPath,
@@ -1483,14 +1452,9 @@ public void endLearnerReconnect() {
                 logger.warn(RECONNECT.getMarker(), "virtual map hashing thread was never started");
             }
             logger.info(RECONNECT.getMarker(), "call postInit()");
-            nodeRemover = null;
             originalMap = null;
             metadata = new VirtualMapMetadata(reconnectState.getSize());
             postInit();
-            if (reconnectCache != null) {
-                reconnectCache.shutdown();
-                reconnectCache = null;
-            }
         } catch (ExecutionException e) {
             throw new MerkleSynchronizationException(e);
         } catch (InterruptedException e) {

platform-sdk/swirlds-virtualmap/src/main/java/com/swirlds/virtualmap/datasource/DataSourceHashChunkPreloader.java

@@ -0,0 +1,57 @@
+// SPDX-License-Identifier: Apache-2.0
+package com.swirlds.virtualmap.datasource;
+
+import edu.umd.cs.findbugs.annotations.NonNull;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.LongFunction;
+
+/**
+ * A {@link LongFunction<VirtualHashChunk>} implementation that preloads hash chunks from a {@link VirtualDataSource} and caches them in memory.
+ * {@link #clearCache(long)} can be used to clear hash chunk cache when chunk is processed.
+ */
+public class DataSourceHashChunkPreloader implements LongFunction<VirtualHashChunk> {
+
+    private final VirtualDataSource dataSource;
+    private final int hashChunkHeight;
+
+    /**
+     * Hash chunk id to chunk map.
+     */
+    private final Map<Long, VirtualHashChunk> hashChunkMap = new ConcurrentHashMap<>();
+
+    public DataSourceHashChunkPreloader(@NonNull VirtualDataSource dataSource) {
+        this.dataSource = Objects.requireNonNull(dataSource, "dataSource must not be null");
+        hashChunkHeight = dataSource.getHashChunkHeight();
+    }
+
+    /**
+     * Clear the cached hash chunk with the given id.
+     * @param hashChunkId hash chunk id
+     */
+    public void clearCache(long hashChunkId) {
+        hashChunkMap.remove(hashChunkId);
+    }
+
+    @Override
+    public VirtualHashChunk apply(long path) {
+        final long hashChunkId = VirtualHashChunk.chunkPathToChunkId(path, hashChunkHeight);
+
+        return hashChunkMap.computeIfAbsent(hashChunkId, id -> {
+            VirtualHashChunk chunk;
+            try {
+                chunk = dataSource.loadHashChunk(id);
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+            if (chunk == null) {
+                final long hashChunkPath = VirtualHashChunk.chunkIdToChunkPath(hashChunkId, hashChunkHeight);
+                chunk = new VirtualHashChunk(hashChunkPath, hashChunkHeight);
+            }
+            return chunk;
+        });
+    }
+}

platform-sdk/swirlds-virtualmap/src/main/java/com/swirlds/virtualmap/internal/reconnect/ReconnectHashListener.java

@@ -3,6 +3,7 @@
 
 import static java.util.Objects.requireNonNull;
 
+import com.swirlds.virtualmap.datasource.DataSourceHashChunkPreloader;
 import com.swirlds.virtualmap.datasource.VirtualHashChunk;
 import com.swirlds.virtualmap.datasource.VirtualLeafBytes;
 import com.swirlds.virtualmap.internal.hash.VirtualHashListener;
@@ -23,14 +24,18 @@
 public class ReconnectHashListener implements VirtualHashListener {
 
     private final ReconnectHashLeafFlusher flusher;
+    private final DataSourceHashChunkPreloader hashChunkPreloader;
 
     /**
      * Create a new {@link ReconnectHashListener}.
      *
      * @param flusher Hash / leaf flusher to use to flush data to disk
      */
-    public ReconnectHashListener(@NonNull final ReconnectHashLeafFlusher flusher) {
+    public ReconnectHashListener(
+            @NonNull final ReconnectHashLeafFlusher flusher,
+            @NonNull final DataSourceHashChunkPreloader hashChunkPreloader) {
         this.flusher = requireNonNull(flusher);
+        this.hashChunkPreloader = requireNonNull(hashChunkPreloader);
     }
 
     /**
@@ -47,6 +52,7 @@ public void onHashingStarted(long firstLeafPath, long lastLeafPath) {
     @Override
     public void onHashChunkHashed(@NonNull final VirtualHashChunk chunk) {
         flusher.updateHashChunk(chunk);
+        hashChunkPreloader.clearCache(chunk.getChunkId());
     }
 
     /**

platform-sdk/swirlds-virtualmap/src/main/java/com/swirlds/virtualmap/internal/reconnect/TeacherPullVirtualTreeView.java

@@ -4,7 +4,6 @@
 import static com.swirlds.logging.legacy.LogMarker.EXCEPTION;
 
 import com.swirlds.base.time.Time;
-import com.swirlds.common.merkle.synchronization.TeachingSynchronizer;
 import com.swirlds.common.merkle.synchronization.streams.AsyncInputStream;
 import com.swirlds.common.merkle.synchronization.streams.AsyncOutputStream;
 import com.swirlds.common.merkle.synchronization.views.TeacherTreeView;
@@ -57,7 +56,6 @@ public TeacherPullVirtualTreeView(final ReconnectConfig reconnectConfig, final V
 
     @Override
     public void startTeacherTasks(
-            final TeachingSynchronizer teachingSynchronizer,
             final Time time,
             final StandardWorkGroup workGroup,
             final AsyncInputStream in,

platform-sdk/swirlds-virtualmap/src/main/java/com/swirlds/virtualmap/internal/reconnect/TeacherPushVirtualTreeView.java

@@ -8,7 +8,6 @@
 import static com.swirlds.virtualmap.internal.Path.getRightChildPath;
 
 import com.swirlds.base.time.Time;
-import com.swirlds.common.merkle.synchronization.TeachingSynchronizer;
 import com.swirlds.common.merkle.synchronization.streams.AsyncInputStream;
 import com.swirlds.common.merkle.synchronization.streams.AsyncOutputStream;
 import com.swirlds.common.merkle.synchronization.task.TeacherPushReceiveTask;
@@ -120,7 +119,6 @@ public TeacherPushVirtualTreeView(final ReconnectConfig reconnectConfig, final V
 
     @Override
     public void startTeacherTasks(
-            final TeachingSynchronizer teachingSynchronizer,
             final Time time,
             final StandardWorkGroup workGroup,
             final AsyncInputStream in,