Keep Non Finalized States (#7742)

* keep non finalized states

* Update beacon-chain/db/kv/state.go

Co-authored-by: terence tsao <terence@prysmaticlabs.com>

beacon-chain/db/kv/state.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/pkg/errors"
 	ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1"
+	"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"
@@ -327,6 +328,7 @@ func createStateIndicesFromStateSlot(ctx context.Context, slot uint64) map[strin
 //   (e.g. archived_interval=2048, states with slots after 1365).
 //   This is to tolerate skip slots. Not every state lays on the boundary.
 // 3.) state with current finalized root
+// 4.) unfinalized States
 func (s *Store) CleanUpDirtyStates(ctx context.Context, slotsPerArchivedPoint uint64) error {
 	ctx, span := trace.StartSpan(ctx, "BeaconDB. CleanUpDirtyStates")
 	defer span.End()
@@ -335,6 +337,10 @@ func (s *Store) CleanUpDirtyStates(ctx context.Context, slotsPerArchivedPoint ui
 	if err != nil {
 		return err
 	}
+	finalizedSlot, err := helpers.StartSlot(f.Epoch)
+	if err != nil {
+		return err
+	}
 	deletedRoots := make([][32]byte, 0)
 
 	err = s.db.View(func(tx *bolt.Tx) error {
@@ -344,11 +350,13 @@ func (s *Store) CleanUpDirtyStates(ctx context.Context, slotsPerArchivedPoint ui
 				return ctx.Err()
 			}
 
-			finalized := bytesutil.ToBytes32(f.Root) == bytesutil.ToBytes32(v)
+			finalizedChkpt := bytesutil.ToBytes32(f.Root) == bytesutil.ToBytes32(v)
 			slot := bytesutil.BytesToUint64BigEndian(k)
 			mod := slot % slotsPerArchivedPoint
-			// The following conditions cover 1, 2, and 3 above.
-			if mod != 0 && mod <= slotsPerArchivedPoint-slotsPerArchivedPoint/3 && !finalized {
+			nonFinalized := slot > finalizedSlot
+
+			// The following conditions cover 1, 2, 3 and 4 above.
+			if mod != 0 && mod <= slotsPerArchivedPoint-slotsPerArchivedPoint/3 && !finalizedChkpt && !nonFinalized {
 				deletedRoots = append(deletedRoots, bytesutil.ToBytes32(v))
 			}
 			return nil

beacon-chain/db/kv/state_test.go

@@ -231,21 +231,30 @@ func TestStore_GenesisState_CanGetHighestBelow(t *testing.T) {
 func TestStore_CleanUpDirtyStates_AboveThreshold(t *testing.T) {
 	db := setupDB(t)
 
+	genesisState := testutil.NewBeaconState()
+	genesisRoot := [32]byte{'a'}
+	require.NoError(t, db.SaveGenesisBlockRoot(context.Background(), genesisRoot))
+	require.NoError(t, db.SaveState(context.Background(), genesisState, genesisRoot))
+
 	bRoots := make([][32]byte, 0)
 	slotsPerArchivedPoint := uint64(128)
+	prevRoot := genesisRoot
 	for i := uint64(1); i <= slotsPerArchivedPoint; i++ {
 		b := testutil.NewBeaconBlock()
 		b.Block.Slot = i
+		b.Block.ParentRoot = prevRoot[:]
 		r, err := b.Block.HashTreeRoot()
 		require.NoError(t, err)
 		require.NoError(t, db.SaveBlock(context.Background(), b))
 		bRoots = append(bRoots, r)
+		prevRoot = r
 
 		st := testutil.NewBeaconState()
 		require.NoError(t, st.SetSlot(i))
 		require.NoError(t, db.SaveState(context.Background(), st, r))
 	}
 
+	require.NoError(t, db.SaveFinalizedCheckpoint(context.Background(), &ethpb.Checkpoint{Root: bRoots[len(bRoots)-1][:], Epoch: slotsPerArchivedPoint / params.BeaconConfig().SlotsPerEpoch}))
 	require.NoError(t, db.CleanUpDirtyStates(context.Background(), slotsPerArchivedPoint))
 
 	for i, root := range bRoots {
@@ -281,3 +290,33 @@ func TestStore_CleanUpDirtyStates_Finalized(t *testing.T) {
 	require.NoError(t, db.CleanUpDirtyStates(context.Background(), params.BeaconConfig().SlotsPerEpoch))
 	require.Equal(t, true, db.HasState(context.Background(), genesisRoot))
 }
+
+func TestStore_CleanUpDirtyStates_DontDeleteNonFinalized(t *testing.T) {
+	db := setupDB(t)
+
+	genesisState := testutil.NewBeaconState()
+	genesisRoot := [32]byte{'a'}
+	require.NoError(t, db.SaveGenesisBlockRoot(context.Background(), genesisRoot))
+	require.NoError(t, db.SaveState(context.Background(), genesisState, genesisRoot))
+
+	unfinalizedRoots := [][32]byte{}
+	for i := uint64(1); i <= params.BeaconConfig().SlotsPerEpoch; i++ {
+		b := testutil.NewBeaconBlock()
+		b.Block.Slot = i
+		r, err := b.Block.HashTreeRoot()
+		require.NoError(t, err)
+		require.NoError(t, db.SaveBlock(context.Background(), b))
+		unfinalizedRoots = append(unfinalizedRoots, r)
+
+		st := testutil.NewBeaconState()
+		require.NoError(t, st.SetSlot(i))
+		require.NoError(t, db.SaveState(context.Background(), st, r))
+	}
+
+	require.NoError(t, db.SaveFinalizedCheckpoint(context.Background(), &ethpb.Checkpoint{Root: genesisRoot[:]}))
+	require.NoError(t, db.CleanUpDirtyStates(context.Background(), params.BeaconConfig().SlotsPerEpoch))
+
+	for _, rt := range unfinalizedRoots {
+		require.Equal(t, true, db.HasState(context.Background(), rt))
+	}
+}