Update delete state(s) functions (#7754)

beacon-chain/db/kv/state.go

@@ -121,23 +121,6 @@ func (s *Store) HasState(ctx context.Context, blockRoot [32]byte) bool {
 func (s *Store) DeleteState(ctx context.Context, blockRoot [32]byte) error {
 	ctx, span := trace.StartSpan(ctx, "BeaconDB.DeleteState")
 	defer span.End()
-	return s.DeleteStates(ctx, [][32]byte{blockRoot})
-}
-
-// DeleteStates by block roots.
-//
-// Note: bkt.Delete(key) uses a binary search to find the item in the database. Iterating with a
-// cursor is faster when there are a large set of keys to delete. This method is O(n) deletion where
-// n is the number of keys in the database. The alternative of calling  bkt.Delete on each key to
-// delete would be O(m*log(n)) which would be much slower given a large set of keys to delete.
-func (s *Store) DeleteStates(ctx context.Context, blockRoots [][32]byte) error {
-	ctx, span := trace.StartSpan(ctx, "BeaconDB.DeleteStates")
-	defer span.End()
-
-	rootMap := make(map[[32]byte]bool, len(blockRoots))
-	for _, blockRoot := range blockRoots {
-		rootMap[blockRoot] = true
-	}
 
 	return s.db.Update(func(tx *bolt.Tx) error {
 		bkt := tx.Bucket(blocksBucket)
@@ -155,32 +138,36 @@ func (s *Store) DeleteStates(ctx context.Context, blockRoots [][32]byte) error {
 		blockBkt := tx.Bucket(blocksBucket)
 		headBlkRoot := blockBkt.Get(headBlockRootKey)
 		bkt = tx.Bucket(stateBucket)
-		c := bkt.Cursor()
+		// Safe guard against deleting genesis, finalized, head state.
+		if bytes.Equal(blockRoot[:], checkpoint.Root) || bytes.Equal(blockRoot[:], genesisBlockRoot) || bytes.Equal(blockRoot[:], headBlkRoot) {
+			return errors.New("cannot delete genesis, finalized, or head state")
+		}
 
-		for blockRoot, _ := c.First(); blockRoot != nil; blockRoot, _ = c.Next() {
-			if !rootMap[bytesutil.ToBytes32(blockRoot)] {
-				continue
-			}
-			// Safe guard against deleting genesis, finalized, head state.
-			if bytes.Equal(blockRoot, checkpoint.Root) || bytes.Equal(blockRoot, genesisBlockRoot) || bytes.Equal(blockRoot, headBlkRoot) {
-				return errors.New("cannot delete genesis, finalized, or head state")
-			}
+		slot, err := slotByBlockRoot(ctx, tx, blockRoot[:])
+		if err != nil {
+			return err
+		}
+		indicesByBucket := createStateIndicesFromStateSlot(ctx, slot)
+		if err := deleteValueForIndices(ctx, indicesByBucket, blockRoot[:], tx); err != nil {
+			return errors.Wrap(err, "could not delete root for DB indices")
+		}
 
-			slot, err := slotByBlockRoot(ctx, tx, blockRoot)
-			if err != nil {
-				return err
-			}
-			indicesByBucket := createStateIndicesFromStateSlot(ctx, slot)
-			if err := deleteValueForIndices(ctx, indicesByBucket, blockRoot, tx); err != nil {
-				return errors.Wrap(err, "could not delete root for DB indices")
-			}
+		return bkt.Delete(blockRoot[:])
+	})
+}
 
-			if err := c.Delete(); err != nil {
-				return err
-			}
+// DeleteStates by block roots.
+func (s *Store) DeleteStates(ctx context.Context, blockRoots [][32]byte) error {
+	ctx, span := trace.StartSpan(ctx, "BeaconDB.DeleteStates")
+	defer span.End()
+
+	for _, r := range blockRoots {
+		if err := s.DeleteState(ctx, r); err != nil {
+			return err
 		}
-		return nil
-	})
+	}
+
+	return nil
 }
 
 // creates state from marshaled proto state bytes.
@@ -366,6 +353,11 @@ func (s *Store) CleanUpDirtyStates(ctx context.Context, slotsPerArchivedPoint ui
 		return err
 	}
 
+	// Length of to be deleted roots is 0. Nothing to do.
+	if len(deletedRoots) == 0 {
+		return nil
+	}
+
 	log.WithField("count", len(deletedRoots)).Info("Cleaning up dirty states")
 	if err := s.DeleteStates(ctx, deletedRoots); err != nil {
 		return err

beacon-chain/state/stategen/setter_test.go

@@ -182,7 +182,11 @@ func TestEnableSaveHotStateToDB_Disabled(t *testing.T) {
 	db, _ := testDB.SetupDB(t)
 	service := New(db, cache.NewStateSummaryCache())
 	service.saveHotStateDB.enabled = true
-	service.saveHotStateDB.savedStateRoots = [][32]byte{{'a'}}
+	b := testutil.NewBeaconBlock()
+	require.NoError(t, db.SaveBlock(ctx, b))
+	r, err := b.Block.HashTreeRoot()
+	require.NoError(t, err)
+	service.saveHotStateDB.savedStateRoots = [][32]byte{r}
 	require.NoError(t, service.DisableSaveHotStateToDB(ctx))
 	require.LogsContain(t, hook, "Exiting mode to save hot states in DB")
 	require.Equal(t, false, service.saveHotStateDB.enabled)