Integrate state-diff into State()

Author: Inspector-Butters

beacon-chain/db/kv/state.go

@@ -28,6 +28,17 @@ func (s *Store) State(ctx context.Context, blockRoot [32]byte) (state.BeaconStat
 	ctx, span := trace.StartSpan(ctx, "BeaconDB.State")
 	defer span.End()
 	startTime := time.Now()
+
+	// If state diff is enabled, we get the state from the state-diff db.
+	if features.Get().EnableStateDiff {
+		st, err := s.getStateUsingStateDiff(ctx, blockRoot)
+		if err != nil {
+			return nil, err
+		}
+		stateReadingTime.Observe(float64(time.Since(startTime).Milliseconds()))
+		return st, nil
+	}
+
 	enc, err := s.stateBytes(ctx, blockRoot)
 	if err != nil {
 		return nil, err
@@ -1031,3 +1042,34 @@ func (s *Store) isStateValidatorMigrationOver() (bool, error) {
 	}
 	return returnFlag, nil
 }
+
+func (s *Store) getStateUsingStateDiff(ctx context.Context, blockRoot [32]byte) (state.BeaconState, error) {
+	var slot primitives.Slot
+
+	stateSummary, err := s.StateSummary(ctx, blockRoot)
+	if err != nil {
+		return nil, err
+	}
+	if stateSummary == nil {
+		blk, err := s.Block(ctx, blockRoot)
+		if err != nil {
+			return nil, err
+		}
+		if blk == nil || blk.IsNil() {
+			return nil, errors.New("neither state summary nor block found")
+		}
+		slot = blk.Block().Slot()
+	} else {
+		slot = stateSummary.Slot
+	}
+
+	st, err := s.stateByDiff(ctx, slot)
+	if err != nil {
+		return nil, err
+	}
+	if st == nil || st.IsNil() {
+		return nil, errors.New("state not found")
+	}
+
+	return st, nil
+}

beacon-chain/db/kv/state_diff_test.go

@@ -20,7 +20,7 @@ import (
 )
 
 func TestStateDiff_LoadOrInitOffset(t *testing.T) {
-	setDefaultExponents()
+	setDefaultStateDiffExponents()
 
 	db := setupDB(t)
 	err := setOffsetInDB(db, 10)
@@ -36,7 +36,7 @@ func TestStateDiff_LoadOrInitOffset(t *testing.T) {
 
 func TestStateDiff_ComputeLevel(t *testing.T) {
 	db := setupDB(t)
-	setDefaultExponents()
+	setDefaultStateDiffExponents()
 
 	err := setOffsetInDB(db, 0)
 	require.NoError(t, err)
@@ -118,7 +118,7 @@ func TestStateDiff_ComputeLevel(t *testing.T) {
 }
 
 func TestStateDiff_SaveFullSnapshot(t *testing.T) {
-	setDefaultExponents()
+	setDefaultStateDiffExponents()
 
 	for v := range version.All() {
 		t.Run(version.String(v), func(t *testing.T) {
@@ -151,7 +151,7 @@ func TestStateDiff_SaveFullSnapshot(t *testing.T) {
 }
 
 func TestStateDiff_SaveAndReadFullSnapshot(t *testing.T) {
-	setDefaultExponents()
+	setDefaultStateDiffExponents()
 
 	for v := range version.All() {
 		t.Run(version.String(v), func(t *testing.T) {
@@ -179,7 +179,7 @@ func TestStateDiff_SaveAndReadFullSnapshot(t *testing.T) {
 }
 
 func TestStateDiff_SaveDiff(t *testing.T) {
-	setDefaultExponents()
+	setDefaultStateDiffExponents()
 
 	for v := range version.All() {
 		t.Run(version.String(v), func(t *testing.T) {
@@ -245,7 +245,7 @@ func TestStateDiff_SaveDiff(t *testing.T) {
 }
 
 func TestStateDiff_SaveAndReadDiff(t *testing.T) {
-	setDefaultExponents()
+	setDefaultStateDiffExponents()
 
 	for v := range version.All() {
 		t.Run(version.String(v), func(t *testing.T) {
@@ -279,7 +279,7 @@ func TestStateDiff_SaveAndReadDiff(t *testing.T) {
 }
 
 func TestStateDiff_SaveAndReadDiff_MultipleLevels(t *testing.T) {
-	setDefaultExponents()
+	setDefaultStateDiffExponents()
 
 	for v := range version.All() {
 		t.Run(version.String(v), func(t *testing.T) {
@@ -345,7 +345,7 @@ func TestStateDiff_SaveAndReadDiff_MultipleLevels(t *testing.T) {
 }
 
 func TestStateDiff_SaveAndReadDiffForkTransition(t *testing.T) {
-	setDefaultExponents()
+	setDefaultStateDiffExponents()
 
 	for v := range version.All()[:len(version.All())-1] {
 		t.Run(version.String(v), func(t *testing.T) {
@@ -379,7 +379,7 @@ func TestStateDiff_SaveAndReadDiffForkTransition(t *testing.T) {
 }
 
 func TestStateDiff_OffsetCache(t *testing.T) {
-	setDefaultExponents()
+	setDefaultStateDiffExponents()
 
 	// test for slot numbers 0 and 1 for every version
 	for slotNum := range 2 {
@@ -410,7 +410,7 @@ func TestStateDiff_OffsetCache(t *testing.T) {
 }
 
 func TestStateDiff_AnchorCache(t *testing.T) {
-	setDefaultExponents()
+	setDefaultStateDiffExponents()
 
 	for v := range version.All() {
 		t.Run(version.String(v), func(t *testing.T) {
@@ -614,7 +614,7 @@ func setOffsetInDB(s *Store, offset uint64) error {
 	return nil
 }
 
-func setDefaultExponents() {
+func setDefaultStateDiffExponents() {
 	globalFlags := flags.GlobalFlags{
 		StateDiffExponents: []int{21, 18, 16, 13, 11, 9, 5},
 	}