From 1da043b99a5d8e3e2432561c8298cf8c3b215a06 Mon Sep 17 00:00:00 2001
From: Stefan Bratanov <stefan.bratanov93@gmail.com>
Date: Wed, 27 Nov 2024 10:24:12 +0000
Subject: [PATCH] Track produced block when builder doesn't reveal the payload

---
 .../publisher/AbstractBlockPublisher.java         | 15 ++++++++++++++-
 .../publisher/AbstractBlockPublisherTest.java     | 14 ++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/beacon/validator/src/main/java/tech/pegasys/teku/validator/coordinator/publisher/AbstractBlockPublisher.java b/beacon/validator/src/main/java/tech/pegasys/teku/validator/coordinator/publisher/AbstractBlockPublisher.java
index 93b39d9bf74..7fb5df7fc4d 100644
--- a/beacon/validator/src/main/java/tech/pegasys/teku/validator/coordinator/publisher/AbstractBlockPublisher.java
+++ b/beacon/validator/src/main/java/tech/pegasys/teku/validator/coordinator/publisher/AbstractBlockPublisher.java
@@ -18,12 +18,14 @@
 
 import com.google.common.base.Suppliers;
 import java.util.List;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Supplier;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import tech.pegasys.teku.ethereum.performance.trackers.BlockPublishingPerformance;
 import tech.pegasys.teku.infrastructure.async.AsyncRunner;
 import tech.pegasys.teku.infrastructure.async.SafeFuture;
+import tech.pegasys.teku.infrastructure.exceptions.ExceptionUtil;
 import tech.pegasys.teku.networking.eth2.gossip.BlockGossipChannel;
 import tech.pegasys.teku.spec.datastructures.blobs.versions.deneb.BlobSidecar;
 import tech.pegasys.teku.spec.datastructures.blocks.SignedBeaconBlock;
@@ -75,7 +77,18 @@ public SafeFuture<SendSignedBlockResult> sendSignedBlock(
       final BlockPublishingPerformance blockPublishingPerformance) {
     return blockFactory
         .unblindSignedBlockIfBlinded(blockContainer.getSignedBlock(), blockPublishingPerformance)
-        .thenPeek(performanceTracker::saveProducedBlock)
+        .whenComplete(
+            (block, ex) -> {
+              if (ex != null) {
+                // in case of relay API timeouts when unblinding, we still want to assume the block
+                // as produced, since the relay could have published it
+                if (ExceptionUtil.hasCause(ex, TimeoutException.class)) {
+                  performanceTracker.saveProducedBlock(blockContainer.getSignedBlock());
+                }
+              } else {
+                performanceTracker.saveProducedBlock(block);
+              }
+            })
         .thenCompose(
             // creating blob sidecars after unblinding the block to ensure in the blinded flow we
             // will have the cached builder payload
diff --git a/beacon/validator/src/test/java/tech/pegasys/teku/validator/coordinator/publisher/AbstractBlockPublisherTest.java b/beacon/validator/src/test/java/tech/pegasys/teku/validator/coordinator/publisher/AbstractBlockPublisherTest.java
index 50391494269..aeb0a152435 100644
--- a/beacon/validator/src/test/java/tech/pegasys/teku/validator/coordinator/publisher/AbstractBlockPublisherTest.java
+++ b/beacon/validator/src/test/java/tech/pegasys/teku/validator/coordinator/publisher/AbstractBlockPublisherTest.java
@@ -23,6 +23,7 @@
 import static tech.pegasys.teku.infrastructure.async.SafeFutureAssert.assertThatSafeFuture;
 
 import java.util.List;
+import java.util.concurrent.TimeoutException;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import tech.pegasys.teku.ethereum.performance.trackers.BlockPublishingPerformance;
@@ -252,6 +253,19 @@ public void sendSignedBlock_shouldReturnNotImportedWhenBlockImportFails() {
     verify(blockPublisher).importBlobSidecars(blobSidecars, BlockPublishingPerformance.NOOP);
   }
 
+  @Test
+  public void sendSignedBlock_shouldTrackBlockAsProducedIfUnblindingTimeouts() {
+    when(blockFactory.unblindSignedBlockIfBlinded(signedBlock, BlockPublishingPerformance.NOOP))
+        .thenReturn(SafeFuture.failedFuture(new TimeoutException()));
+
+    blockPublisher.sendSignedBlock(
+        signedBlockContents,
+        BroadcastValidationLevel.NOT_REQUIRED,
+        BlockPublishingPerformance.NOOP);
+
+    verify(performanceTracker).saveProducedBlock(signedBlockContents.getSignedBlock());
+  }
+
   private SafeFuture<BlockImportAndBroadcastValidationResults> prepareBlockImportResult(
       final BlockImportResult blockImportResult) {
     return SafeFuture.completedFuture(