From e397f8a2bdc6154fb78c3cac5ad68a5020e8db07 Mon Sep 17 00:00:00 2001
From: terence <terence@prysmaticlabs.com>
Date: Wed, 24 Jan 2024 09:22:50 -0800
Subject: [PATCH] Skip origin root when cleaning dirty state (#13521)

* Skip origin root when cleaning dirty state

* Clean up
---
 beacon-chain/db/kv/state.go      | 32 +++++++++++++++++++++++++++-----
 beacon-chain/db/kv/state_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/beacon-chain/db/kv/state.go b/beacon-chain/db/kv/state.go
index 33070f27e0bb..db408d4318af 100644
--- a/beacon-chain/db/kv/state.go
+++ b/beacon-chain/db/kv/state.go
@@ -893,6 +893,7 @@ func createStateIndicesFromStateSlot(ctx context.Context, slot primitives.Slot)
 //
 // 3.) state with current finalized root
 // 4.) unfinalized States
+// 5.) not origin root
 func (s *Store) CleanUpDirtyStates(ctx context.Context, slotsPerArchivedPoint primitives.Slot) error {
 	ctx, span := trace.StartSpan(ctx, "BeaconDB. CleanUpDirtyStates")
 	defer span.End()
@@ -907,6 +908,11 @@ func (s *Store) CleanUpDirtyStates(ctx context.Context, slotsPerArchivedPoint pr
 	}
 	deletedRoots := make([][32]byte, 0)
 
+	oRoot, err := s.OriginCheckpointBlockRoot(ctx)
+	if err != nil {
+		return err
+	}
+
 	err = s.db.View(func(tx *bolt.Tx) error {
 		bkt := tx.Bucket(stateSlotIndicesBucket)
 		return bkt.ForEach(func(k, v []byte) error {
@@ -914,15 +920,31 @@ func (s *Store) CleanUpDirtyStates(ctx context.Context, slotsPerArchivedPoint pr
 				return ctx.Err()
 			}
 
-			finalizedChkpt := bytesutil.ToBytes32(f.Root) == bytesutil.ToBytes32(v)
+			root := bytesutil.ToBytes32(v)
 			slot := bytesutil.BytesToSlotBigEndian(k)
 			mod := slot % slotsPerArchivedPoint
-			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))
+			if mod == 0 {
+				return nil
 			}
+
+			if mod > slotsPerArchivedPoint-slotsPerArchivedPoint/3 {
+				return nil
+			}
+
+			if bytesutil.ToBytes32(f.Root) == root {
+				return nil
+			}
+
+			if slot > finalizedSlot {
+				return nil
+			}
+
+			if oRoot == root {
+				return nil
+			}
+
+			deletedRoots = append(deletedRoots, root)
 			return nil
 		})
 	})
diff --git a/beacon-chain/db/kv/state_test.go b/beacon-chain/db/kv/state_test.go
index f5b04c1a655f..132c6897942e 100644
--- a/beacon-chain/db/kv/state_test.go
+++ b/beacon-chain/db/kv/state_test.go
@@ -675,6 +675,7 @@ func TestStore_CleanUpDirtyStates_AboveThreshold(t *testing.T) {
 	genesisRoot := [32]byte{'a'}
 	require.NoError(t, db.SaveGenesisBlockRoot(context.Background(), genesisRoot))
 	require.NoError(t, db.SaveState(context.Background(), genesisState, genesisRoot))
+	require.NoError(t, db.SaveOriginCheckpointBlockRoot(context.Background(), [32]byte{'a'}))
 
 	bRoots := make([][32]byte, 0)
 	slotsPerArchivedPoint := primitives.Slot(128)
@@ -720,6 +721,7 @@ func TestStore_CleanUpDirtyStates_Finalized(t *testing.T) {
 	genesisRoot := [32]byte{'a'}
 	require.NoError(t, db.SaveGenesisBlockRoot(context.Background(), genesisRoot))
 	require.NoError(t, db.SaveState(context.Background(), genesisState, genesisRoot))
+	require.NoError(t, db.SaveOriginCheckpointBlockRoot(context.Background(), [32]byte{'a'}))
 
 	for i := primitives.Slot(1); i <= params.BeaconConfig().SlotsPerEpoch; i++ {
 		b := util.NewBeaconBlock()
@@ -741,6 +743,35 @@ func TestStore_CleanUpDirtyStates_Finalized(t *testing.T) {
 	require.Equal(t, true, db.HasState(context.Background(), genesisRoot))
 }
 
+func TestStore_CleanUpDirtyStates_OriginRoot(t *testing.T) {
+	db := setupDB(t)
+
+	genesisState, err := util.NewBeaconState()
+	require.NoError(t, err)
+	r := [32]byte{'a'}
+	require.NoError(t, db.SaveGenesisBlockRoot(context.Background(), r))
+	require.NoError(t, db.SaveState(context.Background(), genesisState, r))
+
+	for i := primitives.Slot(1); i <= params.BeaconConfig().SlotsPerEpoch; i++ {
+		b := util.NewBeaconBlock()
+		b.Block.Slot = i
+		r, err := b.Block.HashTreeRoot()
+		require.NoError(t, err)
+		wsb, err := blocks.NewSignedBeaconBlock(b)
+		require.NoError(t, err)
+		require.NoError(t, db.SaveBlock(context.Background(), wsb))
+
+		st, err := util.NewBeaconState()
+		require.NoError(t, err)
+		require.NoError(t, st.SetSlot(i))
+		require.NoError(t, db.SaveState(context.Background(), st, r))
+	}
+
+	require.NoError(t, db.SaveOriginCheckpointBlockRoot(context.Background(), r))
+	require.NoError(t, db.CleanUpDirtyStates(context.Background(), params.BeaconConfig().SlotsPerEpoch))
+	require.Equal(t, true, db.HasState(context.Background(), r))
+}
+
 func TestStore_CleanUpDirtyStates_DontDeleteNonFinalized(t *testing.T) {
 	db := setupDB(t)
 
@@ -749,6 +780,7 @@ func TestStore_CleanUpDirtyStates_DontDeleteNonFinalized(t *testing.T) {
 	genesisRoot := [32]byte{'a'}
 	require.NoError(t, db.SaveGenesisBlockRoot(context.Background(), genesisRoot))
 	require.NoError(t, db.SaveState(context.Background(), genesisState, genesisRoot))
+	require.NoError(t, db.SaveOriginCheckpointBlockRoot(context.Background(), [32]byte{'a'}))
 
 	var unfinalizedRoots [][32]byte
 	for i := primitives.Slot(1); i <= params.BeaconConfig().SlotsPerEpoch; i++ {