Save hot state (#5083)

* loadEpochBoundaryRoot
* Tests
* Span
* Merge branch 'master' of https://github.com/prysmaticlabs/prysm into save-hot-state
* Starting test
* Tests
* Merge refs/heads/master into save-hot-state
* Merge branch 'master' into save-hot-state
* Use copy
* Merge branch 'save-hot-state' of https://github.com/prysmaticlabs/prysm into save-hot-state
* Merge refs/heads/master into save-hot-state

beacon-chain/state/stategen/BUILD.bazel

@@ -21,6 +21,7 @@ go_library(
         "//beacon-chain/db:go_default_library",
         "//beacon-chain/db/filters:go_default_library",
         "//beacon-chain/state:go_default_library",
+        "//proto/beacon/p2p/v1:go_default_library",
         "//shared/bytesutil:go_default_library",
         "//shared/featureconfig:go_default_library",
         "//shared/params:go_default_library",

beacon-chain/state/stategen/hot.go

@@ -2,15 +2,59 @@ package stategen
 
 import (
 	"context"
+	"encoding/hex"
 
 	"github.com/pkg/errors"
 	"github.com/prysmaticlabs/prysm/beacon-chain/core/helpers"
 	"github.com/prysmaticlabs/prysm/beacon-chain/state"
+	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/bytesutil"
 	"github.com/prysmaticlabs/prysm/shared/params"
+	"github.com/sirupsen/logrus"
 	"go.opencensus.io/trace"
 )
 
+// This saves a post finalized beacon state in the hot section of the DB. On the epoch boundary,
+// it saves a full state. On an intermediate slot, it saves a back pointer to the
+// nearest epoch boundary state.
+func (s *State) saveHotState(ctx context.Context, blockRoot [32]byte, state *state.BeaconState) error {
+	ctx, span := trace.StartSpan(ctx, "stateGen.saveHotState")
+	defer span.End()
+
+	// If the hot state is already in cache, one can be sure the state was processed and in the DB.
+	if s.hotStateCache.Has(blockRoot) {
+		return nil
+	}
+
+	// Only on an epoch boundary slot, saves the whole state.
+	if helpers.IsEpochStart(state.Slot()) {
+		if err := s.beaconDB.SaveState(ctx, state, blockRoot); err != nil {
+			return err
+		}
+		log.WithFields(logrus.Fields{
+			"slot":      state.Slot(),
+			"blockRoot": hex.EncodeToString(bytesutil.Trunc(blockRoot[:]))}).Info("Saved full state on epoch boundary")
+	}
+
+	// On an intermediate slots, save the hot state summary.
+	epochRoot, err := s.loadEpochBoundaryRoot(ctx, blockRoot, state)
+	if err != nil {
+		return errors.Wrap(err, "could not get epoch boundary root to save hot state")
+	}
+	if err := s.beaconDB.SaveStateSummary(ctx, &pb.StateSummary{
+		Slot:         state.Slot(),
+		Root:         blockRoot[:],
+		BoundaryRoot: epochRoot[:],
+	}); err != nil {
+		return err
+	}
+
+	// Store the copied state in the cache.
+	s.hotStateCache.Put(blockRoot, state.Copy())
+
+	return nil
+}
+
 // This loads a post finalized beacon state from the hot section of the DB. If necessary it will
 // replay blocks starting from the nearest epoch boundary. It returns the beacon state that
 // corresponds to the input block root.

beacon-chain/state/stategen/hot_test.go

@@ -11,8 +11,94 @@ import (
 	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/params"
 	"github.com/prysmaticlabs/prysm/shared/testutil"
+	logTest "github.com/sirupsen/logrus/hooks/test"
 )
 
+func TestSaveHotState_AlreadyHas(t *testing.T) {
+	hook := logTest.NewGlobal()
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	service := New(db)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch)
+	r := [32]byte{'A'}
+
+	// Pre cache the hot state.
+	service.hotStateCache.Put(r, beaconState)
+	if err := service.saveHotState(ctx, r, beaconState); err != nil {
+		t.Fatal(err)
+	}
+
+	// Should not save the state and state summary.
+	if service.beaconDB.HasState(ctx, r) {
+		t.Error("Should not have saved the state")
+	}
+	if service.beaconDB.HasStateSummary(ctx, r) {
+		t.Error("Should have saved the state summary")
+	}
+	testutil.AssertLogsDoNotContain(t, hook, "Saved full state on epoch boundary")
+}
+
+func TestSaveHotState_CanSaveOnEpochBoundary(t *testing.T) {
+	hook := logTest.NewGlobal()
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	service := New(db)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch)
+	r := [32]byte{'A'}
+
+	if err := service.saveHotState(ctx, r, beaconState); err != nil {
+		t.Fatal(err)
+	}
+
+	// Should save both state and state summary.
+	if !service.beaconDB.HasState(ctx, r) {
+		t.Error("Should have saved the state")
+	}
+	if !service.beaconDB.HasStateSummary(ctx, r) {
+		t.Error("Should have saved the state summary")
+	}
+	testutil.AssertLogsContain(t, hook, "Saved full state on epoch boundary")
+}
+
+func TestSaveHotState_NoSaveNotEpochBoundary(t *testing.T) {
+	hook := logTest.NewGlobal()
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+	service := New(db)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	beaconState.SetSlot(params.BeaconConfig().SlotsPerEpoch - 1)
+	r := [32]byte{'A'}
+	b := &ethpb.SignedBeaconBlock{Block: &ethpb.BeaconBlock{}}
+	if err := db.SaveBlock(ctx, b); err != nil {
+		t.Fatal(err)
+	}
+	gRoot, _ := ssz.HashTreeRoot(b.Block)
+	if err := db.SaveGenesisBlockRoot(ctx, gRoot); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := service.saveHotState(ctx, r, beaconState); err != nil {
+		t.Fatal(err)
+	}
+
+	// Should only save state summary.
+	if service.beaconDB.HasState(ctx, r) {
+		t.Error("Should not have saved the state")
+	}
+	if !service.beaconDB.HasStateSummary(ctx, r) {
+		t.Error("Should have saved the state summary")
+	}
+	testutil.AssertLogsDoNotContain(t, hook, "Saved full state on epoch boundary")
+}
+
 func TestLoadHoteStateByRoot_Cached(t *testing.T) {
 	ctx := context.Background()
 	db := testDB.SetupDB(t)