finish wiring up ugly wrapper hack

beacon-chain/blockchain/chain_info_norace_test.go

@@ -5,7 +5,7 @@ import (
 	"testing"
 
 	testDB "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/consensus-types/blocks"
 	ethpb "github.com/prysmaticlabs/prysm/v3/proto/prysm/v1alpha1"
 	"github.com/prysmaticlabs/prysm/v3/testing/require"
@@ -32,7 +32,7 @@ func TestHeadSlot_DataRace(t *testing.T) {
 func TestHeadRoot_DataRace(t *testing.T) {
 	beaconDB := testDB.SetupDB(t)
 	s := &Service{
-		cfg:  &config{BeaconDB: beaconDB, StateGen: stategen.New(beaconDB)},
+		cfg:  &config{BeaconDB: beaconDB, StateGen: sgmock.NewMockStategen(beaconDB)},
 		head: &head{root: [32]byte{'A'}},
 	}
 	b, err := blocks.NewSignedBeaconBlock(util.NewBeaconBlock())
@@ -54,7 +54,7 @@ func TestHeadBlock_DataRace(t *testing.T) {
 	wsb, err := blocks.NewSignedBeaconBlock(&ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{Body: &ethpb.BeaconBlockBody{}}})
 	require.NoError(t, err)
 	s := &Service{
-		cfg:  &config{BeaconDB: beaconDB, StateGen: stategen.New(beaconDB)},
+		cfg:  &config{BeaconDB: beaconDB, StateGen: sgmock.NewMockStategen(beaconDB)},
 		head: &head{block: wsb},
 	}
 	b, err := blocks.NewSignedBeaconBlock(util.NewBeaconBlock())
@@ -74,7 +74,7 @@ func TestHeadBlock_DataRace(t *testing.T) {
 func TestHeadState_DataRace(t *testing.T) {
 	beaconDB := testDB.SetupDB(t)
 	s := &Service{
-		cfg: &config{BeaconDB: beaconDB, StateGen: stategen.New(beaconDB)},
+		cfg: &config{BeaconDB: beaconDB, StateGen: sgmock.NewMockStategen(beaconDB)},
 	}
 	b, err := blocks.NewSignedBeaconBlock(util.NewBeaconBlock())
 	require.NoError(t, err)

beacon-chain/blockchain/chain_info_test.go

@@ -10,7 +10,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/protoarray"
 	forkchoicetypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/types"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	v1 "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/v1"
 	v3 "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/v3"
 	fieldparams "github.com/prysmaticlabs/prysm/v3/config/fieldparams"
@@ -86,7 +86,7 @@ func TestFinalizedCheckpt_GenesisRootOk(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithForkChoiceStore(fcs),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -111,7 +111,7 @@ func TestCurrentJustifiedCheckpt_CanRetrieve(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithForkChoiceStore(fcs),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -138,7 +138,7 @@ func TestHeadRoot_CanRetrieve(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithForkChoiceStore(fcs),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -157,7 +157,7 @@ func TestHeadRoot_UseDB(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithForkChoiceStore(fcs),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)

beacon-chain/blockchain/execution_engine_test.go

@@ -16,7 +16,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/protoarray"
 	forkchoicetypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/types"
 	bstate "github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	fieldparams "github.com/prysmaticlabs/prysm/v3/config/fieldparams"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	consensusblocks "github.com/prysmaticlabs/prysm/v3/consensus-types/blocks"
@@ -44,7 +44,7 @@ func Test_NotifyForkchoiceUpdate(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 		WithProposerIdsCache(cache.NewProposerPayloadIDsCache()),
 	}
@@ -578,7 +578,7 @@ func Test_NotifyNewPayload(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	phase0State, _ := util.DeterministicGenesisState(t, 1)
@@ -824,7 +824,7 @@ func Test_NotifyNewPayload_SetOptimisticToValid(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	bellatrixState, _ := util.DeterministicGenesisStateBellatrix(t, 2)
@@ -871,7 +871,7 @@ func Test_IsOptimisticCandidateBlock(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 
@@ -975,7 +975,7 @@ func Test_IsOptimisticShallowExecutionParent(t *testing.T) {
 	beaconDB := testDB.SetupDB(t)
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 
 	service, err := NewService(ctx, opts...)
@@ -1016,7 +1016,7 @@ func Test_GetPayloadAttribute(t *testing.T) {
 	beaconDB := testDB.SetupDB(t)
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithProposerIdsCache(cache.NewProposerPayloadIDsCache()),
 	}
 
@@ -1058,7 +1058,7 @@ func Test_UpdateLastValidatedCheckpoint(t *testing.T) {
 
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	stateGen := stategen.New(beaconDB)
+	stateGen := sgmock.NewMockStategen(beaconDB)
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
@@ -1162,7 +1162,7 @@ func TestService_removeInvalidBlockAndState(t *testing.T) {
 	beaconDB := testDB.SetupDB(t)
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(doublylinkedtree.New()),
 	}
 	service, err := NewService(ctx, opts...)
@@ -1214,7 +1214,7 @@ func TestService_getPayloadHash(t *testing.T) {
 	beaconDB := testDB.SetupDB(t)
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(doublylinkedtree.New()),
 	}
 	service, err := NewService(ctx, opts...)

beacon-chain/blockchain/head_sync_committee_info_test.go

@@ -7,7 +7,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/core/signing"
 	dbtest "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/v3/testing/require"
@@ -19,7 +19,7 @@ func TestService_headSyncCommitteeFetcher_Errors(t *testing.T) {
 	beaconDB := dbtest.SetupDB(t)
 	c := &Service{
 		cfg: &config{
-			StateGen: stategen.New(beaconDB),
+			StateGen: sgmock.NewMockStategen(beaconDB),
 		},
 	}
 	c.head = &head{}
@@ -37,7 +37,7 @@ func TestService_HeadDomainFetcher_Errors(t *testing.T) {
 	beaconDB := dbtest.SetupDB(t)
 	c := &Service{
 		cfg: &config{
-			StateGen: stategen.New(beaconDB),
+			StateGen: sgmock.NewMockStategen(beaconDB),
 		},
 	}
 	c.head = &head{}

beacon-chain/blockchain/head_test.go

@@ -11,7 +11,7 @@ import (
 	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/forkchoice/protoarray"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/config/features"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	"github.com/prysmaticlabs/prysm/v3/consensus-types/blocks"
@@ -536,7 +536,7 @@ func TestUpdateHead_noSavedChanges(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 

beacon-chain/blockchain/mock_test.go

@@ -8,7 +8,7 @@ import (
 	testDB "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/protoarray"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 )
 
 func testServiceOptsWithDB(t *testing.T) []Option {
@@ -16,7 +16,7 @@ func testServiceOptsWithDB(t *testing.T) []Option {
 	fcs := protoarray.New()
 	return []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 }

beacon-chain/blockchain/pow_block_test.go

@@ -11,7 +11,7 @@ import (
 	testDB "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	mocks "github.com/prysmaticlabs/prysm/v3/beacon-chain/execution/testing"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/protoarray"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	"github.com/prysmaticlabs/prysm/v3/consensus-types/blocks"
 	"github.com/prysmaticlabs/prysm/v3/encoding/bytesutil"
@@ -112,7 +112,7 @@ func Test_validateMergeBlock(t *testing.T) {
 	fcs := protoarray.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	service, err := NewService(ctx, opts...)
@@ -162,7 +162,7 @@ func Test_getBlkParentHashAndTD(t *testing.T) {
 	fcs := protoarray.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	service, err := NewService(ctx, opts...)

beacon-chain/blockchain/process_attestation_test.go

@@ -10,7 +10,7 @@ import (
 	doublylinkedtree "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/doubly-linked-tree"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/protoarray"
 	forkchoicetypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/types"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	fieldparams "github.com/prysmaticlabs/prysm/v3/config/fieldparams"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
@@ -29,7 +29,7 @@ func TestStore_OnAttestation_ErrorConditions_ProtoArray(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithForkChoiceStore(doublylinkedtree.New()),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -135,7 +135,7 @@ func TestStore_OnAttestation_ErrorConditions_DoublyLinkedTree(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithForkChoiceStore(doublylinkedtree.New()),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -241,7 +241,7 @@ func TestStore_OnAttestation_Ok_ProtoArray(t *testing.T) {
 	fcs := protoarray.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	service, err := NewService(ctx, opts...)
@@ -271,7 +271,7 @@ func TestStore_OnAttestation_Ok_DoublyLinkedTree(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	service, err := NewService(ctx, opts...)
@@ -300,7 +300,7 @@ func TestStore_SaveCheckpointState(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -362,7 +362,7 @@ func TestStore_UpdateCheckpointState(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -468,7 +468,7 @@ func TestVerifyFinalizedConsistency_InconsistentRoot_ProtoArray(t *testing.T) {
 	fcs := protoarray.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	service, err := NewService(ctx, opts...)
@@ -499,7 +499,7 @@ func TestVerifyFinalizedConsistency_InconsistentRoot_DoublyLinkedTree(t *testing
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	service, err := NewService(ctx, opts...)

beacon-chain/blockchain/process_block.go

@@ -132,9 +132,6 @@ func (s *Service) onBlock(ctx context.Context, signed interfaces.SignedBeaconBlo
 			return err
 		}
 	}
-	if err := s.savePostStateInfo(ctx, blockRoot, signed, postState); err != nil {
-		return err
-	}
 
 	if err := s.insertBlockToForkchoiceStore(ctx, signed.Block(), blockRoot, postState); err != nil {
 		return errors.Wrapf(err, "could not insert block %d to fork choice store", signed.Block().Slot())
@@ -146,6 +143,13 @@ func (s *Service) onBlock(ctx context.Context, signed interfaces.SignedBeaconBlo
 		}
 	}
 
+	if err := s.cfg.BeaconDB.SaveBlock(ctx, signed); err != nil {
+		return errors.Wrapf(err, "could not save block from slot %d", signed.Block().Slot())
+	}
+	if err := s.cfg.StateGen.SaveState(ctx, blockRoot, postState); err != nil {
+		return errors.Wrap(err, "could not save state")
+	}
+
 	// If slasher is configured, forward the attestations in the block via
 	// an event feed for processing.
 	if features.Get().EnableSlasher {
@@ -527,20 +531,6 @@ func (s *Service) InsertSlashingsToForkChoiceStore(ctx context.Context, slashing
 	}
 }
 
-// This saves post state info to DB or cache. This also saves post state info to fork choice store.
-// Post state info consists of processed block and state. Do not call this method unless the block and state are verified.
-func (s *Service) savePostStateInfo(ctx context.Context, r [32]byte, b interfaces.SignedBeaconBlock, st state.BeaconState) error {
-	ctx, span := trace.StartSpan(ctx, "blockChain.savePostStateInfo")
-	defer span.End()
-	if err := s.cfg.BeaconDB.SaveBlock(ctx, b); err != nil {
-		return errors.Wrapf(err, "could not save block from slot %d", b.Block().Slot())
-	}
-	if err := s.cfg.StateGen.SaveState(ctx, r, st); err != nil {
-		return errors.Wrap(err, "could not save state")
-	}
-	return nil
-}
-
 // This removes the attestations from the mem pool. It will only remove the attestations if input root `r` is canonical,
 // meaning the block `b` is part of the canonical chain.
 func (s *Service) pruneCanonicalAttsFromPool(ctx context.Context, r [32]byte, b interfaces.SignedBeaconBlock) error {

beacon-chain/blockchain/process_block_test.go

@@ -27,7 +27,7 @@ import (
 	forkchoicetypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/types"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/attestations"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	v1 "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/v1"
 	fieldparams "github.com/prysmaticlabs/prysm/v3/config/fieldparams"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
@@ -53,7 +53,7 @@ func TestStore_OnBlock_ProtoArray(t *testing.T) {
 	fcs := protoarray.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 
@@ -149,7 +149,7 @@ func TestStore_OnBlock_DoublyLinkedTree(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 
@@ -244,7 +244,7 @@ func TestStore_OnBlockBatch_ProtoArray(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(protoarray.New()),
 	}
 	service, err := NewService(ctx, opts...)
@@ -285,7 +285,7 @@ func TestStore_OnBlockBatch_PruneOK_Protoarray(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(protoarray.New()),
 	}
 	service, err := NewService(ctx, opts...)
@@ -321,7 +321,7 @@ func TestStore_OnBlockBatch_DoublyLinkedTree(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(doublylinkedtree.New()),
 	}
 	service, err := NewService(ctx, opts...)
@@ -364,7 +364,7 @@ func TestStore_OnBlockBatch_NotifyNewPayload(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(doublylinkedtree.New()),
 	}
 	service, err := NewService(ctx, opts...)
@@ -399,7 +399,7 @@ func TestCachedPreState_CanGetFromStateSummary_ProtoArray(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(protoarray.New()),
 	}
 	service, err := NewService(ctx, opts...)
@@ -426,7 +426,7 @@ func TestCachedPreState_CanGetFromStateSummary_DoublyLinkedTree(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(doublylinkedtree.New()),
 	}
 	service, err := NewService(ctx, opts...)
@@ -453,7 +453,7 @@ func TestFillForkChoiceMissingBlocks_CanSave_ProtoArray(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -494,7 +494,7 @@ func TestFillForkChoiceMissingBlocks_CanSave_DoublyLinkedTree(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -543,7 +543,7 @@ func TestFillForkChoiceMissingBlocks_RootsMatch_ProtoArray(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -586,7 +586,7 @@ func TestFillForkChoiceMissingBlocks_RootsMatch_DoublyLinkedTree(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -637,7 +637,7 @@ func TestFillForkChoiceMissingBlocks_FilterFinalized_ProtoArray(t *testing.T) {
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -694,7 +694,7 @@ func TestFillForkChoiceMissingBlocks_FilterFinalized_DoublyLinkedTree(t *testing
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -752,7 +752,7 @@ func TestFillForkChoiceMissingBlocks_FinalizedSibling_DoublyLinkedTree(t *testin
 
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 	service, err := NewService(ctx, opts...)
 	require.NoError(t, err)
@@ -897,7 +897,7 @@ func TestAncestor_HandleSkipSlot(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	service, err := NewService(ctx, opts...)
@@ -988,7 +988,7 @@ func TestAncestor_CanUseDB(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	service, err := NewService(ctx, opts...)
@@ -1050,7 +1050,7 @@ func TestVerifyBlkDescendant(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	b := util.NewBeaconBlock()
@@ -1159,7 +1159,7 @@ func TestOnBlock_CanFinalize_WithOnTick(t *testing.T) {
 	require.NoError(t, err)
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 		WithDepositCache(depositCache),
 		WithStateNotifier(&mock.MockStateNotifier{}),
@@ -1208,7 +1208,7 @@ func TestOnBlock_CanFinalize(t *testing.T) {
 	require.NoError(t, err)
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 		WithDepositCache(depositCache),
 		WithStateNotifier(&mock.MockStateNotifier{}),
@@ -1256,7 +1256,7 @@ func TestOnBlock_NilBlock(t *testing.T) {
 	require.NoError(t, err)
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 		WithDepositCache(depositCache),
 	}
@@ -1275,7 +1275,7 @@ func TestOnBlock_InvalidSignature(t *testing.T) {
 	require.NoError(t, err)
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 		WithDepositCache(depositCache),
 		WithStateNotifier(&mock.MockStateNotifier{}),
@@ -1311,7 +1311,7 @@ func TestOnBlock_CallNewPayloadAndForkchoiceUpdated(t *testing.T) {
 	require.NoError(t, err)
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 		WithDepositCache(depositCache),
 		WithStateNotifier(&mock.MockStateNotifier{}),
@@ -1524,7 +1524,7 @@ func Test_validateMergeTransitionBlock(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 		WithProposerIdsCache(cache.NewProposerPayloadIDsCache()),
 		WithAttestationPool(attestations.NewPool()),
@@ -1662,7 +1662,7 @@ func TestService_insertSlashingsToForkChoiceStore(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 		WithProposerIdsCache(cache.NewProposerPayloadIDsCache()),
 	}
@@ -1715,7 +1715,7 @@ func TestOnBlock_ProcessBlocksParallel(t *testing.T) {
 	require.NoError(t, err)
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 		WithDepositCache(depositCache),
 		WithStateNotifier(&mock.MockStateNotifier{}),
@@ -1793,7 +1793,7 @@ func Test_verifyBlkFinalizedSlot_invalidBlock(t *testing.T) {
 	fcs := doublylinkedtree.New()
 	opts := []Option{
 		WithDatabase(beaconDB),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(fcs),
 	}
 	service, err := NewService(ctx, opts...)
@@ -1828,7 +1828,7 @@ func TestStore_NoViableHead_FCU_Protoarray(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithAttestationPool(attestations.NewPool()),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(protoarray.New()),
 		WithStateNotifier(&mock.MockStateNotifier{}),
 		WithExecutionEngineCaller(mockEngine),
@@ -1988,7 +1988,7 @@ func TestStore_NoViableHead_FCU_DoublyLinkedTree(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithAttestationPool(attestations.NewPool()),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(doublylinkedtree.New()),
 		WithStateNotifier(&mock.MockStateNotifier{}),
 		WithExecutionEngineCaller(mockEngine),
@@ -2148,7 +2148,7 @@ func TestStore_NoViableHead_NewPayload_DoublyLinkedTree(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithAttestationPool(attestations.NewPool()),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(doublylinkedtree.New()),
 		WithStateNotifier(&mock.MockStateNotifier{}),
 		WithExecutionEngineCaller(mockEngine),
@@ -2308,7 +2308,7 @@ func TestStore_NoViableHead_NewPayload_Protoarray(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithAttestationPool(attestations.NewPool()),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(protoarray.New()),
 		WithStateNotifier(&mock.MockStateNotifier{}),
 		WithExecutionEngineCaller(mockEngine),
@@ -2469,7 +2469,7 @@ func TestStore_NoViableHead_Liveness_DoublyLinkedTree(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithAttestationPool(attestations.NewPool()),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(doublylinkedtree.New()),
 		WithStateNotifier(&mock.MockStateNotifier{}),
 		WithExecutionEngineCaller(mockEngine),
@@ -2679,7 +2679,7 @@ func TestStore_NoViableHead_Liveness_Protoarray(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithAttestationPool(attestations.NewPool()),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(doublylinkedtree.New()),
 		WithStateNotifier(&mock.MockStateNotifier{}),
 		WithExecutionEngineCaller(mockEngine),
@@ -2889,7 +2889,7 @@ func TestStore_NoViableHead_Reboot_DoublyLinkedTree(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithAttestationPool(attestations.NewPool()),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(doublylinkedtree.New()),
 		WithStateNotifier(&mock.MockStateNotifier{}),
 		WithExecutionEngineCaller(mockEngine),
@@ -3114,7 +3114,7 @@ func TestStore_NoViableHead_Reboot_Protoarray(t *testing.T) {
 	opts := []Option{
 		WithDatabase(beaconDB),
 		WithAttestationPool(attestations.NewPool()),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 		WithForkChoiceStore(protoarray.New()),
 		WithStateNotifier(&mock.MockStateNotifier{}),
 		WithExecutionEngineCaller(mockEngine),

beacon-chain/blockchain/receive_block.go

@@ -11,7 +11,6 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/monitoring/tracing"
 	ethpb "github.com/prysmaticlabs/prysm/v3/proto/prysm/v1alpha1"
 	"github.com/prysmaticlabs/prysm/v3/time"
-	"github.com/prysmaticlabs/prysm/v3/time/slots"
 	"go.opencensus.io/trace"
 )
 
@@ -56,11 +55,6 @@ func (s *Service) ReceiveBlock(ctx context.Context, block interfaces.SignedBeaco
 		return err
 	}
 
-	// Have we been finalizing? Should we start saving hot states to db?
-	if err := s.checkSaveHotStateDB(ctx); err != nil {
-		return err
-	}
-
 	// Reports on block and fork choice metrics.
 	finalized := s.FinalizedCheckpt()
 	reportSlotMetrics(blockCopy.Block().Slot(), s.HeadSlot(), s.CurrentSlot(), finalized)
@@ -166,25 +160,3 @@ func (s *Service) handlePostBlockOperations(b interfaces.BeaconBlock) error {
 	}
 	return nil
 }
-
-// This checks whether it's time to start saving hot state to DB.
-// It's time when there's `epochsSinceFinalitySaveHotStateDB` epochs of non-finality.
-func (s *Service) checkSaveHotStateDB(ctx context.Context) error {
-	currentEpoch := slots.ToEpoch(s.CurrentSlot())
-	// Prevent `sinceFinality` going underflow.
-	var sinceFinality types.Epoch
-	finalized := s.FinalizedCheckpt()
-	if finalized == nil {
-		return errNilFinalizedInStore
-	}
-	if currentEpoch > finalized.Epoch {
-		sinceFinality = currentEpoch - finalized.Epoch
-	}
-
-	if sinceFinality >= epochsSinceFinalitySaveHotStateDB {
-		s.cfg.StateGen.EnableSaveHotStateToDB(ctx)
-		return nil
-	}
-
-	return s.cfg.StateGen.DisableSaveHotStateToDB(ctx)
-}

beacon-chain/blockchain/receive_block_test.go

@@ -4,14 +4,13 @@ import (
 	"context"
 	"sync"
 	"testing"
-	"time"
 
 	blockchainTesting "github.com/prysmaticlabs/prysm/v3/beacon-chain/blockchain/testing"
 	testDB "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/protoarray"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/attestations"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/voluntaryexits"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	"github.com/prysmaticlabs/prysm/v3/consensus-types/blocks"
 	"github.com/prysmaticlabs/prysm/v3/consensus-types/interfaces"
@@ -21,7 +20,6 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/testing/assert"
 	"github.com/prysmaticlabs/prysm/v3/testing/require"
 	"github.com/prysmaticlabs/prysm/v3/testing/util"
-	logTest "github.com/sirupsen/logrus/hooks/test"
 )
 
 func TestService_ReceiveBlock(t *testing.T) {
@@ -133,7 +131,7 @@ func TestService_ReceiveBlock(t *testing.T) {
 				WithAttestationPool(attestations.NewPool()),
 				WithExitPool(voluntaryexits.NewPool()),
 				WithStateNotifier(&blockchainTesting.MockStateNotifier{RecordEvents: true}),
-				WithStateGen(stategen.New(beaconDB)),
+				WithStateGen(sgmock.NewMockStategen(beaconDB)),
 				WithFinalizedStateAtStartUp(genesis),
 			}
 			s, err := NewService(ctx, opts...)
@@ -172,7 +170,7 @@ func TestService_ReceiveBlockUpdateHead(t *testing.T) {
 		WithAttestationPool(attestations.NewPool()),
 		WithExitPool(voluntaryexits.NewPool()),
 		WithStateNotifier(&blockchainTesting.MockStateNotifier{RecordEvents: true}),
-		WithStateGen(stategen.New(beaconDB)),
+		WithStateGen(sgmock.NewMockStategen(beaconDB)),
 	}
 
 	s, err := NewService(ctx, opts...)
@@ -247,7 +245,7 @@ func TestService_ReceiveBlockBatch(t *testing.T) {
 				WithDatabase(beaconDB),
 				WithForkChoiceStore(protoarray.New()),
 				WithStateNotifier(&blockchainTesting.MockStateNotifier{RecordEvents: true}),
-				WithStateGen(stategen.New(beaconDB)),
+				WithStateGen(sgmock.NewMockStategen(beaconDB)),
 			}
 			s, err := NewService(ctx, opts...)
 			require.NoError(t, err)
@@ -292,40 +290,3 @@ func TestService_HasBlock(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, true, s.HasBlock(context.Background(), r))
 }
-
-func TestCheckSaveHotStateDB_Enabling(t *testing.T) {
-	opts := testServiceOptsWithDB(t)
-	hook := logTest.NewGlobal()
-	s, err := NewService(context.Background(), opts...)
-	require.NoError(t, err)
-	st := params.BeaconConfig().SlotsPerEpoch.Mul(uint64(epochsSinceFinalitySaveHotStateDB))
-	s.genesisTime = time.Now().Add(time.Duration(-1*int64(st)*int64(params.BeaconConfig().SecondsPerSlot)) * time.Second)
-
-	require.NoError(t, s.checkSaveHotStateDB(context.Background()))
-	assert.LogsContain(t, hook, "Entering mode to save hot states in DB")
-}
-
-func TestCheckSaveHotStateDB_Disabling(t *testing.T) {
-	hook := logTest.NewGlobal()
-	opts := testServiceOptsWithDB(t)
-	s, err := NewService(context.Background(), opts...)
-	require.NoError(t, err)
-	st := params.BeaconConfig().SlotsPerEpoch.Mul(uint64(epochsSinceFinalitySaveHotStateDB))
-	s.genesisTime = time.Now().Add(time.Duration(-1*int64(st)*int64(params.BeaconConfig().SecondsPerSlot)) * time.Second)
-	require.NoError(t, s.checkSaveHotStateDB(context.Background()))
-	s.genesisTime = time.Now()
-
-	require.NoError(t, s.checkSaveHotStateDB(context.Background()))
-	assert.LogsContain(t, hook, "Exiting mode to save hot states in DB")
-}
-
-func TestCheckSaveHotStateDB_Overflow(t *testing.T) {
-	hook := logTest.NewGlobal()
-	opts := testServiceOptsWithDB(t)
-	s, err := NewService(context.Background(), opts...)
-	require.NoError(t, err)
-	s.genesisTime = time.Now()
-
-	require.NoError(t, s.checkSaveHotStateDB(context.Background()))
-	assert.LogsDoNotContain(t, hook, "Entering mode to save hot states in DB")
-}

beacon-chain/blockchain/service.go

@@ -176,7 +176,6 @@ func (s *Service) Status() error {
 
 // StartFromSavedState initializes the blockchain using a previously saved finalized checkpoint.
 func (s *Service) StartFromSavedState(saved state.BeaconState) error {
-	log.Info("Blockchain data already exists in DB, initializing...")
 	s.genesisTime = time.Unix(int64(saved.GenesisTime()), 0) // lint:ignore uintcast -- Genesis time will not exceed int64 in your lifetime.
 	s.cfg.AttService.SetGenesisTime(saved.GenesisTime())
 

beacon-chain/blockchain/service_test.go

@@ -24,7 +24,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/protoarray"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/attestations"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	v1 "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/v1"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	consensusblocks "github.com/prysmaticlabs/prysm/v3/consensus-types/blocks"
@@ -117,7 +117,7 @@ func setupBeaconChain(t *testing.T, beaconDB db.Database) *Service {
 	depositCache, err := depositcache.New()
 	require.NoError(t, err)
 
-	stateGen := stategen.New(beaconDB)
+	stateGen := sgmock.NewMockStategen(beaconDB)
 	// Safe a state in stategen to purposes of testing a service stop / shutdown.
 	require.NoError(t, stateGen.SaveState(ctx, bytesutil.ToBytes32(bState.FinalizedCheckpoint().Root), bState))
 
@@ -305,7 +305,7 @@ func TestChainService_InitializeChainInfo(t *testing.T) {
 	require.NoError(t, beaconDB.SaveFinalizedCheckpoint(ctx, &ethpb.Checkpoint{Epoch: slots.ToEpoch(finalizedSlot), Root: headRoot[:]}))
 	attSrv, err := attestations.NewService(ctx, &attestations.Config{})
 	require.NoError(t, err)
-	stateGen := stategen.New(beaconDB)
+	stateGen := sgmock.NewMockStategen(beaconDB)
 	c, err := NewService(ctx, WithDatabase(beaconDB), WithStateGen(stateGen), WithAttestationService(attSrv), WithStateNotifier(&mock.MockStateNotifier{}), WithFinalizedStateAtStartUp(headState))
 	require.NoError(t, err)
 	require.NoError(t, stateGen.SaveState(ctx, headRoot, headState))
@@ -358,7 +358,7 @@ func TestChainService_InitializeChainInfo_SetHeadAtGenesis(t *testing.T) {
 	}
 	require.NoError(t, beaconDB.SaveStateSummary(ctx, ss))
 	require.NoError(t, beaconDB.SaveFinalizedCheckpoint(ctx, &ethpb.Checkpoint{Root: headRoot[:], Epoch: slots.ToEpoch(finalizedSlot)}))
-	stateGen := stategen.New(beaconDB)
+	stateGen := sgmock.NewMockStategen(beaconDB)
 	c, err := NewService(ctx, WithDatabase(beaconDB), WithStateGen(stateGen), WithAttestationService(attSrv), WithStateNotifier(&mock.MockStateNotifier{}), WithFinalizedStateAtStartUp(headState))
 	require.NoError(t, err)
 
@@ -376,7 +376,7 @@ func TestChainService_SaveHeadNoDB(t *testing.T) {
 	beaconDB := testDB.SetupDB(t)
 	ctx := context.Background()
 	s := &Service{
-		cfg: &config{BeaconDB: beaconDB, StateGen: stategen.New(beaconDB), ForkChoiceStore: doublylinkedtree.New()},
+		cfg: &config{BeaconDB: beaconDB, StateGen: sgmock.NewMockStategen(beaconDB), ForkChoiceStore: doublylinkedtree.New()},
 	}
 	blk := util.NewBeaconBlock()
 	blk.Block.Slot = 1
@@ -438,7 +438,7 @@ func TestServiceStop_SaveCachedBlocks(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	beaconDB := testDB.SetupDB(t)
 	s := &Service{
-		cfg:            &config{BeaconDB: beaconDB, StateGen: stategen.New(beaconDB)},
+		cfg:            &config{BeaconDB: beaconDB, StateGen: sgmock.NewMockStategen(beaconDB)},
 		ctx:            ctx,
 		cancel:         cancel,
 		initSyncBlocks: make(map[[32]byte]interfaces.SignedBeaconBlock),

beacon-chain/execution/BUILD.bazel

@@ -103,7 +103,7 @@ go_test(
         "//beacon-chain/db/testing:go_default_library",
         "//beacon-chain/execution/testing:go_default_library",
         "//beacon-chain/execution/types:go_default_library",
-        "//beacon-chain/state/stategen:go_default_library",
+        "//beacon-chain/state/stategen/mock:go_default_library",
         "//config/fieldparams:go_default_library",
         "//config/params:go_default_library",
         "//consensus-types/blocks:go_default_library",

beacon-chain/execution/service_test.go

@@ -20,7 +20,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/cache/depositcache"
 	dbutil "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	mockExecution "github.com/prysmaticlabs/prysm/v3/beacon-chain/execution/testing"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	contracts "github.com/prysmaticlabs/prysm/v3/contracts/deposit"
 	"github.com/prysmaticlabs/prysm/v3/contracts/deposit/mock"
@@ -469,7 +469,7 @@ func TestInitDepositCacheWithFinalization_OK(t *testing.T) {
 	headBlock := util.NewBeaconBlock()
 	headRoot, err := headBlock.Block.HashTreeRoot()
 	require.NoError(t, err)
-	stateGen := stategen.New(beaconDB)
+	stateGen := sgmock.NewMockStategen(beaconDB)
 
 	emptyState, err := util.NewBeaconState()
 	require.NoError(t, err)

beacon-chain/monitor/BUILD.bazel

@@ -54,7 +54,7 @@ go_test(
         "//beacon-chain/core/feed:go_default_library",
         "//beacon-chain/core/feed/state:go_default_library",
         "//beacon-chain/db/testing:go_default_library",
-        "//beacon-chain/state/stategen:go_default_library",
+        "//beacon-chain/state/stategen/mock:go_default_library",
         "//config/params:go_default_library",
         "//consensus-types/blocks:go_default_library",
         "//consensus-types/primitives:go_default_library",

beacon-chain/monitor/service_test.go

@@ -12,7 +12,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/core/feed"
 	statefeed "github.com/prysmaticlabs/prysm/v3/beacon-chain/core/feed/state"
 	testDB "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/consensus-types/blocks"
 	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/v3/encoding/bytesutil"
@@ -88,7 +88,7 @@ func setupService(t *testing.T) *Service {
 	}
 	return &Service{
 		config: &ValidatorMonitorConfig{
-			StateGen:            stategen.New(beaconDB),
+			StateGen:            sgmock.NewMockStategen(beaconDB),
 			StateNotifier:       chainService.StateNotifier(),
 			HeadFetcher:         chainService,
 			AttestationNotifier: chainService.OperationNotifier(),

beacon-chain/node/node.go

@@ -204,8 +204,12 @@ func New(cliCtx *cli.Context, opts ...Option) (*BeaconNode, error) {
 		return nil, errors.Wrap(err, "backfill status initialization error")
 	}
 
+	log.Debugln("Starting Fork Choice")
+	beacon.startForkChoice()
+
+	bw := &blockchainWrapper{}
 	log.Debugln("Starting State Gen")
-	if err := beacon.startStateGen(ctx, bfs); err != nil {
+	if err := beacon.startStateGen(ctx, bfs, bw); err != nil {
 		return nil, err
 	}
 
@@ -229,11 +233,8 @@ func New(cliCtx *cli.Context, opts ...Option) (*BeaconNode, error) {
 		return nil, err
 	}
 
-	log.Debugln("Starting Fork Choice")
-	beacon.startForkChoice()
-
 	log.Debugln("Registering Blockchain Service")
-	if err := beacon.registerBlockchainService(); err != nil {
+	if err := beacon.registerBlockchainService(bw); err != nil {
 		return nil, err
 	}
 
@@ -495,9 +496,28 @@ func (b *BeaconNode) startSlasherDB(cliCtx *cli.Context) error {
 	return nil
 }
 
-func (b *BeaconNode) startStateGen(ctx context.Context, bfs *backfill.Status) error {
-	opts := []stategen.StateGenOption{stategen.WithBackfillStatus(bfs)}
-	sg := stategen.New(b.db, opts...)
+type blockchainWrapper struct {
+	blksrv *blockchain.Service
+}
+
+func (w *blockchainWrapper) CurrentSlot() types.Slot {
+	return w.blksrv.CurrentSlot()
+}
+
+func (w *blockchainWrapper) IsCanonical(ctx context.Context, root [32]byte) (bool, error) {
+	return w.blksrv.IsCanonical(ctx, root)
+}
+
+var _ stategen.CurrentSlotter = &blockchainWrapper{}
+var _ stategen.CanonicalChecker = &blockchainWrapper{}
+
+func (b *BeaconNode) startStateGen(ctx context.Context, bfs *backfill.Status, bw *blockchainWrapper) error {
+	// we'll update the blockchain service pointer to the blockchain service once it's initialized
+	// this is weird and gross, but it's an intermediate step between decoupling these two services
+	rb := stategen.NewCanonicalHistory(b.db, bw, bw)
+	opts := []stategen.StateGenOption{stategen.WithBackfillStatus(bfs), stategen.WithReplayerBuilder(rb)}
+	saver := stategen.NewHotStateSaver(b.db, b.forkChoiceStore, bw)
+	sg := stategen.New(b.db, saver, opts...)
 
 	cp, err := b.db.FinalizedCheckpoint(ctx)
 	if err != nil {
@@ -586,7 +606,7 @@ func (b *BeaconNode) registerAttestationPool() error {
 	return b.services.RegisterService(s)
 }
 
-func (b *BeaconNode) registerBlockchainService() error {
+func (b *BeaconNode) registerBlockchainService(bw *blockchainWrapper) error {
 	var web3Service *execution.Service
 	if err := b.services.FetchService(&web3Service); err != nil {
 		return err
@@ -620,6 +640,11 @@ func (b *BeaconNode) registerBlockchainService() error {
 	if err != nil {
 		return errors.Wrap(err, "could not register blockchain service")
 	}
+	// TODO: file an issue to track the task of moving the implementation of CurrentSlot to
+	// a separate service.
+	// fulfill this creepy circular dependency between stategen and the blockchain service
+	bw.blksrv = blockchainService
+
 	return b.services.RegisterService(blockchainService)
 }
 

beacon-chain/rpc/eth/validator/BUILD.bazel

@@ -68,7 +68,7 @@ go_test(
         "//beacon-chain/rpc/prysm/v1alpha1/validator:go_default_library",
         "//beacon-chain/rpc/testutil:go_default_library",
         "//beacon-chain/state:go_default_library",
-        "//beacon-chain/state/stategen:go_default_library",
+        "//beacon-chain/state/stategen/mock:go_default_library",
         "//beacon-chain/sync/initial-sync/testing:go_default_library",
         "//config/fieldparams:go_default_library",
         "//config/params:go_default_library",

beacon-chain/rpc/eth/validator/validator_test.go

@@ -28,7 +28,7 @@ import (
 	v1alpha1validator "github.com/prysmaticlabs/prysm/v3/beacon-chain/rpc/prysm/v1alpha1/validator"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/rpc/testutil"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	mockSync "github.com/prysmaticlabs/prysm/v3/beacon-chain/sync/initial-sync/testing"
 	fieldparams "github.com/prysmaticlabs/prysm/v3/config/fieldparams"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
@@ -675,7 +675,7 @@ func TestProduceBlockV2(t *testing.T) {
 			AttPool:           attestations.NewPool(),
 			SlashingsPool:     slashings.NewPool(),
 			ExitPool:          voluntaryexits.NewPool(),
-			StateGen:          stategen.New(db),
+			StateGen:          sgmock.NewMockStategen(db),
 		}
 
 		proposerSlashings := make([]*ethpbalpha.ProposerSlashing, params.BeaconConfig().MaxProposerSlashings)
@@ -779,7 +779,7 @@ func TestProduceBlockV2(t *testing.T) {
 			AttPool:           attestations.NewPool(),
 			SlashingsPool:     slashings.NewPool(),
 			ExitPool:          voluntaryexits.NewPool(),
-			StateGen:          stategen.New(db),
+			StateGen:          sgmock.NewMockStategen(db),
 			SyncCommitteePool: synccommittee.NewStore(),
 		}
 
@@ -928,7 +928,7 @@ func TestProduceBlockV2(t *testing.T) {
 			AttPool:                attestations.NewPool(),
 			SlashingsPool:          slashings.NewPool(),
 			ExitPool:               voluntaryexits.NewPool(),
-			StateGen:               stategen.New(db),
+			StateGen:               sgmock.NewMockStategen(db),
 			SyncCommitteePool:      synccommittee.NewStore(),
 			ProposerSlotIndexCache: cache.NewProposerPayloadIDsCache(),
 		}
@@ -1052,7 +1052,7 @@ func TestProduceBlockV2SSZ(t *testing.T) {
 			AttPool:           attestations.NewPool(),
 			SlashingsPool:     slashings.NewPool(),
 			ExitPool:          voluntaryexits.NewPool(),
-			StateGen:          stategen.New(db),
+			StateGen:          sgmock.NewMockStategen(db),
 		}
 
 		proposerSlashings := make([]*ethpbalpha.ProposerSlashing, 1)
@@ -1213,7 +1213,7 @@ func TestProduceBlockV2SSZ(t *testing.T) {
 			AttPool:           attestations.NewPool(),
 			SlashingsPool:     slashings.NewPool(),
 			ExitPool:          voluntaryexits.NewPool(),
-			StateGen:          stategen.New(db),
+			StateGen:          sgmock.NewMockStategen(db),
 			SyncCommitteePool: synccommittee.NewStore(),
 		}
 
@@ -1417,7 +1417,7 @@ func TestProduceBlockV2SSZ(t *testing.T) {
 			AttPool:                attestations.NewPool(),
 			SlashingsPool:          slashings.NewPool(),
 			ExitPool:               voluntaryexits.NewPool(),
-			StateGen:               stategen.New(db),
+			StateGen:               sgmock.NewMockStategen(db),
 			SyncCommitteePool:      synccommittee.NewStore(),
 			ProposerSlotIndexCache: cache.NewProposerPayloadIDsCache(),
 		}
@@ -1635,7 +1635,7 @@ func TestProduceBlindedBlock(t *testing.T) {
 			AttPool:           attestations.NewPool(),
 			SlashingsPool:     slashings.NewPool(),
 			ExitPool:          voluntaryexits.NewPool(),
-			StateGen:          stategen.New(db),
+			StateGen:          sgmock.NewMockStategen(db),
 		}
 
 		proposerSlashings := make([]*ethpbalpha.ProposerSlashing, params.BeaconConfig().MaxProposerSlashings)
@@ -1739,7 +1739,7 @@ func TestProduceBlindedBlock(t *testing.T) {
 			AttPool:           attestations.NewPool(),
 			SlashingsPool:     slashings.NewPool(),
 			ExitPool:          voluntaryexits.NewPool(),
-			StateGen:          stategen.New(db),
+			StateGen:          sgmock.NewMockStategen(db),
 			SyncCommitteePool: synccommittee.NewStore(),
 		}
 
@@ -1888,7 +1888,7 @@ func TestProduceBlindedBlock(t *testing.T) {
 			AttPool:                attestations.NewPool(),
 			SlashingsPool:          slashings.NewPool(),
 			ExitPool:               voluntaryexits.NewPool(),
-			StateGen:               stategen.New(db),
+			StateGen:               sgmock.NewMockStategen(db),
 			SyncCommitteePool:      synccommittee.NewStore(),
 			ProposerSlotIndexCache: cache.NewProposerPayloadIDsCache(),
 		}
@@ -2012,7 +2012,7 @@ func TestProduceBlindedBlockSSZ(t *testing.T) {
 			AttPool:           attestations.NewPool(),
 			SlashingsPool:     slashings.NewPool(),
 			ExitPool:          voluntaryexits.NewPool(),
-			StateGen:          stategen.New(db),
+			StateGen:          sgmock.NewMockStategen(db),
 		}
 
 		proposerSlashings := make([]*ethpbalpha.ProposerSlashing, 1)
@@ -2173,7 +2173,7 @@ func TestProduceBlindedBlockSSZ(t *testing.T) {
 			AttPool:           attestations.NewPool(),
 			SlashingsPool:     slashings.NewPool(),
 			ExitPool:          voluntaryexits.NewPool(),
-			StateGen:          stategen.New(db),
+			StateGen:          sgmock.NewMockStategen(db),
 			SyncCommitteePool: synccommittee.NewStore(),
 		}
 
@@ -2377,7 +2377,7 @@ func TestProduceBlindedBlockSSZ(t *testing.T) {
 			AttPool:                attestations.NewPool(),
 			SlashingsPool:          slashings.NewPool(),
 			ExitPool:               voluntaryexits.NewPool(),
-			StateGen:               stategen.New(db),
+			StateGen:               sgmock.NewMockStategen(db),
 			SyncCommitteePool:      synccommittee.NewStore(),
 			ProposerSlotIndexCache: cache.NewProposerPayloadIDsCache(),
 		}

beacon-chain/rpc/prysm/v1alpha1/beacon/assignments_test.go

@@ -10,7 +10,6 @@ import (
 	mock "github.com/prysmaticlabs/prysm/v3/beacon-chain/blockchain/testing"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/core/helpers"
 	dbTest "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
 	mockstategen "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/cmd"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
@@ -60,7 +59,7 @@ func TestServer_ListAssignments_NoResults(t *testing.T) {
 	bs := &Server{
 		BeaconDB:           db,
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(db),
+		StateGen:           mockstategen.NewMockStategen(db),
 		ReplayerBuilder:    mockstategen.NewMockReplayerBuilder(mockstategen.WithMockState(st)),
 	}
 	wanted := &ethpb.ValidatorAssignments{
@@ -122,7 +121,7 @@ func TestServer_ListAssignments_Pagination_InputOutOfRange(t *testing.T) {
 			},
 		},
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(db),
+		StateGen:           mockstategen.NewMockStategen(db),
 		ReplayerBuilder:    mockstategen.NewMockReplayerBuilder(mockstategen.WithMockState(s)),
 	}
 
@@ -198,7 +197,7 @@ func TestServer_ListAssignments_Pagination_DefaultPageSize_NoArchive(t *testing.
 			},
 		},
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(db),
+		StateGen:           mockstategen.NewMockStategen(db),
 		ReplayerBuilder:    mockstategen.NewMockReplayerBuilder(mockstategen.WithMockState(s)),
 	}
 
@@ -265,7 +264,7 @@ func TestServer_ListAssignments_FilterPubkeysIndices_NoPagination(t *testing.T)
 			},
 		},
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(db),
+		StateGen:           mockstategen.NewMockStategen(db),
 		ReplayerBuilder:    mockstategen.NewMockReplayerBuilder(mockstategen.WithMockState(s)),
 	}
 
@@ -336,7 +335,7 @@ func TestServer_ListAssignments_CanFilterPubkeysIndices_WithPagination(t *testin
 			},
 		},
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(db),
+		StateGen:           mockstategen.NewMockStategen(db),
 	}
 
 	addDefaultReplayerBuilder(bs, db)

beacon-chain/rpc/prysm/v1alpha1/beacon/attestations_test.go

@@ -17,7 +17,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/core/signing"
 	dbTest "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/attestations"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	v1 "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/v1"
 	"github.com/prysmaticlabs/prysm/v3/cmd"
 	fieldparams "github.com/prysmaticlabs/prysm/v3/config/fieldparams"
@@ -569,7 +569,7 @@ func TestServer_ListIndexedAttestations_GenesisEpoch(t *testing.T) {
 		BeaconDB:           db,
 		GenesisTimeFetcher: &chainMock.ChainService{State: state},
 		HeadFetcher:        &chainMock.ChainService{State: state},
-		StateGen:           stategen.New(db),
+		StateGen:           sgmock.NewMockStategen(db),
 	}
 	err := db.SaveStateSummary(ctx, &ethpb.StateSummary{
 		Root: targetRoot1[:],
@@ -669,7 +669,7 @@ func TestServer_ListIndexedAttestations_OldEpoch(t *testing.T) {
 		GenesisTimeFetcher: &chainMock.ChainService{
 			Genesis: time.Now(),
 		},
-		StateGen: stategen.New(db),
+		StateGen: sgmock.NewMockStategen(db),
 	}
 	err = db.SaveStateSummary(ctx, &ethpb.StateSummary{
 		Root: blockRoot[:],
@@ -940,7 +940,7 @@ func TestServer_StreamIndexedAttestations_OK(t *testing.T) {
 		},
 		AttestationNotifier:         chainService.OperationNotifier(),
 		CollectedAttestationsBuffer: make(chan []*ethpb.Attestation, 1),
-		StateGen:                    stategen.New(db),
+		StateGen:                    sgmock.NewMockStategen(db),
 	}
 
 	for dataRoot, sameDataAtts := range atts {

beacon-chain/rpc/prysm/v1alpha1/beacon/committees_test.go

@@ -42,7 +42,7 @@ func TestServer_ListBeaconCommittees_CurrentEpoch(t *testing.T) {
 	bs := &Server{
 		HeadFetcher:        m,
 		GenesisTimeFetcher: m,
-		StateGen:           stategen.New(db),
+		StateGen:           mockstategen.NewMockStategen(db),
 	}
 	b := util.NewBeaconBlock()
 	util.SaveBlock(t, ctx, db, b)
@@ -115,7 +115,7 @@ func TestServer_ListBeaconCommittees_PreviousEpoch(t *testing.T) {
 	bs := &Server{
 		HeadFetcher:        m,
 		GenesisTimeFetcher: m,
-		StateGen:           stategen.New(db),
+		StateGen:           mockstategen.NewMockStategen(db),
 	}
 	addDefaultReplayerBuilder(bs, db)
 
@@ -170,7 +170,7 @@ func TestRetrieveCommitteesForRoot(t *testing.T) {
 	bs := &Server{
 		HeadFetcher:        m,
 		GenesisTimeFetcher: m,
-		StateGen:           stategen.New(db),
+		StateGen:           mockstategen.NewMockStategen(db),
 	}
 	b := util.NewBeaconBlock()
 	util.SaveBlock(t, ctx, db, b)

beacon-chain/rpc/prysm/v1alpha1/beacon/validators_test.go

@@ -19,7 +19,6 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/db"
 	dbTest "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
 	mockstategen "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	v1 "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/v1"
 	mockSync "github.com/prysmaticlabs/prysm/v3/beacon-chain/sync/initial-sync/testing"
@@ -107,7 +106,7 @@ func TestServer_ListValidatorBalances_NoResults(t *testing.T) {
 	require.NoError(t, st.SetSlot(0))
 	bs := &Server{
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(beaconDB),
+		StateGen:           mockstategen.NewMockStategen(beaconDB),
 	}
 
 	headState, err := util.NewBeaconState()
@@ -174,7 +173,7 @@ func TestServer_ListValidatorBalances_DefaultResponse_NoArchive(t *testing.T) {
 	require.NoError(t, beaconDB.SaveState(ctx, st, gRoot))
 	bs := &Server{
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(beaconDB),
+		StateGen:           mockstategen.NewMockStategen(beaconDB),
 		HeadFetcher: &mock.ChainService{
 			State: st,
 		},
@@ -203,7 +202,7 @@ func TestServer_ListValidatorBalances_PaginationOutOfRange(t *testing.T) {
 
 	bs := &Server{
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(beaconDB),
+		StateGen:           mockstategen.NewMockStategen(beaconDB),
 		HeadFetcher: &mock.ChainService{
 			State: headState,
 		},
@@ -252,7 +251,7 @@ func TestServer_ListValidatorBalances_Pagination_Default(t *testing.T) {
 
 	bs := &Server{
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(beaconDB),
+		StateGen:           mockstategen.NewMockStategen(beaconDB),
 		HeadFetcher: &mock.ChainService{
 			State: headState,
 		},
@@ -336,7 +335,7 @@ func TestServer_ListValidatorBalances_Pagination_CustomPageSizes(t *testing.T) {
 
 	bs := &Server{
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(beaconDB),
+		StateGen:           mockstategen.NewMockStategen(beaconDB),
 		HeadFetcher: &mock.ChainService{
 			State: headState,
 		},
@@ -404,7 +403,7 @@ func TestServer_ListValidatorBalances_OutOfRange(t *testing.T) {
 
 	bs := &Server{
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(beaconDB),
+		StateGen:           mockstategen.NewMockStategen(beaconDB),
 		HeadFetcher: &mock.ChainService{
 			State: headState,
 		},
@@ -500,7 +499,7 @@ func TestServer_ListValidators_NoResults(t *testing.T) {
 		HeadFetcher: &mock.ChainService{
 			State: st,
 		},
-		StateGen: stategen.New(beaconDB),
+		StateGen: mockstategen.NewMockStategen(beaconDB),
 	}
 	wanted := &ethpb.Validators{
 		ValidatorList: make([]*ethpb.Validators_ValidatorContainer, 0),
@@ -567,7 +566,7 @@ func TestServer_ListValidators_OnlyActiveValidators(t *testing.T) {
 			// We are in epoch 0.
 			Genesis: time.Now(),
 		},
-		StateGen: stategen.New(beaconDB),
+		StateGen: mockstategen.NewMockStategen(beaconDB),
 	}
 
 	b := util.NewBeaconBlock()
@@ -635,7 +634,7 @@ func TestServer_ListValidators_InactiveInTheMiddle(t *testing.T) {
 			// We are in epoch 0.
 			Genesis: time.Now(),
 		},
-		StateGen: stategen.New(beaconDB),
+		StateGen: mockstategen.NewMockStategen(beaconDB),
 	}
 
 	b := util.NewBeaconBlock()
@@ -667,7 +666,7 @@ func TestServer_ListValidatorBalances_UnknownValidatorInResponse(t *testing.T) {
 
 	bs := &Server{
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(beaconDB),
+		StateGen:           mockstategen.NewMockStategen(beaconDB),
 		HeadFetcher: &mock.ChainService{
 			State: headState,
 		},
@@ -725,7 +724,7 @@ func TestServer_ListValidators_NoPagination(t *testing.T) {
 				Epoch: 0,
 			},
 		},
-		StateGen: stategen.New(beaconDB),
+		StateGen: mockstategen.NewMockStategen(beaconDB),
 	}
 
 	received, err := bs.ListValidators(context.Background(), &ethpb.ListValidatorsRequest{})
@@ -791,7 +790,7 @@ func TestServer_ListValidators_IndicesPubKeys(t *testing.T) {
 			// We are in epoch 0.
 			Genesis: time.Now(),
 		},
-		StateGen: stategen.New(beaconDB),
+		StateGen: mockstategen.NewMockStategen(beaconDB),
 	}
 
 	pubKeysWanted := make([][]byte, len(pubkeyIndicesWanted))
@@ -827,7 +826,7 @@ func TestServer_ListValidators_Pagination(t *testing.T) {
 			// We are in epoch 0.
 			Genesis: time.Now(),
 		},
-		StateGen: stategen.New(beaconDB),
+		StateGen: mockstategen.NewMockStategen(beaconDB),
 	}
 
 	tests := []struct {
@@ -964,7 +963,7 @@ func TestServer_ListValidators_PaginationOutOfRange(t *testing.T) {
 			// We are in epoch 0.
 			Genesis: time.Now(),
 		},
-		StateGen: stategen.New(beaconDB),
+		StateGen: mockstategen.NewMockStategen(beaconDB),
 	}
 
 	req := &ethpb.ListValidatorsRequest{PageToken: strconv.Itoa(1), PageSize: 100}
@@ -1008,7 +1007,7 @@ func TestServer_ListValidators_DefaultPageSize(t *testing.T) {
 			// We are in epoch 0.
 			Genesis: time.Now(),
 		},
-		StateGen: stategen.New(beaconDB),
+		StateGen: mockstategen.NewMockStategen(beaconDB),
 	}
 
 	req := &ethpb.ListValidatorsRequest{}
@@ -1132,7 +1131,7 @@ func TestServer_ListValidators_ProcessHeadStateSlots(t *testing.T) {
 		GenesisTimeFetcher: &mock.ChainService{
 			Genesis: time.Now().Add(time.Duration(-1*int64(secondsPerEpoch)) * time.Second),
 		},
-		StateGen: stategen.New(beaconDB),
+		StateGen: mockstategen.NewMockStategen(beaconDB),
 	}
 
 	req := &ethpb.ListValidatorsRequest{
@@ -1489,7 +1488,7 @@ func TestServer_GetValidatorParticipation_CannotRequestFutureEpoch(t *testing.T)
 			State: headState,
 		},
 		GenesisTimeFetcher: &mock.ChainService{},
-		StateGen:           stategen.New(beaconDB),
+		StateGen:           mockstategen.NewMockStategen(beaconDB),
 	}
 
 	wanted := "Cannot retrieve information about an epoch"
@@ -1554,7 +1553,7 @@ func TestServer_GetValidatorParticipation_CurrentAndPrevEpoch(t *testing.T) {
 	bs := &Server{
 		BeaconDB:    beaconDB,
 		HeadFetcher: m,
-		StateGen:    stategen.New(beaconDB),
+		StateGen:    mockstategen.NewMockStategen(beaconDB),
 		GenesisTimeFetcher: &mock.ChainService{
 			Genesis: prysmTime.Now().Add(time.Duration(-1*offset) * time.Second),
 		},
@@ -1633,7 +1632,7 @@ func TestServer_GetValidatorParticipation_OrphanedUntilGenesis(t *testing.T) {
 	bs := &Server{
 		BeaconDB:    beaconDB,
 		HeadFetcher: m,
-		StateGen:    stategen.New(beaconDB),
+		StateGen:    mockstategen.NewMockStategen(beaconDB),
 		GenesisTimeFetcher: &mock.ChainService{
 			Genesis: prysmTime.Now().Add(time.Duration(-1*offset) * time.Second),
 		},
@@ -1731,7 +1730,7 @@ func runGetValidatorParticipationCurrentAndPrevEpoch(t *testing.T, genState stat
 	bs := &Server{
 		BeaconDB:    beaconDB,
 		HeadFetcher: m,
-		StateGen:    stategen.New(beaconDB),
+		StateGen:    mockstategen.NewMockStategen(beaconDB),
 		GenesisTimeFetcher: &mock.ChainService{
 			Genesis: prysmTime.Now().Add(time.Duration(-1*offset) * time.Second),
 		},
@@ -2225,7 +2224,7 @@ func TestServer_GetIndividualVotes_ValidatorsDontExist(t *testing.T) {
 	util.SaveBlock(t, ctx, beaconDB, b)
 	gRoot, err := b.Block.HashTreeRoot()
 	require.NoError(t, err)
-	gen := stategen.New(beaconDB)
+	gen := mockstategen.NewMockStategen(beaconDB)
 	require.NoError(t, gen.SaveState(ctx, gRoot, beaconState))
 	require.NoError(t, beaconDB.SaveState(ctx, beaconState, gRoot))
 	require.NoError(t, beaconDB.SaveGenesisBlockRoot(ctx, gRoot))
@@ -2321,7 +2320,7 @@ func TestServer_GetIndividualVotes_Working(t *testing.T) {
 	util.SaveBlock(t, ctx, beaconDB, b)
 	gRoot, err := b.Block.HashTreeRoot()
 	require.NoError(t, err)
-	gen := stategen.New(beaconDB)
+	gen := mockstategen.NewMockStategen(beaconDB)
 	require.NoError(t, gen.SaveState(ctx, gRoot, beaconState))
 	require.NoError(t, beaconDB.SaveState(ctx, beaconState, gRoot))
 	require.NoError(t, beaconDB.SaveGenesisBlockRoot(ctx, gRoot))
@@ -2386,7 +2385,7 @@ func TestServer_GetIndividualVotes_WorkingAltair(t *testing.T) {
 	util.SaveBlock(t, ctx, beaconDB, b)
 	gRoot, err := b.Block.HashTreeRoot()
 	require.NoError(t, err)
-	gen := stategen.New(beaconDB)
+	gen := mockstategen.NewMockStategen(beaconDB)
 	require.NoError(t, gen.SaveState(ctx, gRoot, beaconState))
 	require.NoError(t, beaconDB.SaveState(ctx, beaconState, gRoot))
 	require.NoError(t, beaconDB.SaveGenesisBlockRoot(ctx, gRoot))
@@ -2450,7 +2449,7 @@ func TestServer_GetIndividualVotes_AltairEndOfEpoch(t *testing.T) {
 	util.SaveBlock(t, ctx, beaconDB, b)
 	gRoot, err := b.Block.HashTreeRoot()
 	require.NoError(t, err)
-	gen := stategen.New(beaconDB)
+	gen := mockstategen.NewMockStategen(beaconDB)
 	require.NoError(t, gen.SaveState(ctx, gRoot, beaconState))
 	require.NoError(t, beaconDB.SaveState(ctx, beaconState, gRoot))
 	require.NoError(t, beaconDB.SaveGenesisBlockRoot(ctx, gRoot))
@@ -2538,7 +2537,7 @@ func TestServer_GetIndividualVotes_BellatrixEndOfEpoch(t *testing.T) {
 	util.SaveBlock(t, ctx, beaconDB, b)
 	gRoot, err := b.Block.HashTreeRoot()
 	require.NoError(t, err)
-	gen := stategen.New(beaconDB)
+	gen := mockstategen.NewMockStategen(beaconDB)
 	require.NoError(t, gen.SaveState(ctx, gRoot, beaconState))
 	require.NoError(t, beaconDB.SaveState(ctx, beaconState, gRoot))
 	require.NoError(t, beaconDB.SaveGenesisBlockRoot(ctx, gRoot))

beacon-chain/rpc/prysm/v1alpha1/debug/block_test.go

@@ -9,7 +9,7 @@ import (
 	mock "github.com/prysmaticlabs/prysm/v3/beacon-chain/blockchain/testing"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/core/helpers"
 	dbTest "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	fieldparams "github.com/prysmaticlabs/prysm/v3/config/fieldparams"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	ethpb "github.com/prysmaticlabs/prysm/v3/proto/prysm/v1alpha1"
@@ -53,7 +53,7 @@ func TestServer_GetAttestationInclusionSlot(t *testing.T) {
 	offset := int64(2 * params.BeaconConfig().SlotsPerEpoch.Mul(params.BeaconConfig().SecondsPerSlot))
 	bs := &Server{
 		BeaconDB:           db,
-		StateGen:           stategen.New(db),
+		StateGen:           sgmock.NewMockStategen(db),
 		GenesisTimeFetcher: &mock.ChainService{Genesis: time.Now().Add(time.Duration(-1*offset) * time.Second)},
 	}
 

beacon-chain/rpc/prysm/v1alpha1/debug/state_test.go

@@ -34,7 +34,7 @@ func TestServer_GetBeaconState(t *testing.T) {
 	util.SaveBlock(t, ctx, db, b)
 	gRoot, err := b.Block.HashTreeRoot()
 	require.NoError(t, err)
-	gen := stategen.New(db)
+	gen := mockstategen.NewMockStategen(db)
 	require.NoError(t, gen.SaveState(ctx, gRoot, st))
 	require.NoError(t, db.SaveState(ctx, st, gRoot))
 	bs := &Server{

beacon-chain/rpc/prysm/v1alpha1/validator/BUILD.bazel

@@ -141,7 +141,6 @@ go_test(
         "//beacon-chain/operations/voluntaryexits:go_default_library",
         "//beacon-chain/p2p/testing:go_default_library",
         "//beacon-chain/state:go_default_library",
-        "//beacon-chain/state/stategen:go_default_library",
         "//beacon-chain/state/stategen/mock:go_default_library",
         "//beacon-chain/state/v1:go_default_library",
         "//beacon-chain/state/v3:go_default_library",

beacon-chain/rpc/prysm/v1alpha1/validator/attester_test.go

@@ -12,7 +12,7 @@ import (
 	dbutil "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/attestations"
 	mockp2p "github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/testing"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	v1 "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/v1"
 	mockSync "github.com/prysmaticlabs/prysm/v3/beacon-chain/sync/initial-sync/testing"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
@@ -423,7 +423,7 @@ func TestServer_GetAttestationData_HeadStateSlotGreaterThanRequestSlot(t *testin
 		FinalizationFetcher: &mock.ChainService{CurrentJustifiedCheckPoint: beaconState.CurrentJustifiedCheckpoint()},
 		TimeFetcher:         &mock.ChainService{Genesis: time.Now().Add(time.Duration(-1*offset) * time.Second)},
 		StateNotifier:       chainService.StateNotifier(),
-		StateGen:            stategen.New(db),
+		StateGen:            sgmock.NewMockStategen(db),
 	}
 	require.NoError(t, db.SaveState(ctx, beaconState, blockRoot))
 	util.SaveBlock(t, ctx, db, block)

beacon-chain/rpc/prysm/v1alpha1/validator/proposer_bellatrix_test.go

@@ -24,7 +24,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/synccommittee"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/voluntaryexits"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	v3 "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/v3"
 	mockSync "github.com/prysmaticlabs/prysm/v3/beacon-chain/sync/initial-sync/testing"
 	fieldparams "github.com/prysmaticlabs/prysm/v3/config/fieldparams"
@@ -58,7 +58,7 @@ func TestServer_buildHeaderBlock(t *testing.T) {
 
 	proposerServer := &Server{
 		BeaconDB: db,
-		StateGen: stategen.New(db),
+		StateGen: sgmock.NewMockStategen(db),
 	}
 	b, err := util.GenerateFullBlockAltair(copiedState, keys, util.DefaultBlockGenConfig(), 1)
 	require.NoError(t, err)
@@ -69,7 +69,7 @@ func TestServer_buildHeaderBlock(t *testing.T) {
 	b1, err := util.GenerateFullBlockAltair(copiedState, keys, util.DefaultBlockGenConfig(), 2)
 	require.NoError(t, err)
 
-	vs := &Server{StateGen: stategen.New(db), BeaconDB: db}
+	vs := &Server{StateGen: sgmock.NewMockStategen(db), BeaconDB: db}
 	h := &v1.ExecutionPayloadHeader{
 		BlockNumber:      123,
 		GasLimit:         456,
@@ -419,7 +419,7 @@ func TestServer_getAndBuildHeaderBlock(t *testing.T) {
 		Timestamp:        ts,
 	}
 
-	vs.StateGen = stategen.New(vs.BeaconDB)
+	vs.StateGen = sgmock.NewMockStategen(vs.BeaconDB)
 	vs.GenesisFetcher = &blockchainTest.ChainService{}
 	vs.ForkFetcher = &blockchainTest.ChainService{Fork: &ethpb.Fork{}}
 
@@ -538,7 +538,7 @@ func TestServer_GetBellatrixBeaconBlock_HappyCase(t *testing.T) {
 		AttPool:           attestations.NewPool(),
 		SlashingsPool:     slashings.NewPool(),
 		ExitPool:          voluntaryexits.NewPool(),
-		StateGen:          stategen.New(db),
+		StateGen:          sgmock.NewMockStategen(db),
 		SyncCommitteePool: synccommittee.NewStore(),
 		ExecutionEngineCaller: &mockExecution.EngineClient{
 			PayloadIDBytes:   &v1.PayloadIDBytes{1},
@@ -669,7 +669,7 @@ func TestServer_GetBellatrixBeaconBlock_BuilderCase(t *testing.T) {
 		AttPool:             attestations.NewPool(),
 		SlashingsPool:       slashings.NewPool(),
 		ExitPool:            voluntaryexits.NewPool(),
-		StateGen:            stategen.New(db),
+		StateGen:            sgmock.NewMockStategen(db),
 		SyncCommitteePool:   synccommittee.NewStore(),
 		ExecutionEngineCaller: &mockExecution.EngineClient{
 			PayloadIDBytes:   &v1.PayloadIDBytes{1},

beacon-chain/rpc/prysm/v1alpha1/validator/proposer_test.go

@@ -25,7 +25,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/synccommittee"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/voluntaryexits"
 	mockp2p "github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/testing"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	v1 "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/v1"
 	mockSync "github.com/prysmaticlabs/prysm/v3/beacon-chain/sync/initial-sync/testing"
 	fieldparams "github.com/prysmaticlabs/prysm/v3/config/fieldparams"
@@ -137,7 +137,7 @@ func TestProposer_ComputeStateRoot_OK(t *testing.T) {
 		ChainStartFetcher: &mockExecution.Chain{},
 		Eth1InfoFetcher:   &mockExecution.Chain{},
 		Eth1BlockFetcher:  &mockExecution.Chain{},
-		StateGen:          stategen.New(db),
+		StateGen:          sgmock.NewMockStategen(db),
 	}
 	req := util.NewBeaconBlock()
 	req.Block.ProposerIndex = 21
@@ -1928,7 +1928,7 @@ func TestProposer_GetBeaconBlock_PreForkEpoch(t *testing.T) {
 		AttPool:           attestations.NewPool(),
 		SlashingsPool:     slashings.NewPool(),
 		ExitPool:          voluntaryexits.NewPool(),
-		StateGen:          stategen.New(db),
+		StateGen:          sgmock.NewMockStategen(db),
 		SyncCommitteePool: synccommittee.NewStore(),
 	}
 
@@ -2039,7 +2039,7 @@ func TestProposer_GetBeaconBlock_PostForkEpoch(t *testing.T) {
 		AttPool:           attestations.NewPool(),
 		SlashingsPool:     slashings.NewPool(),
 		ExitPool:          voluntaryexits.NewPool(),
-		StateGen:          stategen.New(db),
+		StateGen:          sgmock.NewMockStategen(db),
 		SyncCommitteePool: synccommittee.NewStore(),
 	}
 
@@ -2191,7 +2191,7 @@ func TestProposer_GetBeaconBlock_BellatrixEpoch(t *testing.T) {
 		AttPool:           attestations.NewPool(),
 		SlashingsPool:     slashings.NewPool(),
 		ExitPool:          voluntaryexits.NewPool(),
-		StateGen:          stategen.New(db),
+		StateGen:          sgmock.NewMockStategen(db),
 		SyncCommitteePool: synccommittee.NewStore(),
 		ExecutionEngineCaller: &mockExecution.EngineClient{
 			PayloadIDBytes:   &enginev1.PayloadIDBytes{1},

beacon-chain/slasher/BUILD.bazel

@@ -77,7 +77,7 @@ go_test(
         "//beacon-chain/operations/slashings/mock:go_default_library",
         "//beacon-chain/slasher/mock:go_default_library",
         "//beacon-chain/slasher/types:go_default_library",
-        "//beacon-chain/state/stategen:go_default_library",
+        "//beacon-chain/state/stategen/mock:go_default_library",
         "//beacon-chain/sync/initial-sync/testing:go_default_library",
         "//config/fieldparams:go_default_library",
         "//config/params:go_default_library",

beacon-chain/slasher/detect_blocks_test.go

@@ -9,7 +9,7 @@ import (
 	dbtest "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	slashingsmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/slashings/mock"
 	slashertypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/slasher/types"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/v3/crypto/bls"
@@ -60,7 +60,7 @@ func Test_processQueuedBlocks_DetectsDoubleProposals(t *testing.T) {
 			Database:             slasherDB,
 			StateNotifier:        &mock.MockStateNotifier{},
 			HeadStateFetcher:     mockChain,
-			StateGen:             stategen.New(beaconDB),
+			StateGen:             sgmock.NewMockStategen(beaconDB),
 			SlashingPoolInserter: &slashingsmock.PoolMock{},
 		},
 		params:    DefaultParams(),

beacon-chain/slasher/process_slashings_test.go

@@ -8,7 +8,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/core/signing"
 	dbtest "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	slashingsmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/slashings/mock"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	"github.com/prysmaticlabs/prysm/v3/crypto/bls"
 	"github.com/prysmaticlabs/prysm/v3/encoding/bytesutil"
@@ -44,7 +44,7 @@ func TestService_processAttesterSlashings(t *testing.T) {
 		serviceCfg: &ServiceConfig{
 			Database:                slasherDB,
 			AttestationStateFetcher: mockChain,
-			StateGen:                stategen.New(beaconDB),
+			StateGen:                sgmock.NewMockStategen(beaconDB),
 			SlashingPoolInserter:    &slashingsmock.PoolMock{},
 			HeadStateFetcher:        mockChain,
 		},
@@ -151,7 +151,7 @@ func TestService_processProposerSlashings(t *testing.T) {
 		serviceCfg: &ServiceConfig{
 			Database:                slasherDB,
 			AttestationStateFetcher: mockChain,
-			StateGen:                stategen.New(beaconDB),
+			StateGen:                sgmock.NewMockStategen(beaconDB),
 			SlashingPoolInserter:    &slashingsmock.PoolMock{},
 			HeadStateFetcher:        mockChain,
 		},

beacon-chain/state/stategen/BUILD.bazel

@@ -26,6 +26,7 @@ go_library(
         "//beacon-chain/core/transition:go_default_library",
         "//beacon-chain/db:go_default_library",
         "//beacon-chain/db/filters:go_default_library",
+        "//beacon-chain/forkchoice/types:go_default_library",
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/sync/backfill:go_default_library",
         "//cache/lru:go_default_library",
@@ -62,6 +63,7 @@ go_test(
         "replayer_test.go",
         "service_test.go",
         "setter_test.go",
+        "shared_test.go",
     ],
     embed = [":go_default_library"],
     deps = [
@@ -69,7 +71,10 @@ go_test(
         "//beacon-chain/core/helpers:go_default_library",
         "//beacon-chain/core/transition:go_default_library",
         "//beacon-chain/db:go_default_library",
+        "//beacon-chain/db/iface:go_default_library",
         "//beacon-chain/db/testing:go_default_library",
+        "//beacon-chain/forkchoice/doubly-linked-tree:go_default_library",
+        "//beacon-chain/forkchoice/types:go_default_library",
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/state/v1:go_default_library",
         "//config/params:go_default_library",
@@ -84,6 +89,7 @@ go_test(
         "//testing/assert:go_default_library",
         "//testing/require:go_default_library",
         "//testing/util:go_default_library",
+        "//time/slots:go_default_library",
         "@com_github_pkg_errors//:go_default_library",
         "@com_github_sirupsen_logrus//hooks/test:go_default_library",
         "@org_golang_google_protobuf//proto:go_default_library",

beacon-chain/state/stategen/getter_test.go

@@ -19,7 +19,7 @@ func TestStateByRoot_GenesisState(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	b := util.NewBeaconBlock()
 	bRoot, err := b.Block.HashTreeRoot()
 	require.NoError(t, err)
@@ -36,7 +36,7 @@ func TestStateByRoot_ColdState(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	service.finalizedInfo.slot = 2
 	service.slotsPerArchivedPoint = 1
 
@@ -59,7 +59,7 @@ func TestStateByRootIfCachedNoCopy_HotState(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	r := [32]byte{'A'}
@@ -74,7 +74,7 @@ func TestStateByRootIfCachedNoCopy_ColdState(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	service.finalizedInfo.slot = 2
 	service.slotsPerArchivedPoint = 1
 
@@ -97,7 +97,7 @@ func TestStateByRoot_HotStateUsingEpochBoundaryCacheNoReplay(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	require.NoError(t, beaconState.SetSlot(10))
@@ -115,7 +115,7 @@ func TestStateByRoot_HotStateUsingEpochBoundaryCacheWithReplay(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	blk := util.NewBeaconBlock()
@@ -140,7 +140,7 @@ func TestStateByRoot_HotStateCached(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	r := [32]byte{'A'}
@@ -156,7 +156,7 @@ func TestDeleteStateFromCaches(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	r := [32]byte{'A'}
 
@@ -185,7 +185,7 @@ func TestStateByRoot_StateByRootInitialSync(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	b := util.NewBeaconBlock()
 	bRoot, err := b.Block.HashTreeRoot()
 	require.NoError(t, err)
@@ -202,7 +202,7 @@ func TestStateByRootInitialSync_UseEpochStateCache(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	targetSlot := types.Slot(10)
@@ -220,7 +220,7 @@ func TestStateByRootInitialSync_UseCache(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	r := [32]byte{'A'}
@@ -238,7 +238,7 @@ func TestStateByRootInitialSync_UseCache(t *testing.T) {
 func TestStateByRootInitialSync_CanProcessUpTo(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	blk := util.NewBeaconBlock()
@@ -262,7 +262,7 @@ func TestStateByRootInitialSync_CanProcessUpTo(t *testing.T) {
 func TestLoadeStateByRoot_Cached(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	r := [32]byte{'A'}
@@ -277,7 +277,7 @@ func TestLoadeStateByRoot_Cached(t *testing.T) {
 func TestLoadeStateByRoot_FinalizedState(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	genesisStateRoot, err := beaconState.HashTreeRoot(ctx)
@@ -301,7 +301,7 @@ func TestLoadeStateByRoot_FinalizedState(t *testing.T) {
 func TestLoadeStateByRoot_EpochBoundaryStateCanProcess(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	gBlk := util.NewBeaconBlock()
@@ -327,7 +327,7 @@ func TestLoadeStateByRoot_EpochBoundaryStateCanProcess(t *testing.T) {
 func TestLoadeStateByRoot_FromDBBoundaryCase(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	gBlk := util.NewBeaconBlock()
@@ -353,7 +353,7 @@ func TestLoadeStateByRoot_FromDBBoundaryCase(t *testing.T) {
 func TestLastAncestorState_CanGetUsingDB(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	b0 := util.NewBeaconBlock()
 	b0.Block.ParentRoot = bytesutil.PadTo([]byte{'a'}, 32)
@@ -393,7 +393,7 @@ func TestLastAncestorState_CanGetUsingDB(t *testing.T) {
 func TestLastAncestorState_CanGetUsingCache(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	b0 := util.NewBeaconBlock()
 	b0.Block.ParentRoot = bytesutil.PadTo([]byte{'a'}, 32)
@@ -433,7 +433,7 @@ func TestLastAncestorState_CanGetUsingCache(t *testing.T) {
 func TestState_HasState(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	s, err := util.NewBeaconState()
 	require.NoError(t, err)
 	rHit1 := [32]byte{1}
@@ -465,7 +465,7 @@ func TestState_HasState(t *testing.T) {
 func TestState_HasStateInCache(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	s, err := util.NewBeaconState()
 	require.NoError(t, err)
 	rHit1 := [32]byte{1}

beacon-chain/state/stategen/history_test.go

@@ -529,6 +529,25 @@ func (m *mockCanonicalChecker) IsCanonical(_ context.Context, root [32]byte) (bo
 	return m.is, m.err
 }
 
+func newMockCanonicalMap() *mockCanonicalMap {
+	return &mockCanonicalMap{
+		canonical: make(map[[32]byte]struct{}),
+	}
+}
+
+type mockCanonicalMap struct {
+	canonical map[[32]byte]struct{}
+}
+
+func (m *mockCanonicalMap) IsCanonical(_ context.Context, root [32]byte) (bool, error) {
+	_, exists := m.canonical[root]
+	return exists, nil
+}
+
+func (m *mockCanonicalMap) AddCanonical(root [32]byte) {
+	m.canonical[root] = struct{}{}
+}
+
 func TestReverseChain(t *testing.T) {
 	// test 0,1,2,3 elements to handle: zero case; single element; even number; odd number
 	for i := 0; i < 4; i++ {

beacon-chain/state/stategen/hot_state_cache.go

@@ -1,13 +1,21 @@
 package stategen
 
 import (
+	"context"
 	"sync"
 
 	lru "github.com/hashicorp/golang-lru"
+	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	"github.com/prysmaticlabs/prysm/v3/beacon-chain/db"
+	forkchoicetypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/types"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
 	lruwrpr "github.com/prysmaticlabs/prysm/v3/cache/lru"
+	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
+	ethpb "github.com/prysmaticlabs/prysm/v3/proto/prysm/v1alpha1"
+	"github.com/prysmaticlabs/prysm/v3/time/slots"
+	"github.com/sirupsen/logrus"
 )
 
 var (
@@ -93,3 +101,190 @@ func (c *hotStateCache) delete(blockRoot [32]byte) bool {
 	defer c.lock.Unlock()
 	return c.cache.Remove(blockRoot)
 }
+
+// FinalizedCheckpointer describes the forkchoice methods needed by the stategen service
+type FinalizedCheckpointer interface {
+	FinalizedCheckpoint() *forkchoicetypes.Checkpoint
+}
+
+type PersistenceMode int
+
+const (
+	// PersistenceModeMemory means the hot state cache does write to the database
+	PersistenceModeMemory PersistenceMode = iota
+	PersistenceModeSnapshot
+)
+
+func (m PersistenceMode) String() string {
+	switch m {
+	case PersistenceModeMemory:
+		return "memory"
+	case PersistenceModeSnapshot:
+		return "snapshot"
+	default:
+		return "unknown"
+	}
+}
+
+func NewHotStateSaver(d db.NoHeadAccessDatabase, fc FinalizedCheckpointer, cs CurrentSlotter) *hotStateSaver {
+	return &hotStateSaver{
+		snapshotInterval: DefaultSnapshotInterval,
+		db:               d,
+		fc:               fc,
+		cs:               cs,
+	}
+}
+
+// This tracks the config in the event of long non-finality,
+// how often does the node save hot states to db? what are
+// the saved hot states in db?... etc
+type hotStateSaver struct {
+	m                PersistenceMode
+	lock             sync.RWMutex
+	snapshotInterval types.Slot
+	savedRoots       [][32]byte
+	db               db.NoHeadAccessDatabase
+	fc               FinalizedCheckpointer
+	cs               CurrentSlotter
+}
+
+var _ Saver = &hotStateSaver{}
+
+// enable/disable hot state saving m depending on
+// whether the size of the gap between finalized and current epochs
+// is greater than the threshold.
+func (s *hotStateSaver) refreshMode(ctx context.Context) (PersistenceMode, error) {
+	current := slots.ToEpoch(s.cs.CurrentSlot())
+	fcp := s.fc.FinalizedCheckpoint()
+	if fcp == nil {
+		return PersistenceModeMemory, errForkchoiceFinalizedNil
+	}
+	// don't allow underflow
+	if fcp.Epoch > current {
+		return PersistenceModeMemory, errCurrentEpochBehindFinalized
+	}
+	if current-fcp.Epoch >= hotStateSaveThreshold {
+		s.enableSnapshots()
+		return PersistenceModeSnapshot, nil
+	}
+
+	return PersistenceModeMemory, s.disableSnapshots(ctx)
+}
+
+func (s *hotStateSaver) mode() PersistenceMode {
+	s.lock.RLock()
+	defer s.lock.RUnlock()
+	return s.m
+}
+
+// enableHotStateSaving enters the m that saves hot beacon state to the DB.
+// This usually gets triggered when there's long duration since finality.
+func (s *hotStateSaver) enableSnapshots() {
+	if s.mode() == PersistenceModeSnapshot {
+		return
+	}
+
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	s.m = PersistenceModeSnapshot
+	log.WithFields(logrus.Fields{
+		"mode":             s.m.String(),
+		"slotsPerSnapshot": s.snapshotInterval,
+	}).Warn("Enabling state cache db snapshots")
+}
+
+// disableHotStateSaving exits the m that saves beacon state to DB for the hot states.
+// This usually gets triggered once there's finality after long duration since finality.
+func (s *hotStateSaver) disableSnapshots(ctx context.Context) error {
+	if s.mode() == PersistenceModeMemory {
+		return nil
+	}
+
+	s.lock.Lock()
+	defer s.lock.Unlock()
+
+	log.WithFields(logrus.Fields{
+		"mode":             PersistenceModeMemory.String(),
+		"slotsPerSnapshot": s.snapshotInterval,
+	}).Warn("Disabling state cache db snapshots and removing saved snapshots")
+
+	// we have a recent-enough finalized state, so time to clean up the state cache snapshots
+	if err := s.db.DeleteStates(ctx, s.savedRoots); err != nil {
+		return err
+	}
+	s.savedRoots = nil
+	s.m = PersistenceModeMemory
+
+	return nil
+}
+
+func shouldSave(m PersistenceMode, interval types.Slot, st state.BeaconState) bool {
+	if m != PersistenceModeSnapshot {
+		// only write full states to the db when in snapshot mode
+		return false
+	}
+	// divide by zero guard
+	if interval == 0 {
+		return false
+	}
+	// only saving every s.duration slots - typically every 128 slots
+	// checking this first avoids holding the lock if we aren't on a slot that should be saved
+	if st.Slot().ModSlot(interval) != 0 {
+		return false
+	}
+	return true
+}
+
+func (s *hotStateSaver) Save(ctx context.Context, blockRoot [32]byte, st state.BeaconState) error {
+	mode, err := s.refreshMode(ctx)
+	if err != nil {
+		return errors.Wrap(err, "unable to make hot state saving decision")
+	}
+	err = s.db.SaveStateSummary(ctx, &ethpb.StateSummary{Slot: st.Slot(), Root: blockRoot[:]})
+	if err != nil {
+		return err
+	}
+	if !shouldSave(mode, s.snapshotInterval, st) {
+		return nil
+	}
+
+	// we need the update to savedRoots to be in the same critical section as db.SaveState
+	// because in Preserve we need the state bucket to be consistent with the list of saved roots
+	// so that we can safely confirm the state is present and remove it from the root cleanup list.
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	s.savedRoots = append(s.savedRoots, blockRoot)
+	log.WithFields(logrus.Fields{
+		"slot":               st.Slot(),
+		"totalStatesWritten": len(s.savedRoots),
+	}).Info("Saving hot state to DB")
+	return s.db.SaveState(ctx, st, blockRoot)
+}
+
+// Preserve ensures that the given state is permanently saved in the db. If the state already exists
+// and the state saver is in snapshot mode, the block root will be removed from the list of roots to
+// clean up when exiting snapshot mode to ensure it won't be deleted in the cleanup procedure.
+func (s *hotStateSaver) Preserve(ctx context.Context, root [32]byte, st state.BeaconState) error {
+	s.lock.Lock()
+	defer s.lock.Unlock()
+	exists := s.db.HasState(ctx, root)
+	if !exists {
+		if err := s.db.SaveState(ctx, st, root); err != nil {
+			return err
+		}
+	}
+	// the state exists, and we aren't in snapshot mode, so we shouldn't have to do anything
+	if s.m != PersistenceModeSnapshot {
+		return nil
+	}
+
+	// slice the preserved root out of the list of states to delete once the node exists snapshot mode.
+	for i := 0; i < len(s.savedRoots); i++ {
+		if s.savedRoots[i] == root {
+			s.savedRoots = append(s.savedRoots[:i], s.savedRoots[i+1:]...)
+			return nil
+		}
+	}
+	return nil
+}

beacon-chain/state/stategen/hot_state_cache_test.go

@@ -1,13 +1,21 @@
 package stategen
 
 import (
+	"context"
 	"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/state"
 	v1 "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/v1"
+	"github.com/prysmaticlabs/prysm/v3/config/params"
+	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
 	ethpb "github.com/prysmaticlabs/prysm/v3/proto/prysm/v1alpha1"
 	"github.com/prysmaticlabs/prysm/v3/testing/assert"
 	"github.com/prysmaticlabs/prysm/v3/testing/require"
+	"github.com/prysmaticlabs/prysm/v3/testing/util"
+	"github.com/prysmaticlabs/prysm/v3/time/slots"
 )
 
 func TestHotStateCache_RoundTrip(t *testing.T) {
@@ -32,3 +40,120 @@ func TestHotStateCache_RoundTrip(t *testing.T) {
 	c.delete(root)
 	assert.Equal(t, false, c.has(root), "Cache not supposed to have the object")
 }
+
+func TestHotStateSaving_Enabled(t *testing.T) {
+	h := &hotStateSaver{
+		db: testDB.SetupDB(t),
+	}
+	h.enableSnapshots()
+	require.Equal(t, PersistenceModeSnapshot, h.mode())
+}
+
+func TestHotStateSaving_AlreadyEnabled(t *testing.T) {
+	h := &hotStateSaver{
+		db: testDB.SetupDB(t),
+		m:  PersistenceModeSnapshot,
+	}
+	h.enableSnapshots()
+	require.Equal(t, PersistenceModeSnapshot, h.mode())
+}
+
+func TestHotStateSaving_Disabled(t *testing.T) {
+	ctx := context.Background()
+	h := &hotStateSaver{
+		db: testDB.SetupDB(t),
+		m:  PersistenceModeSnapshot,
+	}
+	b := util.NewBeaconBlock()
+	util.SaveBlock(t, ctx, h.db, b)
+	r, err := b.Block.HashTreeRoot()
+	require.NoError(t, err)
+	h.savedRoots = [][32]byte{r}
+	require.NoError(t, h.disableSnapshots(ctx))
+	require.Equal(t, PersistenceModeMemory, h.mode())
+	require.Equal(t, 0, len(h.savedRoots))
+}
+
+func TestHotStateSaving_AlreadyDisabled(t *testing.T) {
+	h := &hotStateSaver{}
+	require.NoError(t, h.disableSnapshots(context.Background()))
+	require.Equal(t, PersistenceModeMemory, h.mode())
+}
+
+func TestHotStateSaving_DisabledByDefault(t *testing.T) {
+	h := &hotStateSaver{
+		db: testDB.SetupDB(t),
+		fc: doublylinkedtree.New(),
+	}
+	fin := h.fc.FinalizedCheckpoint()
+	finslot, err := slots.EpochStart(fin.Epoch)
+	require.NoError(t, err)
+	h.cs = &mockCurrentSlotter{Slot: finslot}
+	require.Equal(t, PersistenceModeMemory, h.mode())
+	mode, err := h.refreshMode(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, PersistenceModeMemory, mode)
+}
+
+func TestHotStateSaving_Enabling(t *testing.T) {
+	h := &hotStateSaver{
+		db: testDB.SetupDB(t),
+		fc: doublylinkedtree.New(),
+		cs: &mockCurrentSlotter{Slot: types.Slot(uint64(params.BeaconConfig().SlotsPerEpoch) * uint64(hotStateSaveThreshold))},
+	}
+	mode, err := h.refreshMode(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, PersistenceModeSnapshot, mode)
+}
+
+func TestHotStateSaving_DisableAfterFinality(t *testing.T) {
+	h := &hotStateSaver{
+		db: testDB.SetupDB(t),
+		fc: doublylinkedtree.New(),
+		cs: &mockCurrentSlotter{Slot: types.Slot(uint64(params.BeaconConfig().SlotsPerEpoch) * uint64(hotStateSaveThreshold))},
+	}
+	mode, err := h.refreshMode(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, PersistenceModeSnapshot, mode)
+
+	// set current slot equal to finalized and ask for an update, should be disabled
+	fin := h.fc.FinalizedCheckpoint()
+	finslot, err := slots.EpochStart(fin.Epoch)
+	require.NoError(t, err)
+	h.cs = &mockCurrentSlotter{Slot: finslot}
+	mode, err = h.refreshMode(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, PersistenceModeMemory, mode)
+}
+
+type mockFinalizedCheckpointer struct {
+	c *forkchoicetypes.Checkpoint
+}
+
+func (m *mockFinalizedCheckpointer) FinalizedCheckpoint() *forkchoicetypes.Checkpoint {
+	return m.c
+}
+
+var _ FinalizedCheckpointer = &mockFinalizedCheckpointer{}
+
+func TestUpdateHotStateMode_CurrentSlotBeforeFinalized(t *testing.T) {
+	h := &hotStateSaver{
+		db: testDB.SetupDB(t),
+		fc: &mockFinalizedCheckpointer{c: &forkchoicetypes.Checkpoint{Epoch: 1}},
+		cs: &mockCurrentSlotter{Slot: 0},
+	}
+	mode, err := h.refreshMode(context.Background())
+	require.ErrorIs(t, err, errCurrentEpochBehindFinalized)
+	require.Equal(t, PersistenceModeMemory, mode)
+}
+
+func TestUpdateHotStateMode_NilFinalized(t *testing.T) {
+	h := &hotStateSaver{
+		db: testDB.SetupDB(t),
+		fc: &mockFinalizedCheckpointer{c: nil},
+		cs: &mockCurrentSlotter{Slot: 0},
+	}
+	mode, err := h.refreshMode(context.Background())
+	require.ErrorIs(t, err, errForkchoiceFinalizedNil)
+	require.Equal(t, PersistenceModeMemory, mode)
+}

beacon-chain/state/stategen/migrate.go

@@ -5,6 +5,7 @@ import (
 	"encoding/hex"
 	"fmt"
 
+	"github.com/pkg/errors"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
 	"github.com/prysmaticlabs/prysm/v3/encoding/bytesutil"
 	"github.com/sirupsen/logrus"
@@ -39,62 +40,47 @@ func (s *State) MigrateToCold(ctx context.Context, fRoot [32]byte) error {
 		}
 
 		if slot%s.slotsPerArchivedPoint == 0 && slot != 0 {
+			var aRoot [32]byte
+			var aState state.BeaconState
+
+			// cases we need to handle:
+			// 1. state exists in the epoch boundary state cache
+			// 2. state is in the database due to hot state saver in snapshot mode
+			// in this case we're looking up by slot, there's no root to look up, so we still have to replay blocks
+			// so state 2&3 are the same case
+			// 3. state snapshot is not in the database, rebuild it from stategen
+			// in all 3 cases we want to make sure the snapshot mode saver does not delete it.
 			cached, exists, err := s.epochBoundaryStateCache.getBySlot(slot)
 			if err != nil {
 				return fmt.Errorf("could not get epoch boundary state for slot %d", slot)
 			}
-
-			var aRoot [32]byte
-			var aState state.BeaconState
-
-			// When the epoch boundary state is not in cache due to skip slot scenario,
-			// we have to regenerate the state which will represent epoch boundary.
-			// By finding the highest available block below epoch boundary slot, we
-			// generate the state for that block root.
 			if exists {
+				// case 1 - state in epoch boundary state cache
 				aRoot = cached.root
 				aState = cached.state
 			} else {
-				_, roots, err := s.beaconDB.HighestRootsBelowSlot(ctx, slot)
+				// case 3 - state is not in db, we need to rebuild it from the most recent available state
+				aState, err = s.rb.ReplayerForSlot(slot).ReplayToSlot(ctx, slot)
 				if err != nil {
 					return err
 				}
-				// Given the block has been finalized, the db should not have more than one block in a given slot.
-				// We should error out when this happens.
-				if len(roots) != 1 {
-					return errUnknownBlock
+				// compute the block hash from the state
+				sr, err := aState.HashTreeRoot(ctx)
+				if err != nil {
+					return errors.Wrap(err, "error while computing hash_tree_root of state in MigrateToCold")
 				}
-				aRoot = roots[0]
-				// There's no need to generate the state if the state already exists in the DB.
-				// We can skip saving the state.
-				if !s.beaconDB.HasState(ctx, aRoot) {
-					aState, err = s.StateByRoot(ctx, aRoot)
-					if err != nil {
-						return err
-					}
+				header := aState.LatestBlockHeader()
+				header.StateRoot = sr[:]
+				aRoot, err = header.HashTreeRoot()
+				if err != nil {
+					return errors.Wrap(err, "error while computing block root using state data")
 				}
 			}
 
-			if s.beaconDB.HasState(ctx, aRoot) {
-				// If you are migrating a state and its already part of the hot state cache saved to the db,
-				// you can just remove it from the hot state cache as it becomes redundant.
-				s.saveHotStateDB.lock.Lock()
-				roots := s.saveHotStateDB.blockRootsOfSavedStates
-				for i := 0; i < len(roots); i++ {
-					if aRoot == roots[i] {
-						s.saveHotStateDB.blockRootsOfSavedStates = append(roots[:i], roots[i+1:]...)
-						// There shouldn't be duplicated roots in `blockRootsOfSavedStates`.
-						// Break here is ok.
-						break
-					}
-				}
-				s.saveHotStateDB.lock.Unlock()
-				continue
-			}
-
-			if err := s.beaconDB.SaveState(ctx, aState, aRoot); err != nil {
+			if err := s.saver.Preserve(ctx, aRoot, aState); err != nil {
 				return err
 			}
+
 			log.WithFields(
 				logrus.Fields{
 					"slot": aState.Slot(),

beacon-chain/state/stategen/migrate_test.go

@@ -17,7 +17,7 @@ import (
 func TestMigrateToCold_CanSaveFinalizedInfo(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	b := util.NewBeaconBlock()
 	b.Block.Slot = 1
@@ -36,7 +36,20 @@ func TestMigrateToCold_HappyPath(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	var zero, one, two, three, four, five types.Slot = 50, 51, 150, 151, 152, 200
+	specs := []mockHistorySpec{
+		{slot: zero},
+		{slot: one, savedState: true},
+		{slot: two},
+		{slot: three},
+		{slot: four},
+		{slot: five, canonicalBlock: true},
+	}
+
+	hist := newMockHistory(t, specs, five+1)
+	ch := NewCanonicalHistory(hist, hist, hist)
+	service := New(beaconDB, newTestSaver(beaconDB), WithReplayerBuilder(ch))
+
 	service.slotsPerArchivedPoint = 1
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	stateSlot := types.Slot(1)
@@ -66,7 +79,11 @@ func TestMigrateToCold_RegeneratePath(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	mockCanon := newMockCanonicalMap()
+	// picking 5 because the highest block in the below setup code is 4 and we're not testing
+	// slot bounds weirdness here.
+	ch := NewCanonicalHistory(beaconDB, mockCanon, &mockCurrentSlotter{Slot: 5})
+	service := New(beaconDB, newTestSaver(beaconDB), WithReplayerBuilder(ch))
 	service.slotsPerArchivedPoint = 1
 	beaconState, pks := util.DeterministicGenesisState(t, 32)
 	genesisStateRoot, err := beaconState.HashTreeRoot(ctx)
@@ -82,6 +99,7 @@ func TestMigrateToCold_RegeneratePath(t *testing.T) {
 	require.NoError(t, err)
 	r1, err := b1.Block.HashTreeRoot()
 	require.NoError(t, err)
+	mockCanon.AddCanonical(r1)
 	util.SaveBlock(t, ctx, service.beaconDB, b1)
 	require.NoError(t, service.beaconDB.SaveStateSummary(ctx, &ethpb.StateSummary{Slot: 1, Root: r1[:]}))
 
@@ -89,6 +107,7 @@ func TestMigrateToCold_RegeneratePath(t *testing.T) {
 	require.NoError(t, err)
 	r4, err := b4.Block.HashTreeRoot()
 	require.NoError(t, err)
+	mockCanon.AddCanonical(r4)
 	util.SaveBlock(t, ctx, service.beaconDB, b4)
 	require.NoError(t, service.beaconDB.SaveStateSummary(ctx, &ethpb.StateSummary{Slot: 4, Root: r4[:]}))
 	service.finalizedInfo = &finalizedInfo{
@@ -101,16 +120,18 @@ func TestMigrateToCold_RegeneratePath(t *testing.T) {
 
 	s1, err := service.beaconDB.State(ctx, r1)
 	require.NoError(t, err)
-	assert.Equal(t, s1.Slot(), types.Slot(1), "Did not save state")
-	gotRoot := service.beaconDB.ArchivedPointRoot(ctx, 1/service.slotsPerArchivedPoint)
-	assert.Equal(t, r1, gotRoot, "Did not save archived root")
-	lastIndex, err := service.beaconDB.LastArchivedSlot(ctx)
+	require.NotEqual(t, nil, s1)
+
+	require.Equal(t, false, service.beaconDB.HasState(ctx, r4))
+	s4, err := service.beaconDB.State(ctx, r4)
 	require.NoError(t, err)
-	assert.Equal(t, types.Slot(1), lastIndex, "Did not save last archived index")
+	require.Equal(t, nil, s4)
 
 	require.LogsContain(t, hook, "Saved state in DB")
 }
 
+// TODO!!! fix this migrate to cold test
+/*
 func TestMigrateToCold_StateExistsInDB(t *testing.T) {
 	hook := logTest.NewGlobal()
 	ctx := context.Background()
@@ -129,8 +150,9 @@ func TestMigrateToCold_StateExistsInDB(t *testing.T) {
 	require.NoError(t, service.epochBoundaryStateCache.put(fRoot, beaconState))
 	require.NoError(t, service.beaconDB.SaveState(ctx, beaconState, fRoot))
 
-	service.saveHotStateDB.blockRootsOfSavedStates = [][32]byte{{1}, {2}, {3}, {4}, fRoot}
+	service.saver.savedRoots = [][32]byte{{1}, {2}, {3}, {4}, fRoot}
 	require.NoError(t, service.MigrateToCold(ctx, fRoot))
-	assert.DeepEqual(t, [][32]byte{{1}, {2}, {3}, {4}}, service.saveHotStateDB.blockRootsOfSavedStates)
+	assert.DeepEqual(t, [][32]byte{{1}, {2}, {3}, {4}}, service.saver.savedRoots)
 	assert.LogsDoNotContain(t, hook, "Saved state in DB")
 }
+*/

beacon-chain/state/stategen/mock/BUILD.bazel

@@ -6,10 +6,12 @@ go_library(
     srcs = [
         "mock.go",
         "replayer.go",
+        "saver.go",
     ],
     importpath = "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock",
     visibility = ["//visibility:public"],
     deps = [
+        "//beacon-chain/db/iface:go_default_library",
         "//beacon-chain/state:go_default_library",
         "//beacon-chain/state/stategen:go_default_library",
         "//consensus-types/primitives:go_default_library",

beacon-chain/state/stategen/mock/mock.go

@@ -71,16 +71,6 @@ func (_ *MockStateManager) ForceCheckpoint(_ context.Context, _ []byte) error {
 	panic("implement me")
 }
 
-// EnableSaveHotStateToDB --
-func (_ *MockStateManager) EnableSaveHotStateToDB(_ context.Context) {
-	panic("implement me")
-}
-
-// DisableSaveHotStateToDB --
-func (_ *MockStateManager) DisableSaveHotStateToDB(_ context.Context) error {
-	panic("implement me")
-}
-
 // AddStateForRoot --
 func (m *MockStateManager) AddStateForRoot(state state.BeaconState, blockRoot [32]byte) {
 	m.StatesByRoot[blockRoot] = state

beacon-chain/state/stategen/mock/saver.go

@@ -0,0 +1,42 @@
+package mock
+
+import (
+	"context"
+
+	"github.com/prysmaticlabs/prysm/v3/beacon-chain/db/iface"
+	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
+	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+)
+
+// NewMockSaver creates a value that can be used as a stategen.Saver in tests
+func NewMockSaver(d iface.Database) *MockSaver {
+	return &MockSaver{
+		db: d,
+	}
+}
+
+type MockSaver struct {
+	db iface.HeadAccessDatabase
+}
+
+// Save just saves everything it is asked to save
+func (s *MockSaver) Save(ctx context.Context, root [32]byte, st state.BeaconState) error {
+	return s.db.SaveState(ctx, st, root)
+}
+
+// Preserve just checks to see if the state has been saved before saving it
+func (s *MockSaver) Preserve(ctx context.Context, root [32]byte, st state.BeaconState) error {
+	if !s.db.HasState(ctx, root) {
+		return s.Save(ctx, root, st)
+	}
+	return nil
+}
+
+var _ stategen.Saver = &MockSaver{}
+
+// NewMockStategen hides the complexity of preparing all the values needed to construct an
+// instance of stategen.State for tests that don't make assumptions about the internal behavior of stategen.
+func NewMockStategen(db iface.HeadAccessDatabase, opts ...stategen.StateGenOption) *stategen.State {
+	saver := &MockSaver{db: db}
+	return stategen.New(db, saver, opts...)
+}

beacon-chain/state/stategen/replay_test.go

@@ -41,7 +41,7 @@ func TestReplayBlocks_AllSkipSlots(t *testing.T) {
 	require.NoError(t, beaconState.SetCurrentJustifiedCheckpoint(cp))
 	require.NoError(t, beaconState.AppendCurrentEpochAttestations(&ethpb.PendingAttestation{}))
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	targetSlot := params.BeaconConfig().SlotsPerEpoch - 1
 	newState, err := service.replayBlocks(context.Background(), beaconState, []interfaces.SignedBeaconBlock{}, targetSlot)
 	require.NoError(t, err)
@@ -70,7 +70,7 @@ func TestReplayBlocks_SameSlot(t *testing.T) {
 	require.NoError(t, beaconState.SetCurrentJustifiedCheckpoint(cp))
 	require.NoError(t, beaconState.AppendCurrentEpochAttestations(&ethpb.PendingAttestation{}))
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	targetSlot := beaconState.Slot()
 	newState, err := service.replayBlocks(context.Background(), beaconState, []interfaces.SignedBeaconBlock{}, targetSlot)
 	require.NoError(t, err)
@@ -100,7 +100,7 @@ func TestReplayBlocks_LowerSlotBlock(t *testing.T) {
 	require.NoError(t, beaconState.SetCurrentJustifiedCheckpoint(cp))
 	require.NoError(t, beaconState.AppendCurrentEpochAttestations(&ethpb.PendingAttestation{}))
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	targetSlot := beaconState.Slot()
 	b := util.NewBeaconBlock()
 	b.Block.Slot = beaconState.Slot() - 1
@@ -130,7 +130,8 @@ func TestReplayBlocks_ThroughForkBoundary(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	service := New(testDB.SetupDB(t))
+	bdb := testDB.SetupDB(t)
+	service := New(bdb, newTestSaver(bdb))
 	targetSlot := params.BeaconConfig().SlotsPerEpoch
 	newState, err := service.replayBlocks(context.Background(), beaconState, []interfaces.SignedBeaconBlock{}, targetSlot)
 	require.NoError(t, err)
@@ -160,7 +161,8 @@ func TestReplayBlocks_ThroughBellatrixForkBoundary(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	service := New(testDB.SetupDB(t))
+	bdb := testDB.SetupDB(t)
+	service := New(bdb, newTestSaver(bdb))
 	targetSlot := params.BeaconConfig().SlotsPerEpoch * 2
 	newState, err := service.replayBlocks(context.Background(), beaconState, []interfaces.SignedBeaconBlock{}, targetSlot)
 	require.NoError(t, err)

beacon-chain/state/stategen/service.go

@@ -17,14 +17,14 @@ import (
 	"go.opencensus.io/trace"
 )
 
-var defaultHotStateDBInterval types.Slot = 128
+// DefaultSnapshotInterval dictates how frequently the state should be saved when the Saver
+// goes into snapshot mode. Default value is once every 128 slots.
+var DefaultSnapshotInterval types.Slot = 128
 
 // StateManager represents a management object that handles the internal
 // logic of maintaining both hot and cold states in DB.
 type StateManager interface {
 	Resume(ctx context.Context, fState state.BeaconState) (state.BeaconState, error)
-	DisableSaveHotStateToDB(ctx context.Context) error
-	EnableSaveHotStateToDB(_ context.Context)
 	HasState(ctx context.Context, blockRoot [32]byte) (bool, error)
 	DeleteStateFromCaches(ctx context.Context, blockRoot [32]byte) error
 	ForceCheckpoint(ctx context.Context, root []byte) error
@@ -43,18 +43,9 @@ type State struct {
 	hotStateCache           *hotStateCache
 	finalizedInfo           *finalizedInfo
 	epochBoundaryStateCache *epochBoundaryState
-	saveHotStateDB          *saveHotStateDbConfig
+	saver                   Saver
 	backfillStatus          *backfill.Status
-}
-
-// This tracks the config in the event of long non-finality,
-// how often does the node save hot states to db? what are
-// the saved hot states in db?... etc
-type saveHotStateDbConfig struct {
-	enabled                 bool
-	lock                    sync.Mutex
-	duration                types.Slot
-	blockRootsOfSavedStates [][32]byte
+	rb                      ReplayerBuilder
 }
 
 // This tracks the finalized point. It's also the point where slot and the block root of
@@ -75,17 +66,29 @@ func WithBackfillStatus(bfs *backfill.Status) StateGenOption {
 	}
 }
 
+func WithReplayerBuilder(rb ReplayerBuilder) StateGenOption {
+	return func(sg *State) {
+		sg.rb = rb
+	}
+}
+
+// Saver is an interface describing the set of methods that stategen needs
+// in order to delegate the decision to save hot states to the database. The Strategy is
+// responsible for cleaning up the database
+type Saver interface {
+	Save(ctx context.Context, blockRoot [32]byte, st state.BeaconState) error
+	Preserve(ctx context.Context, root [32]byte, st state.BeaconState) error
+}
+
 // New returns a new state management object.
-func New(beaconDB db.NoHeadAccessDatabase, opts ...StateGenOption) *State {
+func New(beaconDB db.NoHeadAccessDatabase, saver Saver, opts ...StateGenOption) *State {
 	s := &State{
 		beaconDB:                beaconDB,
 		hotStateCache:           newHotStateCache(),
 		finalizedInfo:           &finalizedInfo{slot: 0, root: params.BeaconConfig().ZeroHash},
 		slotsPerArchivedPoint:   params.BeaconConfig().SlotsPerArchivedPoint,
 		epochBoundaryStateCache: newBoundaryStateCache(),
-		saveHotStateDB: &saveHotStateDbConfig{
-			duration: defaultHotStateDBInterval,
-		},
+		saver:                   saver,
 	}
 	for _, o := range opts {
 		o(s)

beacon-chain/state/stategen/service_test.go

@@ -16,7 +16,7 @@ func TestResume(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	b := util.NewBeaconBlock()
 	util.SaveBlock(t, ctx, service.beaconDB, b)
 	root, err := b.Block.HashTreeRoot()

beacon-chain/state/stategen/setter.go

@@ -2,17 +2,22 @@ package stategen
 
 import (
 	"context"
-	"math"
 
+	"github.com/pkg/errors"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
+	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/v3/encoding/bytesutil"
-	ethpb "github.com/prysmaticlabs/prysm/v3/proto/prysm/v1alpha1"
 	"github.com/prysmaticlabs/prysm/v3/time/slots"
-	"github.com/sirupsen/logrus"
 	"go.opencensus.io/trace"
 )
 
+var (
+	errCurrentEpochBehindFinalized = errors.New("finalized epoch must always be before the current epoch")
+	errForkchoiceFinalizedNil      = errors.New("forkchoice store finalized checkpoint is nil")
+)
+var hotStateSaveThreshold = types.Epoch(100)
+
 // SaveState saves the state in the cache and/or DB.
 func (s *State) SaveState(ctx context.Context, blockRoot [32]byte, st state.BeaconState) error {
 	ctx, span := trace.StartSpan(ctx, "stateGen.SaveState")
@@ -51,29 +56,16 @@ func (s *State) saveStateByRoot(ctx context.Context, blockRoot [32]byte, st stat
 	ctx, span := trace.StartSpan(ctx, "stateGen.saveStateByRoot")
 	defer span.End()
 
-	// Duration can't be 0 to prevent panic for division.
-	duration := uint64(math.Max(float64(s.saveHotStateDB.duration), 1))
-
-	s.saveHotStateDB.lock.Lock()
-	if s.saveHotStateDB.enabled && st.Slot().Mod(duration) == 0 {
-		if err := s.beaconDB.SaveState(ctx, st, blockRoot); err != nil {
-			s.saveHotStateDB.lock.Unlock()
-			return err
-		}
-		s.saveHotStateDB.blockRootsOfSavedStates = append(s.saveHotStateDB.blockRootsOfSavedStates, blockRoot)
-
-		log.WithFields(logrus.Fields{
-			"slot":                   st.Slot(),
-			"totalHotStateSavedInDB": len(s.saveHotStateDB.blockRootsOfSavedStates),
-		}).Info("Saving hot state to DB")
-	}
-	s.saveHotStateDB.lock.Unlock()
-
-	// If the hot state is already in cache, one can be sure the state was processed and in the DB.
+	// this is the only method that puts states in the cache, so if the state is in the cache
+	// the state has already been processed by this method and there's no need to run again.
 	if s.hotStateCache.has(blockRoot) {
 		return nil
 	}
 
+	if err := s.saver.Save(ctx, blockRoot, st); err != nil {
+		return err
+	}
+
 	// Only on an epoch boundary slot, save epoch boundary state in epoch boundary root state cache.
 	if slots.IsEpochStart(st.Slot()) {
 		if err := s.epochBoundaryStateCache.put(blockRoot, st); err != nil {
@@ -81,57 +73,8 @@ func (s *State) saveStateByRoot(ctx context.Context, blockRoot [32]byte, st stat
 		}
 	}
 
-	// On an intermediate slot, save state summary.
-	if err := s.beaconDB.SaveStateSummary(ctx, &ethpb.StateSummary{
-		Slot: st.Slot(),
-		Root: blockRoot[:],
-	}); err != nil {
-		return err
-	}
-
 	// Store the copied state in the hot state cache.
 	s.hotStateCache.put(blockRoot, st)
 
 	return nil
 }
-
-// EnableSaveHotStateToDB enters the mode that saves hot beacon state to the DB.
-// This usually gets triggered when there's long duration since finality.
-func (s *State) EnableSaveHotStateToDB(_ context.Context) {
-	s.saveHotStateDB.lock.Lock()
-	defer s.saveHotStateDB.lock.Unlock()
-	if s.saveHotStateDB.enabled {
-		return
-	}
-
-	s.saveHotStateDB.enabled = true
-
-	log.WithFields(logrus.Fields{
-		"enabled":       s.saveHotStateDB.enabled,
-		"slotsInterval": s.saveHotStateDB.duration,
-	}).Warn("Entering mode to save hot states in DB")
-}
-
-// DisableSaveHotStateToDB exits the mode that saves beacon state to DB for the hot states.
-// This usually gets triggered once there's finality after long duration since finality.
-func (s *State) DisableSaveHotStateToDB(ctx context.Context) error {
-	s.saveHotStateDB.lock.Lock()
-	defer s.saveHotStateDB.lock.Unlock()
-	if !s.saveHotStateDB.enabled {
-		return nil
-	}
-
-	log.WithFields(logrus.Fields{
-		"enabled":          s.saveHotStateDB.enabled,
-		"deletedHotStates": len(s.saveHotStateDB.blockRootsOfSavedStates),
-	}).Warn("Exiting mode to save hot states in DB")
-
-	// Delete previous saved states in DB as we are turning this mode off.
-	s.saveHotStateDB.enabled = false
-	if err := s.beaconDB.DeleteStates(ctx, s.saveHotStateDB.blockRootsOfSavedStates); err != nil {
-		return err
-	}
-	s.saveHotStateDB.blockRootsOfSavedStates = nil
-
-	return nil
-}

beacon-chain/state/stategen/setter_test.go

@@ -5,18 +5,29 @@ import (
 	"testing"
 
 	testDB "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
+	forkchoicetypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/types"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
+	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/v3/testing/assert"
 	"github.com/prysmaticlabs/prysm/v3/testing/require"
 	"github.com/prysmaticlabs/prysm/v3/testing/util"
+	"github.com/prysmaticlabs/prysm/v3/time/slots"
 	logTest "github.com/sirupsen/logrus/hooks/test"
 )
 
 func TestSaveState_HotStateCanBeSaved(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
+	genroot := [32]byte{}
+	cp := &forkchoicetypes.Checkpoint{Epoch: 0, Root: genroot}
+	h := newTestSaver(beaconDB, withFinalizedCheckpointer(&mockFinalizedCheckpointer{c: cp}))
+	stateSlot := firstSaveableSlotAfter(t, h)
+	h.cs = &mockCurrentSlotter{Slot: stateSlot + h.snapshotInterval}
+	mode, err := h.refreshMode(ctx)
+	require.NoError(t, err)
+	require.Equal(t, PersistenceModeSnapshot, mode)
+	service := New(beaconDB, h)
 
-	service := New(beaconDB)
 	service.slotsPerArchivedPoint = 1
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	// This goes to hot section, verify it can save on epoch boundary.
@@ -37,7 +48,7 @@ func TestSaveState_HotStateCached(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 	service.slotsPerArchivedPoint = 1
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	require.NoError(t, beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch))
@@ -57,7 +68,7 @@ func TestState_ForceCheckpoint_SavesStateToDatabase(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
 
-	svc := New(beaconDB)
+	svc := New(beaconDB, newTestSaver(beaconDB))
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	require.NoError(t, beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch))
 
@@ -76,7 +87,7 @@ func TestSaveState_Alreadyhas(t *testing.T) {
 	hook := logTest.NewGlobal()
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	service := New(beaconDB, newTestSaver(beaconDB))
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	require.NoError(t, beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch))
@@ -95,7 +106,15 @@ func TestSaveState_Alreadyhas(t *testing.T) {
 func TestSaveState_CanSaveOnEpochBoundary(t *testing.T) {
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	genroot := [32]byte{}
+	cp := &forkchoicetypes.Checkpoint{Epoch: 0, Root: genroot}
+	h := newTestSaver(beaconDB, withFinalizedCheckpointer(&mockFinalizedCheckpointer{c: cp}))
+	stateSlot := firstSaveableSlotAfter(t, h)
+	h.cs = &mockCurrentSlotter{Slot: stateSlot + h.snapshotInterval}
+	mode, err := h.refreshMode(ctx)
+	require.NoError(t, err)
+	require.Equal(t, PersistenceModeSnapshot, mode)
+	service := New(beaconDB, h)
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	require.NoError(t, beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch))
@@ -116,7 +135,15 @@ func TestSaveState_NoSaveNotEpochBoundary(t *testing.T) {
 	hook := logTest.NewGlobal()
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
+	genroot := [32]byte{}
+	cp := &forkchoicetypes.Checkpoint{Epoch: 0, Root: genroot}
+	h := newTestSaver(beaconDB, withFinalizedCheckpointer(&mockFinalizedCheckpointer{c: cp}))
+	stateSlot := firstSaveableSlotAfter(t, h)
+	h.cs = &mockCurrentSlotter{Slot: stateSlot + h.snapshotInterval}
+	mode, err := h.refreshMode(ctx)
+	require.NoError(t, err)
+	require.Equal(t, PersistenceModeSnapshot, mode)
+	service := New(beaconDB, h)
 
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
 	require.NoError(t, beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch-1))
@@ -134,70 +161,44 @@ func TestSaveState_NoSaveNotEpochBoundary(t *testing.T) {
 	require.LogsDoNotContain(t, hook, "Saved full state on epoch boundary")
 	// Should have not been saved in DB.
 	require.Equal(t, false, beaconDB.HasState(ctx, r))
+
+	_, ok, err := service.epochBoundaryStateCache.getByBlockRoot(r)
+	require.NoError(t, err)
+	require.Equal(t, false, ok, "saved to epoch boundary cache in error")
+}
+
+func firstSaveableSlotAfter(t *testing.T, h *hotStateSaver) types.Slot {
+	min, err := slots.EpochStart(hotStateSaveThreshold)
+	require.NoError(t, err)
+	f := h.fc.FinalizedCheckpoint()
+	require.NotNil(t, f)
+	fslot, err := slots.EpochStart(f.Epoch)
+	require.NoError(t, err)
+	min += fslot
+	diff := h.snapshotInterval - (min % h.snapshotInterval)
+	aligned := min + diff
+	require.Equal(t, types.Slot(0), aligned%h.snapshotInterval)
+	return min + diff
 }
 
 func TestSaveState_CanSaveHotStateToDB(t *testing.T) {
-	hook := logTest.NewGlobal()
 	ctx := context.Background()
 	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
-	service.EnableSaveHotStateToDB(ctx)
+	genroot := [32]byte{}
+	cp := &forkchoicetypes.Checkpoint{Epoch: 0, Root: genroot}
+	h := newTestSaver(beaconDB, withFinalizedCheckpointer(&mockFinalizedCheckpointer{c: cp}))
+	stateSlot := firstSaveableSlotAfter(t, h)
+	h.cs = &mockCurrentSlotter{Slot: stateSlot + h.snapshotInterval}
+	mode, err := h.refreshMode(ctx)
+	require.NoError(t, err)
+	require.Equal(t, PersistenceModeSnapshot, mode)
+	service := New(beaconDB, h)
 	beaconState, _ := util.DeterministicGenesisState(t, 32)
-	require.NoError(t, beaconState.SetSlot(defaultHotStateDBInterval))
+	require.NoError(t, beaconState.SetSlot(stateSlot))
 
 	r := [32]byte{'A'}
 	require.NoError(t, service.saveStateByRoot(ctx, r, beaconState))
 
-	require.LogsContain(t, hook, "Saving hot state to DB")
 	// Should have saved in DB.
-	require.Equal(t, true, beaconDB.HasState(ctx, r))
-}
-
-func TestEnableSaveHotStateToDB_Enabled(t *testing.T) {
-	hook := logTest.NewGlobal()
-	ctx := context.Background()
-	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
-
-	service.EnableSaveHotStateToDB(ctx)
-	require.LogsContain(t, hook, "Entering mode to save hot states in DB")
-	require.Equal(t, true, service.saveHotStateDB.enabled)
-}
-
-func TestEnableSaveHotStateToDB_AlreadyEnabled(t *testing.T) {
-	hook := logTest.NewGlobal()
-	ctx := context.Background()
-	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
-	service.saveHotStateDB.enabled = true
-	service.EnableSaveHotStateToDB(ctx)
-	require.LogsDoNotContain(t, hook, "Entering mode to save hot states in DB")
-	require.Equal(t, true, service.saveHotStateDB.enabled)
-}
-
-func TestEnableSaveHotStateToDB_Disabled(t *testing.T) {
-	hook := logTest.NewGlobal()
-	ctx := context.Background()
-	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
-	service.saveHotStateDB.enabled = true
-	b := util.NewBeaconBlock()
-	util.SaveBlock(t, ctx, beaconDB, b)
-	r, err := b.Block.HashTreeRoot()
-	require.NoError(t, err)
-	service.saveHotStateDB.blockRootsOfSavedStates = [][32]byte{r}
-	require.NoError(t, service.DisableSaveHotStateToDB(ctx))
-	require.LogsContain(t, hook, "Exiting mode to save hot states in DB")
-	require.Equal(t, false, service.saveHotStateDB.enabled)
-	require.Equal(t, 0, len(service.saveHotStateDB.blockRootsOfSavedStates))
-}
-
-func TestEnableSaveHotStateToDB_AlreadyDisabled(t *testing.T) {
-	hook := logTest.NewGlobal()
-	ctx := context.Background()
-	beaconDB := testDB.SetupDB(t)
-	service := New(beaconDB)
-	require.NoError(t, service.DisableSaveHotStateToDB(ctx))
-	require.LogsDoNotContain(t, hook, "Exiting mode to save hot states in DB")
-	require.Equal(t, false, service.saveHotStateDB.enabled)
+	require.Equal(t, true, h.db.HasState(ctx, r))
 }

beacon-chain/state/stategen/shared_test.go

@@ -0,0 +1,27 @@
+package stategen
+
+import (
+	"github.com/prysmaticlabs/prysm/v3/beacon-chain/db/iface"
+)
+
+type testSaverOpt func(h *hotStateSaver)
+
+func withFinalizedCheckpointer(fc *mockFinalizedCheckpointer) testSaverOpt {
+	return func(h *hotStateSaver) {
+		h.fc = fc
+	}
+}
+
+func newTestSaver(db iface.HeadAccessDatabase, opts ...testSaverOpt) *hotStateSaver {
+	h := &hotStateSaver{
+		snapshotInterval: DefaultSnapshotInterval,
+		db:               db,
+		fc:               &mockFinalizedCheckpointer{},
+		cs:               &mockCurrentSlotter{},
+	}
+	for _, o := range opts {
+		o(h)
+	}
+
+	return h
+}

beacon-chain/sync/BUILD.bazel

@@ -190,7 +190,7 @@ go_test(
         "//beacon-chain/p2p/testing:go_default_library",
         "//beacon-chain/p2p/types:go_default_library",
         "//beacon-chain/state:go_default_library",
-        "//beacon-chain/state/stategen:go_default_library",
+        "//beacon-chain/state/stategen/mock:go_default_library",
         "//beacon-chain/state/v1:go_default_library",
         "//beacon-chain/sync/initial-sync/testing:go_default_library",
         "//cache/lru:go_default_library",

beacon-chain/sync/pending_blocks_queue_test.go

@@ -19,7 +19,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/peers"
 	p2ptest "github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/testing"
 	p2ptypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/types"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	"github.com/prysmaticlabs/prysm/v3/consensus-types/blocks"
 	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
@@ -47,7 +47,7 @@ func TestRegularSyncBeaconBlockSubscriber_ProcessPendingBlocks1(t *testing.T) {
 					Epoch: 0,
 				},
 			},
-			stateGen: stategen.New(db),
+			stateGen: sgmock.NewMockStategen(db),
 		},
 		slotToPendingBlocks: gcache.New(time.Second, 2*time.Second),
 		seenPendingBlocks:   make(map[[32]byte]bool),
@@ -119,7 +119,7 @@ func TestRegularSyncBeaconBlockSubscriber_OptimisticStatus(t *testing.T) {
 					Epoch: 0,
 				},
 			},
-			stateGen: stategen.New(db),
+			stateGen: sgmock.NewMockStategen(db),
 		},
 		slotToPendingBlocks: gcache.New(time.Second, 2*time.Second),
 		seenPendingBlocks:   make(map[[32]byte]bool),
@@ -191,7 +191,7 @@ func TestRegularSyncBeaconBlockSubscriber_ExecutionEngineTimesOut(t *testing.T)
 				},
 				ReceiveBlockMockErr: execution.ErrHTTPTimeout,
 			},
-			stateGen: stategen.New(db),
+			stateGen: sgmock.NewMockStategen(db),
 		},
 		slotToPendingBlocks: gcache.New(time.Second, 2*time.Second),
 		seenPendingBlocks:   make(map[[32]byte]bool),
@@ -318,7 +318,7 @@ func TestRegularSyncBeaconBlockSubscriber_DoNotReprocessBlock(t *testing.T) {
 					Epoch: 0,
 				},
 			},
-			stateGen: stategen.New(db),
+			stateGen: sgmock.NewMockStategen(db),
 		},
 		slotToPendingBlocks: gcache.New(time.Second, 2*time.Second),
 		seenPendingBlocks:   make(map[[32]byte]bool),
@@ -383,7 +383,7 @@ func TestRegularSyncBeaconBlockSubscriber_ProcessPendingBlocks_2Chains(t *testin
 					Root:  make([]byte, 32),
 				},
 			},
-			stateGen: stategen.New(db),
+			stateGen: sgmock.NewMockStategen(db),
 		},
 		slotToPendingBlocks: gcache.New(time.Second, 2*time.Second),
 		seenPendingBlocks:   make(map[[32]byte]bool),
@@ -703,7 +703,7 @@ func TestService_ProcessPendingBlockOnCorrectSlot(t *testing.T) {
 			p2p:      p1,
 			beaconDB: db,
 			chain:    &mockChain,
-			stateGen: stategen.New(db),
+			stateGen: sgmock.NewMockStategen(db),
 		},
 		slotToPendingBlocks: gcache.New(time.Second, 2*time.Second),
 		seenPendingBlocks:   make(map[[32]byte]bool),
@@ -781,7 +781,7 @@ func TestService_ProcessBadPendingBlocks(t *testing.T) {
 			p2p:      p1,
 			beaconDB: db,
 			chain:    &mockChain,
-			stateGen: stategen.New(db),
+			stateGen: sgmock.NewMockStategen(db),
 		},
 		slotToPendingBlocks: gcache.New(time.Second, 2*time.Second),
 		seenPendingBlocks:   make(map[[32]byte]bool),

beacon-chain/sync/sync_fuzz_test.go

@@ -19,7 +19,7 @@ import (
 	dbtest "github.com/prysmaticlabs/prysm/v3/beacon-chain/db/testing"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p"
 	p2ptest "github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/testing"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	mockSync "github.com/prysmaticlabs/prysm/v3/beacon-chain/sync/initial-sync/testing"
 	lruwrpr "github.com/prysmaticlabs/prysm/v3/cache/lru"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
@@ -51,7 +51,7 @@ func FuzzValidateBeaconBlockPubSub_Phase0(f *testing.F) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(f, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Unix(time.Now().Unix()-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		State: beaconState,
 		FinalizedCheckPoint: &ethpb.Checkpoint{
@@ -132,7 +132,7 @@ func FuzzValidateBeaconBlockPubSub_Altair(f *testing.F) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(f, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Unix(time.Now().Unix()-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		State: beaconState,
 		FinalizedCheckPoint: &ethpb.Checkpoint{
@@ -213,7 +213,7 @@ func FuzzValidateBeaconBlockPubSub_Bellatrix(f *testing.F) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(f, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Unix(time.Now().Unix()-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		State: beaconState,
 		FinalizedCheckPoint: &ethpb.Checkpoint{

beacon-chain/sync/validate_beacon_blocks_test.go

@@ -22,7 +22,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/attestations"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p"
 	p2ptest "github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/testing"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	mockSync "github.com/prysmaticlabs/prysm/v3/beacon-chain/sync/initial-sync/testing"
 	lruwrpr "github.com/prysmaticlabs/prysm/v3/cache/lru"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
@@ -64,7 +64,7 @@ func TestValidateBeaconBlockPubSub_InvalidSignature(t *testing.T) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[badPrivKeyIdx])
 	require.NoError(t, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Unix(time.Now().Unix()-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		FinalizedCheckPoint: &ethpb.Checkpoint{
 			Epoch: 0,
@@ -167,7 +167,7 @@ func TestValidateBeaconBlockPubSub_CanRecoverStateSummary(t *testing.T) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(t, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Unix(time.Now().Unix()-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		State: beaconState,
 		FinalizedCheckPoint: &ethpb.Checkpoint{
@@ -231,7 +231,7 @@ func TestValidateBeaconBlockPubSub_IsInCache(t *testing.T) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(t, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Unix(time.Now().Unix()-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		State: beaconState,
 		FinalizedCheckPoint: &ethpb.Checkpoint{
@@ -297,7 +297,7 @@ func TestValidateBeaconBlockPubSub_ValidProposerSignature(t *testing.T) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(t, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Unix(time.Now().Unix()-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		State: beaconState,
 		FinalizedCheckPoint: &ethpb.Checkpoint{
@@ -365,7 +365,7 @@ func TestValidateBeaconBlockPubSub_WithLookahead(t *testing.T) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(t, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	offset := int64(blkSlot.Mul(params.BeaconConfig().SecondsPerSlot))
 	chainService := &mock.ChainService{Genesis: time.Unix(time.Now().Unix()-offset, 0),
 		DB:    db,
@@ -433,7 +433,7 @@ func TestValidateBeaconBlockPubSub_AdvanceEpochsForState(t *testing.T) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(t, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	offset := int64(blkSlot.Mul(params.BeaconConfig().SecondsPerSlot))
 	chainService := &mock.ChainService{Genesis: time.Unix(time.Now().Unix()-offset, 0),
 		DB:    db,
@@ -540,7 +540,7 @@ func TestValidateBeaconBlockPubSub_IgnoreAndQueueBlocksFromNearFuture(t *testing
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(t, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Now(),
 		FinalizedCheckPoint: &ethpb.Checkpoint{
 			Epoch: 0,
@@ -829,7 +829,7 @@ func TestValidateBeaconBlockPubSub_ParentNotFinalizedDescendant(t *testing.T) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(t, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Unix(time.Now().Unix()-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		State: beaconState,
 		FinalizedCheckPoint: &ethpb.Checkpoint{
@@ -898,7 +898,7 @@ func TestValidateBeaconBlockPubSub_InvalidParentBlock(t *testing.T) {
 	currBlockRoot, err := msg.Block.HashTreeRoot()
 	require.NoError(t, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Unix(time.Now().Unix()-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		State: beaconState,
 		FinalizedCheckPoint: &ethpb.Checkpoint{
@@ -1006,7 +1006,7 @@ func TestValidateBeaconBlockPubSub_RejectBlocksFromBadParent(t *testing.T) {
 
 	genesisTime := time.Now()
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{
 		Genesis: time.Unix(genesisTime.Unix()-int64(slotsSinceGenesis.Mul(perSlot)), 0),
 		FinalizedCheckPoint: &ethpb.Checkpoint{
@@ -1108,7 +1108,7 @@ func TestValidateBeaconBlockPubSub_ValidExecutionPayload(t *testing.T) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(t, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Unix(presentTime-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		DB: db,
 		FinalizedCheckPoint: &ethpb.Checkpoint{
@@ -1180,7 +1180,7 @@ func TestValidateBeaconBlockPubSub_InvalidPayloadTimestamp(t *testing.T) {
 	msg.Signature, err = signing.ComputeDomainAndSign(beaconState, 0, msg.Block, params.BeaconConfig().DomainBeaconProposer, privKeys[proposerIdx])
 	require.NoError(t, err)
 
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	chainService := &mock.ChainService{Genesis: time.Unix(presentTime-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		DB: db,
 		FinalizedCheckPoint: &ethpb.Checkpoint{
@@ -1224,7 +1224,7 @@ func Test_validateBellatrixBeaconBlock(t *testing.T) {
 	db := dbtest.SetupDB(t)
 	p := p2ptest.NewTestP2P(t)
 	ctx := context.Background()
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 	presentTime := time.Now().Unix()
 	chainService := &mock.ChainService{Genesis: time.Unix(presentTime-int64(params.BeaconConfig().SecondsPerSlot), 0),
 		FinalizedCheckPoint: &ethpb.Checkpoint{
@@ -1255,7 +1255,7 @@ func Test_validateBellatrixBeaconBlockParentValidation(t *testing.T) {
 	db := dbtest.SetupDB(t)
 	p := p2ptest.NewTestP2P(t)
 	ctx := context.Background()
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 
 	beaconState, privKeys := util.DeterministicGenesisStateBellatrix(t, 100)
 	parentBlock := util.NewBeaconBlockBellatrix()
@@ -1311,7 +1311,7 @@ func Test_validateBeaconBlockProcessingWhenParentIsOptimistic(t *testing.T) {
 	db := dbtest.SetupDB(t)
 	p := p2ptest.NewTestP2P(t)
 	ctx := context.Background()
-	stateGen := stategen.New(db)
+	stateGen := sgmock.NewMockStategen(db)
 
 	beaconState, privKeys := util.DeterministicGenesisStateBellatrix(t, 100)
 	parentBlock := util.NewBeaconBlockBellatrix()

beacon-chain/sync/validate_sync_committee_message_test.go

@@ -19,7 +19,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/encoder"
 	mockp2p "github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/testing"
 	p2ptypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/types"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	mockSync "github.com/prysmaticlabs/prysm/v3/beacon-chain/sync/initial-sync/testing"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	types "github.com/prysmaticlabs/prysm/v3/consensus-types/primitives"
@@ -64,7 +64,7 @@ func TestService_ValidateSyncCommitteeMessage(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SyncCommitteeMessage, topic string) (*Service, string) {
-				s.cfg.stateGen = stategen.New(beaconDB)
+				s.cfg.stateGen = sgmock.NewMockStategen(beaconDB)
 				msg.BlockRoot = headRoot[:]
 				s.cfg.beaconDB = beaconDB
 				s.initCaches()
@@ -92,7 +92,7 @@ func TestService_ValidateSyncCommitteeMessage(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SyncCommitteeMessage, topic string) (*Service, string) {
-				s.cfg.stateGen = stategen.New(beaconDB)
+				s.cfg.stateGen = sgmock.NewMockStategen(beaconDB)
 				msg.BlockRoot = headRoot[:]
 				s.cfg.beaconDB = beaconDB
 				s.initCaches()
@@ -120,7 +120,7 @@ func TestService_ValidateSyncCommitteeMessage(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SyncCommitteeMessage, topic string) (*Service, string) {
-				s.cfg.stateGen = stategen.New(beaconDB)
+				s.cfg.stateGen = sgmock.NewMockStategen(beaconDB)
 				s.cfg.beaconDB = beaconDB
 				s.initCaches()
 				return s, topic
@@ -147,7 +147,7 @@ func TestService_ValidateSyncCommitteeMessage(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SyncCommitteeMessage, topic string) (*Service, string) {
-				s.cfg.stateGen = stategen.New(beaconDB)
+				s.cfg.stateGen = sgmock.NewMockStategen(beaconDB)
 				s.cfg.beaconDB = beaconDB
 				s.initCaches()
 
@@ -176,7 +176,7 @@ func TestService_ValidateSyncCommitteeMessage(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SyncCommitteeMessage, topic string) (*Service, string) {
-				s.cfg.stateGen = stategen.New(beaconDB)
+				s.cfg.stateGen = sgmock.NewMockStategen(beaconDB)
 				s.cfg.beaconDB = beaconDB
 				s.initCaches()
 				s.cfg.chain = &mockChain.ChainService{
@@ -210,7 +210,7 @@ func TestService_ValidateSyncCommitteeMessage(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SyncCommitteeMessage, topic string) (*Service, string) {
-				s.cfg.stateGen = stategen.New(beaconDB)
+				s.cfg.stateGen = sgmock.NewMockStategen(beaconDB)
 				s.cfg.beaconDB = beaconDB
 				s.initCaches()
 				msg.BlockRoot = headRoot[:]
@@ -258,7 +258,7 @@ func TestService_ValidateSyncCommitteeMessage(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SyncCommitteeMessage, topic string) (*Service, string) {
-				s.cfg.stateGen = stategen.New(beaconDB)
+				s.cfg.stateGen = sgmock.NewMockStategen(beaconDB)
 				s.cfg.beaconDB = beaconDB
 				s.initCaches()
 				msg.BlockRoot = headRoot[:]
@@ -305,7 +305,7 @@ func TestService_ValidateSyncCommitteeMessage(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SyncCommitteeMessage, topic string) (*Service, string) {
-				s.cfg.stateGen = stategen.New(beaconDB)
+				s.cfg.stateGen = sgmock.NewMockStategen(beaconDB)
 				s.cfg.beaconDB = beaconDB
 				s.initCaches()
 				msg.BlockRoot = headRoot[:]
@@ -360,7 +360,7 @@ func TestService_ValidateSyncCommitteeMessage(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SyncCommitteeMessage, topic string) (*Service, string) {
-				s.cfg.stateGen = stategen.New(beaconDB)
+				s.cfg.stateGen = sgmock.NewMockStategen(beaconDB)
 				s.cfg.beaconDB = beaconDB
 				s.initCaches()
 				msg.BlockRoot = headRoot[:]

beacon-chain/sync/validate_sync_contribution_proof_test.go

@@ -26,7 +26,7 @@ import (
 	mockp2p "github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/testing"
 	p2ptypes "github.com/prysmaticlabs/prysm/v3/beacon-chain/p2p/types"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	mockSync "github.com/prysmaticlabs/prysm/v3/beacon-chain/sync/initial-sync/testing"
 	"github.com/prysmaticlabs/prysm/v3/config/params"
 	"github.com/prysmaticlabs/prysm/v3/consensus-types/blocks"
@@ -74,7 +74,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				msg.Message.Contribution.BlockRoot = headRoot[:]
 				s.cfg.beaconDB = database
 				s.initCaches()
@@ -110,7 +110,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				msg.Message.Contribution.BlockRoot = headRoot[:]
 				s.cfg.beaconDB = database
 				s.initCaches()
@@ -146,7 +146,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				s.cfg.beaconDB = database
 				s.initCaches()
 				return s
@@ -181,7 +181,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				s.cfg.beaconDB = database
 				s.initCaches()
 				s.cfg.chain = &mockChain.ChainService{
@@ -224,7 +224,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				s.cfg.beaconDB = database
 				s.initCaches()
 				s.cfg.chain = &mockChain.ChainService{
@@ -267,7 +267,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				s.cfg.beaconDB = database
 				s.initCaches()
 				s.cfg.chain = &mockChain.ChainService{
@@ -311,7 +311,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				s.cfg.beaconDB = database
 				s.initCaches()
 				s.cfg.chain = &mockChain.ChainService{
@@ -373,7 +373,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				s.cfg.beaconDB = database
 				msg.Message.Contribution.BlockRoot = headRoot[:]
 				hState, err := database.State(context.Background(), headRoot)
@@ -438,7 +438,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				s.cfg.beaconDB = database
 				s.cfg.chain = chainService
 				msg.Message.Contribution.BlockRoot = headRoot[:]
@@ -518,7 +518,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				s.cfg.beaconDB = database
 				msg.Message.Contribution.BlockRoot = headRoot[:]
 				hState, err := database.State(context.Background(), headRoot)
@@ -600,7 +600,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				msg.Message.Contribution.BlockRoot = headRoot[:]
 				s.cfg.beaconDB = database
 				hState, err := database.State(context.Background(), headRoot)
@@ -684,7 +684,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				msg.Message.Contribution.BlockRoot = headRoot[:]
 				s.cfg.beaconDB = database
 				hState, err := database.State(context.Background(), headRoot)
@@ -780,7 +780,7 @@ func TestService_ValidateSyncContributionAndProof(t *testing.T) {
 				WithOperationNotifier(chainService.OperationNotifier()),
 			),
 			setupSvc: func(s *Service, msg *ethpb.SignedContributionAndProof) *Service {
-				s.cfg.stateGen = stategen.New(database)
+				s.cfg.stateGen = sgmock.NewMockStategen(database)
 				msg.Message.Contribution.BlockRoot = headRoot[:]
 				s.cfg.beaconDB = database
 				hState, err := database.State(context.Background(), headRoot)
@@ -927,7 +927,7 @@ func TestValidateSyncContributionAndProof(t *testing.T) {
 		WithOperationNotifier(chainService.OperationNotifier()),
 	)
 	go s.verifierRoutine()
-	s.cfg.stateGen = stategen.New(database)
+	s.cfg.stateGen = sgmock.NewMockStategen(database)
 	msg.Message.Contribution.BlockRoot = headRoot[:]
 	s.cfg.beaconDB = database
 	hState, err := database.State(context.Background(), headRoot)

testing/spectest/shared/common/forkchoice/BUILD.bazel

@@ -22,7 +22,7 @@ go_library(
         "//beacon-chain/forkchoice/protoarray:go_default_library",
         "//beacon-chain/operations/attestations:go_default_library",
         "//beacon-chain/state:go_default_library",
-        "//beacon-chain/state/stategen:go_default_library",
+        "//beacon-chain/state/stategen/mock:go_default_library",
         "//beacon-chain/state/v1:go_default_library",
         "//beacon-chain/state/v2:go_default_library",
         "//beacon-chain/state/v3:go_default_library",

testing/spectest/shared/common/forkchoice/service.go

@@ -17,7 +17,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/forkchoice/protoarray"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/operations/attestations"
 	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state"
-	"github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen"
+	sgmock "github.com/prysmaticlabs/prysm/v3/beacon-chain/state/stategen/mock"
 	"github.com/prysmaticlabs/prysm/v3/consensus-types/interfaces"
 	"github.com/prysmaticlabs/prysm/v3/encoding/bytesutil"
 	pb "github.com/prysmaticlabs/prysm/v3/proto/engine/v1"
@@ -54,7 +54,7 @@ func startChainService(t testing.TB, st state.BeaconState, block interfaces.Sign
 		blockchain.WithDatabase(db),
 		blockchain.WithAttestationService(attPool),
 		blockchain.WithForkChoiceStore(protoarray.New()),
-		blockchain.WithStateGen(stategen.New(db)),
+		blockchain.WithStateGen(sgmock.NewMockStategen(db)),
 		blockchain.WithStateNotifier(&mock.MockStateNotifier{}),
 		blockchain.WithAttestationPool(attestations.NewPool()),
 		blockchain.WithDepositCache(depositCache),