refactor(distribution)!: add cometinfo (#20588)

Author: tac0turtle

simapp/app.go

@@ -357,7 +357,7 @@ func NewSimApp(
 
 	app.PoolKeeper = poolkeeper.NewKeeper(appCodec, runtime.NewEnvironment(runtime.NewKVStoreService(keys[pooltypes.StoreKey]), logger.With(log.ModuleKey, "x/protocolpool")), app.AuthKeeper, app.BankKeeper, app.StakingKeeper, authtypes.NewModuleAddress(govtypes.ModuleName).String())
 
-	app.DistrKeeper = distrkeeper.NewKeeper(appCodec, runtime.NewEnvironment(runtime.NewKVStoreService(keys[distrtypes.StoreKey]), logger.With(log.ModuleKey, "x/distribution")), app.AuthKeeper, app.BankKeeper, app.StakingKeeper, app.PoolKeeper, authtypes.FeeCollectorName, authtypes.NewModuleAddress(govtypes.ModuleName).String())
+	app.DistrKeeper = distrkeeper.NewKeeper(appCodec, runtime.NewEnvironment(runtime.NewKVStoreService(keys[distrtypes.StoreKey]), logger.With(log.ModuleKey, "x/distribution")), app.AuthKeeper, app.BankKeeper, app.StakingKeeper, app.PoolKeeper, cometService, authtypes.FeeCollectorName, authtypes.NewModuleAddress(govtypes.ModuleName).String())
 
 	app.SlashingKeeper = slashingkeeper.NewKeeper(runtime.NewEnvironment(runtime.NewKVStoreService(keys[slashingtypes.StoreKey]), logger.With(log.ModuleKey, "x/slashing")),
 		appCodec, legacyAmino, app.StakingKeeper, authtypes.NewModuleAddress(govtypes.ModuleName).String(),

tests/integration/distribution/keeper/msg_server_test.go

@@ -137,7 +137,7 @@ func initFixture(t *testing.T) *fixture {
 	poolKeeper := poolkeeper.NewKeeper(cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[pooltypes.StoreKey]), log.NewNopLogger()), accountKeeper, bankKeeper, stakingKeeper, authority.String())
 
 	distrKeeper := distrkeeper.NewKeeper(
-		cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[distrtypes.StoreKey]), logger), accountKeeper, bankKeeper, stakingKeeper, poolKeeper, distrtypes.ModuleName, authority.String(),
+		cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[distrtypes.StoreKey]), logger), accountKeeper, bankKeeper, stakingKeeper, poolKeeper, cometService, distrtypes.ModuleName, authority.String(),
 	)
 
 	authModule := auth.NewAppModule(cdc, accountKeeper, acctsModKeeper, authsims.RandomGenesisAccounts)
@@ -163,6 +163,7 @@ func initFixture(t *testing.T) *fixture {
 				},
 			},
 		},
