go/oasis-node/cmd/storage: Add command for offline pruning (POC)

When enabling aggresive pruning node may fall behind.

To prevent it, we should offer validators a maintenance command
that should be called if pruning is enabled later on, before
starting the node.

Author: martintomazic

go/consensus/cometbft/abci/prune.go

@@ -195,6 +195,8 @@ PruneLoop:
 	return nil
 }
 
+// Warning: When registering new handler DO NOT forget to update the logic for
+// "oasis-node storage prune" command as well.
 func (p *genericPruner) RegisterHandler(handler consensus.StatePruneHandler) {
 	p.Lock()
 	defer p.Unlock()

go/oasis-node/cmd/storage/storage.go

@@ -6,11 +6,13 @@ import (
 	"errors"
 	"fmt"
 	"io/fs"
+	"math"
 	"os"
 	"path/filepath"
 	"strings"
 	"time"
 
+	cmtBlockstore "github.com/cometbft/cometbft/store"
 	badgerDB "github.com/dgraph-io/badger/v4"
 	"github.com/spf13/cobra"
 
@@ -20,6 +22,7 @@ import (
 	"github.com/oasisprotocol/oasis-core/go/config"
 	"github.com/oasisprotocol/oasis-core/go/consensus/cometbft/abci"
 	cmtCommon "github.com/oasisprotocol/oasis-core/go/consensus/cometbft/common"
+	cmtConfig "github.com/oasisprotocol/oasis-core/go/consensus/cometbft/config"
 	cometbftBadger "github.com/oasisprotocol/oasis-core/go/consensus/cometbft/db/badger"
 	cmdCommon "github.com/oasisprotocol/oasis-core/go/oasis-node/cmd/common"
 	roothash "github.com/oasisprotocol/oasis-core/go/roothash/api"
@@ -70,6 +73,13 @@ WARNING: Ensure you have at least as much of a free disk as your largest databas
 		RunE: doDBCompactions,
 	}
 
+	pruneCmd = &cobra.Command{
+		Use:   "prune-experimental",
+		Args:  cobra.NoArgs,
+		Short: "EXPERIMENTAL: trigger pruning for all consensus databases",
+		RunE:  doPrune,
+	}
+
 	logger = logging.GetLogger("cmd/storage")
 
 	pretty = cmdCommon.Isatty(1)
@@ -395,7 +405,17 @@ func flattenBadgerDB(db *badgerDB.DB, logger *logging.Logger, path string) error
 }
 
 func compactConsensusNodeDB(dataDir string) error {
-	ldb, ndb, _, err := abci.InitStateStorage(
+	ndb, err := openConsensusStateNodeDB(dataDir)
+	if err != nil {
+		return fmt.Errorf("failed to initialize ABCI storage backend: %w", err)
+	}
+	defer ndb.Close()
+
+	return ndb.Compact()
+}
+
+func openConsensusStateNodeDB(dataDir string) (db.NodeDB, error) {
+	_, ndb, _, err := abci.InitStateStorage(
 		&abci.ApplicationConfig{
 			DataDir:             filepath.Join(dataDir, cmtCommon.StateDir),
 			StorageBackend:      config.GlobalConfig.Storage.Backend,
@@ -405,16 +425,163 @@ func compactConsensusNodeDB(dataDir string) error {
 			// ChainContext:     doc.ChainContext(), TODO: Should we read this from the doc?
 		},
 	)
-	if err != nil {
-		return fmt.Errorf("failed to initialize ABCI storage backend: %w", err)
+
+	return ndb, err
+}
+
+func doPrune(_ *cobra.Command, args []string) error {
+	if err := cmdCommon.Init(); err != nil {
+		cmdCommon.EarlyLogAndExit(err)
+	}
+
+	// TODO consider validating correct mode?
+
+	dataDir := cmdCommon.DataDir()
+	if err := pruneConsensusDBs(dataDir); err != nil {
+		return fmt.Errorf("failed to prune consensus databases: %w", err)
 	}
 
-	// Close the resources. Both Close and Cleanup only close NodeDB.
-	// Closing both here, to prevent resource leaks in things change in the future.
+	return nil
+}
+
+func pruneConsensusDBs(dataDir string) error {
+	if config.GlobalConfig.Consensus.Prune.Strategy == cmtConfig.PruneStrategyNone {
+		logger.Info("skipping consensus pruning: (strategy=%s)", cmtConfig.PruneStrategyNone)
+		return nil
+	}
+
+	ndb, err := openConsensusStateNodeDB(dataDir)
+	if err != nil {
+		return fmt.Errorf("failed to open NodeDB: %w", err)
+	}
 	defer ndb.Close()
-	defer ldb.Cleanup()
 
-	return ndb.Compact()
+	latest, ok := ndb.GetLatestVersion()
+	if !ok {
+		logger.Info("skipping consensus pruning as state db is empty")
+		return nil
+	}
+
+	earliest, err := pruneConsensusState(dataDir, ndb, latest)
+	if err != nil {
+		return fmt.Errorf("failed to prune application state: %w", err)
+	}
+
+	if err := pruneCometDBs(dataDir, int64(earliest)); err != nil {
+		return fmt.Errorf("failed to prune CometBFT managed databases: %w", err)
+	}
+
+	return nil
+}
+
+func pruneConsensusState(dataDir string, ndb db.NodeDB, latest uint64) (uint64, error) {
+	if latest < config.GlobalConfig.Consensus.Prune.NumKept {
+		logger.Info("consensus state pruning skipped: latest version is smaller than the number of versions to keep")
+		return latest, nil
+	}
+
+	// In case of configured runtimes, we should not prune past the latest reindexed
+	// consensus height, so that light history can be populated correctly.
+	minReindexed, err := minReindexedHeight(dataDir)
+	if err != nil {
+		return 0, fmt.Errorf("failed to fetch minimum reindexed consensus height: %w", err)
+	}
+
+	start := ndb.GetEarliestVersion()
+	end := min(
+		latest-config.GlobalConfig.Consensus.Prune.NumKept, // does not underflow due to if at the top.
+		uint64(minReindexed),
+	)
+
+	if end <= start {
+		logger.Info("consensus state already pruned")
+		return end, nil
+	}
+
+	logger.Info("pruning consensus state", "start", start, "end", end)
+	for i := start; i < end; i++ {
+		if err := ndb.Prune(i); err != nil {
+			return 0, fmt.Errorf("failed to prune version %d: %w", i, err)
+		}
+
+		if i%10_000 == 0 { // TODO not sure this is even needed.
+			if err := ndb.Sync(); err != nil {
+				return 0, fmt.Errorf("failed to sync NodeDB: %w", err)
+			}
+			logger.Debug("forcing NodeDB disk sync during pruning", "version", i)
+		}
+	}
+
+	if err := ndb.Sync(); err != nil {
+		return 0, fmt.Errorf("failed to sync NodeDB: %w", err)
+	}
+
+	return end, nil
+}
+
+// minReindexedHeight returns the smallest consensus height reindexed by any
+// of the configured runtimes.
+//
+// In case of no configured runtimes it returns max int64.
+func minReindexedHeight(dataDir string) (int64, error) {
+	fetchLastReindexedHeight := func(runtimeID common.Namespace) (int64, error) {
+		rtDir := runtimeConfig.GetRuntimeStateDir(dataDir, runtimeID)
+		mode := config.GlobalConfig.Mode
+		hasLocalStorage := mode.HasLocalStorage() && !mode.IsArchive()
+
+		// TODO ideally we would not start whole light history with all background workers, but this would
+		// require as to refactor existing code...
+		history, err := history.New(runtimeID, rtDir, history.NewNonePrunerFactory(), hasLocalStorage)
+		if err != nil {
+			return 0, fmt.Errorf("failed to open new light history: %w", err)
+		}
+		defer history.Close()
+
+		h, err := history.LastConsensusHeight()
+		if err != nil {
+			return 0, fmt.Errorf("failed to get last consensus height: %w", err)
+		}
+
+		return h, nil
+	}
+
+	var minH int64 = math.MaxInt64
+	for _, rt := range config.GlobalConfig.Runtime.Runtimes {
+		h, err := fetchLastReindexedHeight(rt.ID)
+		if err != nil {
+			return 0, fmt.Errorf("failed to fetch last reindexed height for %s: %w", rt.ID, err)
+		}
+
+		if h < minH {
+			minH = h
+		}
+	}
+
+	return minH, nil
+}
+
+func pruneCometDBs(dataDir string, height int64) error {
+	// TODO: This is a hack. In fact even if we manage to get this right via
+	// BadgerDBProvider and somehow pass correct config via context, this will
+	// still not be intended way to use it. I believe this hack is worth it, but
+	// we should definitely release this command as experimental first.
+	blockstorePath := fmt.Sprintf("%s/consensus/data/blockstore.badger.db", dataDir)
+	blockDB, err := cometbftBadger.New(blockstorePath, false)
+	if err != nil {
+		return fmt.Errorf("failed to open blockstore: %w", err)
+	}
+	blockstore := cmtBlockstore.NewBlockStore(blockDB)
+
+	logger.Info("pruning consensus blockstore", "target_height", height)
+	n, err := blockstore.PruneBlocks(height)
+	if err != nil {
+		return fmt.Errorf("failed to prune blocks: %w", err)
+	}
+	logger.Info("consensus blockstore finished", "pruned", n)
+
+	// TODO add pruning of state.badger.db
+
+	return nil
 }
 
 // Register registers the client sub-command and all of its children.
@@ -425,5 +592,6 @@ func Register(parentCmd *cobra.Command) {
 	storageCmd.AddCommand(storageCheckCmd)
 	storageCmd.AddCommand(storageRenameNsCmd)
 	storageCmd.AddCommand(storageCompactCmd)
+	storageCmd.AddCommand(pruneCmd)
 	parentCmd.AddCommand(storageCmd)
 }