cometbft · thanethomson · Aug 22, 2023 · Aug 18, 2023 · Aug 18, 2023 · Aug 18, 2023
diff --git a/node/node.go b/node/node.go
@@ -241,26 +241,16 @@ func NewNode(ctx context.Context,
 		return nil, err
 	}
 
-	err = initApplicationRetainHeight(stateStore)
-	if err != nil {
-		return nil, err
-	}
-
-	err = initCompanionBlockRetainHeight(
+	pruner, err := createPruner(
+		config,
 		stateStore,
-		config.Storage.Pruning.DataCompanion.Enabled,
-		config.Storage.Pruning.DataCompanion.InitialBlockRetainHeight,
+		blockStore,
+		smMetrics,
+		logger.With("module", "state"),
 	)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to create pruner: %w", err)
 	}
-	pruner := sm.NewPruner(
-		stateStore,
-		blockStore,
-		logger,
-		sm.WithPrunerInterval(config.Storage.Pruning.Interval),
-		sm.WithPrunerMetrics(smMetrics),
-	)
 
 	// make block executor for consensus and blocksync reactors to execute blocks
 	blockExec := sm.NewBlockExecutor(
@@ -871,9 +861,40 @@ func makeNodeInfo(
 	return nodeInfo, err
 }
 
-// Set the initial application retain height to 0 to avoid the data companion pruning blocks
-// before the application indicates it is ok
-// We set this to 0 only if the retain height was not set before by the application
+func createPruner(
+	config *cfg.Config,
+	stateStore sm.Store,
+	blockStore *store.BlockStore,
+	metrics *sm.Metrics,
+	logger log.Logger,
+) (*sm.Pruner, error) {
+	if err := initApplicationRetainHeight(stateStore); err != nil {
+		return nil, err
+	}
+
+	prunerOpts := []sm.PrunerOption{
+		sm.WithPrunerInterval(config.Storage.Pruning.Interval),
+		sm.WithPrunerMetrics(metrics),
+	}
+
+	if config.Storage.Pruning.DataCompanion.Enabled {
+		err := initCompanionRetainHeights(
+			stateStore,
+			config.Storage.Pruning.DataCompanion.InitialBlockRetainHeight,
+			config.Storage.Pruning.DataCompanion.InitialBlockResultsRetainHeight,
+		)
+		if err != nil {
+			return nil, err
+		}
+		prunerOpts = append(prunerOpts, sm.WithPrunerCompanionEnabled())
+	}
+
+	return sm.NewPruner(stateStore, blockStore, logger, prunerOpts...), nil
+}
+
+// Set the initial application retain height to 0 to avoid the data companion
+// pruning blocks before the application indicates it is OK. We set this to 0
+// only if the retain height was not set before by the application.
 func initApplicationRetainHeight(stateStore sm.Store) error {
 	if _, err := stateStore.GetApplicationRetainHeight(); err != nil {
 		if errors.Is(err, sm.ErrKeyNotFound) {
@@ -884,27 +905,27 @@ func initApplicationRetainHeight(stateStore sm.Store) error {
 	return nil
 }
 
-func initCompanionBlockRetainHeight(stateStore sm.Store, companionEnabled bool, initialRetainHeight int64) error {
-	if _, err := stateStore.GetCompanionBlockRetainHeight(); err != nil {
-		// If the data companion block retain height has not yet been set in
-		// the database
-		if errors.Is(err, sm.ErrKeyNotFound) {
-			if companionEnabled && initialRetainHeight > 0 {
-				// This will set the data companion retain height into the
-				// database. We bypass the sanity checks by
-				// pruner.SetCompanionBlockRetainHeight. These checks do not
-				// allow a retain height below the current blockstore height or
-				// above the blockstore height  to be set. But this is a retain
-				// height that can be set before the chain starts to indicate
-				// potentially that no pruning should be done before the data
-				// companion comes online.
-				err = stateStore.SaveCompanionBlockRetainHeight(initialRetainHeight)
-				if err != nil {
-					return fmt.Errorf("failed to set initial data companion block retain height: %w", err)
-				}
-			}
-		} else {
-			return fmt.Errorf("failed to obtain companion retain height: %w", err)
+// Sets the data companion retain heights if one of two possible conditions is
+// met:
+// 1. One or more of the retain heights has not yet been set.
+// 2. One or more of the retain heights is currently 0.
+func initCompanionRetainHeights(stateStore sm.Store, initBlockRH, initBlockResultsRH int64) error {
+	curBlockRH, err := stateStore.GetCompanionBlockRetainHeight()
+	if err != nil && !errors.Is(err, sm.ErrKeyNotFound) {
+		return fmt.Errorf("failed to obtain companion block retain height: %w", err)
+	}
+	if curBlockRH == 0 {
+		if err := stateStore.SaveCompanionBlockRetainHeight(initBlockRH); err != nil {
+			return fmt.Errorf("failed to set initial data companion block retain height: %w", err)
+		}
+	}
+	curBlockResultsRH, err := stateStore.GetABCIResRetainHeight()
+	if err != nil && !errors.Is(err, sm.ErrKeyNotFound) {
+		return fmt.Errorf("failed to obtain companion block results retain height: %w", err)
+	}
+	if curBlockResultsRH == 0 {
+		if err := stateStore.SaveABCIResRetainHeight(initBlockResultsRH); err != nil {
+			return fmt.Errorf("failed to set initial data companion block results retain height: %w", err)
 		}
 	}
 	return nil

diff --git a/rpc/grpc/server/services/pruningservice/service.go b/rpc/grpc/server/services/pruningservice/service.go
@@ -40,9 +40,9 @@ func (s *pruningServiceServer) SetBlockRetainHeight(_ context.Context, req *v1.S
 		logger.Error("Error generating RPC trace ID", "err", err)
 		return nil, status.Error(codes.Internal, "Internal server error - see logs for details")
 	}
-	if err := s.pruner.SetCompanionRetainHeight(int64(height)); err != nil {
+	if err := s.pruner.SetCompanionBlockRetainHeight(int64(height)); err != nil {
 		logger.Error("Cannot set block retain height", "err", err, "traceID", traceID)
-		return nil, status.Error(codes.Internal, "Failed to set block retain height")
+		return nil, status.Errorf(codes.Internal, "Failed to set block retain height (see logs for trace ID: %s)", traceID)
 	}
 	return &v1.SetBlockRetainHeightResponse{}, nil
 }
@@ -58,12 +58,12 @@ func (s *pruningServiceServer) GetBlockRetainHeight(_ context.Context, _ *v1.Get
 	svcHeight, err := s.pruner.GetCompanionBlockRetainHeight()
 	if err != nil {
 		logger.Error("Cannot get block retain height stored by companion", "err", err, "traceID", traceID)
-		return nil, status.Error(codes.Internal, "Failed to get companion block retain height")
+		return nil, status.Errorf(codes.Internal, "Failed to get companion block retain height (see logs for trace ID: %s)", traceID)
 	}
 	appHeight, err := s.pruner.GetApplicationRetainHeight()
 	if err != nil {
 		logger.Error("Cannot get block retain height specified by application", "err", err, "traceID", traceID)
-		return nil, status.Error(codes.Internal, "Failed to get app block retain height")
+		return nil, status.Errorf(codes.Internal, "Failed to get app block retain height (see logs for trace ID: %s)", traceID)
 	}
 	return &v1.GetBlockRetainHeightResponse{
 		PruningServiceRetainHeight: uint64(svcHeight),
@@ -86,7 +86,7 @@ func (s *pruningServiceServer) SetBlockResultsRetainHeight(_ context.Context, re
 	}
 	if err := s.pruner.SetABCIResRetainHeight(int64(height)); err != nil {
 		logger.Error("Cannot set block results retain height", "err", err, "traceID", traceID)
-		return nil, status.Error(codes.Internal, "Failed to set block results retain height")
+		return nil, status.Errorf(codes.Internal, "Failed to set block results retain height (see logs for trace ID: %s)", traceID)
 	}
 	return &v1.SetBlockResultsRetainHeightResponse{}, nil
 }
@@ -102,7 +102,7 @@ func (s *pruningServiceServer) GetBlockResultsRetainHeight(_ context.Context, _
 	height, err := s.pruner.GetABCIResRetainHeight()
 	if err != nil {
 		logger.Error("Cannot get block results retain height", "err", err, "traceID", traceID)
-		return nil, status.Errorf(codes.Internal, "Failed to get block results retain height")
+		return nil, status.Errorf(codes.Internal, "Failed to get block results retain height (see logs for trace ID: %s)", traceID)
 	}
 	return &v1.GetBlockResultsRetainHeightResponse{PruningServiceRetainHeight: uint64(height)}, nil
 }
diff --git a/state/errors.go b/state/errors.go
@@ -52,6 +52,11 @@ type (
 		Height int64
 	}
 
+	ErrPrunerFailedToGetRetainHeight struct {
+		Which string
+		Err   error
+	}
+
 	ErrPrunerFailedToLoadState struct {
 		Err error
 	}
@@ -117,6 +122,14 @@ func (e ErrNoABCIResponsesForHeight) Error() string {
 	return fmt.Sprintf("could not find results for height #%d", e.Height)
 }
 
+func (e ErrPrunerFailedToGetRetainHeight) Error() string {
+	return fmt.Sprintf("pruner failed to get existing %s retain height: %s", e.Which, e.Err.Error())
+}
+
+func (e ErrPrunerFailedToGetRetainHeight) Unwrap() error {
+	return e.Err
+}
+
 func (e ErrPrunerFailedToLoadState) Error() string {
 	return fmt.Sprintf("failed to load state, cannot prune: %s", e.Err.Error())
 }

diff --git a/state/execution.go b/state/execution.go
@@ -304,9 +304,9 @@ func (blockExec *BlockExecutor) ApplyBlock(
 
 	// Prune old heights, if requested by ABCI app.
 	if retainHeight > 0 && blockExec.pruner != nil {
-		err := blockExec.pruner.SetApplicationRetainHeight(retainHeight)
+		err := blockExec.pruner.SetApplicationBlockRetainHeight(retainHeight)
 		if err != nil {
-			blockExec.logger.Error("Failed to set application retain height", "retainHeight", retainHeight, "err", err)
+			blockExec.logger.Error("Failed to set application block retain height", "retainHeight", retainHeight, "err", err)
 		}
 	}
 

diff --git a/state/export_test.go b/state/export_test.go
@@ -45,10 +45,11 @@ func SaveValidatorsInfo(db dbm.DB, height, lastHeightChanged int64, valSet *type
 	return stateStore.saveValidatorsInfo(height, lastHeightChanged, valSet)
 }
 
-// FindMinRetainHeight is an alias for the private findMinRetainHeight method
-// in pruner.go, exported exclusively and expicitly for testing.
+// FindMinBlockRetainHeight is an alias for the private
+// findMinBlockRetainHeight method in pruner.go, exported exclusively and
+// expicitly for testing.
 func (p *Pruner) FindMinRetainHeight() int64 {
-	return p.findMinRetainHeight()
+	return p.findMinBlockRetainHeight()
 }
 
 func (p *Pruner) PruneABCIResToRetainHeight(lastRetainHeight int64) int64 {