New state getter (#5101)

* getter.go
* getter_test.go
* fixed a cold bug
* fmt gaz
* All tests pass
* Merge branch 'master' into new-state-getter
* Merge refs/heads/master into new-state-getter

beacon-chain/state/stategen/BUILD.bazel

@@ -6,6 +6,7 @@ go_library(
         "cold.go",
         "epoch_boundary_root.go",
         "errors.go",
+        "getter.go",
         "hot.go",
         "log.go",
         "migrate.go",
@@ -39,6 +40,7 @@ go_test(
     srcs = [
         "cold_test.go",
         "epoch_boundary_root_test.go",
+        "getter_test.go",
         "hot_test.go",
         "migrate_test.go",
         "replay_test.go",

beacon-chain/state/stategen/cold.go

@@ -124,7 +124,7 @@ func (s *State) loadColdIntermediateStateBySlot(ctx context.Context, slot uint64
 		highArchivedPointRoot = s.splitInfo.root
 		highArchivedPointSlot = s.splitInfo.slot
 	} else {
-		if _, err := s.archivedPointByIndex(ctx, highArchivedPointSlot); err != nil {
+		if _, err := s.archivedPointByIndex(ctx, highArchivedPointIdx); err != nil {
 			return nil, errors.Wrap(err, "could not get upper bound archived state using index")
 		}
 		highArchivedPointRoot = s.beaconDB.ArchivedPointRoot(ctx, highArchivedPointIdx)

beacon-chain/state/stategen/cold_test.go

@@ -192,7 +192,6 @@ func TestLoadColdStateBySlotIntermediatePlayback_AfterCutoff(t *testing.T) {
 	}
 }
 
-
 func TestLoadColdStateByRoot_UnknownArchivedState(t *testing.T) {
 	ctx := context.Background()
 	db := testDB.SetupDB(t)
@@ -200,8 +199,7 @@ func TestLoadColdStateByRoot_UnknownArchivedState(t *testing.T) {
 
 	service := New(db)
 	service.slotsPerArchivedPoint = 1
-	if _, err := service.loadColdIntermediateStateBySlot(ctx, 0);
-		!strings.Contains(err.Error(), errUnknownArchivedState.Error()) {
+	if _, err := service.loadColdIntermediateStateBySlot(ctx, 0); !strings.Contains(err.Error(), errUnknownArchivedState.Error()) {
 		t.Log(err)
 		t.Error("Did not get wanted error")
 	}

beacon-chain/state/stategen/getter.go

@@ -0,0 +1,44 @@
+package stategen
+
+import (
+	"context"
+
+	"github.com/pkg/errors"
+	"github.com/prysmaticlabs/prysm/beacon-chain/state"
+	"go.opencensus.io/trace"
+)
+
+// StateByRoot retrieves the state from DB using input block root.
+// It retrieves state from the hot section if the state summary slot
+// is below the split point cut off.
+func (s *State) StateByRoot(ctx context.Context, blockRoot [32]byte) (*state.BeaconState, error) {
+	ctx, span := trace.StartSpan(ctx, "stateGen.StateByRoot")
+	defer span.End()
+
+	summary, err := s.beaconDB.StateSummary(ctx, blockRoot)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not get state summary")
+	}
+
+	if summary.Slot < s.splitInfo.slot {
+		return s.loadColdStateByRoot(ctx, blockRoot)
+	}
+
+	return s.loadHotStateByRoot(ctx, blockRoot)
+}
+
+// StateBySlot retrieves the state from DB using input slot.
+// It retrieves state from the cold section if the input slot
+// is below the split point cut off.
+// Note: `StateByRoot` is preferred over this. Retrieving state
+// by root `StateByRoot` is more performant than retrieving by slot.
+func (s *State) StateBySlot(ctx context.Context, slot uint64) (*state.BeaconState, error) {
+	ctx, span := trace.StartSpan(ctx, "stateGen.StateBySlot")
+	defer span.End()
+
+	if slot < s.splitInfo.slot {
+		return s.loadColdIntermediateStateBySlot(ctx, slot)
+	}
+
+	return s.loadHotStateBySlot(ctx, slot)
+}

beacon-chain/state/stategen/getter_test.go

@@ -0,0 +1,185 @@
+package stategen
+
+import (
+	"context"
+	"testing"
+
+	"github.com/gogo/protobuf/proto"
+	testDB "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
+	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
+	"github.com/prysmaticlabs/prysm/shared/params"
+	"github.com/prysmaticlabs/prysm/shared/testutil"
+)
+
+func TestStateByRoot_ColdState(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+	service.splitInfo.slot = 2
+	service.slotsPerArchivedPoint = 1
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	if err := service.beaconDB.SaveArchivedPointState(ctx, beaconState, 1); err != nil {
+		t.Fatal(err)
+	}
+	r := [32]byte{'a'}
+	if err := service.beaconDB.SaveStateSummary(ctx, &pb.StateSummary{
+		Root: r[:],
+		Slot: 1,
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	loadedState, err := service.StateByRoot(ctx, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !proto.Equal(loadedState.InnerStateUnsafe(), beaconState.InnerStateUnsafe()) {
+		t.Error("Did not correctly save state")
+	}
+}
+
+func TestStateByRoot_HotStateDB(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	boundaryRoot := [32]byte{'A'}
+	blkRoot := [32]byte{'B'}
+	if err := service.beaconDB.SaveState(ctx, beaconState, boundaryRoot); err != nil {
+		t.Fatal(err)
+	}
+	targetSlot := uint64(10)
+	if err := service.beaconDB.SaveStateSummary(ctx, &pb.StateSummary{
+		Slot:         targetSlot,
+		Root:         blkRoot[:],
+		BoundaryRoot: boundaryRoot[:],
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	loadedState, err := service.StateByRoot(ctx, blkRoot)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if loadedState.Slot() != targetSlot {
+		t.Error("Did not correctly load state")
+	}
+}
+
+func TestStateByRoot_HotStateCached(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	r := [32]byte{'A'}
+	if err := service.beaconDB.SaveStateSummary(ctx, &pb.StateSummary{
+		Root:         r[:],
+		BoundaryRoot: r[:],
+	}); err != nil {
+		t.Fatal(err)
+	}
+	service.hotStateCache.Put(r, beaconState)
+
+	loadedState, err := service.StateByRoot(ctx, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !proto.Equal(loadedState.InnerStateUnsafe(), beaconState.InnerStateUnsafe()) {
+		t.Error("Did not correctly cache state")
+	}
+}
+
+func TestStateBySlot_ColdState(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+	service.slotsPerArchivedPoint = params.BeaconConfig().SlotsPerEpoch * 2
+	service.splitInfo.slot = service.slotsPerArchivedPoint + 1
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	r := [32]byte{}
+	if err := service.beaconDB.SaveArchivedPointState(ctx, beaconState, 0); err != nil {
+		t.Fatal(err)
+	}
+	if err := service.beaconDB.SaveArchivedPointRoot(ctx, r, 0); err != nil {
+		t.Fatal(err)
+	}
+	if err := service.beaconDB.SaveArchivedPointState(ctx, beaconState, 1); err != nil {
+		t.Fatal(err)
+	}
+	if err := service.beaconDB.SaveArchivedPointRoot(ctx, r, 1); err != nil {
+		t.Fatal(err)
+	}
+	if err := service.beaconDB.SaveStateSummary(ctx, &pb.StateSummary{
+		Slot: service.slotsPerArchivedPoint,
+		Root: r[:],
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	slot := uint64(20)
+	loadedState, err := service.StateBySlot(ctx, slot)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if loadedState.Slot() != slot {
+		t.Error("Did not correctly save state")
+	}
+}
+
+func TestStateBySlot_HotStateCached(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	r := [32]byte{'A'}
+	service.hotStateCache.Put(r, beaconState)
+	service.setEpochBoundaryRoot(0, r)
+
+	slot := uint64(10)
+	loadedState, err := service.StateBySlot(ctx, slot)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if loadedState.Slot() != slot {
+		t.Error("Did not correctly load state")
+	}
+}
+
+func TestStateBySlot_HotStateDB(t *testing.T) {
+	ctx := context.Background()
+	db := testDB.SetupDB(t)
+	defer testDB.TeardownDB(t, db)
+
+	service := New(db)
+
+	beaconState, _ := testutil.DeterministicGenesisState(t, 32)
+	r := [32]byte{'A'}
+	service.setEpochBoundaryRoot(0, r)
+	if err := service.beaconDB.SaveState(ctx, beaconState, r); err != nil {
+		t.Fatal(err)
+	}
+
+	slot := uint64(10)
+	loadedState, err := service.StateBySlot(ctx, slot)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if loadedState.Slot() != slot {
+		t.Error("Did not correctly load state")
+	}
+}