implement getStateForBlockProduction for late block reorg (#7929)

Currently late block reorg is disabled so this defers to `getStateAtSlotExact`

This functionality diverges from `getStateAtSlotExact` with late block reorg, as sometimes the state we're wanting to use is the parent of the head. In this scenario we actually need to check the proposer head, and get the appropriate state based on that instead.

The fallback given all things being equal currently is to call getStateAtSlotExact, which is basically falling back to the old functionality.
 - if late block reorg is disabled
 - if the head root can't be read
 - if the headRoot is the same as the root we're told to use from getProposerHead.

 In late block reorg case where we're using parent, we need to call `regenerateBeaconState` to apply empty slots.

 partially addresses #6595

Signed-off-by: Paul Harris <paul.harris@consensys.net>

beacon/validator/src/main/java/tech/pegasys/teku/validator/coordinator/ValidatorApiHandler.java

@@ -323,7 +323,10 @@ public SafeFuture<Optional<BlockContainer>> createUnsignedBlock(
         blockProductionPerformanceFactory.create(slot);
     return forkChoiceTrigger
         .prepareForBlockProduction(slot, blockProductionPerformance)
-        .thenCompose(__ -> combinedChainDataClient.getStateAtSlotExact(slot))
+        .thenCompose(
+            __ ->
+                combinedChainDataClient.getStateForBlockProduction(
+                    slot, forkChoiceTrigger.isForkChoiceOverrideLateBlockEnabled()))
         .thenPeek(
             maybeState -> {
               maybeState.ifPresent(

beacon/validator/src/test/java/tech/pegasys/teku/validator/coordinator/ValidatorApiHandlerTest.java

@@ -500,7 +500,7 @@ public void createUnsignedBlock_shouldFailWhenNodeIsSyncing() {
   public void createUnsignedBlock_shouldFailWhenParentBlockIsOptimistic() {
     final UInt64 newSlot = UInt64.valueOf(25);
     final BeaconState blockSlotState = dataStructureUtil.randomBeaconState(newSlot);
-    when(chainDataClient.getStateAtSlotExact(newSlot))
+    when(chainDataClient.getStateForBlockProduction(newSlot, false))
         .thenReturn(SafeFuture.completedFuture(Optional.of(blockSlotState)));
     final Bytes32 parentRoot = spec.getBlockRootAtSlot(blockSlotState, newSlot.minus(1));
     when(chainDataClient.isOptimisticBlock(parentRoot)).thenReturn(true);
@@ -521,7 +521,7 @@ public void createUnsignedBlock_shouldCreateBlock() {
     final BLSSignature randaoReveal = dataStructureUtil.randomSignature();
     final BeaconBlock createdBlock = dataStructureUtil.randomBeaconBlock(newSlot.longValue());
 
-    when(chainDataClient.getStateAtSlotExact(newSlot))
+    when(chainDataClient.getStateForBlockProduction(newSlot, false))
         .thenReturn(SafeFuture.completedFuture(Optional.of(blockSlotState)));
     when(blockFactory.createUnsignedBlock(
             blockSlotState,

ethereum/statetransition/src/main/java/tech/pegasys/teku/statetransition/forkchoice/ForkChoiceTrigger.java

@@ -40,6 +40,10 @@ public SafeFuture<Void> prepareForBlockProduction(
     return forkChoice.prepareForBlockProduction(slot, blockProductionPerformance);
   }
 
+  public boolean isForkChoiceOverrideLateBlockEnabled() {
+    return forkChoice.isForkChoiceLateBlockReorgEnabled();
+  }
+
   public SafeFuture<Void> prepareForAttestationProduction(final UInt64 slot) {
     return forkChoiceRatchet.ensureForkChoiceCompleteForSlot(slot);
   }

storage/src/main/java/tech/pegasys/teku/storage/client/CombinedChainDataClient.java

@@ -240,6 +240,30 @@ public SafeFuture<Optional<BeaconState>> getStateAtSlotExact(final UInt64 slot)
     return regenerateStateAndSlotExact(slot);
   }
 
+  public SafeFuture<Optional<BeaconState>> getStateForBlockProduction(
+      final UInt64 slot, final boolean isForkChoiceLateBlockReorgEnabled) {
+    if (!isForkChoiceLateBlockReorgEnabled) {
+      return getStateAtSlotExact(slot);
+    }
+    final Optional<Bytes32> headRoot = getBestBlockRoot();
+    if (headRoot.isEmpty()) {
+      return getStateAtSlotExact(slot);
+    }
+    final Bytes32 root = recentChainData.getProposerHead(headRoot.get(), slot);
+    if (root.equals(headRoot.get())) {
+      return getStateAtSlotExact(slot);
+    }
+    // otherwise we're looking for the parent slot
+    return getStateByBlockRoot(root)
+        .thenCompose(
+            maybeState ->
+                maybeState
+                    .map(
+                        beaconState ->
+                            SafeFuture.completedFuture(regenerateBeaconState(beaconState, slot)))
+                    .orElseGet(() -> getStateAtSlotExact(slot)));
+  }
+
   public SafeFuture<Optional<BeaconState>> getStateAtSlotExact(
       final UInt64 slot, final Bytes32 chainHead) {
     final Optional<Bytes32> recentBlockRoot =

storage/src/test/java/tech/pegasys/teku/storage/client/CombinedChainDataClientTest.java

@@ -16,11 +16,14 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.ExecutionException;
 import org.apache.tuweni.bytes.Bytes32;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -32,6 +35,7 @@
 import tech.pegasys.teku.spec.TestSpecFactory;
 import tech.pegasys.teku.spec.datastructures.blobs.versions.deneb.BlobSidecar;
 import tech.pegasys.teku.spec.datastructures.blocks.SignedBeaconBlock;
+import tech.pegasys.teku.spec.datastructures.blocks.SlotAndBlockRoot;
 import tech.pegasys.teku.spec.datastructures.metadata.BlockAndMetaData;
 import tech.pegasys.teku.spec.datastructures.state.Checkpoint;
 import tech.pegasys.teku.spec.datastructures.state.CommitteeAssignment;
@@ -40,6 +44,7 @@
 import tech.pegasys.teku.spec.util.DataStructureUtil;
 import tech.pegasys.teku.storage.api.StorageQueryChannel;
 import tech.pegasys.teku.storage.protoarray.ForkChoiceStrategy;
+import tech.pegasys.teku.storage.store.UpdatableStore;
 
 /** Note: Most tests should be added to the integration-test directory */
 class CombinedChainDataClientTest {
@@ -48,6 +53,8 @@ class CombinedChainDataClientTest {
   private final RecentChainData recentChainData = mock(RecentChainData.class);
   private final ForkChoiceStrategy forkChoiceStrategy = mock(ForkChoiceStrategy.class);
   private final StorageQueryChannel historicalChainData = mock(StorageQueryChannel.class);
+
+  private final UpdatableStore store = mock(UpdatableStore.class);
   private final CombinedChainDataClient client =
       new CombinedChainDataClient(
           recentChainData,
@@ -146,6 +153,98 @@ void getsEarliestAvailableBlobSidecarSlot() {
     assertThat(result).hasValue(UInt64.ONE);
   }
 
+  @Test
+  void getStateForBlockProduction_directsToStateAtSlotExact()
+      throws ExecutionException, InterruptedException {
+    final BeaconState state = dataStructureUtil.randomBeaconState(UInt64.valueOf(2));
+    final Optional<Bytes32> recentBlockRoot =
+        Optional.of(spec.getBlockRootAtSlot(state, UInt64.ONE));
+    final SlotAndBlockRoot slotAndBlockRoot =
+        new SlotAndBlockRoot(UInt64.ONE, recentBlockRoot.get());
+    when(recentChainData.getBlockRootInEffectBySlot(UInt64.ONE)).thenReturn(recentBlockRoot);
+    when(recentChainData.getStore()).thenReturn(store);
+    when(store.retrieveStateAtSlot(slotAndBlockRoot))
+        .thenReturn(SafeFuture.completedFuture(Optional.of(state)));
+    final SafeFuture<Optional<BeaconState>> future =
+        client.getStateForBlockProduction(UInt64.ONE, false);
+    assertThat(future.get()).contains(state);
+    // getStateAtSlotExact
+    verify(recentChainData).getBlockRootInEffectBySlot(UInt64.ONE);
+    verify(store).retrieveStateAtSlot(slotAndBlockRoot);
+  }
+
+  @Test
+  void getStateForBlockProduction_whenEnabledAndHaveNoBestBlockRoot()
+      throws ExecutionException, InterruptedException {
+    final BeaconState state = dataStructureUtil.randomBeaconState(UInt64.valueOf(2));
+    final Optional<Bytes32> recentBlockRoot =
+        Optional.of(spec.getBlockRootAtSlot(state, UInt64.ONE));
+    final SlotAndBlockRoot slotAndBlockRoot =
+        new SlotAndBlockRoot(UInt64.ONE, recentBlockRoot.get());
+    when(recentChainData.getStore()).thenReturn(store);
+
+    when(recentChainData.getBestBlockRoot()).thenReturn(Optional.empty());
+    when(recentChainData.getBlockRootInEffectBySlot(UInt64.ONE)).thenReturn(recentBlockRoot);
+    when(store.retrieveStateAtSlot(slotAndBlockRoot))
+        .thenReturn(SafeFuture.completedFuture(Optional.of(state)));
+
+    final SafeFuture<Optional<BeaconState>> future =
+        client.getStateForBlockProduction(UInt64.ONE, true);
+    assertThat(future.get()).contains(state);
+    // getStateAtSlotExact
+    verify(recentChainData).getBlockRootInEffectBySlot(UInt64.ONE);
+    verify(store).retrieveStateAtSlot(slotAndBlockRoot);
+  }
+
+  @Test
+  void getStateForBlockProduction_whenEnabledAndBestBlockRootMatches()
+      throws ExecutionException, InterruptedException {
+    final BeaconState state = dataStructureUtil.randomBeaconState(UInt64.valueOf(2));
+    final Optional<Bytes32> recentBlockRoot =
+        Optional.of(spec.getBlockRootAtSlot(state, UInt64.ONE));
+    final SlotAndBlockRoot slotAndBlockRoot =
+        new SlotAndBlockRoot(UInt64.ONE, recentBlockRoot.get());
+    when(recentChainData.getStore()).thenReturn(store);
+
+    when(recentChainData.getBestBlockRoot()).thenReturn(recentBlockRoot);
+    when(recentChainData.getProposerHead(any(), any())).thenReturn(recentBlockRoot.get());
+    when(recentChainData.getBlockRootInEffectBySlot(UInt64.ONE)).thenReturn(recentBlockRoot);
+    when(store.retrieveStateAtSlot(slotAndBlockRoot))
+        .thenReturn(SafeFuture.completedFuture(Optional.of(state)));
+
+    final SafeFuture<Optional<BeaconState>> future =
+        client.getStateForBlockProduction(UInt64.ONE, true);
+    assertThat(future.get()).contains(state);
+    // getStateAtSlotExact
+    verify(recentChainData).getBlockRootInEffectBySlot(UInt64.ONE);
+    verify(store).retrieveStateAtSlot(slotAndBlockRoot);
+  }
+
+  @Test
+  void getStateForBlockProduction_whenEnabledAndBestBlockRootDifferent()
+      throws ExecutionException, InterruptedException {
+    final BeaconState state = dataStructureUtil.randomBeaconState(UInt64.valueOf(2));
+    final Bytes32 proposerHead = dataStructureUtil.randomBytes32();
+    final BeaconState proposerState = dataStructureUtil.randomBeaconState(UInt64.ONE);
+    final Optional<Bytes32> recentBlockRoot =
+        Optional.of(spec.getBlockRootAtSlot(state, UInt64.ONE));
+    final SlotAndBlockRoot slotAndBlockRoot =
+        new SlotAndBlockRoot(UInt64.ONE, recentBlockRoot.get());
+    when(recentChainData.getStore()).thenReturn(store);
+
+    when(recentChainData.getBestBlockRoot()).thenReturn(recentBlockRoot);
+    when(recentChainData.getProposerHead(any(), any())).thenReturn(proposerHead);
+    when(recentChainData.getBlockRootInEffectBySlot(UInt64.ONE)).thenReturn(recentBlockRoot);
+    when(store.retrieveBlockState(proposerHead))
+        .thenReturn(SafeFuture.completedFuture(Optional.of(proposerState)));
+
+    final SafeFuture<Optional<BeaconState>> future =
+        client.getStateForBlockProduction(UInt64.ONE, true);
+    assertThat(future.get()).contains(proposerState);
+    verify(store).retrieveBlockState(proposerHead);
+    verify(store, never()).retrieveStateAtSlot(slotAndBlockRoot);
+  }
+
   @Test
   void getsBlobSidecarBySlotAndBlockRootAndBlobIndex() {
     final SlotAndBlockRootAndBlobIndex correctKey =