Merge pull request #6355 from oasisprotocol/martin/buxfix/runtime-pruning

go: Fix runtime pruning

Author: martintomazic

.changelog/6352.bugfix.md

@@ -0,0 +1,7 @@
+Fix runtime state pruning performance
+
+Previously, pruning the runtime state with many local versions could
+significantly slow down state synchronization.
+
+This has been fixed by moving state pruning out of the BadgerDB
+transaction responsible for pruning the runtime’s light history.

go/consensus/api/api.go

@@ -428,18 +428,15 @@ type StatePruner interface {
 	RegisterHandler(handler StatePruneHandler)
 }
 
-// StatePruneHandler is a handler that is called when heights are pruned
+// StatePruneHandler is a handler that is called before heights are pruned
 // from history.
 type StatePruneHandler interface {
-	// Prune is called before the specified height is pruned.
-	//
-	// If an error is returned, pruning is aborted and the height is
-	// not pruned from history.
+	// CanPruneConsensus is called to check if the specified height can be pruned.
 	//
 	// Note that this can be called for the same height multiple
 	// times (e.g., if one of the handlers fails but others succeed
 	// and pruning is later retried).
-	Prune(height int64) error
+	CanPruneConsensus(height int64) error
 }
 
 // EstimateGasRequest is a EstimateGas request.

go/consensus/cometbft/abci/prune.go

@@ -136,25 +136,20 @@ func (p *genericPruner) Prune(latestVersion uint64) error {
 	)
 
 	preserveFrom := latestVersion - p.keepN
