fix(distr)!: use wrapped keeper for distribution begin blocker (#2234)

x/distribution/abci.go

@@ -0,0 +1,46 @@
+/*
+This file is identical to the Cosmos SDK distribution ABCI file (cosmos-sdk/x/distribution/abci.go).
+It is duplicated here because the BeginBlocker function accepts a keeper struct as parameter.
+
+Since we have our own keeper overrides the token allocation, we need to copy this file and modify the keeper type
+to use our custom keeper instead of the SDK's keeper.
+*/
+
+package distribution
+
+import (
+	"time"
+
+	"github.com/cosmos/cosmos-sdk/telemetry"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/x/distribution/types"
+	abci "github.com/tendermint/tendermint/abci/types"
+
+	"github.com/axelarnetwork/axelar-core/x/distribution/keeper"
+)
+
+// BeginBlocker sets the proposer for determining distribution during endblock
+// and distribute rewards for the previous block
+func BeginBlocker(ctx sdk.Context, req abci.RequestBeginBlock, k keeper.Keeper) {
+	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyBeginBlocker)
+
+	// determine the total power signing the block
+	var previousTotalPower, sumPreviousPrecommitPower int64
+	for _, voteInfo := range req.LastCommitInfo.GetVotes() {
+		previousTotalPower += voteInfo.Validator.Power
+		if voteInfo.SignedLastBlock {
+			sumPreviousPrecommitPower += voteInfo.Validator.Power
+		}
+	}
+
+	// TODO this is Tendermint-dependent
+	// ref https://github.com/cosmos/cosmos-sdk/issues/3095
+	if ctx.BlockHeight() > 1 {
+		previousProposer := k.GetPreviousProposerConsAddr(ctx)
+		k.AllocateTokens(ctx, sumPreviousPrecommitPower, previousTotalPower, previousProposer, req.LastCommitInfo.GetVotes())
+	}
+
+	// record the proposer for when we payout on the next block
+	consAddr := sdk.ConsAddress(req.Header.ProposerAddress)
+	k.SetPreviousProposerConsAddr(ctx, consAddr)
+}

x/distribution/abic_test.go

@@ -0,0 +1,51 @@
+package distribution_test
+
+import (
+	"fmt"
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"os"
+	"reflect"
+	"runtime"
+	"testing"
+
+	distr "github.com/cosmos/cosmos-sdk/x/distribution"
+	"github.com/stretchr/testify/assert"
+
+	"github.com/axelarnetwork/axelar-core/x/distribution"
+	"github.com/axelarnetwork/utils/funcs"
+)
+
+// TestEnsureIdentical ensure that our BeginBlocker implementation exactly matches the cosmos-sdk version.
+func TestEnsureIdentical(t *testing.T) {
+	customFn := runtime.FuncForPC(reflect.ValueOf(distribution.BeginBlocker).Pointer())
+	sdkFn := runtime.FuncForPC(reflect.ValueOf(distr.BeginBlocker).Pointer())
+
+	localFile, _ := customFn.FileLine(customFn.Entry())
+	sdkFile, _ := sdkFn.FileLine(sdkFn.Entry())
+
+	customContent := funcs.Must(os.ReadFile(localFile))
+	sdkContent := funcs.Must(os.ReadFile(sdkFile))
+
+	funcName := "BeginBlocker"
+	customImpl := extractFunctionBody(string(customContent), funcName)
+	sdkImpl := extractFunctionBody(string(sdkContent), funcName)
+
+	assert.Equal(t, customImpl, sdkImpl, fmt.Sprintf("%s implementation differs from SDK, please update abci.go to match the SDK version", funcName))
+}
+
+func extractFunctionBody(content, funcName string) string {
+	fset := token.NewFileSet()
+	file := funcs.Must(parser.ParseFile(fset, "", content, parser.ParseComments))
+
+	var functionBody string
+	ast.Inspect(file, func(n ast.Node) bool {
+		if fn, ok := n.(*ast.FuncDecl); ok && fn.Name.Name == funcName {
+			functionBody = content[fn.Body.Pos()-1 : fn.Body.End()]
+			return false
+		}
+		return true
+	})
+	return functionBody
+}

x/distribution/module.go

@@ -1,8 +1,10 @@
 package distribution
 
 import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	distr "github.com/cosmos/cosmos-sdk/x/distribution"
+	abci "github.com/tendermint/tendermint/abci/types"
 
 	"github.com/axelarnetwork/axelar-core/x/distribution/keeper"
 )
@@ -22,3 +24,7 @@ func NewAppModule(distrAppModule distr.AppModule, keeper keeper.Keeper) AppModul
 		keeper:    keeper,
 	}
 }
+
+func (am AppModule) BeginBlock(ctx sdk.Context, req abci.RequestBeginBlock) {
+	BeginBlocker(ctx, req, am.keeper)
+}