Add in progress handler to committee cache (#9664)

* Add in progress handler for committee cache

* Remove debug print

* Update validators.go

* Fix all the tests

* More tests

* Update committee_disabled.go

* Update committee_disabled.go

* Update testing util

* Update main.go

Co-authored-by: Nishant Das <nishdas93@gmail.com>

beacon-chain/core/helpers/BUILD.bazel

@@ -51,6 +51,8 @@ go_library(
         "//time:go_default_library",
         "@com_github_ferranbt_fastssz//:go_default_library",
         "@com_github_pkg_errors//:go_default_library",
+        "@com_github_prometheus_client_golang//prometheus:go_default_library",
+        "@com_github_prometheus_client_golang//prometheus/promauto:go_default_library",
         "@com_github_prysmaticlabs_eth2_types//:go_default_library",
         "@com_github_prysmaticlabs_go_bitfield//:go_default_library",
         "@com_github_sirupsen_logrus//:go_default_library",

beacon-chain/core/helpers/attestation_test.go

@@ -1,6 +1,7 @@
 package helpers_test
 
 import (
+	"context"
 	"strconv"
 	"testing"
 	"time"
@@ -22,7 +23,7 @@ import (
 func TestAttestation_IsAggregator(t *testing.T) {
 	t.Run("aggregator", func(t *testing.T) {
 		beaconState, privKeys := util.DeterministicGenesisState(t, 100)
-		committee, err := helpers.BeaconCommitteeFromState(beaconState, 0, 0)
+		committee, err := helpers.BeaconCommitteeFromState(context.Background(), beaconState, 0, 0)
 		require.NoError(t, err)
 		sig := privKeys[0].Sign([]byte{'A'})
 		agg, err := helpers.IsAggregator(uint64(len(committee)), sig.Marshal())
@@ -35,7 +36,7 @@ func TestAttestation_IsAggregator(t *testing.T) {
 		defer params.UseMainnetConfig()
 		beaconState, privKeys := util.DeterministicGenesisState(t, 2048)
 
-		committee, err := helpers.BeaconCommitteeFromState(beaconState, 0, 0)
+		committee, err := helpers.BeaconCommitteeFromState(context.Background(), beaconState, 0, 0)
 		require.NoError(t, err)
 		sig := privKeys[0].Sign([]byte{'A'})
 		agg, err := helpers.IsAggregator(uint64(len(committee)), sig.Marshal())
@@ -117,7 +118,7 @@ func TestAttestation_ComputeSubnetForAttestation(t *testing.T) {
 		},
 		Signature: []byte{'B'},
 	}
-	valCount, err := helpers.ActiveValidatorCount(state, core.SlotToEpoch(att.Data.Slot))
+	valCount, err := helpers.ActiveValidatorCount(context.Background(), state, core.SlotToEpoch(att.Data.Slot))
 	require.NoError(t, err)
 	sub := helpers.ComputeSubnetForAttestation(valCount, att)
 	assert.Equal(t, uint64(6), sub, "Did not get correct subnet for attestation")

beacon-chain/core/helpers/beacon_committee.go

@@ -4,6 +4,7 @@ package helpers
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"sort"
 
@@ -71,14 +72,14 @@ func SlotCommitteeCount(activeValidatorCount uint64) uint64 {
 //        index=(slot % SLOTS_PER_EPOCH) * committees_per_slot + index,
 //        count=committees_per_slot * SLOTS_PER_EPOCH,
 //    )
-func BeaconCommitteeFromState(state state.ReadOnlyBeaconState, slot types.Slot, committeeIndex types.CommitteeIndex) ([]types.ValidatorIndex, error) {
+func BeaconCommitteeFromState(ctx context.Context, state state.ReadOnlyBeaconState, slot types.Slot, committeeIndex types.CommitteeIndex) ([]types.ValidatorIndex, error) {
 	epoch := core.SlotToEpoch(slot)
 	seed, err := Seed(state, epoch, params.BeaconConfig().DomainBeaconAttester)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not get seed")
 	}
 
-	committee, err := committeeCache.Committee(slot, seed, committeeIndex)
+	committee, err := committeeCache.Committee(ctx, slot, seed, committeeIndex)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not interface with committee cache")
 	}
@@ -86,12 +87,12 @@ func BeaconCommitteeFromState(state state.ReadOnlyBeaconState, slot types.Slot,
 		return committee, nil
 	}
 
-	activeIndices, err := ActiveValidatorIndices(state, epoch)
+	activeIndices, err := ActiveValidatorIndices(ctx, state, epoch)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not get active indices")
 	}
 
-	return BeaconCommittee(activeIndices, seed, slot, committeeIndex)
+	return BeaconCommittee(ctx, activeIndices, seed, slot, committeeIndex)
 }
 
 // BeaconCommittee returns the beacon committee of a given slot and committee index. The
@@ -112,12 +113,13 @@ func BeaconCommitteeFromState(state state.ReadOnlyBeaconState, slot types.Slot,
 //        count=committees_per_slot * SLOTS_PER_EPOCH,
 //    )
 func BeaconCommittee(
+	ctx context.Context,
 	validatorIndices []types.ValidatorIndex,
 	seed [32]byte,
 	slot types.Slot,
 	committeeIndex types.CommitteeIndex,
 ) ([]types.ValidatorIndex, error) {
-	committee, err := committeeCache.Committee(slot, seed, committeeIndex)
+	committee, err := committeeCache.Committee(ctx, slot, seed, committeeIndex)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not interface with committee cache")
 	}
@@ -151,6 +153,7 @@ type CommitteeAssignmentContainer struct {
 // 3. Determine the attesting slot for each committee.
 // 4. Construct a map of validator indices pointing to the respective committees.
 func CommitteeAssignments(
+	ctx context.Context,
 	state state.BeaconState,
 	epoch types.Epoch,
 ) (map[types.ValidatorIndex]*CommitteeAssignmentContainer, map[types.ValidatorIndex][]types.Slot, error) {
@@ -181,14 +184,14 @@ func CommitteeAssignments(
 		if err := state.SetSlot(slot); err != nil {
 			return nil, nil, err
 		}
-		i, err := BeaconProposerIndex(state)
+		i, err := BeaconProposerIndex(ctx, state)
 		if err != nil {
 			return nil, nil, errors.Wrapf(err, "could not check proposer at slot %d", state.Slot())
 		}
 		proposerIndexToSlots[i] = append(proposerIndexToSlots[i], slot)
 	}
 
-	activeValidatorIndices, err := ActiveValidatorIndices(state, epoch)
+	activeValidatorIndices, err := ActiveValidatorIndices(ctx, state, epoch)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -202,7 +205,7 @@ func CommitteeAssignments(
 		// Compute committees.
 		for j := uint64(0); j < numCommitteesPerSlot; j++ {
 			slot := startSlot + i
-			committee, err := BeaconCommitteeFromState(state, slot, types.CommitteeIndex(j) /*committee index*/)
+			committee, err := BeaconCommitteeFromState(ctx, state, slot, types.CommitteeIndex(j) /*committee index*/)
 			if err != nil {
 				return nil, nil, err
 			}
@@ -234,8 +237,8 @@ func VerifyBitfieldLength(bf bitfield.Bitfield, committeeSize uint64) error {
 
 // VerifyAttestationBitfieldLengths verifies that an attestations aggregation bitfields is
 // a valid length matching the size of the committee.
-func VerifyAttestationBitfieldLengths(state state.ReadOnlyBeaconState, att *ethpb.Attestation) error {
-	committee, err := BeaconCommitteeFromState(state, att.Data.Slot, att.Data.CommitteeIndex)
+func VerifyAttestationBitfieldLengths(ctx context.Context, state state.ReadOnlyBeaconState, att *ethpb.Attestation) error {
+	committee, err := BeaconCommitteeFromState(ctx, state, att.Data.Slot, att.Data.CommitteeIndex)
 	if err != nil {
 		return errors.Wrap(err, "could not retrieve beacon committees")
 	}
@@ -280,7 +283,6 @@ func UpdateCommitteeCache(state state.ReadOnlyBeaconState, epoch types.Epoch) er
 		if err != nil {
 			return err
 		}
-
 		if committeeCache.HasEntry(string(seed[:])) {
 			return nil
 		}
@@ -315,7 +317,7 @@ func UpdateCommitteeCache(state state.ReadOnlyBeaconState, epoch types.Epoch) er
 }
 
 // UpdateProposerIndicesInCache updates proposer indices entry of the committee cache.
-func UpdateProposerIndicesInCache(state state.ReadOnlyBeaconState) error {
+func UpdateProposerIndicesInCache(ctx context.Context, state state.ReadOnlyBeaconState) error {
 	// The cache uses the state root at the (current epoch - 2)'s slot as key. (e.g. for epoch 2, the key is root at slot 31)
 	// Which is the reason why we skip genesis epoch.
 	if core.CurrentEpoch(state) <= params.BeaconConfig().GenesisEpoch+params.BeaconConfig().MinSeedLookahead {
@@ -348,7 +350,7 @@ func UpdateProposerIndicesInCache(state state.ReadOnlyBeaconState) error {
 		return nil
 	}
 
-	indices, err := ActiveValidatorIndices(state, core.CurrentEpoch(state))
+	indices, err := ActiveValidatorIndices(ctx, state, core.CurrentEpoch(state))
 	if err != nil {
 		return err
 	}

beacon-chain/core/helpers/beacon_committee_test.go

@@ -1,6 +1,7 @@
 package helpers
 
 import (
+	"context"
 	"fmt"
 	"strconv"
 	"testing"
@@ -44,7 +45,7 @@ func TestComputeCommittee_WithoutCache(t *testing.T) {
 	require.NoError(t, err)
 
 	epoch := core.CurrentEpoch(state)
-	indices, err := ActiveValidatorIndices(state, epoch)
+	indices, err := ActiveValidatorIndices(context.Background(), state, epoch)
 	require.NoError(t, err)
 	seed, err := Seed(state, epoch, params.BeaconConfig().DomainBeaconAttester)
 	require.NoError(t, err)
@@ -94,7 +95,7 @@ func TestCommitteeAssignments_CannotRetrieveFutureEpoch(t *testing.T) {
 		Slot: 0, // Epoch 0.
 	})
 	require.NoError(t, err)
-	_, _, err = CommitteeAssignments(state, epoch+1)
+	_, _, err = CommitteeAssignments(context.Background(), state, epoch+1)
 	assert.ErrorContains(t, "can't be greater than next epoch", err)
 }
 
@@ -117,7 +118,7 @@ func TestCommitteeAssignments_NoProposerForSlot0(t *testing.T) {
 	})
 	require.NoError(t, err)
 	ClearCache()
-	_, proposerIndexToSlots, err := CommitteeAssignments(state, 0)
+	_, proposerIndexToSlots, err := CommitteeAssignments(context.Background(), state, 0)
 	require.NoError(t, err, "Failed to determine CommitteeAssignments")
 	for _, slots := range proposerIndexToSlots {
 		for _, s := range slots {
@@ -189,7 +190,7 @@ func TestCommitteeAssignments_CanRetrieve(t *testing.T) {
 	for i, tt := range tests {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
 			ClearCache()
-			validatorIndexToCommittee, proposerIndexToSlots, err := CommitteeAssignments(state, core.SlotToEpoch(tt.slot))
+			validatorIndexToCommittee, proposerIndexToSlots, err := CommitteeAssignments(context.Background(), state, core.SlotToEpoch(tt.slot))
 			require.NoError(t, err, "Failed to determine CommitteeAssignments")
 			cac := validatorIndexToCommittee[tt.index]
 			assert.Equal(t, tt.committeeIndex, cac.CommitteeIndex, "Unexpected committeeIndex for validator index %d", tt.index)
@@ -224,11 +225,11 @@ func TestCommitteeAssignments_CannotRetrieveFuture(t *testing.T) {
 		RandaoMixes: make([][]byte, params.BeaconConfig().EpochsPerHistoricalVector),
 	})
 	require.NoError(t, err)
-	_, proposerIndxs, err := CommitteeAssignments(state, core.CurrentEpoch(state))
+	_, proposerIndxs, err := CommitteeAssignments(context.Background(), state, core.CurrentEpoch(state))
 	require.NoError(t, err)
 	require.NotEqual(t, 0, len(proposerIndxs), "wanted non-zero proposer index set")
 
-	_, proposerIndxs, err = CommitteeAssignments(state, core.CurrentEpoch(state)+1)
+	_, proposerIndxs, err = CommitteeAssignments(context.Background(), state, core.CurrentEpoch(state)+1)
 	require.NoError(t, err)
 	require.Equal(t, 0, len(proposerIndxs), "wanted empty proposer index set")
 }
@@ -250,7 +251,7 @@ func TestCommitteeAssignments_EverySlotHasMin1Proposer(t *testing.T) {
 	require.NoError(t, err)
 	ClearCache()
 	epoch := types.Epoch(1)
-	_, proposerIndexToSlots, err := CommitteeAssignments(state, epoch)
+	_, proposerIndexToSlots, err := CommitteeAssignments(context.Background(), state, epoch)
 	require.NoError(t, err, "Failed to determine CommitteeAssignments")
 
 	slotsWithProposers := make(map[types.Slot]bool)
@@ -358,7 +359,7 @@ func TestVerifyAttestationBitfieldLengths_OK(t *testing.T) {
 	for i, tt := range tests {
 		ClearCache()
 		require.NoError(t, state.SetSlot(tt.stateSlot))
-		err := VerifyAttestationBitfieldLengths(state, tt.attestation)
+		err := VerifyAttestationBitfieldLengths(context.Background(), state, tt.attestation)
 		if tt.verificationFailure {
 			assert.NotNil(t, err, "Verification succeeded when it was supposed to fail")
 		} else {
@@ -391,7 +392,7 @@ func TestUpdateCommitteeCache_CanUpdate(t *testing.T) {
 	seed, err := Seed(state, epoch, params.BeaconConfig().DomainBeaconAttester)
 	require.NoError(t, err)
 
-	indices, err = committeeCache.Committee(params.BeaconConfig().SlotsPerEpoch.Mul(uint64(epoch)), seed, idx)
+	indices, err = committeeCache.Committee(context.Background(), params.BeaconConfig().SlotsPerEpoch.Mul(uint64(epoch)), seed, idx)
 	require.NoError(t, err)
 	assert.Equal(t, params.BeaconConfig().TargetCommitteeSize, uint64(len(indices)), "Did not save correct indices lengths")
 }
@@ -410,7 +411,7 @@ func BenchmarkComputeCommittee300000_WithPreCache(b *testing.B) {
 	require.NoError(b, err)
 
 	epoch := core.CurrentEpoch(state)
-	indices, err := ActiveValidatorIndices(state, epoch)
+	indices, err := ActiveValidatorIndices(context.Background(), state, epoch)
 	require.NoError(b, err)
 	seed, err := Seed(state, epoch, params.BeaconConfig().DomainBeaconAttester)
 	require.NoError(b, err)
@@ -444,7 +445,7 @@ func BenchmarkComputeCommittee3000000_WithPreCache(b *testing.B) {
 	require.NoError(b, err)
 
 	epoch := core.CurrentEpoch(state)
-	indices, err := ActiveValidatorIndices(state, epoch)
+	indices, err := ActiveValidatorIndices(context.Background(), state, epoch)
 	require.NoError(b, err)
 	seed, err := Seed(state, epoch, params.BeaconConfig().DomainBeaconAttester)
 	require.NoError(b, err)
@@ -478,7 +479,7 @@ func BenchmarkComputeCommittee128000_WithOutPreCache(b *testing.B) {
 	require.NoError(b, err)
 
 	epoch := core.CurrentEpoch(state)
-	indices, err := ActiveValidatorIndices(state, epoch)
+	indices, err := ActiveValidatorIndices(context.Background(), state, epoch)
 	require.NoError(b, err)
 	seed, err := Seed(state, epoch, params.BeaconConfig().DomainBeaconAttester)
 	require.NoError(b, err)
@@ -513,7 +514,7 @@ func BenchmarkComputeCommittee1000000_WithOutCache(b *testing.B) {
 	require.NoError(b, err)
 
 	epoch := core.CurrentEpoch(state)
-	indices, err := ActiveValidatorIndices(state, epoch)
+	indices, err := ActiveValidatorIndices(context.Background(), state, epoch)
 	require.NoError(b, err)
 	seed, err := Seed(state, epoch, params.BeaconConfig().DomainBeaconAttester)
 	require.NoError(b, err)
@@ -548,7 +549,7 @@ func BenchmarkComputeCommittee4000000_WithOutCache(b *testing.B) {
 	require.NoError(b, err)
 
 	epoch := core.CurrentEpoch(state)
-	indices, err := ActiveValidatorIndices(state, epoch)
+	indices, err := ActiveValidatorIndices(context.Background(), state, epoch)
 	require.NoError(b, err)
 	seed, err := Seed(state, epoch, params.BeaconConfig().DomainBeaconAttester)
 	require.NoError(b, err)
@@ -584,13 +585,13 @@ func TestBeaconCommitteeFromState_UpdateCacheForPreviousEpoch(t *testing.T) {
 		RandaoMixes: make([][]byte, params.BeaconConfig().EpochsPerHistoricalVector),
 	})
 	require.NoError(t, err)
-	_, err = BeaconCommitteeFromState(state, 1 /* previous epoch */, 0)
+	_, err = BeaconCommitteeFromState(context.Background(), state, 1 /* previous epoch */, 0)
 	require.NoError(t, err)
 
 	// Verify previous epoch is cached
 	seed, err := Seed(state, 0, params.BeaconConfig().DomainBeaconAttester)
 	require.NoError(t, err)
-	activeIndices, err := committeeCache.ActiveIndices(seed)
+	activeIndices, err := committeeCache.ActiveIndices(context.Background(), seed)
 	require.NoError(t, err)
 	assert.NotNil(t, activeIndices, "Did not cache active indices")
 }
@@ -609,7 +610,7 @@ func TestPrecomputeProposerIndices_Ok(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	indices, err := ActiveValidatorIndices(state, 0)
+	indices, err := ActiveValidatorIndices(context.Background(), state, 0)
 	require.NoError(t, err)
 
 	proposerIndices, err := precomputeProposerIndices(state, indices)

beacon-chain/core/helpers/signing_root_test.go

@@ -2,6 +2,7 @@ package helpers_test
 
 import (
 	"bytes"
+	"context"
 	"testing"
 
 	fuzz "github.com/google/gofuzz"
@@ -80,7 +81,7 @@ func TestSigningRoot_ComputeDomainAndSign(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			beaconState, privKeys := tt.genState(t)
-			idx, err := helpers.BeaconProposerIndex(beaconState)
+			idx, err := helpers.BeaconProposerIndex(context.Background(), beaconState)
 			require.NoError(t, err)
 			block := tt.genBlock(t, beaconState, privKeys)
 			got, err := helpers.ComputeDomainAndSign(

beacon-chain/core/helpers/validators.go

@@ -2,9 +2,13 @@ package helpers
 
 import (
 	"bytes"
+	"context"
 
 	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
 	types "github.com/prysmaticlabs/eth2-types"
+	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state"
 	"github.com/prysmaticlabs/prysm/config/params"
@@ -12,8 +16,14 @@ import (
 	"github.com/prysmaticlabs/prysm/crypto/hash"
 	"github.com/prysmaticlabs/prysm/encoding/bytesutil"
 	ethpb "github.com/prysmaticlabs/prysm/proto/prysm/v1alpha1"
+	log "github.com/sirupsen/logrus"
 )
 
+var CommitteeCacheInProgressHit = promauto.NewCounter(prometheus.CounterOpts{
+	Name: "committee_cache_in_progress_hit",
+	Help: "The number of committee requests that are present in the cache.",
+})
+
 // IsActiveValidator returns the boolean value on whether the validator
 // is active or not.
 //
@@ -73,18 +83,39 @@ func checkValidatorSlashable(activationEpoch, withdrawableEpoch types.Epoch, sla
 //    Return the sequence of active validator indices at ``epoch``.
 //    """
 //    return [ValidatorIndex(i) for i, v in enumerate(state.validators) if is_active_validator(v, epoch)]
-func ActiveValidatorIndices(s state.ReadOnlyBeaconState, epoch types.Epoch) ([]types.ValidatorIndex, error) {
+func ActiveValidatorIndices(ctx context.Context, s state.ReadOnlyBeaconState, epoch types.Epoch) ([]types.ValidatorIndex, error) {
 	seed, err := Seed(s, epoch, params.BeaconConfig().DomainBeaconAttester)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not get seed")
 	}
-	activeIndices, err := committeeCache.ActiveIndices(seed)
+	activeIndices, err := committeeCache.ActiveIndices(ctx, seed)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not interface with committee cache")
 	}
 	if activeIndices != nil {
 		return activeIndices, nil
 	}
+
+	if err := committeeCache.MarkInProgress(seed); err != nil {
+		if errors.Is(err, cache.ErrAlreadyInProgress) {
+			activeIndices, err := committeeCache.ActiveIndices(ctx, seed)
+			if err != nil {
+				return nil, err
+			}
+			if activeIndices == nil {
+				return nil, errors.New("nil active indices")
+			}
+			CommitteeCacheInProgressHit.Inc()
+			return activeIndices, nil
+		}
+		return nil, errors.Wrap(err, "could not mark committee cache as in progress")
+	}
+	defer func() {
+		if err := committeeCache.MarkNotInProgress(seed); err != nil {
+			log.WithError(err).Error("Could not mark cache not in progress")
+		}
+	}()
+
 	var indices []types.ValidatorIndex
 	if err := s.ReadFromEveryValidator(func(idx int, val state.ReadOnlyValidator) error {
 		if IsActiveValidatorUsingTrie(val, epoch) {
@@ -104,12 +135,12 @@ func ActiveValidatorIndices(s state.ReadOnlyBeaconState, epoch types.Epoch) ([]t
 
 // ActiveValidatorCount returns the number of active validators in the state
 // at the given epoch.
-func ActiveValidatorCount(s state.ReadOnlyBeaconState, epoch types.Epoch) (uint64, error) {
+func ActiveValidatorCount(ctx context.Context, s state.ReadOnlyBeaconState, epoch types.Epoch) (uint64, error) {
 	seed, err := Seed(s, epoch, params.BeaconConfig().DomainBeaconAttester)
 	if err != nil {
 		return 0, errors.Wrap(err, "could not get seed")
 	}
-	activeCount, err := committeeCache.ActiveIndicesCount(seed)
+	activeCount, err := committeeCache.ActiveIndicesCount(ctx, seed)
 	if err != nil {
 		return 0, errors.Wrap(err, "could not interface with committee cache")
 	}
@@ -117,6 +148,23 @@ func ActiveValidatorCount(s state.ReadOnlyBeaconState, epoch types.Epoch) (uint6
 		return uint64(activeCount), nil
 	}
 
+	if err := committeeCache.MarkInProgress(seed); err != nil {
+		if errors.Is(err, cache.ErrAlreadyInProgress) {
+			activeCount, err := committeeCache.ActiveIndicesCount(ctx, seed)
+			if err != nil {
+				return 0, err
+			}
+			CommitteeCacheInProgressHit.Inc()
+			return uint64(activeCount), nil
+		}
+		return 0, errors.Wrap(err, "could not mark committee cache as in progress")
+	}
+	defer func() {
+		if err := committeeCache.MarkNotInProgress(seed); err != nil {
+			log.WithError(err).Error("Could not mark cache not in progress")
+		}
+	}()
+
 	count := uint64(0)
 	if err := s.ReadFromEveryValidator(func(idx int, val state.ReadOnlyValidator) error {
 		if IsActiveValidatorUsingTrie(val, epoch) {
@@ -176,7 +224,7 @@ func ValidatorChurnLimit(activeValidatorCount uint64) (uint64, error) {
 //    seed = hash(get_seed(state, epoch, DOMAIN_BEACON_PROPOSER) + uint_to_bytes(state.slot))
 //    indices = get_active_validator_indices(state, epoch)
 //    return compute_proposer_index(state, indices, seed)
-func BeaconProposerIndex(state state.ReadOnlyBeaconState) (types.ValidatorIndex, error) {
+func BeaconProposerIndex(ctx context.Context, state state.ReadOnlyBeaconState) (types.ValidatorIndex, error) {
 	e := core.CurrentEpoch(state)
 	// The cache uses the state root of the previous epoch - minimum_seed_lookahead last slot as key. (e.g. Starting epoch 1, slot 32, the key would be block root at slot 31)
 	// For simplicity, the node will skip caching of genesis epoch.
@@ -204,7 +252,7 @@ func BeaconProposerIndex(state state.ReadOnlyBeaconState) (types.ValidatorIndex,
 				}
 				return proposerIndices[state.Slot()%params.BeaconConfig().SlotsPerEpoch], nil
 			}
-			if err := UpdateProposerIndicesInCache(state); err != nil {
+			if err := UpdateProposerIndicesInCache(ctx, state); err != nil {
 				return 0, errors.Wrap(err, "could not update committee cache")
 			}
 		}
@@ -218,7 +266,7 @@ func BeaconProposerIndex(state state.ReadOnlyBeaconState) (types.ValidatorIndex,
 	seedWithSlot := append(seed[:], bytesutil.Bytes8(uint64(state.Slot()))...)
 	seedWithSlotHash := hash.Hash(seedWithSlot)
 
-	indices, err := ActiveValidatorIndices(state, e)
+	indices, err := ActiveValidatorIndices(ctx, state, e)
 	if err != nil {
 		return 0, errors.Wrap(err, "could not get active indices")
 	}

beacon-chain/core/helpers/validators_test.go

@@ -1,6 +1,7 @@
 package helpers
 
 import (
+	"context"
 	"errors"
 	"testing"
 
@@ -270,7 +271,7 @@ func TestBeaconProposerIndex_OK(t *testing.T) {
 	for _, tt := range tests {
 		ClearCache()
 		require.NoError(t, state.SetSlot(tt.slot))
-		result, err := BeaconProposerIndex(state)
+		result, err := BeaconProposerIndex(context.Background(), state)
 		require.NoError(t, err, "Failed to get shard and committees at slot")
 		assert.Equal(t, tt.index, result, "Result index was an unexpected value")
 	}
@@ -304,7 +305,7 @@ func TestBeaconProposerIndex_BadState(t *testing.T) {
 	// Set a very high slot, so that retrieved block root will be
 	// non existent for the proposer cache.
 	require.NoError(t, state.SetSlot(100))
-	_, err = BeaconProposerIndex(state)
+	_, err = BeaconProposerIndex(context.Background(), state)
 	require.NoError(t, err)
 	assert.Equal(t, 0, len(proposerIndicesCache.ProposerIndicesCache.ListKeys()))
 }
@@ -323,7 +324,7 @@ func TestComputeProposerIndex_Compatibility(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	indices, err := ActiveValidatorIndices(state, 0)
+	indices, err := ActiveValidatorIndices(context.Background(), state, 0)
 	require.NoError(t, err)
 
 	var proposerIndices []types.ValidatorIndex
@@ -375,7 +376,7 @@ func TestActiveValidatorCount_Genesis(t *testing.T) {
 	seed, err := Seed(beaconState, 0, params.BeaconConfig().DomainBeaconAttester)
 	require.NoError(t, err)
 	require.NoError(t, committeeCache.AddCommitteeShuffledList(&cache.Committees{Seed: seed, ShuffledIndices: []types.ValidatorIndex{1, 2, 3}}))
-	validatorCount, err := ActiveValidatorCount(beaconState, core.CurrentEpoch(beaconState))
+	validatorCount, err := ActiveValidatorCount(context.Background(), beaconState, core.CurrentEpoch(beaconState))
 	require.NoError(t, err)
 	assert.Equal(t, uint64(c), validatorCount, "Did not get the correct validator count")
 }
@@ -406,7 +407,7 @@ func TestChurnLimit_OK(t *testing.T) {
 			RandaoMixes: make([][]byte, params.BeaconConfig().EpochsPerHistoricalVector),
 		})
 		require.NoError(t, err)
-		validatorCount, err := ActiveValidatorCount(beaconState, core.CurrentEpoch(beaconState))
+		validatorCount, err := ActiveValidatorCount(context.Background(), beaconState, core.CurrentEpoch(beaconState))
 		require.NoError(t, err)
 		resultChurn, err := ValidatorChurnLimit(validatorCount)
 		require.NoError(t, err)
@@ -591,7 +592,7 @@ func TestActiveValidatorIndices(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			s, err := v1.InitializeFromProto(tt.args.state)
 			require.NoError(t, err)
-			got, err := ActiveValidatorIndices(s, tt.args.epoch)
+			got, err := ActiveValidatorIndices(context.Background(), s, tt.args.epoch)
 			if tt.wantedErr != "" {
 				assert.ErrorContains(t, tt.wantedErr, err)
 				return

beacon-chain/core/helpers/weak_subjectivity.go

@@ -2,6 +2,7 @@ package helpers
 
 import (
 	"bytes"
+	"context"
 	"encoding/hex"
 	"errors"
 	"fmt"
@@ -52,12 +53,12 @@ import (
 //        )
 //
 //    return ws_period
-func ComputeWeakSubjectivityPeriod(st state.ReadOnlyBeaconState) (types.Epoch, error) {
+func ComputeWeakSubjectivityPeriod(ctx context.Context, st state.ReadOnlyBeaconState) (types.Epoch, error) {
 	// Weak subjectivity period cannot be smaller than withdrawal delay.
 	wsp := uint64(params.BeaconConfig().MinValidatorWithdrawabilityDelay)
 
 	// Cardinality of active validator set.
-	N, err := ActiveValidatorCount(st, core.CurrentEpoch(st))
+	N, err := ActiveValidatorCount(ctx, st, core.CurrentEpoch(st))
 	if err != nil {
 		return 0, fmt.Errorf("cannot obtain active valiadtor count: %w", err)
 	}
@@ -120,7 +121,7 @@ func ComputeWeakSubjectivityPeriod(st state.ReadOnlyBeaconState) (types.Epoch, e
 //    current_epoch = compute_epoch_at_slot(get_current_slot(store))
 //    return current_epoch <= ws_state_epoch + ws_period
 func IsWithinWeakSubjectivityPeriod(
-	currentEpoch types.Epoch, wsState state.ReadOnlyBeaconState, wsCheckpoint *eth.WeakSubjectivityCheckpoint) (bool, error) {
+	ctx context.Context, currentEpoch types.Epoch, wsState state.ReadOnlyBeaconState, wsCheckpoint *eth.WeakSubjectivityCheckpoint) (bool, error) {
 	// Make sure that incoming objects are not nil.
 	if wsState == nil || wsState.IsNil() || wsState.LatestBlockHeader() == nil || wsCheckpoint == nil {
 		return false, errors.New("invalid weak subjectivity state or checkpoint")
@@ -137,7 +138,7 @@ func IsWithinWeakSubjectivityPeriod(
 	}
 
 	// Compare given epoch to state epoch + weak subjectivity period.
-	wsPeriod, err := ComputeWeakSubjectivityPeriod(wsState)
+	wsPeriod, err := ComputeWeakSubjectivityPeriod(ctx, wsState)
 	if err != nil {
 		return false, fmt.Errorf("cannot compute weak subjectivity period: %w", err)
 	}
@@ -151,8 +152,8 @@ func IsWithinWeakSubjectivityPeriod(
 // Within the weak subjectivity period, if two conflicting blocks are finalized, 1/3 - D (D := safety decay)
 // of validators will get slashed. Therefore, it is safe to assume that any finalized checkpoint within that
 // period is protected by this safety margin.
-func LatestWeakSubjectivityEpoch(st state.ReadOnlyBeaconState) (types.Epoch, error) {
-	wsPeriod, err := ComputeWeakSubjectivityPeriod(st)
+func LatestWeakSubjectivityEpoch(ctx context.Context, st state.ReadOnlyBeaconState) (types.Epoch, error) {
+	wsPeriod, err := ComputeWeakSubjectivityPeriod(ctx, st)
 	if err != nil {
 		return 0, err
 	}

beacon-chain/core/helpers/weak_subjectivity_test.go

@@ -1,6 +1,7 @@
 package helpers_test
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
@@ -47,7 +48,7 @@ func TestWeakSubjectivity_ComputeWeakSubjectivityPeriod(t *testing.T) {
 		t.Run(fmt.Sprintf("valCount: %d, avgBalance: %d", tt.valCount, tt.avgBalance), func(t *testing.T) {
 			// Reset committee cache - as we need to recalculate active validator set for each test.
 			helpers.ClearCache()
-			got, err := helpers.ComputeWeakSubjectivityPeriod(genState(t, tt.valCount, tt.avgBalance))
+			got, err := helpers.ComputeWeakSubjectivityPeriod(context.Background(), genState(t, tt.valCount, tt.avgBalance))
 			require.NoError(t, err)
 			assert.Equal(t, tt.want, got, "valCount: %v, avgBalance: %v", tt.valCount, tt.avgBalance)
 		})
@@ -192,7 +193,7 @@ func TestWeakSubjectivity_IsWithinWeakSubjectivityPeriod(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := helpers.IsWithinWeakSubjectivityPeriod(tt.epoch, tt.genWsState(), tt.genWsCheckpoint())
+			got, err := helpers.IsWithinWeakSubjectivityPeriod(context.Background(), tt.epoch, tt.genWsState(), tt.genWsCheckpoint())
 			if tt.wantedErr != "" {
 				assert.Equal(t, false, got)
 				assert.ErrorContains(t, tt.wantedErr, err)