diff --git a/beacon-chain/forkchoice/protoarray/BUILD.bazel b/beacon-chain/forkchoice/protoarray/BUILD.bazel index 32b2ba169..6031cee5a 100644 --- a/beacon-chain/forkchoice/protoarray/BUILD.bazel +++ b/beacon-chain/forkchoice/protoarray/BUILD.bazel @@ -64,6 +64,7 @@ go_test( "//beacon-chain/forkchoice/types:go_default_library", "//beacon-chain/state:go_default_library", "//beacon-chain/state/v3:go_default_library", + "//config/features:go_default_library", "//config/params:go_default_library", "//consensus-types/blocks:go_default_library", "//consensus-types/primitives:go_default_library", diff --git a/beacon-chain/forkchoice/protoarray/store.go b/beacon-chain/forkchoice/protoarray/store.go index cffa6170a..9ba3b317b 100644 --- a/beacon-chain/forkchoice/protoarray/store.go +++ b/beacon-chain/forkchoice/protoarray/store.go @@ -885,6 +885,15 @@ func (s *Store) viableForHead(node *Node) bool { // It's also viable if we are in genesis epoch. justified := s.justifiedCheckpoint.Epoch == node.justifiedEpoch || s.justifiedCheckpoint.Epoch == 0 finalized := s.finalizedCheckpoint.Epoch == node.finalizedEpoch || s.finalizedCheckpoint.Epoch == 0 + if features.Get().EnableDefensivePull { + currentEpoch := slots.EpochsSinceGenesis(time.Unix(int64(s.genesisTime), 0)) + if !justified && s.justifiedCheckpoint.Epoch+1 == currentEpoch { + if node.unrealizedJustifiedEpoch+1 >= currentEpoch { + justified = true + finalized = true + } + } + } return justified && finalized } diff --git a/beacon-chain/forkchoice/protoarray/store_test.go b/beacon-chain/forkchoice/protoarray/store_test.go index 657b187e6..f0ba515e6 100644 --- a/beacon-chain/forkchoice/protoarray/store_test.go +++ b/beacon-chain/forkchoice/protoarray/store_test.go @@ -7,6 +7,7 @@ import ( "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice" forkchoicetypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/types" + "github.com/prysmaticlabs/prysm/v3/config/features" "github.com/prysmaticlabs/prysm/v3/config/params" "github.com/prysmaticlabs/prysm/v3/consensus-types/blocks" types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives" @@ -887,6 +888,43 @@ func TestStore_ViableForHead(t *testing.T) { } } +func TestStore_ViableForHead_DefensivePull(t *testing.T) { + resetCfg := features.InitWithReset(&features.Flags{ + EnableDefensivePull: true, + }) + defer resetCfg() + + tests := []struct { + n *Node + justifiedEpoch types.Epoch + finalizedEpoch types.Epoch + currentEpoch types.Epoch + want bool + }{ + {&Node{}, 0, 0, 0, true}, + {&Node{}, 1, 0, 1, false}, + {&Node{}, 0, 1, 1, false}, + {&Node{finalizedEpoch: 1, justifiedEpoch: 1}, 1, 1, 1, true}, + {&Node{finalizedEpoch: 1, justifiedEpoch: 1}, 2, 2, 2, false}, + {&Node{finalizedEpoch: 3, justifiedEpoch: 4}, 4, 3, 3, true}, + {&Node{unrealizedFinalizedEpoch: 3, unrealizedJustifiedEpoch: 4}, 3, 2, 4, true}, + {&Node{unrealizedFinalizedEpoch: 2, unrealizedJustifiedEpoch: 3}, 3, 2, 4, true}, + {&Node{unrealizedFinalizedEpoch: 1, unrealizedJustifiedEpoch: 2}, 3, 2, 4, false}, + } + for _, tc := range tests { + jc := &forkchoicetypes.Checkpoint{Epoch: tc.justifiedEpoch} + fc := &forkchoicetypes.Checkpoint{Epoch: tc.finalizedEpoch} + currentTime := uint64(time.Now().Unix()) + driftSeconds := uint64(params.BeaconConfig().SlotsPerEpoch) * params.BeaconConfig().SecondsPerSlot + s := &Store{ + justifiedCheckpoint: jc, + finalizedCheckpoint: fc, + genesisTime: currentTime - driftSeconds*uint64(tc.currentEpoch), + } + assert.Equal(t, tc.want, s.viableForHead(tc.n)) + } +} + func TestStore_HasParent(t *testing.T) { tests := []struct { m map[[32]byte]uint64