Add and verify context in AddCommitteeShuffledList

beacon-chain/blockchain/process_block.go

@@ -537,7 +537,7 @@ func (s *Service) handleEpochBoundary(ctx context.Context, postState state.Beaco
 			return err
 		}
 		// Update caches for the next epoch at epoch boundary slot - 1.
-		if err := helpers.UpdateCommitteeCache(copied, coreTime.CurrentEpoch(copied)); err != nil {
+		if err := helpers.UpdateCommitteeCache(ctx, copied, coreTime.CurrentEpoch(copied)); err != nil {
 			return err
 		}
 		if err := helpers.UpdateProposerIndicesInCache(ctx, copied); err != nil {
@@ -559,7 +559,7 @@ func (s *Service) handleEpochBoundary(ctx context.Context, postState state.Beaco
 
 		// Update caches at epoch boundary slot.
 		// The following updates have short cut to return nil cheaply if fulfilled during boundary slot - 1.
-		if err := helpers.UpdateCommitteeCache(postState, coreTime.CurrentEpoch(postState)); err != nil {
+		if err := helpers.UpdateCommitteeCache(ctx, postState, coreTime.CurrentEpoch(postState)); err != nil {
 			return err
 		}
 		if err := helpers.UpdateProposerIndicesInCache(ctx, postState); err != nil {

beacon-chain/blockchain/service.go

@@ -451,7 +451,7 @@ func (s *Service) initializeBeaconChain(
 	s.cfg.ChainStartFetcher.ClearPreGenesisData()
 
 	// Update committee shuffled indices for genesis epoch.
-	if err := helpers.UpdateCommitteeCache(genesisState, 0 /* genesis epoch */); err != nil {
+	if err := helpers.UpdateCommitteeCache(ctx, genesisState, 0); err != nil {
 		return nil, err
 	}
 	if err := helpers.UpdateProposerIndicesInCache(ctx, genesisState); err != nil {

beacon-chain/cache/committee.go

@@ -103,9 +103,12 @@ func (c *CommitteeCache) Committee(ctx context.Context, slot types.Slot, seed [3
 
 // AddCommitteeShuffledList adds Committee shuffled list object to the cache. T
 // his method also trims the least recently list if the cache size has ready the max cache size limit.
-func (c *CommitteeCache) AddCommitteeShuffledList(committees *Committees) error {
+func (c *CommitteeCache) AddCommitteeShuffledList(ctx context.Context, committees *Committees) error {
 	c.lock.Lock()
 	defer c.lock.Unlock()
+	if err := ctx.Err(); err != nil {
+		return err
+	}
 	key, err := committeeKeyFn(committees)
 	if err != nil {
 		return err

beacon-chain/cache/committee_fuzz_test.go

@@ -31,7 +31,7 @@ func TestCommitteeCache_FuzzCommitteesByEpoch(t *testing.T) {
 
 	for i := 0; i < 100000; i++ {
 		fuzzer.Fuzz(c)
-		require.NoError(t, cache.AddCommitteeShuffledList(c))
+		require.NoError(t, cache.AddCommitteeShuffledList(context.Background(), c))
 		_, err := cache.Committee(context.Background(), 0, c.Seed, 0)
 		require.NoError(t, err)
 	}
@@ -46,7 +46,7 @@ func TestCommitteeCache_FuzzActiveIndices(t *testing.T) {
 
 	for i := 0; i < 100000; i++ {
 		fuzzer.Fuzz(c)
-		require.NoError(t, cache.AddCommitteeShuffledList(c))
+		require.NoError(t, cache.AddCommitteeShuffledList(context.Background(), c))
 
 		indices, err := cache.ActiveIndices(context.Background(), c.Seed)
 		require.NoError(t, err)

beacon-chain/cache/committee_test.go

@@ -50,7 +50,7 @@ func TestCommitteeCache_CommitteesByEpoch(t *testing.T) {
 	if indices != nil {
 		t.Error("Expected committee not to exist in empty cache")
 	}
-	require.NoError(t, cache.AddCommitteeShuffledList(item))
+	require.NoError(t, cache.AddCommitteeShuffledList(context.Background(), item))
 
 	wantedIndex := types.CommitteeIndex(0)
 	indices, err = cache.Committee(context.Background(), slot, item.Seed, wantedIndex)
@@ -70,7 +70,7 @@ func TestCommitteeCache_ActiveIndices(t *testing.T) {
 		t.Error("Expected committee not to exist in empty cache")
 	}
 
-	require.NoError(t, cache.AddCommitteeShuffledList(item))
+	require.NoError(t, cache.AddCommitteeShuffledList(context.Background(), item))
 
 	indices, err = cache.ActiveIndices(context.Background(), item.Seed)
 	require.NoError(t, err)
@@ -85,7 +85,7 @@ func TestCommitteeCache_ActiveCount(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, 0, count, "Expected active count not to exist in empty cache")
 
-	require.NoError(t, cache.AddCommitteeShuffledList(item))
+	require.NoError(t, cache.AddCommitteeShuffledList(context.Background(), item))
 
 	count, err = cache.ActiveIndicesCount(context.Background(), item.Seed)
 	require.NoError(t, err)
@@ -101,7 +101,7 @@ func TestCommitteeCache_CanRotate(t *testing.T) {
 	for i := start; i < end; i++ {
 		s := []byte(strconv.Itoa(i))
 		item := &Committees{Seed: bytesutil.ToBytes32(s)}
-		require.NoError(t, cache.AddCommitteeShuffledList(item))
+		require.NoError(t, cache.AddCommitteeShuffledList(context.Background(), item))
 	}
 
 	k := cache.CommitteeCache.Keys()
@@ -134,3 +134,20 @@ func TestCommitteeCacheOutOfRange(t *testing.T) {
 	_, err = cache.Committee(context.Background(), 0, seed, math.MaxUint64) // Overflow!
 	require.NotNil(t, err, "Did not fail as expected")
 }
+
+func TestCommitteeCache_DoesNothingWhenCancelledContext(t *testing.T) {
+	cache := NewCommitteesCache()
+
+	item := &Committees{Seed: [32]byte{'A'}, SortedIndices: []types.ValidatorIndex{1, 2, 3, 4, 5, 6}}
+	count, err := cache.ActiveIndicesCount(context.Background(), item.Seed)
+	require.NoError(t, err)
+	assert.Equal(t, 0, count, "Expected active count not to exist in empty cache")
+
+	cancelled, cancel := context.WithCancel(context.Background())
+	cancel()
+	require.ErrorIs(t, cache.AddCommitteeShuffledList(cancelled, item), context.Canceled)
+
+	count, err = cache.ActiveIndicesCount(context.Background(), item.Seed)
+	require.NoError(t, err)
+	assert.Equal(t, 0, count)
+}

beacon-chain/core/helpers/beacon_committee.go

@@ -285,7 +285,7 @@ func ShuffledIndices(s state.ReadOnlyBeaconState, epoch types.Epoch) ([]types.Va
 
 // UpdateCommitteeCache gets called at the beginning of every epoch to cache the committee shuffled indices
 // list with committee index and epoch number. It caches the shuffled indices for current epoch and next epoch.
-func UpdateCommitteeCache(state state.ReadOnlyBeaconState, epoch types.Epoch) error {
+func UpdateCommitteeCache(ctx context.Context, state state.ReadOnlyBeaconState, epoch types.Epoch) error {
 	for _, e := range []types.Epoch{epoch, epoch + 1} {
 		seed, err := Seed(state, e, params.BeaconConfig().DomainBeaconAttester)
 		if err != nil {
@@ -311,7 +311,7 @@ func UpdateCommitteeCache(state state.ReadOnlyBeaconState, epoch types.Epoch) er
 			return sortedIndices[i] < sortedIndices[j]
 		})
 
-		if err := committeeCache.AddCommitteeShuffledList(&cache.Committees{
+		if err := committeeCache.AddCommitteeShuffledList(ctx, &cache.Committees{
 			ShuffledIndices: shuffledIndices,
 			CommitteeCount:  uint64(params.BeaconConfig().SlotsPerEpoch.Mul(count)),
 			Seed:            seed,

beacon-chain/core/helpers/beacon_committee_test.go

@@ -386,7 +386,7 @@ func TestUpdateCommitteeCache_CanUpdate(t *testing.T) {
 		RandaoMixes: make([][]byte, params.BeaconConfig().EpochsPerHistoricalVector),
 	})
 	require.NoError(t, err)
-	require.NoError(t, UpdateCommitteeCache(state, time.CurrentEpoch(state)))
+	require.NoError(t, UpdateCommitteeCache(context.Background(), state, time.CurrentEpoch(state)))
 
 	epoch := types.Epoch(1)
 	idx := types.CommitteeIndex(1)

beacon-chain/core/helpers/validators.go

@@ -126,7 +126,7 @@ func ActiveValidatorIndices(ctx context.Context, s state.ReadOnlyBeaconState, ep
 		return nil, err
 	}
 
-	if err := UpdateCommitteeCache(s, epoch); err != nil {
+	if err := UpdateCommitteeCache(ctx, s, epoch); err != nil {
 		return nil, errors.Wrap(err, "could not update committee cache")
 	}
 
@@ -175,7 +175,7 @@ func ActiveValidatorCount(ctx context.Context, s state.ReadOnlyBeaconState, epoc
 		return 0, err
 	}
 
-	if err := UpdateCommitteeCache(s, epoch); err != nil {
+	if err := UpdateCommitteeCache(ctx, s, epoch); err != nil {
 		return 0, errors.Wrap(err, "could not update committee cache")
 	}
 

beacon-chain/core/helpers/validators_test.go

@@ -301,7 +301,7 @@ func TestActiveValidatorCount_Genesis(t *testing.T) {
 	// Preset cache to a bad count.
 	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}}))
+	require.NoError(t, committeeCache.AddCommitteeShuffledList(context.Background(), &cache.Committees{Seed: seed, ShuffledIndices: []types.ValidatorIndex{1, 2, 3}}))
 	validatorCount, err := ActiveValidatorCount(context.Background(), beaconState, time.CurrentEpoch(beaconState))
 	require.NoError(t, err)
 	assert.Equal(t, uint64(c), validatorCount, "Did not get the correct validator count")

beacon-chain/core/transition/benchmarks_test.go

@@ -53,7 +53,7 @@ func BenchmarkExecuteStateTransition_WithCache(b *testing.B) {
 	// some attestations in block are from previous epoch
 	currentSlot := beaconState.Slot()
 	require.NoError(b, beaconState.SetSlot(beaconState.Slot()-params.BeaconConfig().SlotsPerEpoch))
-	require.NoError(b, helpers.UpdateCommitteeCache(beaconState, time.CurrentEpoch(beaconState)))
+	require.NoError(b, helpers.UpdateCommitteeCache(context.Background(), beaconState, time.CurrentEpoch(beaconState)))
 	require.NoError(b, beaconState.SetSlot(currentSlot))
 	// Run the state transition once to populate the cache.
 	wsb, err := wrapper.WrappedSignedBeaconBlock(block)
@@ -81,7 +81,7 @@ func BenchmarkProcessEpoch_2FullEpochs(b *testing.B) {
 	// some attestations in block are from previous epoch
 	currentSlot := beaconState.Slot()
 	require.NoError(b, beaconState.SetSlot(beaconState.Slot()-params.BeaconConfig().SlotsPerEpoch))
-	require.NoError(b, helpers.UpdateCommitteeCache(beaconState, time.CurrentEpoch(beaconState)))
+	require.NoError(b, helpers.UpdateCommitteeCache(context.Background(), beaconState, time.CurrentEpoch(beaconState)))
 	require.NoError(b, beaconState.SetSlot(currentSlot))
 
 	b.ResetTimer()