fix(x/epochs): di wiring (#25425)

Co-authored-by: Alex | Cosmos Labs <alex@cosmoslabs.io>

CHANGELOG.md

@@ -92,6 +92,7 @@ This patch update also includes minor dependency bumps.
 
 ### Bug Fixes
 
+* (x/epochs) [#25425](https://github.com/cosmos/cosmos-sdk/pull/25425) Fix `InvokeSetHooks` being called with a nil keeper and `AppModule` containing a copy instead of a pointer (hooks set post creating the `AppModule` like with depinject didn't apply because it's a different instance).
 * (client, client/rpc, x/auth/tx) [#24551](https://github.com/cosmos/cosmos-sdk/pull/24551) Handle cancellation properly when supplying context to client methods.
 * (x/authz) [#24638](https://github.com/cosmos/cosmos-sdk/pull/24638) Fixed a minor bug where the grant key was cast as a string and dumped directly into the error message leading to an error string possibly containing invalid UTF-8.
 * (client, client/rpc, x/auth/tx) [#24551](https://github.com/cosmos/cosmos-sdk/pull/24551) Handle cancellation properly when supplying context to client methods.

simapp/app.go

@@ -151,7 +151,7 @@ type SimApp struct {
 	// supplementary keepers
 	FeeGrantKeeper     feegrantkeeper.Keeper
 	AuthzKeeper        authzkeeper.Keeper
-	EpochsKeeper       epochskeeper.Keeper
+	EpochsKeeper       *epochskeeper.Keeper
 	ProtocolPoolKeeper protocolpoolkeeper.Keeper
 
 	// the module manager
@@ -449,11 +449,13 @@ func NewSimApp(
 	// If evidence needs to be handled for the app, set routes in router here and seal
 	app.EvidenceKeeper = *evidenceKeeper
 
-	app.EpochsKeeper = epochskeeper.NewKeeper(
+	epochsKeeper := epochskeeper.NewKeeper(
 		runtime.NewKVStoreService(keys[epochstypes.StoreKey]),
 		appCodec,
 	)
 
+	app.EpochsKeeper = &epochsKeeper
+
 	app.EpochsKeeper.SetHooks(
 		epochstypes.NewMultiEpochHooks(
 		// insert epoch hooks receivers here

simapp/app_di.go

@@ -76,7 +76,7 @@ type SimApp struct {
 	// supplementary keepers
 	FeeGrantKeeper     feegrantkeeper.Keeper
 	AuthzKeeper        authzkeeper.Keeper
-	EpochsKeeper       epochskeeper.Keeper
+	EpochsKeeper       *epochskeeper.Keeper
 	ProtocolPoolKeeper protocolpoolkeeper.Keeper
 
 	// simulation manager

x/epochs/README.md

@@ -95,6 +95,176 @@ func (k MyModuleKeeper) AfterEpochEnd(ctx sdk.Context, epochIdentifier string, e
 }
 ```
 
+### Wiring Hooks
+
+**Manual Wiring:**
+
+Import the following:
+
+```go
+import (
+    // ...
+    "github.com/cosmos/cosmos-sdk/x/epochs"
+    epochskeeper "github.com/cosmos/cosmos-sdk/x/epochs/keeper"
+    epochstypes "github.com/cosmos/cosmos-sdk/x/epochs/types"
+)
+```
+
+Add the epochs keeper to your application struct:
+
+```go
+EpochsKeeper *epochskeeper.Keeper
+```
+
+Add the store key:
+
+```go
+keys := storetypes.NewKVStoreKeys(
+    // ...
+    epochstypes.StoreKey,
+)
+```
+
+Instantiate the keeper:
+
+```go
+epochsKeeper := epochskeeper.NewKeeper(
+    runtime.NewKVStoreService(keys[epochstypes.StoreKey]),
+    appCodec,
+)
+
+app.EpochsKeeper = &epochsKeeper
+```
+
+Set up hooks for the epochs keeper:
+
+```go
+app.EpochsKeeper.SetHooks(
+    epochstypes.NewMultiEpochHooks(
+        // insert epoch hooks receivers here
+        app.SomeOtherModule
+    ),
+)
+```
+
+Add the epochs module to the module manager:
+
+```go
+app.ModuleManager = module.NewManager(
+    // ...
+    epochs.NewAppModule(appCodec, app.EpochsKeeper),
+)
+```
+
+Add entries for SetOrderBeginBlockers and SetOrderInitGenesis:
+
+```go
+app.ModuleManager.SetOrderBeginBlockers(
+    // ...
+    epochstypes.ModuleName,
+)
+```
+
+```go
+app.ModuleManager.SetOrderInitGenesis(
+    // ...
+    epochstypes.ModuleName,
+)
+```
+
+**DI Wiring:**
+
+First, set up the keeper for the application.
+
+Import the epochs keeper:
+
+```go
+epochskeeper "github.com/cosmos/cosmos-sdk/x/epochs/keeper"
+```
+
+Add the keeper to your application struct:
+
+```go
+EpochsKeeper *epochskeeper.Keeper
+```
+
+Add the keeper to the depinject system:
+
+```go
+depinject.Inject(
+    appConfig,
+    &appBuilder,
+    &app.appCodec,
+    &app.legacyAmino,
+    &app.txConfig,
+    &app.interfaceRegistry,
+    // ... other modules
+    &app.EpochsKeeper, // NEW MODULE!
+)
+```
+
+Next, set up configuration for the module.
+
+Import the following:
+
+```go
+import (
+    epochsmodulev1 "cosmossdk.io/api/cosmos/epochs/module/v1"
+    
+    _ "github.com/cosmos/cosmos-sdk/x/epochs" // import for side-effects
+    epochstypes "github.com/cosmos/cosmos-sdk/x/epochs/types"
+)
+```
+
+Add an entry for BeginBlockers and InitGenesis:
+
+```go
+BeginBlockers: []string{
+    // ...
+    epochstypes.ModuleName,
+},
+```
+
+```go
+InitGenesis: []string{
+    // ...
+    epochstypes.ModuleName,
+},
+```
+
+Add an entry for epochs in the ModuleConfig:
+
+```go
+{
+    Name:   epochstypes.ModuleName,
+    Config: appconfig.WrapAny(&epochsmodulev1.Module{}),
+},
+```
+
+depinject can automatically add your hooks to the epochs `Keeper`. For it do so, specify an output of your module with the type `epochtypes.EpochHooksWrapper`, ie:
+
+```go
+type TestInputs struct {
+	depinject.In
+}
+
+type TestOutputs struct {
+	depinject.Out
+
+	Hooks types.EpochHooksWrapper
+}
+
+func DummyProvider(in TestInputs) TestOutputs {
+	return TestOutputs{
+		Hooks: types.EpochHooksWrapper{
+			EpochHooks: testEpochHooks{},
+		},
+	}
+}
+```
+
+for an example see [`depinject_test.go`](https://github.com/cosmos/cosmos-sdk/tree/main/x/epochs/depinject_test.go)
+
 ### Panic isolation
 
 If a given epoch hook panics, its state update is reverted, but we keep

x/epochs/depinject.go

@@ -39,14 +39,14 @@ type ModuleInputs struct {
 type ModuleOutputs struct {
 	depinject.Out
 
-	EpochKeeper keeper.Keeper
+	EpochKeeper *keeper.Keeper
 	Module      appmodule.AppModule
 }
 
 func ProvideModule(in ModuleInputs) ModuleOutputs {
 	k := keeper.NewKeeper(in.StoreService, in.Cdc)
-	m := NewAppModule(k)
-	return ModuleOutputs{EpochKeeper: k, Module: m}
+	m := NewAppModule(&k)
+	return ModuleOutputs{EpochKeeper: m.keeper, Module: m}
 }
 
 func InvokeSetHooks(keeper *keeper.Keeper, hooks map[string]types.EpochHooksWrapper) error {

x/epochs/depinject_test.go

@@ -6,22 +6,41 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	appv1alpha1 "cosmossdk.io/api/cosmos/app/v1alpha1"
+	bankmodulev1 "cosmossdk.io/api/cosmos/bank/module/v1"
+	modulev1 "cosmossdk.io/api/cosmos/epochs/module/v1"
+	"cosmossdk.io/core/appmodule"
+	"cosmossdk.io/core/store"
+	"cosmossdk.io/depinject"
+	"cosmossdk.io/depinject/appconfig"
 	storetypes "cosmossdk.io/store/types"
 
+	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/runtime"
 	"github.com/cosmos/cosmos-sdk/types/module/testutil"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
 	"github.com/cosmos/cosmos-sdk/x/epochs"
 	"github.com/cosmos/cosmos-sdk/x/epochs/keeper"
 	"github.com/cosmos/cosmos-sdk/x/epochs/types"
 )
 
+var _ types.EpochHooks = testEpochHooks{}
+
 type testEpochHooks struct{}
 
-func (h testEpochHooks) AfterEpochEnd(ctx context.Context, epochIdentifier string, epochNumber int64) error {
+func (h testEpochHooks) AfterEpochEnd(
+	ctx context.Context,
+	epochIdentifier string,
+	epochNumber int64,
+) error {
 	return nil
 }
 
-func (h testEpochHooks) BeforeEpochStart(ctx context.Context, epochIdentifier string, epochNumber int64) error {
+func (h testEpochHooks) BeforeEpochStart(
+	ctx context.Context,
+	epochIdentifier string,
+	epochNumber int64,
+) error {
 	return nil
 }
 
@@ -58,3 +77,88 @@ func TestInvokeSetHooks(t *testing.T) {
 	require.Equal(t, hook1, multiHooks[0])
 	require.Equal(t, hook2, multiHooks[1])
 }
+
+type TestInputs struct {
+	depinject.In
+}
+
+type TestOutputs struct {
+	depinject.Out
+
+	Hooks types.EpochHooksWrapper
+}
+
+func DummyProvider(in TestInputs) TestOutputs {
+	return TestOutputs{
+		Hooks: types.EpochHooksWrapper{
+			EpochHooks: testEpochHooks{},
+		},
+	}
+}
+
+func ProvideDeps(depinject.In) struct {
+	depinject.Out
+	Cdc          codec.Codec
+	StoreService store.KVStoreService
+} {
+	encCfg := testutil.MakeTestEncodingConfig()
+
+	key := storetypes.NewKVStoreKey(types.StoreKey)
+	return struct {
+		depinject.Out
+		Cdc          codec.Codec
+		StoreService store.KVStoreService
+	}{
+		Cdc:          encCfg.Codec,
+		StoreService: runtime.NewKVStoreService(key),
+	}
+}
+
+func TestDepinject(t *testing.T) {
+	/// we just need any module's proto to register the provider here, no specific reason to use bank
+	appconfig.RegisterModule(&bankmodulev1.Module{}, appconfig.Provide(DummyProvider))
+	var appModules map[string]appmodule.AppModule
+	keeper := new(keeper.Keeper)
+	require.NoError(t,
+		depinject.Inject(
+			depinject.Configs(
+				appconfig.Compose(
+					&appv1alpha1.Config{
+						Modules: []*appv1alpha1.ModuleConfig{
+							{
+								Name:   banktypes.ModuleName,
+								Config: appconfig.WrapAny(&bankmodulev1.Module{}),
+							},
+							{
+								Name:   types.ModuleName,
+								Config: appconfig.WrapAny(&modulev1.Module{}),
+							},
+						},
+					},
+				),
+				depinject.Provide(ProvideDeps),
+			),
+			&appModules,
+			&keeper,
+		),
+	)
+
+	require.NotNil(t, keeper, "expected keeper to not be nil after depinject")
+	multihooks, ok := keeper.Hooks().(types.MultiEpochHooks)
+	require.True(t, ok, "expected keeper to have MultiEpochHooks after depinject")
+	require.Len(t, multihooks, 1, "expected MultiEpochHooks to have 1 element after depinject")
+	require.Equal(
+		t,
+		types.EpochHooksWrapper{EpochHooks: testEpochHooks{}},
+		multihooks[0],
+		"expected the only hook in MultiEpochHooks to be the test hook",
+	)
+	module, ok := appModules[types.ModuleName].(epochs.AppModule)
+	require.True(t, ok, "expected depinject to fill map with the epochs AppModule")
+	require.Equal(
+		t,
+		types.MultiEpochHooks{types.EpochHooksWrapper{EpochHooks: testEpochHooks{}}},
+		module.Keeper().Hooks(),
+	)
+	require.Same(t, keeper, module.Keeper()) // pointers pointing to the same instance
+}

x/epochs/export_test.go

@@ -0,0 +1,7 @@
+package epochs
+
+import "github.com/cosmos/cosmos-sdk/x/epochs/keeper"
+
+func (am AppModule) Keeper() *keeper.Keeper {
+	return am.keeper
+}

x/epochs/module.go

@@ -33,11 +33,11 @@ const ConsensusVersion = 1
 
 // AppModule implements the AppModule interface for the epochs module.
 type AppModule struct {
-	keeper keeper.Keeper
+	keeper *keeper.Keeper
 }
 
 // NewAppModule creates a new AppModule object.
-func NewAppModule(keeper keeper.Keeper) AppModule {
+func NewAppModule(keeper *keeper.Keeper) AppModule {
 	return AppModule{
 		keeper: keeper,
 	}
@@ -66,7 +66,7 @@ func (AppModule) RegisterGRPCGatewayRoutes(clientCtx client.Context, mux *gwrunt
 
 // RegisterServices registers module services.
 func (am AppModule) RegisterServices(registrar grpc.ServiceRegistrar) error {
-	types.RegisterQueryServer(registrar, keeper.NewQuerier(am.keeper))
+	types.RegisterQueryServer(registrar, keeper.NewQuerier(*am.keeper))
 	return nil
 }