diff --git a/beacon/validator/src/main/java/tech/pegasys/teku/validator/coordinator/ValidatorApiHandler.java b/beacon/validator/src/main/java/tech/pegasys/teku/validator/coordinator/ValidatorApiHandler.java index db7e71f5508..db56b343852 100644 --- a/beacon/validator/src/main/java/tech/pegasys/teku/validator/coordinator/ValidatorApiHandler.java +++ b/beacon/validator/src/main/java/tech/pegasys/teku/validator/coordinator/ValidatorApiHandler.java @@ -323,7 +323,10 @@ public SafeFuture> createUnsignedBlock( blockProductionPerformanceFactory.create(slot); return forkChoiceTrigger .prepareForBlockProduction(slot, blockProductionPerformance) - .thenCompose(__ -> combinedChainDataClient.getStateAtSlotExact(slot)) + .thenCompose( + __ -> + combinedChainDataClient.getStateForBlockProduction( + slot, forkChoiceTrigger.isForkChoiceOverrideLateBlockEnabled())) .thenPeek( maybeState -> { maybeState.ifPresent( diff --git a/beacon/validator/src/test/java/tech/pegasys/teku/validator/coordinator/ValidatorApiHandlerTest.java b/beacon/validator/src/test/java/tech/pegasys/teku/validator/coordinator/ValidatorApiHandlerTest.java index b9c941e7dd9..0edbbf40252 100644 --- a/beacon/validator/src/test/java/tech/pegasys/teku/validator/coordinator/ValidatorApiHandlerTest.java +++ b/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, diff --git a/ethereum/statetransition/src/main/java/tech/pegasys/teku/statetransition/forkchoice/ForkChoiceTrigger.java b/ethereum/statetransition/src/main/java/tech/pegasys/teku/statetransition/forkchoice/ForkChoiceTrigger.java index 9536490a024..204dfaa1e35 100644 --- a/ethereum/statetransition/src/main/java/tech/pegasys/teku/statetransition/forkchoice/ForkChoiceTrigger.java +++ b/ethereum/statetransition/src/main/java/tech/pegasys/teku/statetransition/forkchoice/ForkChoiceTrigger.java @@ -40,6 +40,10 @@ public SafeFuture prepareForBlockProduction( return forkChoice.prepareForBlockProduction(slot, blockProductionPerformance); } + public boolean isForkChoiceOverrideLateBlockEnabled() { + return forkChoice.isForkChoiceLateBlockReorgEnabled(); + } + public SafeFuture prepareForAttestationProduction(final UInt64 slot) { return forkChoiceRatchet.ensureForkChoiceCompleteForSlot(slot); } diff --git a/storage/src/main/java/tech/pegasys/teku/storage/client/CombinedChainDataClient.java b/storage/src/main/java/tech/pegasys/teku/storage/client/CombinedChainDataClient.java index 2e2d0007ac2..6d328b35bc3 100644 --- a/storage/src/main/java/tech/pegasys/teku/storage/client/CombinedChainDataClient.java +++ b/storage/src/main/java/tech/pegasys/teku/storage/client/CombinedChainDataClient.java @@ -240,6 +240,30 @@ public SafeFuture> getStateAtSlotExact(final UInt64 slot) return regenerateStateAndSlotExact(slot); } + public SafeFuture> getStateForBlockProduction( + final UInt64 slot, final boolean isForkChoiceLateBlockReorgEnabled) { + if (!isForkChoiceLateBlockReorgEnabled) { + return getStateAtSlotExact(slot); + } + final Optional 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> getStateAtSlotExact( final UInt64 slot, final Bytes32 chainHead) { final Optional recentBlockRoot = diff --git a/storage/src/test/java/tech/pegasys/teku/storage/client/CombinedChainDataClientTest.java b/storage/src/test/java/tech/pegasys/teku/storage/client/CombinedChainDataClientTest.java index 8a09eeafcb4..f9169921db7 100644 --- a/storage/src/test/java/tech/pegasys/teku/storage/client/CombinedChainDataClientTest.java +++ b/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 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> 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 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> 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 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> 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 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> 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 =