From 5066667b32fc31a18a181934e24ce8dc599c9d5e Mon Sep 17 00:00:00 2001
From: Bryce Neal <bryce@dydx.exchange>
Date: Tue, 31 Jan 2023 17:20:00 -0500
Subject: [PATCH] Add Commit callback

---
 baseapp/abci.go        |  4 ++++
 baseapp/baseapp.go     |  1 +
 baseapp/options.go     |  8 ++++++++
 telemetry/wrapper.go   |  1 +
 types/abci.go          |  4 ++++
 types/module/module.go | 24 ++++++++++++++++++++++++
 6 files changed, 42 insertions(+)

diff --git a/baseapp/abci.go b/baseapp/abci.go
index 310d5b4012d97..41e7137bc863b 100644
--- a/baseapp/abci.go
+++ b/baseapp/abci.go
@@ -458,6 +458,10 @@ func (app *BaseApp) Commit() abci.ResponseCommit {
 	// empty/reset the deliver state
 	app.deliverState = nil
 
+	if app.commiter != nil {
+		app.commiter(app.checkState.ctx)
+	}
+
 	var halt bool
 
 	switch {
diff --git a/baseapp/baseapp.go b/baseapp/baseapp.go
index 391dac8c5e2cb..3c0ad082c3150 100644
--- a/baseapp/baseapp.go
+++ b/baseapp/baseapp.go
@@ -70,6 +70,7 @@ type BaseApp struct { //nolint: maligned
 	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
+	commiter        sdk.Commiter               // logic to run during commit
 	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.
diff --git a/baseapp/options.go b/baseapp/options.go
index 39398826f60ef..812d58bb5ebb5 100644
--- a/baseapp/options.go
+++ b/baseapp/options.go
@@ -166,6 +166,14 @@ func (app *BaseApp) SetEndBlocker(endBlocker sdk.EndBlocker) {
 	app.endBlocker = endBlocker
 }
 
+func (app *BaseApp) SetCommiter(commiter sdk.Commiter) {
+	if app.sealed {
+		panic("SetCommiter() on sealed BaseApp")
+	}
+
+	app.commiter = commiter
+}
+
 func (app *BaseApp) SetAnteHandler(ah sdk.AnteHandler) {
 	if app.sealed {
 		panic("SetAnteHandler() on sealed BaseApp")
diff --git a/telemetry/wrapper.go b/telemetry/wrapper.go
index e6542df86e19d..d354a7c8f2cc2 100644
--- a/telemetry/wrapper.go
+++ b/telemetry/wrapper.go
@@ -10,6 +10,7 @@ import (
 const (
 	MetricKeyBeginBlocker = "begin_blocker"
 	MetricKeyEndBlocker   = "end_blocker"
+	MetricKeyCommiter     = "commiter"
 	MetricLabelNameModule = "module"
 )
 
diff --git a/types/abci.go b/types/abci.go
index 2ea35e62c298e..9b582b20c1eec 100644
--- a/types/abci.go
+++ b/types/abci.go
@@ -19,6 +19,10 @@ type BeginBlocker func(ctx Context, req abci.RequestBeginBlock) (abci.ResponseBe
 // e.g. BFT timestamps rather than block height for any periodic EndBlock logic
 type EndBlocker func(ctx Context, req abci.RequestEndBlock) (abci.ResponseEndBlock, error)
 
+// Commiter runs code during commit after the block has been committed, and the `checkState` has been
+// branched for the new block.
+type Commiter func(ctx Context)
+
 // PeerFilter responds to p2p filtering queries from Tendermint
 type PeerFilter func(info string) abci.ResponseQuery
 
diff --git a/types/module/module.go b/types/module/module.go
index 2b9ab3eb8424b..f692a46fd05e1 100644
--- a/types/module/module.go
+++ b/types/module/module.go
@@ -210,6 +210,12 @@ type EndBlockAppModule interface {
 	EndBlock(sdk.Context, abci.RequestEndBlock) []abci.ValidatorUpdate
 }
 
+// CommitAppModule is an extension interface that contains information about the AppModule and Commit.
+type CommitAppModule interface {
+	AppModule
+	Commit(sdk.Context)
+}
+
 // GenesisOnlyAppModule is an AppModule that only has import/export functionality
 type GenesisOnlyAppModule struct {
 	AppModuleGenesis
@@ -256,6 +262,7 @@ type Manager struct {
 	OrderExportGenesis []string
 	OrderBeginBlockers []string
 	OrderEndBlockers   []string
+	OrderCommiters     []string
 	OrderMigrations    []string
 }
 
@@ -274,6 +281,7 @@ func NewManager(modules ...AppModule) *Manager {
 		OrderExportGenesis: modulesStr,
 		OrderBeginBlockers: modulesStr,
 		OrderEndBlockers:   modulesStr,
+		OrderCommiters:     modulesStr,
 	}
 }
 
@@ -349,6 +357,11 @@ func (m *Manager) SetOrderEndBlockers(moduleNames ...string) {
 	m.OrderEndBlockers = moduleNames
 }
 
+// SetOrderCommiters sets the order of set commiter calls
+func (m *Manager) SetOrderCommiters(moduleNames ...string) {
+	m.OrderCommiters = moduleNames
+}
+
 // SetOrderMigrations sets the order of migrations to be run. If not set
 // then migrations will be run with an order defined in `DefaultMigrationsOrder`.
 func (m *Manager) SetOrderMigrations(moduleNames ...string) {
@@ -682,6 +695,17 @@ func (m *Manager) EndBlock(ctx sdk.Context, req abci.RequestEndBlock) (abci.Resp
 	}, nil
 }
 
+// Commit performs commit functionality for all modules.
+func (m *Manager) Commit(ctx sdk.Context) {
+	for _, moduleName := range m.OrderCommiters {
+		module, ok := m.Modules[moduleName].(CommitAppModule)
+		if !ok {
+			continue
+		}
+		module.Commit(ctx)
+	}
+}
+
 // GetVersionMap gets consensus version from all modules
 func (m *Manager) GetVersionMap() VersionMap {
 	vermap := make(VersionMap)