From 71dd7f13ebd04dc46d37201e53d59a661ac29886 Mon Sep 17 00:00:00 2001
From: codchen <codchen03@gmail.com>
Date: Thu, 25 Jul 2024 11:53:50 +0800
Subject: [PATCH] Add DeliverTx callback (#529)

## Describe your changes and provide context
Add a field in context to set callbacks that would be run at the end of
DeliverTx on the base store, regardless of the outcome of the
transaction (e.g. bump nonce)

## Testing performed to validate your change
integration test in sei-test
---
 baseapp/baseapp.go |  5 +++++
 types/context.go   | 54 +++++++++++++++++++++++++++-------------------
 2 files changed, 37 insertions(+), 22 deletions(-)

diff --git a/baseapp/baseapp.go b/baseapp/baseapp.go
index 1b75579dd..5e3f1cf7f 100644
--- a/baseapp/baseapp.go
+++ b/baseapp/baseapp.go
@@ -945,6 +945,11 @@ func (app *BaseApp) runTx(ctx sdk.Context, mode runTxMode, tx sdk.Tx, checksum [
 			// gas there too.
 			ctx = newCtx.WithMultiStore(ms)
 		}
+		defer func() {
+			if newCtx.DeliverTxCallback() != nil {
+				newCtx.DeliverTxCallback()(ctx.WithGasMeter(sdk.NewInfiniteGasMeterWithMultiplier(ctx)))
+			}
+		}()
 
 		events := ctx.EventManager().Events()
 
diff --git a/types/context.go b/types/context.go
index 48d85cdcd..7cf4b2a59 100644
--- a/types/context.go
+++ b/types/context.go
@@ -24,28 +24,29 @@ but please do not over-use it. We try to keep all data structured
 and standard additions here would be better just to add to the Context struct
 */
 type Context struct {
-	ctx              context.Context
-	ms               MultiStore
-	header           tmproto.Header
-	headerHash       tmbytes.HexBytes
-	chainID          string
-	txBytes          []byte
-	txSum            [32]byte
-	logger           log.Logger
-	voteInfo         []abci.VoteInfo
-	gasMeter         GasMeter
-	occEnabled       bool
-	blockGasMeter    GasMeter
-	checkTx          bool
-	recheckTx        bool // if recheckTx == true, then checkTx must also be true
-	minGasPrice      DecCoins
-	consParams       *tmproto.ConsensusParams
-	eventManager     *EventManager
-	evmEventManager  *EVMEventManager
-	priority         int64                 // The tx priority, only relevant in CheckTx
-	pendingTxChecker abci.PendingTxChecker // Checker for pending transaction, only relevant in CheckTx
-	checkTxCallback  func(Context, error)  // callback to make at the end of CheckTx. Input param is the error (nil-able) of `runMsgs`
-	expireTxHandler  func()                // callback that the mempool invokes when a tx is expired
+	ctx               context.Context
+	ms                MultiStore
+	header            tmproto.Header
+	headerHash        tmbytes.HexBytes
+	chainID           string
+	txBytes           []byte
+	txSum             [32]byte
+	logger            log.Logger
+	voteInfo          []abci.VoteInfo
+	gasMeter          GasMeter
+	occEnabled        bool
+	blockGasMeter     GasMeter
+	checkTx           bool
+	recheckTx         bool // if recheckTx == true, then checkTx must also be true
+	minGasPrice       DecCoins
+	consParams        *tmproto.ConsensusParams
+	eventManager      *EventManager
+	evmEventManager   *EVMEventManager
+	priority          int64                 // The tx priority, only relevant in CheckTx
+	pendingTxChecker  abci.PendingTxChecker // Checker for pending transaction, only relevant in CheckTx
+	checkTxCallback   func(Context, error)  // callback to make at the end of CheckTx. Input param is the error (nil-able) of `runMsgs`
+	deliverTxCallback func(Context)         // callback to make at the end of DeliverTx.
+	expireTxHandler   func()                // callback that the mempool invokes when a tx is expired
 
 	txBlockingChannels   acltypes.MessageAccessOpsChannelMapping
 	txCompletionChannels acltypes.MessageAccessOpsChannelMapping
@@ -169,6 +170,10 @@ func (c Context) CheckTxCallback() func(Context, error) {
 	return c.checkTxCallback
 }
 
+func (c Context) DeliverTxCallback() func(Context) {
+	return c.deliverTxCallback
+}
+
 func (c Context) TxCompletionChannels() acltypes.MessageAccessOpsChannelMapping {
 	return c.txCompletionChannels
 }
@@ -448,6 +453,11 @@ func (c Context) WithCheckTxCallback(checkTxCallback func(Context, error)) Conte
 	return c
 }
 
+func (c Context) WithDeliverTxCallback(deliverTxCallback func(Context)) Context {
+	c.deliverTxCallback = deliverTxCallback
+	return c
+}
+
 func (c Context) WithExpireTxHandler(expireTxHandler func()) Context {
 	c.expireTxHandler = expireTxHandler
 	return c