Make blocks locked when used + minor cleanup

Consensys · Oct 13, 2023 · a871db0 · a871db0
1 parent 3e48739
commit a871db0
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 8 deletions.
diff --git a/...eum/dataproviders/src/main/java/tech/pegasys/teku/dataproviders/lookup/BlockProvider.java b/...eum/dataproviders/src/main/java/tech/pegasys/teku/dataproviders/lookup/BlockProvider.java
@@ -21,6 +21,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.locks.Lock;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -37,6 +38,22 @@ static BlockProvider fromDynamicMap(Supplier<Map<Bytes32, SignedBeaconBlock>> ma
     return (roots) -> fromMap(mapSupplier.get()).getBlocks(roots);
   }
 
+  static BlockProvider fromMapWithLock(
+      final Map<Bytes32, SignedBeaconBlock> blockMap, final Lock readLock) {
+    return (roots) -> {
+      readLock.lock();
+      try {
+        return SafeFuture.completedFuture(
+            roots.stream()
+                .map(blockMap::get)
+                .filter(Objects::nonNull)
+                .collect(Collectors.toMap(SignedBeaconBlock::getRoot, Function.identity())));
+      } finally {
+        readLock.unlock();
+      }
+    };
+  }
+
   static BlockProvider fromMap(final Map<Bytes32, SignedBeaconBlock> blockMap) {
     return (roots) ->
         SafeFuture.completedFuture(

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
@@ -16,7 +16,7 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static tech.pegasys.teku.dataproviders.generators.StateAtSlotTask.AsyncStateProvider.fromAnchor;
 import static tech.pegasys.teku.dataproviders.lookup.BlockProvider.fromDynamicMap;
-import static tech.pegasys.teku.dataproviders.lookup.BlockProvider.fromMap;
+import static tech.pegasys.teku.dataproviders.lookup.BlockProvider.fromMapWithLock;
 import static tech.pegasys.teku.infrastructure.time.TimeUtilities.secondsToMillis;
 
 import com.google.common.annotations.VisibleForTesting;
@@ -178,7 +178,7 @@ private Store(
                         .getSignedBeaconBlock()
                         .map((b) -> Map.of(b.getRoot(), b))
                         .orElseGet(Collections::emptyMap)),
-            fromMap(this.blocks),
+            fromMapWithLock(this.blocks, readLock),
             blockProvider);
     this.blobSidecarsProvider = blobSidecarsProvider;
     this.earliestBlobSidecarSlotProvider = earliestBlobSidecarSlotProvider;
@@ -332,6 +332,7 @@ public ForkChoiceStrategy getForkChoiceStrategy() {
   }
 
   @Override
+  @VisibleForTesting
   public void clearCaches() {
     states.clear();
     checkpointStates.clear();
@@ -492,12 +493,6 @@ public SafeFuture<Optional<SignedBeaconBlock>> retrieveSignedBlock(final Bytes32
     if (!containsBlock(blockRoot)) {
       return EmptyStoreResults.EMPTY_SIGNED_BLOCK_FUTURE;
     }
-    final Optional<SignedBeaconBlock> inMemoryBlock = getBlockIfAvailable(blockRoot);
-    if (inMemoryBlock.isPresent()) {
-      return SafeFuture.completedFuture(inMemoryBlock);
-    }
-
-    // Retrieve and cache block
     return blockProvider.getBlock(blockRoot);
   }