edge cases fix

Consensys · May 29, 2024 · c8054b0 · c8054b0
1 parent a78130c
commit c8054b0
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/...gasys/teku/networking/eth2/rpc/beaconchain/methods/BeaconBlocksByRangeMessageHandler.java b/...gasys/teku/networking/eth2/rpc/beaconchain/methods/BeaconBlocksByRangeMessageHandler.java
@@ -223,7 +223,11 @@ private SafeFuture<Boolean> processNextBlock(final RequestState requestState) {
     // Ensure blocks are loaded off of the event thread
     return requestState
         .loadNextBlock()
-        .thenCompose(block -> handleLoadedBlock(requestState, block));
+        .thenCompose(
+            block -> {
+              requestState.decrementRemainingBlocks();
+              return handleLoadedBlock(requestState, block);
+            });
   }
 
   /** Sends the block and returns true if the request is now complete. */
@@ -287,13 +291,11 @@ SafeFuture<Void> sendBlock(final SignedBeaconBlock block) {
       if (step.isGreaterThan(1L)) {
         remainingBlocks = ZERO;
       }
-      return callback
-          .respond(block)
-          .thenRun(
-              () -> {
-                sentBlocks.incrementAndGet();
-                remainingBlocks = remainingBlocks.minusMinZero(1);
-              });
+      return callback.respond(block).thenRun(sentBlocks::incrementAndGet);
+    }
+
+    void decrementRemainingBlocks() {
+      remainingBlocks = remainingBlocks.minusMinZero(1);
     }
 
     void incrementCurrentSlot() {