+		ProposerAddress: valConsAddr,
 	})
 
 	integrationApp := integration.NewIntegrationApp(ctx, logger, keys, cdc,

x/distribution/CHANGELOG.md

@@ -31,6 +31,8 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+
+* [#20588](https://github.com/cosmos/cosmos-sdk/pull/20588) `x/distribution` now takes cometService in order to get consensus related information.
 * [#19868](https://github.com/cosmos/cosmos-sdk/pull/19868) Removes Accounts String method
     * `NewMsgSetWithdrawAddress` now takes strings as argument instead of `sdk.AccAddress`.
     * `NewGenesisState` now takes a string as argument instead of `sdk.ConsAddress`.
@@ -74,4 +76,4 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### Bug Fixes
 
-* [#19301](https://github.com/cosmos/cosmos-sdk/pull/19301) Fix vulnerability in `incrementReferenceCount` in distribution.
+* [#19301](https://github.com/cosmos/cosmos-sdk/pull/19301) Fix vulnerability in `incrementReferenceCount` in distribution.

x/distribution/depinject.go

@@ -3,6 +3,7 @@ package distribution
 import (
 	modulev1 "cosmossdk.io/api/cosmos/distribution/module/v1"
 	"cosmossdk.io/core/appmodule"
+	"cosmossdk.io/core/comet"
 	"cosmossdk.io/depinject"
 	"cosmossdk.io/depinject/appconfig"
 	authtypes "cosmossdk.io/x/auth/types"
@@ -27,9 +28,10 @@ func init() {
 type ModuleInputs struct {
 	depinject.In
 
-	Config      *modulev1.Module
-	Environment appmodule.Environment
-	Cdc         codec.Codec
+	Config       *modulev1.Module
+	Environment  appmodule.Environment
+	Cdc          codec.Codec
+	CometService comet.Service
 
 	AccountKeeper types.AccountKeeper
 	BankKeeper    types.BankKeeper
@@ -69,6 +71,7 @@ func ProvideModule(in ModuleInputs) ModuleOutputs {
 		in.BankKeeper,
 		in.StakingKeeper,
 		in.PoolKeeper,
+		in.CometService,
 		feeCollectorName,
 		authorityAddr,
 	)

x/distribution/keeper/abci.go

@@ -1,6 +1,8 @@
 package keeper
 
 import (
+	"context"
+
 	"cosmossdk.io/x/distribution/types"
 
 	"github.com/cosmos/cosmos-sdk/telemetry"
@@ -10,31 +12,33 @@ import (
 // BeginBlocker sets the proposer for determining distribution during endblock
 // and distribute rewards for the previous block.
 // TODO: use context.Context after including the comet service
-func (k Keeper) BeginBlocker(ctx sdk.Context) error {
+func (k Keeper) BeginBlocker(ctx context.Context) error {
 	defer telemetry.ModuleMeasureSince(types.ModuleName, telemetry.Now(), telemetry.MetricKeyBeginBlocker)
 
 	// determine the total power signing the block
 	var previousTotalPower int64
-	for _, vote := range ctx.CometInfo().LastCommit.Votes {
+	header := k.HeaderService.HeaderInfo(ctx)
+	ci := k.cometService.CometInfo(ctx)
+	for _, vote := range ci.LastCommit.Votes {
 		previousTotalPower += vote.Validator.Power
 	}
 
 	// TODO this is Tendermint-dependent
 	// ref https://github.com/cosmos/cosmos-sdk/issues/3095
-	if ctx.BlockHeight() > 1 {
-		if err := k.AllocateTokens(ctx, previousTotalPower, ctx.CometInfo().LastCommit.Votes); err != nil {
+	if header.Height > 1 {
+		if err := k.AllocateTokens(ctx, previousTotalPower, ci.LastCommit.Votes); err != nil {
 			return err
 		}
 
 		// every 1000 blocks send whole coins from decimal pool to community pool
-		if ctx.BlockHeight()%1000 == 0 {
+		if header.Height%1000 == 0 {
 			if err := k.sendDecimalPoolToCommunityPool(ctx); err != nil {
 				return err
 			}
 		}
 	}
 
 	// record the proposer for when we payout on the next block
-	consAddr := sdk.ConsAddress(ctx.BlockHeader().ProposerAddress)
+	consAddr := sdk.ConsAddress(ci.ProposerAddress)
 	return k.PreviousProposer.Set(ctx, consAddr)
 }

x/distribution/keeper/allocation_test.go

@@ -1,6 +1,7 @@
 package keeper_test
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -28,6 +29,17 @@ import (
 	moduletestutil "github.com/cosmos/cosmos-sdk/types/module/testutil"
 )
 
+var _ comet.Service = (*emptyCometService)(nil)
+
+type emptyCometService struct{}
+
+// CometInfo implements comet.Service.
+func (e *emptyCometService) CometInfo(context.Context) comet.Info {
+	return comet.Info{}
+}
+
+var testCometService = &emptyCometService{}
+
 func TestAllocateTokensToValidatorWithCommission(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	key := storetypes.NewKVStoreKey(disttypes.StoreKey)
@@ -58,6 +70,7 @@ func TestAllocateTokensToValidatorWithCommission(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)
@@ -124,6 +137,7 @@ func TestAllocateTokensToManyValidators(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)
@@ -264,6 +278,7 @@ func TestAllocateTokensTruncation(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)

x/distribution/keeper/delegation_test.go

@@ -54,6 +54,7 @@ func TestCalculateRewardsBasic(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)
@@ -167,6 +168,7 @@ func TestCalculateRewardsAfterSlash(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)
@@ -283,6 +285,7 @@ func TestCalculateRewardsAfterManySlashes(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)
@@ -420,6 +423,7 @@ func TestCalculateRewardsMultiDelegator(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)
@@ -530,6 +534,7 @@ func TestWithdrawDelegationRewardsBasic(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)
@@ -618,6 +623,7 @@ func TestCalculateRewardsAfterManySlashesInSameBlock(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)
@@ -747,6 +753,7 @@ func TestCalculateRewardsMultiDelegatorMultiSlash(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)
@@ -900,6 +907,7 @@ func TestCalculateRewardsMultiDelegatorMultWithdraw(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)
@@ -1114,6 +1122,7 @@ func Test100PercentCommissionReward(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)

x/distribution/keeper/keeper.go

@@ -8,6 +8,7 @@ import (
 	"cosmossdk.io/collections"
 	collcodec "cosmossdk.io/collections/codec"
 	"cosmossdk.io/core/appmodule"
+	"cosmossdk.io/core/comet"
 	"cosmossdk.io/core/event"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/x/distribution/types"
@@ -21,6 +22,8 @@ import (
 type Keeper struct {
 	appmodule.Environment
 
+	cometService comet.Service
+
 	cdc           codec.BinaryCodec
 	authKeeper    types.AccountKeeper
 	bankKeeper    types.BankKeeper
@@ -57,8 +60,13 @@ type Keeper struct {
 
 // NewKeeper creates a new distribution Keeper instance
 func NewKeeper(
-	cdc codec.BinaryCodec, env appmodule.Environment,
-	ak types.AccountKeeper, bk types.BankKeeper, sk types.StakingKeeper, pk types.PoolKeeper,
+	cdc codec.BinaryCodec,
+	env appmodule.Environment,
+	ak types.AccountKeeper,
+	bk types.BankKeeper,
+	sk types.StakingKeeper,
+	pk types.PoolKeeper,
+	cometService comet.Service,
 	feeCollectorName, authority string,
 ) Keeper {
 	// ensure distribution module account is set
@@ -69,6 +77,7 @@ func NewKeeper(
 	sb := collections.NewSchemaBuilder(env.KVStoreService)
 	k := Keeper{
 		Environment:      env,
+		cometService:     cometService,
 		cdc:              cdc,
 		authKeeper:       ak,
 		bankKeeper:       bk,

x/distribution/keeper/keeper_test.go

@@ -70,6 +70,7 @@ func initFixture(t *testing.T) (sdk.Context, []sdk.AccAddress, keeper.Keeper, de
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		testCometService,
 		"fee_collector",
 		authorityAddr,
 	)

x/distribution/migrations/v4/migrate_funds_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
 
+	"cosmossdk.io/core/comet"
 	"cosmossdk.io/log"
 	storetypes "cosmossdk.io/store/types"
 	"cosmossdk.io/x/auth"
@@ -31,6 +32,13 @@ import (
 	moduletestutil "github.com/cosmos/cosmos-sdk/types/module/testutil"
 )
 
+type emptyCometService struct{}
+
+// CometInfo implements comet.Service.
+func (e *emptyCometService) CometInfo(context.Context) comet.Info {
+	return comet.Info{}
+}
+
 func TestFundsMigration(t *testing.T) {
 	keys := storetypes.NewKVStoreKeys(
 		authtypes.StoreKey, banktypes.StoreKey, disttypes.StoreKey,
@@ -90,6 +98,7 @@ func TestFundsMigration(t *testing.T) {
 		bankKeeper,
 		stakingKeeper,
 		poolKeeper,
+		&emptyCometService{},
 		disttypes.ModuleName,
 		authority,
 	)