Rewrite process_registry_updates not to use st.Validators()

prysmaticlabs · Jun 27, 2024 · 0a1500b · 0a1500b
1 parent d42415e
commit 0a1500b
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 30 deletions.
diff --git a/beacon-chain/core/electra/registry_updates.go b/beacon-chain/core/electra/registry_updates.go
@@ -34,35 +34,70 @@ 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) error {
-	currentEpoch := time.CurrentEpoch(state)
+func ProcessRegistryUpdates(ctx context.Context, st state.BeaconState) error {
+	currentEpoch := time.CurrentEpoch(st)
 	ejectionBal := params.BeaconConfig().EjectionBalance
 	activationEpoch := helpers.ActivationExitEpoch(currentEpoch)
-	vals := state.Validators()
-	for idx, val := range vals {
-		// Handle validators eligible to join the activation queue.
+
+	// To avoid copying the state validator set via st.Validators(), we will perform a read only pass
+	// over the validator set while collecting validator indices where the validator copy is actually
+	// necessary, then we will process these operations.
+	eligibleForActivationQ := make([]primitives.ValidatorIndex, 0)
+	eligibleForEjection := make([]primitives.ValidatorIndex, 0)
+	eligibleForActivation := make([]primitives.ValidatorIndex, 0)
+
+	if err := st.ReadFromEveryValidator(func(idx int, val state.ReadOnlyValidator) error {
+		// Collect validators eligible to enter the activation queue.
 		if helpers.IsEligibleForActivationQueue(val, currentEpoch) {
-			val.ActivationEligibilityEpoch = currentEpoch + 1
-			if err := state.UpdateValidatorAtIndex(primitives.ValidatorIndex(idx), val); err != nil {
-				return fmt.Errorf("failed to update eligible validator at index %d: %w", idx, err)
-			}
+			eligibleForActivationQ = append(eligibleForActivationQ, primitives.ValidatorIndex(idx))
+		}
+
+		// Collect validators to eject.
+		if val.EffectiveBalance() <= ejectionBal && helpers.IsActiveValidatorUsingTrie(val, currentEpoch) {
+			eligibleForEjection = append(eligibleForEjection, primitives.ValidatorIndex(idx))
 		}
-		// Handle validator ejections.
-		if val.EffectiveBalance <= ejectionBal && helpers.IsActiveValidator(val, currentEpoch) {
-			var err error
-			// 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 fmt.Errorf("failed to initiate validator exit at index %d: %w", idx, err)
-			}
+
+		// Collect validators eligible for activation and not yet dequeued for activation.
+		if helpers.IsEligibleForActivationUsingTrie(st, val) {
+			eligibleForActivation = append(eligibleForActivation, primitives.ValidatorIndex(idx))
 		}
 
+		return nil
+	}); err != nil {
+		return fmt.Errorf("failed to read validators: %w", err)
+	}
+
+	// Handle validators eligible to join the activation queue.
+	for _, idx := range eligibleForActivationQ {
+		v, err := st.ValidatorAtIndex(idx)
+		if err != nil {
+			return err
+		}
+		v.ActivationEligibilityEpoch = currentEpoch + 1
+		if err := st.UpdateValidatorAtIndex(idx, v); err != nil {
+			return fmt.Errorf("failed to updated eligible validator at index %d: %w", idx, err)
+		}
+	}
+
+	// Handle validator ejections.
+	for _, idx := range eligibleForEjection {
+		var err error
+		// exitQueueEpoch and churn arguments are not used in electra.
+		st, _, err = validators.InitiateValidatorExit(ctx, st, primitives.ValidatorIndex(idx), 0 /*exitQueueEpoch*/, 0 /*churn*/)
+		if err != nil {
+			return fmt.Errorf("failed to initiate validator exit at index %d: %w", idx, err)
+		}
+	}
+
+	for _, idx := range eligibleForActivation {
 		// Activate all eligible validators.
-		if helpers.IsEligibleForActivation(state, val) {
-			val.ActivationEpoch = activationEpoch
-			if err := state.UpdateValidatorAtIndex(primitives.ValidatorIndex(idx), val); err != nil {
-				return fmt.Errorf("failed to activate validator at index %d: %w", idx, err)
-			}
+		v, err := st.ValidatorAtIndex(idx)
+		if err != nil {
+			return err
+		}
+		v.ActivationEpoch = activationEpoch
+		if err := st.UpdateValidatorAtIndex(idx, v); err != nil {
+			return fmt.Errorf("failed to activate validator at index %d: %w", idx, err)
 		}
 	}
 

diff --git a/beacon-chain/core/electra/registry_updates_test.go b/beacon-chain/core/electra/registry_updates_test.go
@@ -10,16 +10,14 @@ import (
 	state_native "github.com/prysmaticlabs/prysm/v5/beacon-chain/state/state-native"
 	fieldparams "github.com/prysmaticlabs/prysm/v5/config/fieldparams"
 	"github.com/prysmaticlabs/prysm/v5/config/params"
+	"github.com/prysmaticlabs/prysm/v5/consensus-types/primitives"
 	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) {
-	const electraEpoch = 3
-	cfg := params.BeaconConfig()
-	cfg.ElectraForkEpoch = electraEpoch
-	params.SetActiveTestCleanup(t, cfg)
+	finalizedEpoch := primitives.Epoch(4)
 
 	tests := []struct {
 		name  string
@@ -56,11 +54,11 @@ func TestProcessRegistryUpdates(t *testing.T) {
 			state: func() state.BeaconState {
 				base := &eth.BeaconStateElectra{
 					Slot:                5 * params.BeaconConfig().SlotsPerEpoch,
-					FinalizedCheckpoint: &eth.Checkpoint{Epoch: 6, Root: make([]byte, fieldparams.RootLength)},
+					FinalizedCheckpoint: &eth.Checkpoint{Epoch: finalizedEpoch, Root: make([]byte, fieldparams.RootLength)},
 				}
 				for i := uint64(0); i < 10; i++ {
 					base.Validators = append(base.Validators, &eth.Validator{
-						ActivationEligibilityEpoch: params.BeaconConfig().FarFutureEpoch,
+						ActivationEligibilityEpoch: finalizedEpoch,
 						EffectiveBalance:           params.BeaconConfig().MaxEffectiveBalance,
 						ActivationEpoch:            params.BeaconConfig().FarFutureEpoch,
 					})
@@ -82,7 +80,7 @@ func TestProcessRegistryUpdates(t *testing.T) {
 			state: func() state.BeaconState {
 				base := &eth.BeaconStateElectra{
 					Slot:                5 * params.BeaconConfig().SlotsPerEpoch,
-					FinalizedCheckpoint: &eth.Checkpoint{Epoch: 6, Root: make([]byte, fieldparams.RootLength)},
+					FinalizedCheckpoint: &eth.Checkpoint{Epoch: finalizedEpoch, Root: make([]byte, fieldparams.RootLength)},
 				}
 				for i := uint64(0); i < 10; i++ {
 					base.Validators = append(base.Validators, &eth.Validator{
@@ -144,5 +142,4 @@ func Benchmark_ProcessRegistryUpdates_MassEjection(b *testing.B) {
 			panic(err)
 		}
 	}
-
 }