Remove state balance cache (#11964)

* Remove state balance cache

* remove unused metrics

* add missing balance update

* add comment about locks

* update balance only when updataing the justified checkpoint on unrealization

* update checkpoint balances at launch

* fix unit tests

* fix remaining test

* review

* fix build file

* fix test

---------

Co-authored-by: terencechain <[email protected]>
Co-authored-by: Radosław Kapka <[email protected]>

beacon-chain/blockchain/BUILD.bazel

@@ -22,7 +22,6 @@ go_library(
         "receive_attestation.go",
         "receive_block.go",
         "service.go",
-        "state_balance_cache.go",
         "weak_subjectivity_checks.go",
     ],
     importpath = "github.com/prysmaticlabs/prysm/v3/beacon-chain/blockchain",
@@ -118,7 +117,6 @@ go_test(
         "receive_attestation_test.go",
         "receive_block_test.go",
         "service_test.go",
-        "state_balance_cache_test.go",
         "weak_subjectivity_checks_test.go",
     ],
     embed = [":go_default_library"],

beacon-chain/blockchain/chain_info_test.go

@@ -96,7 +96,7 @@ func TestFinalizedCheckpt_GenesisRootOk(t *testing.T) {
 	cp = service.CurrentJustifiedCheckpt()
 	assert.DeepEqual(t, [32]byte{}, bytesutil.ToBytes32(cp.Root))
 	// check that forkchoice has the right genesis root as the node root
-	root, err := fcs.Head(ctx, []uint64{})
+	root, err := fcs.Head(ctx)
 	require.NoError(t, err)
 	require.Equal(t, service.originBlockRoot, root)
 
@@ -114,8 +114,12 @@ func TestCurrentJustifiedCheckpt_CanRetrieve(t *testing.T) {
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
 
-	cp := &forkchoicetypes.Checkpoint{Epoch: 6, Root: [32]byte{'j'}}
-	require.NoError(t, fcs.UpdateJustifiedCheckpoint(cp))
+	jroot := [32]byte{'j'}
+	cp := &forkchoicetypes.Checkpoint{Epoch: 6, Root: jroot}
+	bState, _ := util.DeterministicGenesisState(t, 10)
+	require.NoError(t, beaconDB.SaveState(ctx, bState, jroot))
+
+	require.NoError(t, fcs.UpdateJustifiedCheckpoint(ctx, cp))
 	jp := service.CurrentJustifiedCheckpt()
 	assert.Equal(t, cp.Epoch, jp.Epoch, "Unexpected justified epoch")
 	require.Equal(t, cp.Root, bytesutil.ToBytes32(jp.Root))

beacon-chain/blockchain/execution_engine.go

@@ -98,7 +98,7 @@ func (s *Service) notifyForkchoiceUpdate(ctx context.Context, arg *notifyForkcho
 				return nil, nil
 			}
 
-			r, err := s.cfg.ForkChoiceStore.Head(ctx, s.justifiedBalances.balances)
+			r, err := s.cfg.ForkChoiceStore.Head(ctx)
 			if err != nil {
 				log.WithFields(logrus.Fields{
 					"slot":                 headBlk.Slot(),

beacon-chain/blockchain/execution_engine_test.go

@@ -302,8 +302,8 @@ func Test_NotifyForkchoiceUpdate_NIlLVH(t *testing.T) {
 	service.cfg.ForkChoiceStore = fcs
 	service.cfg.ProposerSlotIndexCache = cache.NewProposerPayloadIDsCache()
 
-	service.justifiedBalances.balances = []uint64{50, 100, 200}
-	require.NoError(t, err)
+	fcs.SetBalancesByRooter(func(context.Context, [32]byte) ([]uint64, error) { return []uint64{50, 100, 200}, nil })
+	require.NoError(t, fcs.UpdateJustifiedCheckpoint(ctx, &forkchoicetypes.Checkpoint{}))
 	ojc := &ethpb.Checkpoint{Root: params.BeaconConfig().ZeroHash[:]}
 	ofc := &ethpb.Checkpoint{Root: params.BeaconConfig().ZeroHash[:]}
 	state, blkRoot, err := prepareForkchoiceState(ctx, 1, bra, [32]byte{}, [32]byte{'A'}, ojc, ofc)
@@ -419,12 +419,15 @@ func Test_NotifyForkchoiceUpdateRecursive_DoublyLinkedTree(t *testing.T) {
 	service.cfg.ForkChoiceStore = fcs
 	service.cfg.ProposerSlotIndexCache = cache.NewProposerPayloadIDsCache()
 
-	service.justifiedBalances.balances = []uint64{50, 100, 200}
-	require.NoError(t, err)
+	fcs.SetBalancesByRooter(func(context.Context, [32]byte) ([]uint64, error) { return []uint64{50, 100, 200}, nil })
+	require.NoError(t, fcs.UpdateJustifiedCheckpoint(ctx, &forkchoicetypes.Checkpoint{}))
 	ojc := &ethpb.Checkpoint{Root: params.BeaconConfig().ZeroHash[:]}
 	ofc := &ethpb.Checkpoint{Root: params.BeaconConfig().ZeroHash[:]}
 	state, blkRoot, err := prepareForkchoiceState(ctx, 1, bra, [32]byte{}, [32]byte{'A'}, ojc, ofc)
 	require.NoError(t, err)
+
+	bState, _ := util.DeterministicGenesisState(t, 10)
+	require.NoError(t, beaconDB.SaveState(ctx, bState, bra))
 	require.NoError(t, fcs.InsertNode(ctx, state, blkRoot))
 	state, blkRoot, err = prepareForkchoiceState(ctx, 2, brb, bra, [32]byte{'B'}, ojc, ofc)
 	require.NoError(t, err)
@@ -450,9 +453,12 @@ func Test_NotifyForkchoiceUpdateRecursive_DoublyLinkedTree(t *testing.T) {
 	fcs.ProcessAttestation(ctx, []uint64{0}, brd, 1)
 	fcs.ProcessAttestation(ctx, []uint64{1}, brf, 1)
 	fcs.ProcessAttestation(ctx, []uint64{2}, brg, 1)
+	fcs.SetBalancesByRooter(service.cfg.StateGen.ActiveNonSlashedBalancesByRoot)
 	jc := &forkchoicetypes.Checkpoint{Epoch: 0, Root: bra}
-	require.NoError(t, fcs.UpdateJustifiedCheckpoint(jc))
-	headRoot, err := fcs.Head(ctx, []uint64{50, 100, 200})
+	require.NoError(t, fcs.UpdateJustifiedCheckpoint(ctx, jc))
+	fcs.SetBalancesByRooter(func(context.Context, [32]byte) ([]uint64, error) { return []uint64{50, 100, 200}, nil })
+	require.NoError(t, fcs.UpdateJustifiedCheckpoint(ctx, &forkchoicetypes.Checkpoint{}))
+	headRoot, err := fcs.Head(ctx)
 	require.NoError(t, err)
 	require.Equal(t, brg, headRoot)
 
@@ -476,7 +482,7 @@ func Test_NotifyForkchoiceUpdateRecursive_DoublyLinkedTree(t *testing.T) {
 	require.Equal(t, brf, InvalidBlockRoot(err))
 
 	// Ensure Head is D
-	headRoot, err = fcs.Head(ctx, service.justifiedBalances.balances)
+	headRoot, err = fcs.Head(ctx)
 	require.NoError(t, err)
 	require.Equal(t, brd, headRoot)
 
@@ -886,7 +892,7 @@ func Test_UpdateLastValidatedCheckpoint(t *testing.T) {
 	ojc := &ethpb.Checkpoint{Root: params.BeaconConfig().ZeroHash[:]}
 	ofc := &ethpb.Checkpoint{Root: params.BeaconConfig().ZeroHash[:]}
 	fjc := &forkchoicetypes.Checkpoint{Epoch: 0, Root: params.BeaconConfig().ZeroHash}
-	require.NoError(t, fcs.UpdateJustifiedCheckpoint(fjc))
+	require.NoError(t, fcs.UpdateJustifiedCheckpoint(ctx, fjc))
 	require.NoError(t, fcs.UpdateFinalizedCheckpoint(fjc))
 	state, blkRoot, err := prepareForkchoiceState(ctx, 0, genesisRoot, params.BeaconConfig().ZeroHash, params.BeaconConfig().ZeroHash, ojc, ofc)
 	require.NoError(t, err)
@@ -953,6 +959,7 @@ func Test_UpdateLastValidatedCheckpoint(t *testing.T) {
 	twentyfc := &ethpb.Checkpoint{Epoch: 20, Root: validRoot[:]}
 	state, blkRoot, err = prepareForkchoiceState(ctx, 640, validRoot, genesisRoot, params.BeaconConfig().ZeroHash, twentyjc, twentyfc)
 	require.NoError(t, err)
+	fcs.SetBalancesByRooter(func(_ context.Context, _ [32]byte) ([]uint64, error) { return []uint64{}, nil })
 	require.NoError(t, fcs.InsertNode(ctx, state, blkRoot))
 	require.NoError(t, fcs.SetOptimisticToValid(ctx, validRoot))
 	assert.NoError(t, beaconDB.SaveGenesisBlockRoot(ctx, validRoot))

beacon-chain/blockchain/head.go

@@ -28,14 +28,7 @@ import (
 // UpdateAndSaveHeadWithBalances updates the beacon state head after getting justified balanced from cache.
 // This function is only used in spec-tests, it does save the head after updating it.
 func (s *Service) UpdateAndSaveHeadWithBalances(ctx context.Context) error {
-	jp := s.CurrentJustifiedCheckpt()
-
-	balances, err := s.justifiedBalances.get(ctx, bytesutil.ToBytes32(jp.Root))
-	if err != nil {
-		msg := fmt.Sprintf("could not read balances for state w/ justified checkpoint %#x", jp.Root)
-		return errors.Wrap(err, msg)
-	}
-	headRoot, err := s.cfg.ForkChoiceStore.Head(ctx, balances)
+	headRoot, err := s.cfg.ForkChoiceStore.Head(ctx)
 	if err != nil {
 		return errors.Wrap(err, "could not update head")
 	}

beacon-chain/blockchain/head_test.go

@@ -10,6 +10,7 @@ import (
 	mock "github.com/prysmaticlabs/prysm/v3/beacon-chain/blockchain/testing"
 	testDB "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	doublylinkedtree "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/doubly-linked-tree"
+	forkchoicetypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/types"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/blstoexec"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
 	"github.com/prysmaticlabs/prysm/v3/config/features"
@@ -152,20 +153,6 @@ func TestSaveHead_Different_Reorg(t *testing.T) {
 	require.LogsContain(t, hook, "depth=1")
 }
 
-func TestCacheJustifiedStateBalances_CanCache(t *testing.T) {
-	beaconDB := testDB.SetupDB(t)
-	service := setupBeaconChain(t, beaconDB)
-	ctx := context.Background()
-
-	state, _ := util.DeterministicGenesisState(t, 100)
-	r := [32]byte{'a'}
-	require.NoError(t, service.cfg.BeaconDB.SaveStateSummary(context.Background(), &ethpb.StateSummary{Root: r[:]}))
-	require.NoError(t, service.cfg.BeaconDB.SaveState(context.Background(), state, r))
-	balances, err := service.justifiedBalances.get(ctx, r)
-	require.NoError(t, err)
-	require.DeepEqual(t, balances, state.Balances(), "Incorrect justified balances")
-}
-
 func Test_notifyNewHeadEvent(t *testing.T) {
 	t.Run("genesis_state_root", func(t *testing.T) {
 		bState, _ := util.DeterministicGenesisState(t, 10)
@@ -559,7 +546,9 @@ func TestUpdateHead_noSavedChanges(t *testing.T) {
 	st, blkRoot, err = prepareForkchoiceState(ctx, 0, bellatrixBlkRoot, [32]byte{}, [32]byte{}, fcp, fcp)
 	require.NoError(t, err)
 	require.NoError(t, fcs.InsertNode(ctx, st, blkRoot))
-	newRoot, err := service.cfg.ForkChoiceStore.Head(ctx, []uint64{1, 2})
+	fcs.SetBalancesByRooter(func(context.Context, [32]byte) ([]uint64, error) { return []uint64{1, 2}, nil })
+	require.NoError(t, fcs.UpdateJustifiedCheckpoint(ctx, &forkchoicetypes.Checkpoint{}))
+	newRoot, err := service.cfg.ForkChoiceStore.Head(ctx)
 	require.NoError(t, err)
 	require.NotEqual(t, headRoot, newRoot)
 	require.Equal(t, headRoot, service.headRoot())

beacon-chain/blockchain/metrics.go

@@ -130,14 +130,6 @@ var (
 		Name: "sync_head_state_hit",
 		Help: "The number of sync head state requests that are present in the cache.",
 	})
-	stateBalanceCacheHit = promauto.NewCounter(prometheus.CounterOpts{
-		Name: "state_balance_cache_hit",
-		Help: "Count the number of state balance cache hits.",
-	})
-	stateBalanceCacheMiss = promauto.NewCounter(prometheus.CounterOpts{
-		Name: "state_balance_cache_miss",
-		Help: "Count the number of state balance cache hits.",
-	})
 	newPayloadValidNodeCount = promauto.NewCounter(prometheus.CounterOpts{
 		Name: "new_payload_valid_node_count",
 		Help: "Count the number of valid nodes after newPayload EE call",

beacon-chain/blockchain/mock_test.go

@@ -1,15 +1,10 @@
 package blockchain
 
 import (
-	"context"
-	"errors"
 	"testing"
 
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/core/helpers"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/core/time"
 	testDB "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	doublylinkedtree "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/doubly-linked-tree"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
 )
 
@@ -27,44 +22,5 @@ func testServiceOptsWithDB(t *testing.T) []Option {
 // in your code path. this is a lightweight way to satisfy the stategen/beacondb
 // initialization requirements w/o the overhead of db init.
 func testServiceOptsNoDB() []Option {
-	return []Option{
-		withStateBalanceCache(satisfactoryStateBalanceCache()),
-	}
-}
-
-type mockBalanceByRooter struct {
-	state state.BeaconState
-	err   error
-}
-
-var _ balanceByRooter = &mockBalanceByRooter{}
-
-func (m mockBalanceByRooter) ActiveNonSlashedBalancesByRoot(_ context.Context, _ [32]byte) ([]uint64, error) {
-	st := m.state
-	if st == nil || st.IsNil() {
-		return nil, errors.New("nil state")
-	}
-	epoch := time.CurrentEpoch(st)
-
-	balances := make([]uint64, st.NumValidators())
-	var balanceAccretor = func(idx int, val state.ReadOnlyValidator) error {
-		if helpers.IsActiveValidatorUsingTrie(val, epoch) {
-			balances[idx] = val.EffectiveBalance()
-		} else {
-			balances[idx] = 0
-		}
-		return nil
-	}
-	if err := st.ReadFromEveryValidator(balanceAccretor); err != nil {
-		return nil, err
-	}
-	return balances, nil
-}
-
-// returns an instance of the state balance cache that can be used
-// to satisfy the requirement for one in NewService, but which will
-// always return an error if used.
-func satisfactoryStateBalanceCache() *stateBalanceCache {
-	err := errors.New("satisfactoryStateBalanceCache doesn't perform real caching")
-	return &stateBalanceCache{stateGen: mockBalanceByRooter{err: err}}
+	return []Option{}
 }

beacon-chain/blockchain/options.go

@@ -156,13 +156,6 @@ func WithSlasherAttestationsFeed(f *event.Feed) Option {
 	}
 }
 
-func withStateBalanceCache(c *stateBalanceCache) Option {
-	return func(s *Service) error {
-		s.justifiedBalances = c
-		return nil
-	}
-}
-
 // WithFinalizedStateAtStartUp to store finalized state at start up.
 func WithFinalizedStateAtStartUp(st state.BeaconState) Option {
 	return func(s *Service) error {

beacon-chain/blockchain/process_attestation_test.go

@@ -409,8 +409,10 @@ func TestVerifyFinalizedConsistency_IsCanonical(t *testing.T) {
 	require.NoError(t, service.cfg.ForkChoiceStore.InsertNode(ctx, state, blkRoot))
 
 	jc := &forkchoicetypes.Checkpoint{Epoch: 0, Root: r32}
-	require.NoError(t, service.cfg.ForkChoiceStore.UpdateJustifiedCheckpoint(jc))
-	_, err = service.cfg.ForkChoiceStore.Head(ctx, []uint64{})
+	bState, _ := util.DeterministicGenesisState(t, 10)
+	require.NoError(t, service.cfg.BeaconDB.SaveState(ctx, bState, r32))
+	require.NoError(t, service.cfg.ForkChoiceStore.UpdateJustifiedCheckpoint(ctx, jc))
+	_, err = service.cfg.ForkChoiceStore.Head(ctx)
 	require.NoError(t, err)
 	err = service.VerifyFinalizedConsistency(context.Background(), r33[:])
 	require.NoError(t, err)

beacon-chain/blockchain/process_block.go

@@ -186,14 +186,8 @@ func (s *Service) onBlock(ctx context.Context, signed interfaces.ReadOnlySignedB
 	}
 
 	justified := s.ForkChoicer().JustifiedCheckpoint()
-	balances, err := s.justifiedBalances.get(ctx, justified.Root)
-	if err != nil {
-		msg := fmt.Sprintf("could not read balances for state w/ justified checkpoint %#x", justified.Root)
-		return errors.Wrap(err, msg)
-	}
-
 	start := time.Now()
-	headRoot, err := s.cfg.ForkChoiceStore.Head(ctx, balances)
+	headRoot, err := s.cfg.ForkChoiceStore.Head(ctx)
 	if err != nil {
 		log.WithError(err).Warn("Could not update head")
 	}

beacon-chain/blockchain/process_block_test.go

@@ -1992,10 +1992,12 @@ func TestNoViableHead_Reboot(t *testing.T) {
 	attSrv, err := attestations.NewService(ctx, &attestations.Config{})
 	require.NoError(t, err)
 	newfc := doublylinkedtree.New()
+	newStateGen := stategen.New(beaconDB, newfc)
+	newfc.SetBalancesByRooter(newStateGen.ActiveNonSlashedBalancesByRoot)
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithAttestationPool(attestations.NewPool()),
-		WithStateGen(stategen.New(beaconDB, newfc)),
+		WithStateGen(newStateGen),
 		WithForkChoiceStore(newfc),
 		WithStateNotifier(&mock.MockStateNotifier{}),
 		WithExecutionEngineCaller(mockEngine),
@@ -2105,6 +2107,12 @@ func TestNoViableHead_Reboot(t *testing.T) {
 	require.Equal(t, blk, nil)
 
 	service.cfg.ForkChoiceStore = doublylinkedtree.New()
+	justified, err := service.cfg.BeaconDB.JustifiedCheckpoint(ctx)
+	require.NoError(t, err)
+
+	jroot := bytesutil.ToBytes32(justified.Root)
+	require.NoError(t, service.cfg.BeaconDB.SaveState(ctx, genesisState, jroot))
+	service.cfg.ForkChoiceStore.SetBalancesByRooter(service.cfg.StateGen.ActiveNonSlashedBalancesByRoot)
 	require.NoError(t, service.StartFromSavedState(genesisState))
 
 	// Forkchoice has the genesisRoot loaded at startup
@@ -2182,6 +2190,7 @@ func TestNoViableHead_Reboot(t *testing.T) {
 	require.NoError(t, err)
 	root, err = b.Block.HashTreeRoot()
 	require.NoError(t, err)
+	service.ForkChoicer().SetBalancesByRooter(service.cfg.StateGen.ActiveNonSlashedBalancesByRoot)
 	err = service.onBlock(ctx, wsb, root)
 	require.NoError(t, err)
 	require.Equal(t, root, service.ForkChoicer().CachedHeadRoot())

beacon-chain/blockchain/receive_attestation.go

@@ -143,13 +143,8 @@ func (s *Service) UpdateHead(ctx context.Context) error {
 	s.processAttestations(ctx)
 	processAttsElapsedTime.Observe(float64(time.Since(start).Milliseconds()))
 
-	justified := s.ForkChoicer().JustifiedCheckpoint()
-	balances, err := s.justifiedBalances.get(ctx, justified.Root)
-	if err != nil {
-		return err
-	}
 	start = time.Now()
-	newHeadRoot, err := s.cfg.ForkChoiceStore.Head(ctx, balances)
+	newHeadRoot, err := s.cfg.ForkChoiceStore.Head(ctx)
 	if err != nil {
 		log.WithError(err).Error("Could not compute head from new attestations")
 	}