diff --git a/beacon-chain/db/kv/state.go b/beacon-chain/db/kv/state.go index d9e83e77ef..ce073a2ddd 100644 --- a/beacon-chain/db/kv/state.go +++ b/beacon-chain/db/kv/state.go @@ -9,6 +9,7 @@ import ( ethpb "github.com/prysmaticlabs/ethereumapis/eth/v1alpha1" "github.com/prysmaticlabs/prysm/beacon-chain/state" pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1" + "github.com/prysmaticlabs/prysm/shared/bytesutil" "go.opencensus.io/trace" ) @@ -198,6 +199,11 @@ func (k *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) + for _, blockRoot := range blockRoots { + rootMap[blockRoot] = true + } + return k.db.Update(func(tx *bolt.Tx) error { bkt := tx.Bucket(blocksBucket) genesisBlockRoot := bkt.Get(genesisBlockRootKey) @@ -213,16 +219,18 @@ func (k *Store) DeleteStates(ctx context.Context, blockRoots [][32]byte) error { bkt = tx.Bucket(blocksBucket) headBlkRoot := bkt.Get(headBlockRootKey) + bkt = tx.Bucket(stateBucket) + c := bkt.Cursor() - for _, blockRoot := range blockRoots { - // Safe guard against deleting genesis, finalized, or head state. - if bytes.Equal(blockRoot[:], checkpoint.Root) || bytes.Equal(blockRoot[:], genesisBlockRoot) || bytes.Equal(blockRoot[:], headBlkRoot) { - return errors.New("could not delete genesis, finalized, or head state") - } - - bkt = tx.Bucket(stateBucket) - if err := bkt.Delete(blockRoot[:]); err != nil { - return err + for blockRoot, _ := c.First(); blockRoot != nil; blockRoot, _ = c.Next() { + if rootMap[bytesutil.ToBytes32(blockRoot)] { + // Safe guard against deleting genesis, finalized, or head state. + if bytes.Equal(blockRoot[:], checkpoint.Root) || bytes.Equal(blockRoot[:], genesisBlockRoot) || bytes.Equal(blockRoot[:], headBlkRoot) { + return errors.New("could not delete genesis, finalized, or head state") + } + if err := c.Delete(); err != nil { + return err + } } } return nil