fix(abci): ignore comet ctx and use app ctx (#2568)

Co-authored-by: aBear <[email protected]> Co-authored-by: Cal Bera <[email protected]> Co-authored-by: Fridrik Asmundsson <[email protected]> Co-authored-by: Rez <[email protected]>
berachain · Mar 6, 2025 · 6699ea9 · 6699ea9
1 parent de59b46
commit 6699ea9
Show file tree

Hide file tree

Showing 10 changed files with 78 additions and 29 deletions.
diff --git a/.gitignore b/.gitignore
@@ -231,6 +231,7 @@ coverage-test-unit-cover.txt
 coverage-merged.txt
 test-simulated.txt
 test-unit-cover.txt
+temp-test-simulated.txt
 .vercel
 
 # server dev env for Air

diff --git a/consensus/cometbft/service/abci.go b/consensus/cometbft/service/abci.go
@@ -37,19 +37,33 @@ var (
 )
 
 func (s *Service) InitChain(
-	ctx context.Context,
+	_ context.Context,
 	req *cmtabci.InitChainRequest,
 ) (*cmtabci.InitChainResponse, error) {
-	return s.initChain(ctx, req)
+	// Check if ctx is still good. CometBFT does not check this.
+	if s.ctx.Err() != nil {
+		// If the context is getting cancelled, we are shutting down.
+		return &cmtabci.InitChainResponse{}, s.ctx.Err()
+	}
+	//nolint:contextcheck // see s.ctx comment for more details
+	return s.initChain(req) // internally this uses s.ctx
 }
 
 // PrepareProposal implements the PrepareProposal ABCI method and returns a
 // ResponsePrepareProposal object to the client.
 func (s *Service) PrepareProposal(
-	ctx context.Context,
+	_ context.Context,
 	req *cmtabci.PrepareProposalRequest,
 ) (*cmtabci.PrepareProposalResponse, error) {
-	return s.prepareProposal(ctx, req)
+	// Check if ctx is still good. CometBFT does not check this.
+	if s.ctx.Err() != nil {
+		// If the context is getting cancelled, we are shutting down.
+		// It is ok returning an empty proposal.
+		//nolint:nilerr // explicitly allowing this case
+		return &cmtabci.PrepareProposalResponse{Txs: req.Txs}, nil
+	}
+	//nolint:contextcheck // see s.ctx comment for more details
+	return s.prepareProposal(s.ctx, req)
 }
 
 func (s *Service) Info(context.Context,
@@ -77,17 +91,31 @@ func (s *Service) Info(context.Context,
 // ProcessProposal implements the ProcessProposal ABCI method and returns a
 // ResponseProcessProposal object to the client.
 func (s *Service) ProcessProposal(
-	ctx context.Context,
+	_ context.Context,
 	req *cmtabci.ProcessProposalRequest,
 ) (*cmtabci.ProcessProposalResponse, error) {
-	return s.processProposal(ctx, req)
+	// Check if ctx is still good. CometBFT does not check this.
+	if s.ctx.Err() != nil {
+		// Node will panic on context cancel with "CONSENSUS FAILURE!!!" due to
+		// returning an error. This is expected. We do not want to accept or
+		// reject a proposal based on incomplete data.
+		return nil, s.ctx.Err()
+	}
+	//nolint:contextcheck // see s.ctx comment for more details
+	return s.processProposal(req) // internally this uses s.ctx
 }
 
 func (s *Service) FinalizeBlock(
-	ctx context.Context,
+	_ context.Context,
 	req *cmtabci.FinalizeBlockRequest,
 ) (*cmtabci.FinalizeBlockResponse, error) {
-	return s.finalizeBlock(ctx, req)
+	// Check if ctx is still good. CometBFT does not check this.
+	if s.ctx.Err() != nil {
+		// Node will panic on context cancel with "CONSENSUS FAILURE!!!" due to error.
+		// We expect this to happen and do not want to finalize any incomplete or invalid state.
+		return nil, s.ctx.Err()
+	}
+	return s.finalizeBlock(req) // internally this uses s.ctx
 }
 
 // Commit implements the ABCI interface. It will commit all state that exists in
@@ -98,9 +126,16 @@ func (s *Service) FinalizeBlock(
 // against that height and gracefully halt if it matches the latest committed
 // height.
 func (s *Service) Commit(
-	ctx context.Context, req *cmtabci.CommitRequest,
+	_ context.Context, req *cmtabci.CommitRequest,
 ) (*cmtabci.CommitResponse, error) {
-	return s.commit(ctx, req)
+	// Check if ctx is still good. CometBFT does not check this.
+	if s.ctx.Err() != nil {
+		// Node will panic on context cancel with "CONSENSUS FAILURE!!!" due to error.
+		// We expect this to happen and do not want to commit any incomplete or invalid state.
+		return nil, s.ctx.Err()
+	}
+
+	return s.commit(req)
 }
 
 //

diff --git a/consensus/cometbft/service/commit.go b/consensus/cometbft/service/commit.go
@@ -22,15 +22,14 @@
 package cometbft
 
 import (
-	"context"
 	"fmt"
 
 	"cosmossdk.io/store/rootmulti"
 	cmtabci "github.com/cometbft/cometbft/abci/types"
 )
 
 func (s *Service) commit(
-	context.Context, *cmtabci.CommitRequest,
+	*cmtabci.CommitRequest,
 ) (*cmtabci.CommitResponse, error) {
 	if s.finalizeBlockState == nil {
 		// This is unexpected since CometBFT should call Commit only

diff --git a/consensus/cometbft/service/finalize_block.go b/consensus/cometbft/service/finalize_block.go
@@ -21,18 +21,16 @@
 package cometbft
 
 import (
-	"context"
 	"fmt"
 
 	cmtabci "github.com/cometbft/cometbft/abci/types"
 	"github.com/sourcegraph/conc/iter"
 )
 
 func (s *Service) finalizeBlock(
-	ctx context.Context,
 	req *cmtabci.FinalizeBlockRequest,
 ) (*cmtabci.FinalizeBlockResponse, error) {
-	res, err := s.finalizeBlockInternal(ctx, req)
+	res, err := s.finalizeBlockInternal(req)
 	if res != nil {
 		res.AppHash = s.workingHash()
 	}
@@ -41,7 +39,6 @@ func (s *Service) finalizeBlock(
 }
 
 func (s *Service) finalizeBlockInternal(
-	ctx context.Context,
 	req *cmtabci.FinalizeBlockRequest,
 ) (*cmtabci.FinalizeBlockResponse, error) {
 	if err := s.validateFinalizeBlockHeight(req); err != nil {
@@ -53,7 +50,10 @@ func (s *Service) finalizeBlockInternal(
 	// here given that during block replay ProcessProposal is not executed by
 	// CometBFT.
 	if s.finalizeBlockState == nil {
-		s.finalizeBlockState = s.resetState(ctx)
+		s.finalizeBlockState = s.resetState(s.ctx)
+	} else {
+		// Preserve the CosmosSDK context while using the correct base ctx.
+		s.finalizeBlockState.SetContext(s.finalizeBlockState.Context().WithContext(s.ctx))
 	}
 
 	// Iterate over all raw transactions in the proposal and attempt to execute

diff --git a/consensus/cometbft/service/init_chain.go b/consensus/cometbft/service/init_chain.go
@@ -21,7 +21,6 @@
 package cometbft
 
 import (
-	"context"
 	"fmt"
 
 	"github.com/berachain/beacon-kit/primitives/encoding/json"
@@ -31,7 +30,6 @@ import (
 )
 
 func (s *Service) initChain(
-	ctx context.Context,
 	req *cmtabci.InitChainRequest,
 ) (*cmtabci.InitChainResponse, error) {
 	if req.ChainId != s.chainID {
@@ -83,9 +81,8 @@ func (s *Service) initChain(
 		}
 	}
 
-	s.finalizeBlockState = s.resetState(ctx)
+	s.finalizeBlockState = s.resetState(s.ctx)
 
-	//nolint:contextcheck // ctx already passed via resetState
 	resValidators, err := s.initChainer(
 		s.finalizeBlockState.Context(),
 		req.AppStateBytes,

diff --git a/consensus/cometbft/service/prepare_proposal.go b/consensus/cometbft/service/prepare_proposal.go
@@ -50,6 +50,7 @@ func (s *Service) prepareProposal(
 	// Always reset state given that PrepareProposal can timeout
 	// and be called again in a subsequent round.
 	s.prepareProposalState = s.resetState(ctx)
+	//nolint:contextcheck // ctx already passed via resetState
 	s.prepareProposalState.SetContext(
 		s.getContextForProposal(
 			s.prepareProposalState.Context(),

diff --git a/consensus/cometbft/service/process_proposal.go b/consensus/cometbft/service/process_proposal.go
@@ -21,15 +21,13 @@
 package cometbft
 
 import (
-	"context"
 	"fmt"
 	"time"
 
 	cmtabci "github.com/cometbft/cometbft/abci/types"
 )
 
 func (s *Service) processProposal(
-	ctx context.Context,
 	req *cmtabci.ProcessProposalRequest,
 ) (*cmtabci.ProcessProposalResponse, error) {
 	startTime := time.Now()
@@ -52,9 +50,9 @@ func (s *Service) processProposal(
 	// processed the first block, as we want to avoid overwriting the
 	// finalizeState
 	// after state changes during InitChain.
-	s.processProposalState = s.resetState(ctx)
+	s.processProposalState = s.resetState(s.ctx)
 	if req.Height > s.initialHeight {
-		s.finalizeBlockState = s.resetState(ctx)
+		s.finalizeBlockState = s.resetState(s.ctx)
 	}
 
 	s.processProposalState.SetContext(

diff --git a/consensus/cometbft/service/service.go b/consensus/cometbft/service/service.go
@@ -96,6 +96,15 @@ type Service struct {
 	minRetainBlocks uint64
 
 	chainID string
+
+	// ctx is the context passed in for the service. CometBFT currently does
+	// not support context usage. It passes "context.TODO()" to apps that
+	// implement the ABCI++ interface, and does not provide a context that is
+	// a child context of the one the node originally provides to comet.
+	// Thus the app cannot tell when the context as been cancelled or not.
+	// TODO: We must use this as a workaround for now until CometBFT properly
+	// generates contexts that inherit from the parent context we provide.
+	ctx context.Context
 }
 
 func NewService(
@@ -165,6 +174,7 @@ func (s *Service) Start(
 		return err
 	}
 
+	s.ctx = ctx
 	s.node, err = node.NewNode(
 		ctx,
 		cfg,
@@ -220,6 +230,12 @@ func (s *Service) Stop() error {
 	return errors.Join(errs...)
 }
 
+// ResetAppCtx sets the app ctx for the service. This is used
+// primarily for the mock service.
+func (s *Service) ResetAppCtx(ctx context.Context) {
+	s.ctx = ctx
+}
+
 // Name returns the name of the cometbft.
 func (s *Service) Name() string {
 	return AppName
@@ -320,8 +336,9 @@ func (s *Service) getContextForProposal(
 		// on initialHeight. Panic appeases nilaway.
 		panic(fmt.Errorf("getContextForProposal: %w", errNilFinalizeBlockState))
 	}
-	ctx, _ = s.finalizeBlockState.Context().CacheContext()
-	return ctx
+	newCtx, _ := s.finalizeBlockState.Context().CacheContext()
+	// Preserve the CosmosSDK context while using the correct base ctx.
+	return newCtx.WithContext(ctx.Context())
 }
 
 // CreateQueryContext creates a new sdk.Context for a query, taking as args

diff --git a/testing/simulated/simcomet.go b/testing/simulated/simcomet.go
@@ -64,7 +64,8 @@ func ProvideSimComet(
 		)}
 }
 
-func (s *SimComet) Start(_ context.Context) error {
+func (s *SimComet) Start(ctx context.Context) error {
+	s.Comet.ResetAppCtx(ctx)
 	return nil
 }
 

diff --git a/testing/simulated/simulated_test.go b/testing/simulated/simulated_test.go
@@ -95,7 +95,7 @@ func (s *SimulatedSuite) SetupTest() {
 	}()
 
 	// Allow a short period for services to fully initialize.
-	time.Sleep(2 * time.Second)
+	time.Sleep(5 * time.Second)
 }
 
 // TearDownTest cleans up the test environment.