Minor refactorings and benchmark

prysmaticlabs · Jun 18, 2024 · c4eda0b · c4eda0b
1 parent 8a6b73c
commit c4eda0b
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 22 deletions.
diff --git a/beacon-chain/core/electra/registry_updates.go b/beacon-chain/core/electra/registry_updates.go
@@ -11,7 +11,8 @@ import (
 	"github.com/prysmaticlabs/prysm/v5/consensus-types/primitives"
 )
 
-// ProcessRegistryUpdates rotates validators in and out of active pool.
+// ProcessRegistryUpdates processes all validators eligible for the activation queue, all validators
+// which should be ejected, and all validators which are eligible for activation from then queue.
 //
 // Spec pseudocode definition:
 //
@@ -32,34 +33,37 @@ import (
 //	         for validator in state.validators:
 //	             if is_eligible_for_activation(state, validator):
 //	                 validator.activation_epoch = activation_epoch
-func ProcessRegistryUpdates(ctx context.Context, state state.BeaconState) (state.BeaconState, error) {
+func ProcessRegistryUpdates(ctx context.Context, state state.BeaconState) error {
 	currentEpoch := time.CurrentEpoch(state)
 	ejectionBal := params.BeaconConfig().EjectionBalance
 	activationEpoch := helpers.ActivationExitEpoch(currentEpoch)
 	vals := state.Validators()
 	for idx, val := range vals {
+		// Handle validators eligible to join the activation queue.
 		if helpers.IsEligibleForActivationQueue(val, currentEpoch) {
 			val.ActivationEligibilityEpoch = currentEpoch + 1
 			if err := state.UpdateValidatorAtIndex(primitives.ValidatorIndex(idx), val); err != nil {
-				return nil, err
+				return err
 			}
 		}
+		// Handle validator ejections.
 		if helpers.IsActiveValidator(val, currentEpoch) && val.EffectiveBalance <= ejectionBal {
 			var err error
-			maxExitEpoch, churn := validators.MaxExitEpochAndChurn(state) 
-			state, _, err = validators.InitiateValidatorExit(ctx, state, primitives.ValidatorIndex(idx), maxExitEpoch, churn)
+			// exitQueueEpoch and churn arguments are not used in electra.
+			state, _, err = validators.InitiateValidatorExit(ctx, state, primitives.ValidatorIndex(idx), 0 /*exitQueueEpoch*/, 0 /*churn*/)
 			if err != nil {
-				return nil, err
+				return err
 			}
 		}
 
+		// Activate all eligible validators.
 		if helpers.IsEligibleForActivation(state, val) {
 			val.ActivationEpoch = activationEpoch
 			if err := state.UpdateValidatorAtIndex(primitives.ValidatorIndex(idx), val); err != nil {
-				return nil, err
+				return err
 			}
 		}
 	}
 
-	return state, nil
+	return nil
 }
diff --git a/beacon-chain/core/electra/registry_updates_test.go b/beacon-chain/core/electra/registry_updates_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v5/config/params"
 	eth "github.com/prysmaticlabs/prysm/v5/proto/prysm/v1alpha1"
 	"github.com/prysmaticlabs/prysm/v5/testing/require"
+	"github.com/prysmaticlabs/prysm/v5/testing/util"
 )
 
 func TestProcessRegistryUpdates(t *testing.T) {
@@ -26,9 +27,9 @@ func TestProcessRegistryUpdates(t *testing.T) {
 		check func(*testing.T, state.BeaconState)
 	}{
 		{
-			name: "No rotation", // No validators exited.
+			name: "No rotation",
 			state: func() state.BeaconState {
-				base := &eth.BeaconState{
+				base := &eth.BeaconStateElectra{
 					Slot: 5 * params.BeaconConfig().SlotsPerEpoch,
 					Validators: []*eth.Validator{
 						{ExitEpoch: params.BeaconConfig().MaxSeedLookahead},
@@ -40,7 +41,7 @@ func TestProcessRegistryUpdates(t *testing.T) {
 					},
 					FinalizedCheckpoint: &eth.Checkpoint{Root: make([]byte, fieldparams.RootLength)},
 				}
-				st, err := state_native.InitializeFromProtoPhase0(base)
+				st, err := state_native.InitializeFromProtoElectra(base)
 				require.NoError(t, err)
 				return st
 			}(),
@@ -53,8 +54,7 @@ func TestProcessRegistryUpdates(t *testing.T) {
 		{
 			name: "Validators are activated",
 			state: func() state.BeaconState {
-				// Construct a state which
-				base := &eth.BeaconState{
+				base := &eth.BeaconStateElectra{
 					Slot:                5 * params.BeaconConfig().SlotsPerEpoch,
 					FinalizedCheckpoint: &eth.Checkpoint{Epoch: 6, Root: make([]byte, fieldparams.RootLength)},
 				}
@@ -65,7 +65,7 @@ func TestProcessRegistryUpdates(t *testing.T) {
 						ActivationEpoch:            params.BeaconConfig().FarFutureEpoch,
 					})
 				}
-				st, err := state_native.InitializeFromProtoPhase0(base)
+				st, err := state_native.InitializeFromProtoElectra(base)
 				require.NoError(t, err)
 				return st
 			}(),
@@ -80,8 +80,7 @@ func TestProcessRegistryUpdates(t *testing.T) {
 		{
 			name: "Validators are exited",
 			state: func() state.BeaconState {
-				// Construct a state which
-				base := &eth.BeaconState{
+				base := &eth.BeaconStateElectra{
 					Slot:                5 * params.BeaconConfig().SlotsPerEpoch,
 					FinalizedCheckpoint: &eth.Checkpoint{Epoch: 6, Root: make([]byte, fieldparams.RootLength)},
 				}
@@ -92,7 +91,7 @@ func TestProcessRegistryUpdates(t *testing.T) {
 						WithdrawableEpoch: params.BeaconConfig().FarFutureEpoch,
 					})
 				}
-				st, err := state_native.InitializeFromProtoPhase0(base)
+				st, err := state_native.InitializeFromProtoElectra(base)
 				require.NoError(t, err)
 				return st
 			}(),
@@ -108,11 +107,42 @@ func TestProcessRegistryUpdates(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			res, err := electra.ProcessRegistryUpdates(context.TODO(), tt.state)
+			err := electra.ProcessRegistryUpdates(context.TODO(), tt.state)
 			require.NoError(t, err)
 			if tt.check != nil {
-				tt.check(t, res)
+				tt.check(t, tt.state)
 			}
 		})
 	}
 }
+
+func Benchmark_ProcessRegistryUpdates_MassEjection(b *testing.B) {
+	bal := params.BeaconConfig().EjectionBalance - 1
+	ffe := params.BeaconConfig().FarFutureEpoch
+	genValidators := func(num uint64) []*eth.Validator {
+		vals := make([]*eth.Validator, num)
+		for i := range vals {
+			vals[i] = &eth.Validator{
+				EffectiveBalance: bal,
+				ExitEpoch:        ffe,
+			}
+		}
+		return vals
+	}
+
+	st, err := util.NewBeaconStateElectra()
+	require.NoError(b, err)
+
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		if err := st.SetValidators(genValidators(100000)); err != nil {
+			panic(err)
+		}
+		b.StartTimer()
+
+		if err := electra.ProcessRegistryUpdates(context.TODO(), st); err != nil {
+			panic(err)
+		}
+	}
+
+}
diff --git a/beacon-chain/core/electra/transition.go b/beacon-chain/core/electra/transition.go
@@ -74,8 +74,7 @@ func ProcessEpoch(ctx context.Context, state state.BeaconState) error {
 		return errors.Wrap(err, "could not process rewards and penalties")
 	}
 
-	state, err = ProcessRegistryUpdates(ctx, state)
-	if err != nil {
+	if err := ProcessRegistryUpdates(ctx, state); err != nil {
 		return errors.Wrap(err, "could not process registry updates")
 	}
 

diff --git a/testing/spectest/shared/electra/epoch_processing/registry_updates.go b/testing/spectest/shared/electra/epoch_processing/registry_updates.go
@@ -28,5 +28,5 @@ func RunRegistryUpdatesTests(t *testing.T, config string) {
 }
 
 func processRegistryUpdatesWrapper(_ *testing.T, state state.BeaconState) (state.BeaconState, error) {
-	return electra.ProcessRegistryUpdates(context.Background(), state)
+	return state, electra.ProcessRegistryUpdates(context.Background(), state)
 }