feat(store): rollback --store flag to to include block store in rollback

cmd/tenderdash/commands/rollback.go

@@ -10,7 +10,7 @@ import (
 )
 
 func MakeRollbackStateCommand(conf *config.Config) *cobra.Command {
-	return &cobra.Command{
+	cmd := &cobra.Command{
 		Use:   "rollback",
 		Short: "rollback tendermint state by one height",
 		Long: `
@@ -20,9 +20,12 @@ progress. Rollback overwrites a state at height n with the state at height n - 1
 The application should also roll back to height n - 1. No blocks are removed, so upon
 restarting Tendermint the transactions in block n will be re-executed against the
 application.
+
+If the --store flag is set, the block store will also be rolled back to match the state height.
 `,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			height, hash, err := RollbackState(conf)
+			storeRollback, _ := cmd.Flags().GetBool("store")
+			height, hash, err := RollbackState(conf, storeRollback)
 			if err != nil {
 				return fmt.Errorf("failed to rollback state: %w", err)
 			}
@@ -32,12 +35,14 @@ application.
 		},
 	}
 
+	cmd.Flags().Bool("store", false, "also roll back the block store to match the state height")
+	return cmd
 }
 
 // RollbackState takes the state at the current height n and overwrites it with the state
 // at height n - 1. Note state here refers to tendermint state not application state.
 // Returns the latest state height and app hash alongside an error if there was one.
-func RollbackState(config *config.Config) (int64, []byte, error) {
+func RollbackState(config *config.Config, rollbackStore bool) (int64, []byte, error) {
 	// use the parsed config to load the block and state store
 	blockStore, stateStore, err := loadStateAndBlockStore(config)
 	if err != nil {
@@ -49,5 +54,20 @@ func RollbackState(config *config.Config) (int64, []byte, error) {
 	}()
 
 	// rollback the last state
-	return state.Rollback(blockStore, stateStore)
+	height, hash, err := state.Rollback(blockStore, stateStore)
+	if err != nil {
+		return -1, nil, err
+	}
+
+	if rollbackStore {
+		// Rollback the block store to match the state height
+		for currentHeight := blockStore.Height(); currentHeight > height; currentHeight-- {
+			_, err := blockStore.DeleteBlock(currentHeight)
+			if err != nil {
+				return -1, nil, fmt.Errorf("failed to delete block at height %d: %w", currentHeight, err)
+			}
+		}
+	}
+
+	return height, hash, nil
 }

cmd/tenderdash/commands/rollback_test.go

@@ -48,7 +48,7 @@ func TestRollbackIntegration(t *testing.T) {
 	t.Run("Rollback", func(t *testing.T) {
 		time.Sleep(time.Second)
 		require.NoError(t, app.Rollback())
-		height, _, err = commands.RollbackState(cfg)
+		height, _, err = commands.RollbackState(cfg, false)
 		require.NoError(t, err, "%d", height)
 	})
 	t.Run("Restart", func(t *testing.T) {

internal/store/store.go

@@ -358,6 +358,57 @@ func (bs *BlockStore) PruneBlocks(height int64) (uint64, error) {
 	return pruned, nil
 }
 
+// DeleteBlock removes the block at the given height, including its parts and commit.
+// It returns the number of blocks removed, which is always 1 in this case.
+func (bs *BlockStore) DeleteBlock(height int64) (uint64, error) {
+	if height <= 0 {
+		return 0, fmt.Errorf("height must be greater than 0")
+	}
+
+	// Check if block exists at this height
+	blockMeta := bs.LoadBlockMeta(height)
+	if blockMeta == nil {
+		return 0, fmt.Errorf("block at height %d does not exist", height)
+	}
+
+	batch := bs.db.NewBatch()
+	defer batch.Close()
+
+	// Delete the hash key corresponding to the block meta's hash
+	if err := batch.Delete(blockHashKey(blockMeta.BlockID.Hash)); err != nil {
+		return 0, fmt.Errorf("failed to delete hash key: %X: %w", blockHashKey(blockMeta.BlockID.Hash), err)
+	}
+
+	// Delete block meta
+	if err := batch.Delete(blockMetaKey(height)); err != nil {
+		return 0, fmt.Errorf("failed to delete block meta at height %d: %w", height, err)
+	}
+
+	// Delete block parts
+	for i := 0; i < int(blockMeta.BlockID.PartSetHeader.Total); i++ {
+		if err := batch.Delete(blockPartKey(height, i)); err != nil {
+			return 0, fmt.Errorf("failed to delete block part at height %d, index %d: %w", height, i, err)
+		}
+	}
+
+	// Delete block commit
+	if err := batch.Delete(blockCommitKey(height)); err != nil {
+		return 0, fmt.Errorf("failed to delete block commit at height %d: %w", height, err)
+	}
+
+	// Delete seen commit at this height if it exists
+	if err := batch.Delete(seenCommitAtKey(height)); err != nil {
+		return 0, fmt.Errorf("failed to delete seen commit at height %d: %w", height, err)
+	}
+
+	// Write all deletions atomically
+	if err := batch.WriteSync(); err != nil {
+		return 0, fmt.Errorf("failed to write deletions for height %d: %w", height, err)
+	}
+
+	return 1, nil
+}
+
 // pruneRange is a generic function for deleting a range of values based on the lowest
 // height up to but excluding retainHeight. For each key/value pair, an optional hook can be
 // executed before the deletion itself is made. pruneRange will use batch delete to delete