feat: Add PrepareCheckState and Precommit callbacks (#14860)

Co-authored-by: dydxwill <119354122+dydxwill@users.noreply.github.com>

baseapp/abci.go

@@ -457,6 +457,10 @@ func (app *BaseApp) Commit() abci.ResponseCommit {
 	header := app.deliverState.ctx.BlockHeader()
 	retainHeight := app.GetBlockRetentionHeight(header.Height)
 
+	if app.precommiter != nil {
+		app.precommiter(app.deliverState.ctx)
+	}
+
 	rms, ok := app.cms.(*rootmulti.Store)
 	if ok {
 		rms.SetCommitHeader(header)
@@ -464,7 +468,7 @@ func (app *BaseApp) Commit() abci.ResponseCommit {
 
 	// Write the DeliverTx state into branched storage and commit the MultiStore.
 	// The write to the DeliverTx state writes all state transitions to the root
-	// MultiStore (app.cms) so when Commit() is called is persists those values.
+	// MultiStore (app.cms) so when Commit() is called it persists those values.
 	app.deliverState.ms.Write()
 	commitID := app.cms.Commit()
 
@@ -497,6 +501,10 @@ func (app *BaseApp) Commit() abci.ResponseCommit {
 	// empty/reset the deliver state
 	app.deliverState = nil
 
+	if app.prepareCheckStater != nil {
+		app.prepareCheckStater(app.checkState.ctx)
+	}
+
 	var halt bool
 
 	switch {

baseapp/abci_test.go

@@ -598,6 +598,60 @@ func TestABCI_EndBlock(t *testing.T) {
 	require.Equal(t, cp.Block.MaxGas, res.ConsensusParamUpdates.Block.MaxGas)
 }
 
+func TestBaseApp_PrepareCheckState(t *testing.T) {
+	db := dbm.NewMemDB()
+	name := t.Name()
+	logger := log.NewTestLogger(t)
+
+	cp := &cmtproto.ConsensusParams{
+		Block: &cmtproto.BlockParams{
+			MaxGas: 5000000,
+		},
+	}
+
+	app := baseapp.NewBaseApp(name, logger, db, nil)
+	app.SetParamStore(&paramStore{db: dbm.NewMemDB()})
+	app.InitChain(abci.RequestInitChain{
+		ConsensusParams: cp,
+	})
+
+	wasPrepareCheckStateCalled := false
+	app.SetPrepareCheckStater(func(ctx sdk.Context) {
+		wasPrepareCheckStateCalled = true
+	})
+	app.Seal()
+
+	app.Commit()
+	require.Equal(t, true, wasPrepareCheckStateCalled)
+}
+
+func TestBaseApp_Precommit(t *testing.T) {
+	db := dbm.NewMemDB()
+	name := t.Name()
+	logger := log.NewTestLogger(t)
+
+	cp := &cmtproto.ConsensusParams{
+		Block: &cmtproto.BlockParams{
+			MaxGas: 5000000,
+		},
+	}
+
+	app := baseapp.NewBaseApp(name, logger, db, nil)
+	app.SetParamStore(&paramStore{db: dbm.NewMemDB()})
+	app.InitChain(abci.RequestInitChain{
+		ConsensusParams: cp,
+	})
+
+	wasPrecommiterCalled := false
+	app.SetPrecommiter(func(ctx sdk.Context) {
+		wasPrecommiterCalled = true
+	})
+	app.Seal()
+
+	app.Commit()
+	require.Equal(t, true, wasPrecommiterCalled)
+}
+
 func TestABCI_CheckTx(t *testing.T) {
 	// This ante handler reads the key and checks that the value matches the
 	// current counter. This ensures changes to the KVStore persist across
@@ -1322,6 +1376,49 @@ func TestABCI_GetBlockRetentionHeight(t *testing.T) {
 	}
 }
 
+// Verifies that PrepareCheckState is called with the checkState.
+func TestPrepareCheckStateCalledWithCheckState(t *testing.T) {
+	t.Parallel()
+
+	logger := log.NewTestLogger(t)
+	db := dbm.NewMemDB()
+	name := t.Name()
+	app := baseapp.NewBaseApp(name, logger, db, nil)
+
+	wasPrepareCheckStateCalled := false
+	app.SetPrepareCheckStater(func(ctx sdk.Context) {
+		require.Equal(t, true, ctx.IsCheckTx())
+		wasPrepareCheckStateCalled = true
+	})
+
+	app.BeginBlock(abci.RequestBeginBlock{Header: cmtproto.Header{Height: 1}})
+	app.Commit()
+
+	require.Equal(t, true, wasPrepareCheckStateCalled)
+}
+
+// Verifies that the Precommiter is called with the deliverState.
+func TestPrecommiterCalledWithDeliverState(t *testing.T) {
+	t.Parallel()
+
+	logger := log.NewTestLogger(t)
+	db := dbm.NewMemDB()
+	name := t.Name()
+	app := baseapp.NewBaseApp(name, logger, db, nil)
+
+	wasPrecommiterCalled := false
+	app.SetPrecommiter(func(ctx sdk.Context) {
+		require.Equal(t, false, ctx.IsCheckTx())
+		require.Equal(t, false, ctx.IsReCheckTx())
+		wasPrecommiterCalled = true
+	})
+
+	app.BeginBlock(abci.RequestBeginBlock{Header: cmtproto.Header{Height: 1}})
+	app.Commit()
+
+	require.Equal(t, true, wasPrecommiterCalled)
+}
+
 func TestABCI_Proposal_HappyPath(t *testing.T) {
 	anteKey := []byte("ante-key")
 	pool := mempool.NewSenderNonceMempool()

baseapp/baseapp.go

@@ -62,17 +62,19 @@ type BaseApp struct {
 	txDecoder         sdk.TxDecoder // unmarshal []byte into sdk.Tx
 	txEncoder         sdk.TxEncoder // marshal sdk.Tx into []byte
 
-	mempool         mempool.Mempool            // application side mempool
-	anteHandler     sdk.AnteHandler            // ante handler for fee and auth
-	postHandler     sdk.PostHandler            // post handler, optional, e.g. for tips
-	initChainer     sdk.InitChainer            // initialize state with validators and state blob
-	beginBlocker    sdk.BeginBlocker           // logic to run before any txs
-	processProposal sdk.ProcessProposalHandler // the handler which runs on ABCI ProcessProposal
-	prepareProposal sdk.PrepareProposalHandler // the handler which runs on ABCI PrepareProposal
-	endBlocker      sdk.EndBlocker             // logic to run after all txs, and to determine valset changes
-	addrPeerFilter  sdk.PeerFilter             // filter peers by address and port
-	idPeerFilter    sdk.PeerFilter             // filter peers by node ID
-	fauxMerkleMode  bool                       // if true, IAVL MountStores uses MountStoresDB for simulation speed.
+	mempool            mempool.Mempool            // application side mempool
+	anteHandler        sdk.AnteHandler            // ante handler for fee and auth
+	postHandler        sdk.PostHandler            // post handler, optional, e.g. for tips
+	initChainer        sdk.InitChainer            // initialize state with validators and state blob
+	beginBlocker       sdk.BeginBlocker           // logic to run before any txs
+	processProposal    sdk.ProcessProposalHandler // the handler which runs on ABCI ProcessProposal
+	prepareProposal    sdk.PrepareProposalHandler // the handler which runs on ABCI PrepareProposal
+	endBlocker         sdk.EndBlocker             // logic to run after all txs, and to determine valset changes
+	prepareCheckStater sdk.PrepareCheckStater     // logic to run during commit using the checkState
+	precommiter        sdk.Precommiter            // logic to run during commit using the deliverState
+	addrPeerFilter     sdk.PeerFilter             // filter peers by address and port
+	idPeerFilter       sdk.PeerFilter             // filter peers by node ID
+	fauxMerkleMode     bool                       // if true, IAVL MountStores uses MountStoresDB for simulation speed.
 
 	// manages snapshots, i.e. dumps of app state at certain intervals
 	snapshotManager *snapshots.Manager

baseapp/baseapp_test.go

@@ -388,6 +388,12 @@ func TestBaseAppOptionSeal(t *testing.T) {
 	require.Panics(t, func() {
 		suite.baseApp.SetEndBlocker(nil)
 	})
+	require.Panics(t, func() {
+		suite.baseApp.SetPrepareCheckStater(nil)
+	})
+	require.Panics(t, func() {
+		suite.baseApp.SetPrecommiter(nil)
+	})
 	require.Panics(t, func() {
 		suite.baseApp.SetAnteHandler(nil)
 	})

baseapp/options.go

@@ -167,6 +167,22 @@ func (app *BaseApp) SetEndBlocker(endBlocker sdk.EndBlocker) {
 	app.endBlocker = endBlocker
 }
 
+func (app *BaseApp) SetPrepareCheckStater(prepareCheckStater sdk.PrepareCheckStater) {
+	if app.sealed {
+		panic("SetPrepareCheckStater() on sealed BaseApp")
+	}
+
+	app.prepareCheckStater = prepareCheckStater
+}
+
+func (app *BaseApp) SetPrecommiter(precommiter sdk.Precommiter) {
+	if app.sealed {
+		panic("SetPrecommiter() on sealed BaseApp")
+	}
+
+	app.precommiter = precommiter
+}
+
 func (app *BaseApp) SetAnteHandler(ah sdk.AnteHandler) {
 	if app.sealed {
 		panic("SetAnteHandler() on sealed BaseApp")