fix: context

omni-network · Dec 11, 2024 · c77570e · c77570e
1 parent c7ae35c
commit c77570e
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 8 deletions.
diff --git a/halo/app/abci.go b/halo/app/abci.go
@@ -21,21 +21,25 @@ type abciWrapper struct {
 	abci.Application
 	postFinalize       postFinalizeCallback
 	multiStoreProvider multiStoreProvider
+	withContext        contextWrapper
 }
 
 func newABCIWrapper(
 	app abci.Application,
 	finaliseCallback postFinalizeCallback,
 	multiStoreProvider multiStoreProvider,
+	ctxWrapper contextWrapper,
 ) *abciWrapper {
 	return &abciWrapper{
 		Application:        app,
 		postFinalize:       finaliseCallback,
 		multiStoreProvider: multiStoreProvider,
+		withContext:        ctxWrapper,
 	}
 }
 
 func (l abciWrapper) Info(ctx context.Context, info *abci.RequestInfo) (*abci.ResponseInfo, error) {
+	ctx = l.withContext(ctx)
 	log.Debug(ctx, "👾 ABCI call: Info")
 	resp, err := l.Application.Info(ctx, info)
 	if err != nil {
@@ -46,15 +50,20 @@ func (l abciWrapper) Info(ctx context.Context, info *abci.RequestInfo) (*abci.Re
 }
 
 func (l abciWrapper) Query(ctx context.Context, query *abci.RequestQuery) (*abci.ResponseQuery, error) {
+	ctx = l.withContext(ctx)
+
 	return l.Application.Query(ctx, query) // No log here since this can be very noisy
 }
 
 func (l abciWrapper) CheckTx(ctx context.Context, tx *abci.RequestCheckTx) (*abci.ResponseCheckTx, error) {
+	ctx = l.withContext(ctx)
 	log.Debug(ctx, "👾 ABCI call: CheckTx")
+
 	return l.Application.CheckTx(ctx, tx)
 }
 
 func (l abciWrapper) InitChain(ctx context.Context, chain *abci.RequestInitChain) (*abci.ResponseInitChain, error) {
+	ctx = l.withContext(ctx)
 	log.Debug(ctx, "👾 ABCI call: InitChain")
 	resp, err := l.Application.InitChain(ctx, chain)
 	if err != nil {
@@ -65,7 +74,7 @@ func (l abciWrapper) InitChain(ctx context.Context, chain *abci.RequestInitChain
 }
 
 func (l abciWrapper) PrepareProposal(ctx context.Context, proposal *abci.RequestPrepareProposal) (*abci.ResponsePrepareProposal, error) {
-	ctx = log.WithCtx(ctx, "height", proposal.Height)
+	ctx = log.WithCtx(l.withContext(ctx), "height", proposal.Height)
 	log.Debug(ctx, "👾 ABCI call: PrepareProposal",
 		log.Hex7("proposer", proposal.ProposerAddress),
 	)
@@ -78,7 +87,7 @@ func (l abciWrapper) PrepareProposal(ctx context.Context, proposal *abci.Request
 }
 
 func (l abciWrapper) ProcessProposal(ctx context.Context, proposal *abci.RequestProcessProposal) (*abci.ResponseProcessProposal, error) {
-	ctx = log.WithCtx(ctx, "height", proposal.Height)
+	ctx = log.WithCtx(l.withContext(ctx), "height", proposal.Height)
 	log.Debug(ctx, "👾 ABCI call: ProcessProposal",
 		log.Hex7("proposer", proposal.ProposerAddress),
 	)
@@ -91,7 +100,7 @@ func (l abciWrapper) ProcessProposal(ctx context.Context, proposal *abci.Request
 }
 
 func (l abciWrapper) FinalizeBlock(ctx context.Context, req *abci.RequestFinalizeBlock) (*abci.ResponseFinalizeBlock, error) {
-	ctx = log.WithCtx(ctx, "height", req.Height)
+	ctx = log.WithCtx(l.withContext(ctx), "height", req.Height)
 	resp, err := l.Application.FinalizeBlock(ctx, req)
 	if err != nil {
 		log.Error(ctx, "Finalize req failed [BUG]", err)
@@ -134,7 +143,7 @@ func (l abciWrapper) FinalizeBlock(ctx context.Context, req *abci.RequestFinaliz
 }
 
 func (l abciWrapper) ExtendVote(ctx context.Context, vote *abci.RequestExtendVote) (*abci.ResponseExtendVote, error) {
-	ctx = log.WithCtx(ctx, "height", vote.Height)
+	ctx = log.WithCtx(l.withContext(ctx), "height", vote.Height)
 	log.Debug(ctx, "👾 ABCI call: ExtendVote")
 	resp, err := l.Application.ExtendVote(ctx, vote)
 	if err != nil {
@@ -145,7 +154,7 @@ func (l abciWrapper) ExtendVote(ctx context.Context, vote *abci.RequestExtendVot
 }
 
 func (l abciWrapper) VerifyVoteExtension(ctx context.Context, extension *abci.RequestVerifyVoteExtension) (*abci.ResponseVerifyVoteExtension, error) {
-	ctx = log.WithCtx(ctx, "height", extension.Height)
+	ctx = log.WithCtx(l.withContext(ctx), "height", extension.Height)
 	log.Debug(ctx, "👾 ABCI call: VerifyVoteExtension")
 	resp, err := l.Application.VerifyVoteExtension(ctx, extension)
 	if err != nil {
@@ -156,26 +165,36 @@ func (l abciWrapper) VerifyVoteExtension(ctx context.Context, extension *abci.Re
 }
 
 func (l abciWrapper) Commit(ctx context.Context, commit *abci.RequestCommit) (*abci.ResponseCommit, error) {
+	ctx = l.withContext(ctx)
 	log.Debug(ctx, "👾 ABCI call: Commit")
+
 	return l.Application.Commit(ctx, commit)
 }
 
 func (l abciWrapper) ListSnapshots(ctx context.Context, listSnapshots *abci.RequestListSnapshots) (*abci.ResponseListSnapshots, error) {
+	ctx = l.withContext(ctx)
 	log.Debug(ctx, "👾 ABCI call: ListSnapshots")
+
 	return l.Application.ListSnapshots(ctx, listSnapshots)
 }
 
 func (l abciWrapper) OfferSnapshot(ctx context.Context, snapshot *abci.RequestOfferSnapshot) (*abci.ResponseOfferSnapshot, error) {
+	ctx = l.withContext(ctx)
 	log.Debug(ctx, "👾 ABCI call: OfferSnapshot")
+
 	return l.Application.OfferSnapshot(ctx, snapshot)
 }
 
 func (l abciWrapper) LoadSnapshotChunk(ctx context.Context, chunk *abci.RequestLoadSnapshotChunk) (*abci.ResponseLoadSnapshotChunk, error) {
+	ctx = l.withContext(ctx)
 	log.Debug(ctx, "👾 ABCI call: LoadSnapshotChunk")
+
 	return l.Application.LoadSnapshotChunk(ctx, chunk)
 }
 
 func (l abciWrapper) ApplySnapshotChunk(ctx context.Context, chunk *abci.RequestApplySnapshotChunk) (*abci.ResponseApplySnapshotChunk, error) {
+	ctx = l.withContext(ctx)
 	log.Debug(ctx, "👾 ABCI call: ApplySnapshotChunk")
+
 	return l.Application.ApplySnapshotChunk(ctx, chunk)
 }
diff --git a/halo/app/start.go b/halo/app/start.go
@@ -51,6 +51,9 @@ import (
 	grpc1 "github.com/cosmos/gogoproto/grpc"
 )
 
+// contextWrapper wraps given context with more context.
+type contextWrapper = func(context.Context) context.Context
+
 // Config wraps the halo (app) and comet (client) configurations.
 type Config struct {
 	halocfg.Config
@@ -103,7 +106,11 @@ func Start(ctx context.Context, cfg Config) (<-chan error, func(context.Context)
 
 	buildinfo.Instrument(ctx)
 
-	ctx = feature.WithFlags(ctx, cfg.FeatureFlags)
+	ctxWrapper := func(ctx context.Context) context.Context {
+		return feature.WithFlags(ctx, cfg.FeatureFlags)
+	}
+
+	ctx = ctxWrapper(ctx)
 
 	tracerIDs := tracer.Identifiers{Network: cfg.Network, Service: "halo", Instance: cfg.Comet.Moniker}
 	stopTracer, err := tracer.Init(ctx, tracerIDs, cfg.Tracer)
@@ -170,7 +177,7 @@ func Start(ctx context.Context, cfg Config) (<-chan error, func(context.Context)
 	app.EVMEngKeeper.SetBuildDelay(cfg.EVMBuildDelay)
 	app.EVMEngKeeper.SetBuildOptimistic(cfg.EVMBuildOptimistic)
 
-	cmtNode, err := newCometNode(ctx, &cfg.Comet, app, privVal)
+	cmtNode, err := newCometNode(ctx, &cfg.Comet, app, privVal, ctxWrapper)
 	if err != nil {
 		return nil, nil, errors.Wrap(err, "create comet node")
 	}
@@ -298,7 +305,7 @@ func startRPCServers(
 	return nil
 }
 
-func newCometNode(ctx context.Context, cfg *cmtcfg.Config, app *App, privVal cmttypes.PrivValidator) (
+func newCometNode(ctx context.Context, cfg *cmtcfg.Config, app *App, privVal cmttypes.PrivValidator, ctxWrapper contextWrapper) (
 	*node.Node, error) {
 	nodeKey, err := p2p.LoadOrGenNodeKey(cfg.NodeKeyFile())
 	if err != nil {
@@ -316,6 +323,7 @@ func newCometNode(ctx context.Context, cfg *cmtcfg.Config, app *App, privVal cmt
 		func() storetypes.CacheMultiStore {
 			return app.CommitMultiStore().CacheMultiStore()
 		},
+		ctxWrapper,
 	)
 
 	// Configure CometBFT prometheus metrics as per provided config