Store only most recent blocks in cache

zilm13 · Sep 19, 2023 · 668ce1e · 668ce1e
1 parent db94bcc
commit 668ce1e
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 25 deletions.
diff --git a/storage/src/main/java/tech/pegasys/teku/storage/store/Store.java b/storage/src/main/java/tech/pegasys/teku/storage/store/Store.java
@@ -83,7 +83,7 @@ class Store implements UpdatableStore {
   private final Lock readLock = lock.readLock();
 
   private final MetricsSystem metricsSystem;
-  private Optional<SettableGauge> blockCountGauge = Optional.empty();
+  Optional<SettableGauge> blockCountGauge = Optional.empty();
 
   private Optional<SettableGauge> epochStatesCountGauge = Optional.empty();
 
@@ -495,13 +495,7 @@ public SafeFuture<Optional<SignedBeaconBlock>> retrieveSignedBlock(final Bytes32
     }
 
     // Retrieve and cache block
-    return blockProvider
-        .getBlock(blockRoot)
-        .thenApply(
-            block -> {
-              block.ifPresent(this::putBlock);
-              return block;
-            });
+    return blockProvider.getBlock(blockRoot);
   }
 
   @Override
@@ -617,7 +611,6 @@ private SafeFuture<Optional<SignedBlockAndState>> getAndCacheBlockAndState(
                       signedBeaconBlock ->
                           SafeFuture.completedFuture(Optional.of(signedBeaconBlock)))
                   .orElseGet(() -> blockProvider.getBlock(blockRoot))
-                  .thenPeek(block -> block.ifPresent(this::putBlock))
                   .thenApply(
                       block -> block.map(b -> new SignedBlockAndState(b, res.get().getState())));
             });
@@ -815,19 +808,6 @@ private boolean isSlotAtNthEpochBoundary(
         .orElse(false);
   }
 
-  private void putBlock(final SignedBeaconBlock block) {
-    final Lock writeLock = lock.writeLock();
-    writeLock.lock();
-    try {
-      if (containsBlock(block.getRoot())) {
-        blocks.put(block.getRoot(), block);
-        blockCountGauge.ifPresent(gauge -> gauge.set(blocks.size()));
-      }
-    } finally {
-      writeLock.unlock();
-    }
-  }
-
   @VisibleForTesting
   Optional<Map<Bytes32, StateAndBlockSummary>> getEpochStates() {
     return maybeEpochStates;

diff --git a/storage/src/main/java/tech/pegasys/teku/storage/store/StoreTransactionUpdates.java b/storage/src/main/java/tech/pegasys/teku/storage/store/StoreTransactionUpdates.java
@@ -16,6 +16,7 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.common.collect.Maps;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -109,7 +110,11 @@ public void applyToStore(final Store store, final UpdateResult updateResult) {
     tx.genesisTime.ifPresent(value -> store.genesisTime = value);
     tx.justifiedCheckpoint.ifPresent(store::updateJustifiedCheckpoint);
     tx.bestJustifiedCheckpoint.ifPresent(store::updateBestJustifiedCheckpoint);
-    hotBlocks.forEach((root, value) -> store.blocks.put(root, value.getBlock()));
+    hotBlocks.values().stream()
+        .sorted(Comparator.comparing(BlockAndCheckpoints::getSlot))
+        .map(BlockAndCheckpoints::getBlock)
+        .forEach(block -> store.blocks.put(block.getRoot(), block));
+    store.blockCountGauge.ifPresent(gauge -> gauge.set(store.blocks.size()));
     store.states.cacheAll(Maps.transformValues(hotBlockAndStates, this::blockAndStateAsSummary));
     if (optimisticTransitionBlockRootSet) {
       store.finalizedOptimisticTransitionPayload =

diff --git a/storage/src/test/java/tech/pegasys/teku/storage/store/AbstractStoreTest.java b/storage/src/test/java/tech/pegasys/teku/storage/store/AbstractStoreTest.java
@@ -47,6 +47,7 @@ public abstract class AbstractStoreTest {
   protected final Spec spec = TestSpecFactory.createMinimalDeneb();
   protected final StorageUpdateChannel storageUpdateChannel = new StubStorageUpdateChannel();
   protected final ChainBuilder chainBuilder = ChainBuilder.create(spec);
+  protected final StoreConfig defaultStoreConfig = StoreConfig.createDefault();
 
   protected void processChainWithLimitedCache(
       BiConsumer<UpdatableStore, SignedBlockAndState> chainProcessor) {
@@ -131,7 +132,7 @@ protected void addBlocks(final UpdatableStore store, final List<SignedBlockAndSt
   }
 
   protected UpdatableStore createGenesisStore() {
-    return createGenesisStore(StoreConfig.createDefault());
+    return createGenesisStore(defaultStoreConfig);
   }
 
   protected UpdatableStore createGenesisStore(final StoreConfig pruningOptions) {

diff --git a/storage/src/test/java/tech/pegasys/teku/storage/store/StoreTest.java b/storage/src/test/java/tech/pegasys/teku/storage/store/StoreTest.java
@@ -73,7 +73,7 @@ public void create_timeLessThanGenesisTime() {
                     genesisCheckpoint,
                     Collections.emptyMap(),
                     Collections.emptyMap(),
-                    StoreConfig.createDefault()))
+                    defaultStoreConfig))
         .isInstanceOf(IllegalArgumentException.class)
         .hasMessageContaining("Time must be greater than or equal to genesisTime");
   }
@@ -298,6 +298,40 @@ public void retrieveFinalizedCheckpointAndState() {
     assertThatThrownBy(result::get).hasCauseInstanceOf(InvalidCheckpointException.class);
   }
 
+  @Test
+  public void shouldKeepOnlyMostRecentBlocksInBlockCache() {
+    final UpdatableStore store = createGenesisStore();
+    final UInt64 epoch = UInt64.valueOf(5);
+    final UInt64 startSlot = spec.computeStartSlotAtEpoch(epoch);
+    chainBuilder.generateBlocksUpToSlot(startSlot);
+
+    // Add blocks
+    final StoreTransaction tx = store.startTransaction(new StubStorageUpdateChannel());
+    chainBuilder
+        .streamBlocksAndStates()
+        .forEach(
+            blockAndState ->
+                tx.putBlockAndState(
+                    blockAndState, spec.calculateBlockCheckpoints(blockAndState.getState())));
+    tx.commit().join();
+    final List<SignedBlockAndState> last32 =
+        chainBuilder
+            .streamBlocksAndStates()
+            .dropWhile(
+                signedBlockAndState ->
+                    signedBlockAndState
+                        .getSlot()
+                        .isLessThanOrEqualTo(
+                            chainBuilder
+                                .getLatestBlockAndState()
+                                .getSlot()
+                                .minus(defaultStoreConfig.getBlockCacheSize())))
+            .toList();
+    for (final SignedBlockAndState signedBlockAndState : last32) {
+      assertThat(store.getBlockIfAvailable(signedBlockAndState.getRoot())).isPresent();
+    }
+  }
+
   private void testApplyChangesWhenTransactionCommits(final boolean withInterleavedTransaction) {
     final UpdatableStore store = createGenesisStore();
     final UInt64 epoch3 = UInt64.valueOf(4);