refactor: simplify optional tendermint pruning migrations (cosmos#2862)

docs/migrations/v6-to-v7.md

@@ -21,7 +21,7 @@ Add the following to the function call to the upgrade handler in `app/app.go`, t
 ```go
 import (
     // ...
-    ibctm "github.com/cosmos/ibc-go/v6/modules/light-clients/07-tendermint"
+    ibctmmigrations "github.com/cosmos/ibc-go/v6/modules/light-clients/07-tendermint/migrations"
 )
 
 // ...
@@ -30,7 +30,10 @@ app.UpgradeKeeper.SetUpgradeHandler(
     upgradeName,
     func(ctx sdk.Context, _ upgradetypes.Plan, _ module.VersionMap) (module.VersionMap, error) {
         // prune expired tendermint consensus states to save storage space
-        ibctm.PruneTendermintConsensusStates(ctx, app.Codec, appCodec, keys[ibchost.StoreKey])
+        _, err := ibctmmigrations.PruneExpiredConsensusStates(ctx, app.Codec, app.IBCKeeper.ClientKeeper)
+        if err != nil {
+            return nil, err
+        }
 
         return app.mm.RunMigrations(ctx, app.configurator, fromVM)
     },

modules/light-clients/07-tendermint/migrations/expected_keepers.go

@@ -0,0 +1,16 @@
+package migrations
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/tendermint/tendermint/libs/log"
+
+	"github.com/cosmos/ibc-go/v6/modules/core/exported"
+)
+
+// ClientKeeper expected account IBC client keeper
+type ClientKeeper interface {
+	GetClientState(ctx sdk.Context, clientID string) (exported.ClientState, bool)
+	IterateClientStates(ctx sdk.Context, prefix []byte, cb func(string, exported.ClientState) bool)
+	ClientStore(ctx sdk.Context, clientID string) sdk.KVStore
+	Logger(ctx sdk.Context) log.Logger
+}

modules/light-clients/07-tendermint/migrations/migrations.go

@@ -0,0 +1,46 @@
+package migrations
+
+import (
+	"github.com/cosmos/cosmos-sdk/codec"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+
+	clienttypes "github.com/cosmos/ibc-go/v6/modules/core/02-client/types"
+	"github.com/cosmos/ibc-go/v6/modules/core/exported"
+	ibctm "github.com/cosmos/ibc-go/v6/modules/light-clients/07-tendermint"
+)
+
+// PruneExpiredConsensusStates prunes all expired tendermint consensus states. This function
+// may optionally be called during in-place store migrations. The ibc store key must be provided.
+func PruneExpiredConsensusStates(ctx sdk.Context, cdc codec.BinaryCodec, clientKeeper ClientKeeper) (int, error) {
+	var clientIDs []string
+	clientKeeper.IterateClientStates(ctx, []byte(exported.Tendermint), func(clientID string, _ exported.ClientState) bool {
+		clientIDs = append(clientIDs, clientID)
+		return false
+	})
+
+	// keep track of the total consensus states pruned so chains can
+	// understand how much space is saved when the migration is run
+	var totalPruned int
+
+	for _, clientID := range clientIDs {
+		clientStore := clientKeeper.ClientStore(ctx, clientID)
+
+		clientState, ok := clientKeeper.GetClientState(ctx, clientID)
+		if !ok {
+			return 0, sdkerrors.Wrapf(clienttypes.ErrClientNotFound, "clientID %s", clientID)
+		}
+
+		tmClientState, ok := clientState.(*ibctm.ClientState)
+		if !ok {
+			return 0, sdkerrors.Wrap(clienttypes.ErrInvalidClient, "client state is not tendermint even though client id contains 07-tendermint")
+		}
+
+		totalPruned += ibctm.PruneAllExpiredConsensusStates(ctx, clientStore, cdc, tmClientState)
+	}
+
+	clientLogger := clientKeeper.Logger(ctx)
+	clientLogger.Info("pruned expired tendermint consensus states", "total", totalPruned)
+
+	return totalPruned, nil
+}

modules/light-clients/07-tendermint/migrations_test.go → modules/light-clients/07-tendermint/migrations/migrations_test.go

@@ -1,17 +1,41 @@
-package tendermint_test
+package migrations_test
 
 import (
+	"testing"
 	"time"
 
+	"github.com/stretchr/testify/suite"
+
 	clienttypes "github.com/cosmos/ibc-go/v6/modules/core/02-client/types"
 	host "github.com/cosmos/ibc-go/v6/modules/core/24-host"
 	"github.com/cosmos/ibc-go/v6/modules/core/exported"
 	ibctm "github.com/cosmos/ibc-go/v6/modules/light-clients/07-tendermint"
+	ibctmmigrations "github.com/cosmos/ibc-go/v6/modules/light-clients/07-tendermint/migrations"
 	ibctesting "github.com/cosmos/ibc-go/v6/testing"
 )
 
+type MigrationsTestSuite struct {
+	suite.Suite
+
+	coordinator *ibctesting.Coordinator
+
+	// testing chains used for convenience and readability
+	chainA *ibctesting.TestChain
+	chainB *ibctesting.TestChain
+}
+
+func (suite *MigrationsTestSuite) SetupTest() {
+	suite.coordinator = ibctesting.NewCoordinator(suite.T(), 2)
+	suite.chainA = suite.coordinator.GetChain(ibctesting.GetChainID(1))
+	suite.chainB = suite.coordinator.GetChain(ibctesting.GetChainID(2))
+}
+
+func TestTendermintTestSuite(t *testing.T) {
+	suite.Run(t, new(MigrationsTestSuite))
+}
+
 // test pruning of multiple expired tendermint consensus states
-func (suite *TendermintTestSuite) TestPruneTendermintConsensusStates() {
+func (suite *MigrationsTestSuite) TestPruneExpiredConsensusStates() {
 	// create multiple tendermint clients and a solo machine client
 	// the solo machine is used to verify this pruning function only modifies
 	// the tendermint store.
@@ -99,8 +123,9 @@ func (suite *TendermintTestSuite) TestPruneTendermintConsensusStates() {
 	// This will cause the consensus states created before the first time increment
 	// to be expired
 	suite.coordinator.IncrementTimeBy(7 * 24 * time.Hour)
-	err = ibctm.PruneTendermintConsensusStates(suite.chainA.GetContext(), suite.chainA.App.AppCodec(), suite.chainA.GetSimApp().GetKey(host.StoreKey))
+	totalPruned, err := ibctmmigrations.PruneExpiredConsensusStates(suite.chainA.GetContext(), suite.chainA.App.AppCodec(), suite.chainA.GetSimApp().IBCKeeper.ClientKeeper)
 	suite.Require().NoError(err)
+	suite.Require().NotZero(totalPruned)
 
 	for _, path := range paths {
 		ctx := suite.chainA.GetContext()

testing/simapp/app.go

@@ -899,7 +899,7 @@ func (app *SimApp) setupUpgradeHandlers() {
 			app.mm,
 			app.configurator,
 			app.appCodec,
-			app.keys[ibchost.StoreKey],
+			app.IBCKeeper.ClientKeeper,
 		),
 	)
 }

testing/simapp/upgrades/v7/upgrades.go

@@ -2,12 +2,12 @@ package v7
 
 import (
 	"github.com/cosmos/cosmos-sdk/codec"
-	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
 
-	ibctm "github.com/cosmos/ibc-go/v6/modules/light-clients/07-tendermint"
+	clientkeeper "github.com/cosmos/ibc-go/v6/modules/core/02-client/keeper"
+	ibctmmigrations "github.com/cosmos/ibc-go/v6/modules/light-clients/07-tendermint/migrations"
 )
 
 const (
@@ -20,13 +20,14 @@ func CreateUpgradeHandler(
 	mm *module.Manager,
 	configurator module.Configurator,
 	cdc codec.BinaryCodec,
-	hostStoreKey *storetypes.KVStoreKey,
+	clientKeeper clientkeeper.Keeper,
 ) upgradetypes.UpgradeHandler {
 	return func(ctx sdk.Context, _ upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
 		// OPTIONAL: prune expired tendermint consensus states to save storage space
-		if err := ibctm.PruneTendermintConsensusStates(ctx, cdc, hostStoreKey); err != nil {
+		if _, err := ibctmmigrations.PruneExpiredConsensusStates(ctx, cdc, clientKeeper); err != nil {
 			return nil, err
 		}
+
 		return mm.RunMigrations(ctx, configurator, vm)
 	}
 }