-PruneLoop:
 	for i := p.earliestVersion; i <= latestVersion; i++ {
 		if i >= preserveFrom {
 			p.earliestVersion = i
 			break
 		}
 
-		// Before pruning anything, run all prune handlers. If any of them
-		// fails we abort the prune.
-		for _, ph := range p.handlers {
-			if err := ph.Prune(int64(i)); err != nil {
-				p.logger.Debug("prune handler blocked pruning version",
-					"err", err,
-					"latest_version", latestVersion,
-					"version", i,
-				)
-				p.earliestVersion = i
-				break PruneLoop
-			}
+		if err := p.canPrune(int64(i)); err != nil {
+			p.logger.Debug("prune handler blocked pruning version",
+				"err", err,
+				"latest_version", latestVersion,
+				"version", i,
+			)
+			p.earliestVersion = i
+			break
 		}
 
 		p.logger.Debug("Prune: Delete",
@@ -195,6 +190,16 @@ PruneLoop:
 	return nil
 }
 
+// canPrune checks if all prune handlers allow pruning the given version.
+func (p *genericPruner) canPrune(v int64) error {
+	for _, ph := range p.handlers {
+		if err := ph.CanPruneConsensus(v); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (p *genericPruner) RegisterHandler(handler consensus.StatePruneHandler) {
 	p.Lock()
 	defer p.Unlock()

go/runtime/history/history.go

@@ -287,12 +287,13 @@ func (h *runtimeHistory) GetEarliestBlock(ctx context.Context) (*block.Block, er
 	return annBlk.Block, nil
 }
 
-func (h *runtimeHistory) Prune(height int64) error {
+// CanPruneConsenus returns no error when the specified consensus height has been already reindexed.
+//
+// Implements consensus.api.StatePruneHandler
+func (h *runtimeHistory) CanPruneConsensus(height int64) error {
 	lastHeight, err := h.LastConsensusHeight()
 	if err != nil {
-		h.logger.Warn("failed to fetch last consensus height for tracked runtime",
-			"err", err,
-		)
+		h.logger.Warn("failed to fetch last consensus height for tracked runtime", "err", err)
 		// We can't be sure if it is ok to prune this version, so prevent pruning to be safe.
 		return fmt.Errorf("failed to fetch last consensus height for tracked runtime: %w", err)
 	}
@@ -334,14 +335,10 @@ func (h *runtimeHistory) pruneWorker() {
 				return
 			}
 
-			h.logger.Debug("pruning runtime history",
-				"round", round.(uint64),
-			)
+			h.logger.Debug("pruning runtime history", "round", round.(uint64))
 
 			if err := h.pruner.Prune(round.(uint64)); err != nil {
-				h.logger.Error("failed to prune",
-					"err", err,
-				)
+				h.logger.Debug("failed to prune", "err", err)
 				continue
 			}
 		case <-h.stopCh:

go/runtime/history/history_test.go

@@ -372,7 +372,7 @@ type testPruneHandler struct {
 	batches      []int
 }
 
-func (h *testPruneHandler) Prune(rounds []uint64) error {
+func (h *testPruneHandler) CanPruneRuntime(rounds []uint64) error {
 	// NOTE: Users must ensure that accessing prunedRounds is safe (e.g., that
 	//       no more pruning happens using this handler before prunedRounds is
 	//       accessed from a different goroutine).
@@ -486,7 +486,7 @@ func TestHistoryPrune(t *testing.T) {
 
 type testPruneFailingHandler struct{}
 
-func (h *testPruneFailingHandler) Prune([]uint64) error {
+func (h *testPruneFailingHandler) CanPruneRuntime([]uint64) error {
 	return fmt.Errorf("thou shall not pass")
 }
 

go/runtime/history/prune.go

@@ -24,18 +24,15 @@ const (
 // PrunerFactory is the runtime history pruner factory interface.
 type PrunerFactory func(runtimeID common.Namespace, db *DB) (Pruner, error)
 
-// PruneHandler is a handler that is called when rounds are pruned
+// PruneHandler is a handler that is called before rounds are pruned
 // from history.
 type PruneHandler interface {
-	// Prune is called before the specified rounds are pruned.
-	//
-	// If an error is returned, pruning is aborted and the rounds are
-	// not pruned from history.
+	// CanPruneRuntime is called before to check if the specified round can be pruned.
 	//
 	// Note that this can be called for the same round multiple
 	// times (e.g., if one of the handlers fails but others succeed
 	// and pruning is later retried).
-	Prune(rounds []uint64) error
+	CanPruneRuntime(rounds []uint64) error
 }
 
 // Pruner is the runtime history pruner interface.
@@ -131,32 +128,38 @@ func (p *keepLastPruner) Prune(latestRound uint64) error {
 			pruned = append(pruned, round)
 		}
 
-		// If there is nothing to prune, do not call any handlers.
-		if len(pruned) == 0 {
+		if len(pruned) == 0 { // nothing to prune
 			return nil
 		}
 
-		// Before pruning anything, run all prune handlers. If any of them
-		// fails we abort the prune.
-		p.mu.RLock()
-		defer p.mu.RUnlock()
-
-		for _, ph := range p.handlers {
-			if err := ph.Prune(pruned); err != nil {
-				p.logger.Error("prune handler failed, aborting prune",
-					"err", err,
-					"round_count", len(pruned),
-					"round_min", pruned[0],
-					"round_max", pruned[len(pruned)-1],
-				)
-				return fmt.Errorf("runtime/history: prune handler failed: %w", err)
-			}
+		if err := p.canPrune(pruned); err != nil {
+			p.logger.Debug("prune handler blocked pruning, aborting prune",
+				"err", err,
+				"round_count", len(pruned),
+				"round_min", pruned[0],
+				"round_max", pruned[len(pruned)-1], // safe due to length check above
+			)
+			return fmt.Errorf("prune handler blocked pruning: %w", err)
 		}
 
 		return nil
 	})
 }
 
+// canPrune checks if all prune handlers allow pruning specified rounds.
+func (p *keepLastPruner) canPrune(rounds []uint64) error {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+
+	for _, ph := range p.handlers {
+		if err := ph.CanPruneRuntime(rounds); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // PruneInterval implements Pruner.
 func (p *keepLastPruner) PruneInterval() time.Duration {
 	return p.pruneInterval

go/worker/compute/executor/committee/prune.go

@@ -23,11 +23,14 @@ func NewPruneHandler(runtimeID common.Namespace, consensus consensus.Service) *P
 	}
 }
 
-// Prune verifies that no rounds beyond the last normal round are pruned.
-func (p *PruneHandler) Prune(rounds []uint64) error {
-	// Make sure we never prune past the last normal round, as some runtimes will do historic queries
-	// for things that are not available in the last consensus state (e.g. delegation/undelegation
-	// events that happened while the runtime was suspended or not producing blocks).
+// CanPruneRuntime returns no error when pruning runtime rounds would not go past last normal round.
+//
+// This is important as some runtimes will do historic queries for things that are not available
+// in the last consensus state (e.g. delegation/undelegation events that happened while the runtime
+// was suspended or not producing blocks).
+//
+// Implements runtime.history.PruneHandler.
+func (p *PruneHandler) CanPruneRuntime(rounds []uint64) error {
 	state, err := p.consensus.RootHash().GetRuntimeState(context.TODO(), &roothash.RuntimeRequest{
 		RuntimeID: p.runtimeID,
 		Height:    consensus.HeightLatest,

go/worker/storage/committee/prune.go

@@ -1,9 +1,12 @@
 package committee
 
 import (
+	"context"
 	"fmt"
+	"time"
 
 	"github.com/oasisprotocol/oasis-core/go/common/logging"
+	"github.com/oasisprotocol/oasis-core/go/runtime/history"
 	mkvsDB "github.com/oasisprotocol/oasis-core/go/storage/mkvs/db/api"
 )
 
@@ -12,8 +15,10 @@ type pruneHandler struct {
 	worker *Worker
 }
 
-func (p *pruneHandler) Prune(rounds []uint64) error {
-	// Make sure we never prune past what was synced.
+// CanPruneRuntime returns no error when pruning would not go past last synced round.
+//
+// Implements runtime.history.PruneHandler.
+func (p *pruneHandler) CanPruneRuntime(rounds []uint64) error {
 	lastSycnedRound, _, _ := p.worker.GetLastSynced()
 
 	for _, round := range rounds {
@@ -24,23 +29,68 @@ func (p *pruneHandler) Prune(rounds []uint64) error {
 		}
 
 		// Old suggestion: Make sure we don't prune rounds that need to be checkpointed but haven't been yet.
+	}
 
-		p.logger.Debug("pruning storage for round", "round", round)
+	return nil
+}
 
-		// Prune given block.
-		err := p.worker.localStorage.NodeDB().Prune(round)
-		switch err {
-		case nil:
-		case mkvsDB.ErrNotEarliest:
-			p.logger.Debug("skipping non-earliest round",
-				"round", round,
-			)
-			continue
-		default:
-			p.logger.Error("failed to prune block",
-				"err", err,
-			)
-			return err
+// statePruner handles pruning of the runtime state.
+//
+// Everytime pruning is triggered, the pruner removes rounds that are older than the earliest
+// round in the runtime’s history.
+//
+// TODO: Pruning logic is not robust as developer changing the pruning of the runtime history
+// may also unexpectedly change the pruning behavior of the state db. This could be fixed
+// by making the storage committee worker responsible for both syncing and pruning of the
+// history and state DB. See https://github.com/oasisprotocol/oasis-core/issues/6400.
+type statePruner struct {
+	state    mkvsDB.NodeDB
+	history  history.History
+	interval time.Duration
+	logger   *logging.Logger
+}
+
+// newPruner creates new runtime state pruner.
+func newPruner(ndb mkvsDB.NodeDB, history history.History, interval time.Duration) *statePruner {
+	return &statePruner{
+		state:    ndb,
+		history:  history,
+		interval: max(interval, time.Second),
+		logger:   logging.GetLogger("/worker/storage/state-pruner").With("runtime_id", history.RuntimeID()),
+	}
+}
+
+// serve periodically triggers the pruning of the runtime state db.
+func (p *statePruner) serve(ctx context.Context) error {
+	p.logger.Info("starting")
+	defer p.logger.Info("stopped")
+
+	ticker := time.NewTicker(p.interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-ticker.C:
+		}
+
+		if err := p.prune(ctx); err != nil {
+			p.logger.Warn("failed to prune", "err", err)
+		}
+	}
+}
+
+func (p *statePruner) prune(ctx context.Context) error {
+	blk, err := p.history.GetEarliestBlock(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get earliest block from runtime history: %w", err)
+	}
+
+	for round := p.state.GetEarliestVersion(); round < blk.Header.Round; round++ {
+		p.logger.Debug("pruning", "round", round)
+		if err := p.state.Prune(round); err != nil {
+			return fmt.Errorf("failed to prune round %d: %w", round, err)
 		}
 	}
 

go/worker/storage/committee/worker.go

@@ -166,6 +166,8 @@ type Worker struct {
 	diffCh     chan *fetchedDiff
 	finalizeCh chan finalizeResult
 
+	pruneInterval time.Duration
+
 	initCh chan struct{}
 }
 
@@ -177,6 +179,7 @@ func New(
 	workerCommonCfg workerCommon.Config,
 	localStorage storageApi.LocalBackend,
 	checkpointSyncCfg *CheckpointSyncConfig,
+	pruneInterval time.Duration,
 ) (*Worker, error) {
 	initMetrics()
 
@@ -200,6 +203,8 @@ func New(
 		diffCh:     make(chan *fetchedDiff),
 		finalizeCh: make(chan finalizeResult),
 
+		pruneInterval: pruneInterval,
+
 		initCh: make(chan struct{}),
 	}
 
@@ -1014,6 +1019,13 @@ func (w *Worker) Serve(ctx context.Context) error { // nolint: gocyclo
 	close(w.initCh)
 	w.logger.Info("initialized")
 
+	statePruner := newPruner(w.localStorage.NodeDB(), w.commonNode.Runtime.History(), w.pruneInterval)
+	wg.Go(func() {
+		if err := statePruner.serve(ctx); err != nil && !errors.Is(err, context.Canceled) {
+			w.logger.Error("state pruner failed: %w", err)
+		}
+	})
+
 	// Notify the checkpointer of the genesis round so it can be checkpointed.
 	w.checkpointer.ForceCheckpoint(genesisBlock.Header.Round)
 	w.checkpointer.Flush()

go/worker/storage/worker.go

@@ -109,6 +109,7 @@ func (w *Worker) registerRuntime(commonNode *committeeCommon.Node) error {
 			Disabled:          config.GlobalConfig.Storage.CheckpointSyncDisabled,
 			ChunkFetcherCount: config.GlobalConfig.Storage.FetcherCount,
 		},
+		config.GlobalConfig.Runtime.Prune.Interval,
 	)
 	if err != nil {
 		return err