Micro optimizations on new-state-mgmt service for initial syncing

beacon-chain/blockchain/init_sync_process_block_test.go

@@ -6,6 +6,7 @@ import (
 
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
 	"github.com/prysmaticlabs/go-ssz"
+	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/blocks"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/state"
@@ -210,7 +211,7 @@ func TestPruneNonBoundary_CanPrune(t *testing.T) {
 func TestGenerateState_CorrectlyGenerated(t *testing.T) {
 	db := testDB.SetupDB(t)
 	defer testDB.TeardownDB(t, db)
-	cfg := &Config{BeaconDB: db, StateGen: stategen.New(db)}
+	cfg := &Config{BeaconDB: db, StateGen: stategen.New(db, cache.NewStateSummaryCache())}
 	service, err := NewService(context.Background(), cfg)
 	if err != nil {
 		t.Fatal(err)

beacon-chain/blockchain/process_block.go

@@ -140,12 +140,13 @@ func (s *Service) onBlock(ctx context.Context, signed *ethpb.SignedBeaconBlock)
 		}
 
 		if featureconfig.Get().NewStateMgmt {
-			finalizedState, err := s.stateGen.StateByRoot(ctx, fRoot)
+			fRoot := bytesutil.ToBytes32(postState.FinalizedCheckpoint().Root)
+			fBlock, err := s.beaconDB.Block(ctx, fRoot)
 			if err != nil {
-				return nil, err
+				return nil, errors.Wrap(err, "could not get finalized block to migrate")
 			}
-			if err := s.stateGen.MigrateToCold(ctx, finalizedState, fRoot); err != nil {
-				return nil, err
+			if err := s.stateGen.MigrateToCold(ctx, fBlock.Block.Slot, fRoot); err != nil {
+				return nil, errors.Wrap(err, "could not migrate to cold")
 			}
 		}
 	}
@@ -300,13 +301,12 @@ func (s *Service) onBlockInitialSyncStateTransition(ctx context.Context, signed
 
 		if featureconfig.Get().NewStateMgmt {
 			fRoot := bytesutil.ToBytes32(postState.FinalizedCheckpoint().Root)
-			finalizedState, err := s.stateGen.StateByRoot(ctx, fRoot)
+			fBlock, err := s.beaconDB.Block(ctx, fRoot)
 			if err != nil {
-				return errors.Wrap(err, "could not get state by root for migration")
+				return errors.Wrap(err, "could not get finalized block to migrate")
 			}
-			if err := s.stateGen.MigrateToCold(ctx, finalizedState, fRoot); err != nil {
-				return errors.Wrap(err, "could not migrate with new finalized root")
-
+			if err := s.stateGen.MigrateToCold(ctx, fBlock.Block.Slot, fRoot); err != nil {
+				return errors.Wrap(err, "could not migrate to cold")
 			}
 		}
 	}

beacon-chain/cache/BUILD.bazel

@@ -11,11 +11,16 @@ go_library(
         "eth1_data.go",
         "hot_state_cache.go",
         "skip_slot_cache.go",
+        "state_summary.go",
     ],
     importpath = "github.com/prysmaticlabs/prysm/beacon-chain/cache",
-    visibility = ["//beacon-chain:__subpackages__"],
+    visibility = [
+        "//beacon-chain:__subpackages__",
+        "//tools:__subpackages__",
+    ],
     deps = [
         "//beacon-chain/state:go_default_library",
+        "//proto/beacon/p2p/v1:go_default_library",
         "//shared/featureconfig:go_default_library",
         "//shared/hashutil:go_default_library",
         "//shared/params:go_default_library",

beacon-chain/cache/state_summary.go

@@ -0,0 +1,65 @@
+package cache
+
+import (
+	"sync"
+
+	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+)
+
+// StateSummaryCache caches state summary object.
+type StateSummaryCache struct {
+	initSyncStateSummaries     map[[32]byte]*pb.StateSummary
+	initSyncStateSummariesLock sync.RWMutex
+}
+
+// NewStateSummaryCache creates a new state summary cache.
+func NewStateSummaryCache() *StateSummaryCache {
+	return &StateSummaryCache{
+		initSyncStateSummaries: make(map[[32]byte]*pb.StateSummary),
+	}
+}
+
+// Put saves a state summary to the initial sync state summaries cache.
+func (s *StateSummaryCache) Put(r [32]byte, b *pb.StateSummary) {
+	s.initSyncStateSummariesLock.Lock()
+	defer s.initSyncStateSummariesLock.Unlock()
+	s.initSyncStateSummaries[r] = b
+}
+
+// Has checks if a state summary exists in the initial sync state summaries cache using the root
+// of the block.
+func (s *StateSummaryCache) Has(r [32]byte) bool {
+	s.initSyncStateSummariesLock.RLock()
+	defer s.initSyncStateSummariesLock.RUnlock()
+	_, ok := s.initSyncStateSummaries[r]
+	return ok
+}
+
+// Get retrieves a state summary from the initial sync state summaries cache using the root of
+// the block.
+func (s *StateSummaryCache) Get(r [32]byte) *pb.StateSummary {
+	s.initSyncStateSummariesLock.RLock()
+	defer s.initSyncStateSummariesLock.RUnlock()
+	b := s.initSyncStateSummaries[r]
+	return b
+}
+
+// GetAll retrieves all the beacon state summaries from the initial sync state summaries cache, the returned
+// state summaries are unordered.
+func (s *StateSummaryCache) GetAll() []*pb.StateSummary {
+	s.initSyncStateSummariesLock.RLock()
+	defer s.initSyncStateSummariesLock.RUnlock()
+
+	blks := make([]*pb.StateSummary, 0, len(s.initSyncStateSummaries))
+	for _, b := range s.initSyncStateSummaries {
+		blks = append(blks, b)
+	}
+	return blks
+}
+
+// Clear clears out the initial sync state summaries cache.
+func (s *StateSummaryCache) Clear() {
+	s.initSyncStateSummariesLock.Lock()
+	defer s.initSyncStateSummariesLock.Unlock()
+	s.initSyncStateSummaries = make(map[[32]byte]*pb.StateSummary)
+}

beacon-chain/db/BUILD.bazel

@@ -26,6 +26,7 @@ go_library(
         "//tools:__subpackages__",
     ],
     deps = [
+        "//beacon-chain/cache:go_default_library",
         "//beacon-chain/db/iface:go_default_library",
         "//beacon-chain/db/kv:go_default_library",
         "@com_github_sirupsen_logrus//:go_default_library",

beacon-chain/db/db.go

@@ -1,8 +1,11 @@
 package db
 
-import "github.com/prysmaticlabs/prysm/beacon-chain/db/kv"
+import (
+	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
+	"github.com/prysmaticlabs/prysm/beacon-chain/db/kv"
+)
 
 // NewDB initializes a new DB.
-func NewDB(dirPath string) (Database, error) {
-	return kv.NewKVStore(dirPath)
+func NewDB(dirPath string, stateSummaryCache *cache.StateSummaryCache) (Database, error) {
+	return kv.NewKVStore(dirPath, stateSummaryCache)
 }

beacon-chain/db/db_kafka_wrapped.go

@@ -1,13 +1,14 @@
 package db
 
 import (
+	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/beacon-chain/db/kafka"
 	"github.com/prysmaticlabs/prysm/beacon-chain/db/kv"
 )
 
 // NewDB initializes a new DB with kafka wrapper.
-func NewDB(dirPath string) (Database, error) {
-	db, err := kv.NewKVStore(dirPath)
+func NewDB(dirPath string, stateSummaryCache *cache.StateSummaryCache) (Database, error) {
+	db, err := kv.NewKVStore(dirPath, stateSummaryCache)
 	if err != nil {
 		return nil, err
 	}

beacon-chain/db/iface/interface.go

@@ -90,6 +90,7 @@ type NoHeadAccessDatabase interface {
 	DeleteState(ctx context.Context, blockRoot [32]byte) error
 	DeleteStates(ctx context.Context, blockRoots [][32]byte) error
 	SaveStateSummary(ctx context.Context, summary *ethereum_beacon_p2p_v1.StateSummary) error
+	SaveStateSummaries(ctx context.Context, summaries []*ethereum_beacon_p2p_v1.StateSummary) error
 	// Slashing operations.
 	SaveProposerSlashing(ctx context.Context, slashing *eth.ProposerSlashing) error
 	SaveAttesterSlashing(ctx context.Context, slashing *eth.AttesterSlashing) error

beacon-chain/db/kafka/passthrough.go

@@ -238,6 +238,11 @@ func (e Exporter) SaveStateSummary(ctx context.Context, summary *pb.StateSummary
 	return e.db.SaveStateSummary(ctx, summary)
 }
 
+// SaveStateSummaries -- passthrough.
+func (e Exporter) SaveStateSummaries(ctx context.Context, summaries []*pb.StateSummary) error {
+	return e.db.SaveStateSummaries(ctx, summaries)
+}
+
 // SaveStates -- passthrough.
 func (e Exporter) SaveStates(ctx context.Context, states []*state.BeaconState, blockRoots [][32]byte) error {
 	return e.db.SaveStates(ctx, states, blockRoots)

beacon-chain/db/kv/BUILD.bazel

@@ -26,6 +26,7 @@ go_library(
     importpath = "github.com/prysmaticlabs/prysm/beacon-chain/db/kv",
     visibility = ["//beacon-chain:__subpackages__"],
     deps = [
+        "//beacon-chain/cache:go_default_library",
         "//beacon-chain/core/helpers:go_default_library",
         "//beacon-chain/db/filters:go_default_library",
         "//beacon-chain/db/iface:go_default_library",
@@ -74,6 +75,7 @@ go_test(
     ],
     embed = [":go_default_library"],
     deps = [
+        "//beacon-chain/cache:go_default_library",
         "//beacon-chain/db/filters:go_default_library",
         "//beacon-chain/state:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",

beacon-chain/db/kv/blocks.go

@@ -263,7 +263,7 @@ func (k *Store) SaveHeadBlockRoot(ctx context.Context, blockRoot [32]byte) error
 	defer span.End()
 	return k.db.Update(func(tx *bolt.Tx) error {
 		if featureconfig.Get().NewStateMgmt {
-			if tx.Bucket(stateSummaryBucket).Get(blockRoot[:]) == nil {
+			if tx.Bucket(stateSummaryBucket).Get(blockRoot[:]) == nil && !k.stateSummaryCache.Has(blockRoot) {
 				return errors.New("no state summary found with head block root")
 			}
 		} else {

beacon-chain/db/kv/checkpoint.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"github.com/prysmaticlabs/prysm/shared/bytesutil"
 	"github.com/prysmaticlabs/prysm/shared/featureconfig"
 	"github.com/prysmaticlabs/prysm/shared/traceutil"
 	bolt "go.etcd.io/bbolt"
@@ -65,7 +66,7 @@ func (k *Store) SaveJustifiedCheckpoint(ctx context.Context, checkpoint *ethpb.C
 	return k.db.Update(func(tx *bolt.Tx) error {
 		bucket := tx.Bucket(checkpointBucket)
 		if featureconfig.Get().NewStateMgmt {
-			if tx.Bucket(stateSummaryBucket).Get(checkpoint.Root) == nil {
+			if tx.Bucket(stateSummaryBucket).Get(checkpoint.Root) == nil && !k.stateSummaryCache.Has(bytesutil.ToBytes32(checkpoint.Root)) {
 				return errors.New("missing state summary for finalized root")
 			}
 		} else {
@@ -93,7 +94,7 @@ func (k *Store) SaveFinalizedCheckpoint(ctx context.Context, checkpoint *ethpb.C
 	return k.db.Update(func(tx *bolt.Tx) error {
 		bucket := tx.Bucket(checkpointBucket)
 		if featureconfig.Get().NewStateMgmt {
-			if tx.Bucket(stateSummaryBucket).Get(checkpoint.Root) == nil {
+			if tx.Bucket(stateSummaryBucket).Get(checkpoint.Root) == nil && !k.stateSummaryCache.Has(bytesutil.ToBytes32(checkpoint.Root)) {
 				return errors.New("missing state summary for finalized root")
 			}
 		} else {

beacon-chain/db/kv/kv.go

@@ -10,6 +10,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	prombolt "github.com/prysmaticlabs/prombbolt"
+	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/beacon-chain/db/iface"
 	bolt "go.etcd.io/bbolt"
 )
@@ -38,12 +39,13 @@ type Store struct {
 	validatorIndexCache *ristretto.Cache
 	stateSlotBitLock    sync.Mutex
 	blockSlotBitLock    sync.Mutex
+	stateSummaryCache   *cache.StateSummaryCache
 }
 
 // NewKVStore initializes a new boltDB key-value store at the directory
 // path specified, creates the kv-buckets based on the schema, and stores
 // an open connection db object as a property of the Store struct.
-func NewKVStore(dirPath string) (*Store, error) {
+func NewKVStore(dirPath string, stateSummaryCache *cache.StateSummaryCache) (*Store, error) {
 	if err := os.MkdirAll(dirPath, 0700); err != nil {
 		return nil, err
 	}
@@ -79,6 +81,7 @@ func NewKVStore(dirPath string) (*Store, error) {
 		databasePath:        dirPath,
 		blockCache:          blockCache,
 		validatorIndexCache: validatorCache,
+		stateSummaryCache:   stateSummaryCache,
 	}
 
 	if err := kv.db.Update(func(tx *bolt.Tx) error {

beacon-chain/db/kv/kv_test.go

@@ -8,6 +8,7 @@ import (
 	"path"
 	"testing"
 
+	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/shared/testutil"
 )
 
@@ -21,7 +22,7 @@ func setupDB(t testing.TB) *Store {
 	if err := os.RemoveAll(path); err != nil {
 		t.Fatalf("Failed to remove directory: %v", err)
 	}
-	db, err := NewKVStore(path)
+	db, err := NewKVStore(path, cache.NewStateSummaryCache())
 	if err != nil {
 		t.Fatalf("Failed to instantiate DB: %v", err)
 	}

beacon-chain/db/kv/state_summary.go

@@ -23,6 +23,26 @@ func (k *Store) SaveStateSummary(ctx context.Context, summary *pb.StateSummary)
 	})
 }
 
+// SaveStateSummaries saves state summary objects to the DB.
+func (k *Store) SaveStateSummaries(ctx context.Context, summaries []*pb.StateSummary) error {
+	ctx, span := trace.StartSpan(ctx, "BeaconDB.SaveStateSummaries")
+	defer span.End()
+
+	return k.db.Update(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket(stateSummaryBucket)
+		for _, summary := range summaries {
+			enc, err := encode(summary)
+			if err != nil {
+				return err
+			}
+			if err := bucket.Put(summary.Root, enc); err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+}
+
 // StateSummary returns the state summary object from the db using input block root.
 func (k *Store) StateSummary(ctx context.Context, blockRoot [32]byte) (*pb.StateSummary, error) {
 	ctx, span := trace.StartSpan(ctx, "BeaconDB.StateSummary")

beacon-chain/db/testing/BUILD.bazel

@@ -7,6 +7,7 @@ go_library(
     importpath = "github.com/prysmaticlabs/prysm/beacon-chain/db/testing",
     visibility = ["//beacon-chain:__subpackages__"],
     deps = [
+        "//beacon-chain/cache:go_default_library",
         "//beacon-chain/db:go_default_library",
         "//beacon-chain/db/kv:go_default_library",
         "//shared/testutil:go_default_library",

beacon-chain/db/testing/setup_db.go

@@ -8,6 +8,7 @@ import (
 	"path"
 	"testing"
 
+	"github.com/prysmaticlabs/prysm/beacon-chain/cache"
 	"github.com/prysmaticlabs/prysm/beacon-chain/db"
 	"github.com/prysmaticlabs/prysm/beacon-chain/db/kv"
 	"github.com/prysmaticlabs/prysm/shared/testutil"
@@ -23,7 +24,7 @@ func SetupDB(t testing.TB) db.Database {
 	if err := os.RemoveAll(p); err != nil {
 		t.Fatalf("failed to remove directory: %v", err)
 	}
-	s, err := kv.NewKVStore(p)
+	s, err := kv.NewKVStore(p, cache.NewStateSummaryCache())
 	if err != nil {
 		t.Fatal(err)
 	}

beacon-chain/node/BUILD.bazel

@@ -8,6 +8,7 @@ go_library(
     deps = [
         "//beacon-chain/archiver:go_default_library",
         "//beacon-chain/blockchain:go_default_library",
+        "//beacon-chain/cache:go_default_library",
         "//beacon-chain/cache/depositcache:go_default_library",
         "//beacon-chain/db:go_default_library",
         "//beacon-chain/flags:go_default_library",