feat(runtime/v2): store loader on simappv2 (#21704)

Co-authored-by: marbar3778 <marbar3778@yahoo.com>

CHANGELOG.md

@@ -43,6 +43,7 @@ Every module contains its own CHANGELOG.md. Please refer to the module you are i
 ### Features
 
 * (baseapp) [#20291](https://github.com/cosmos/cosmos-sdk/pull/20291) Simulate nested messages.
+* (runtime) [#21704](https://github.com/cosmos/cosmos-sdk/pull/21704) Add StoreLoader in simappv2.
 
 ### Improvements
 

runtime/store.go

@@ -7,6 +7,7 @@ import (
 	"cosmossdk.io/core/store"
 	storetypes "cosmossdk.io/store/types"
 
+	"github.com/cosmos/cosmos-sdk/baseapp"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -179,3 +180,23 @@ func (s kvStoreAdapter) ReverseIterator(start, end []byte) store.Iterator {
 func KVStoreAdapter(store store.KVStore) storetypes.KVStore {
 	return &kvStoreAdapter{store}
 }
+
+// UpgradeStoreLoader is used to prepare baseapp with a fixed StoreLoader
+// pattern. This is useful for custom upgrade loading logic.
+func UpgradeStoreLoader(upgradeHeight int64, storeUpgrades *store.StoreUpgrades) baseapp.StoreLoader {
+	return func(ms storetypes.CommitMultiStore) error {
+		if upgradeHeight == ms.LastCommitID().Version+1 {
+			// Check if the current commit version and upgrade height matches
+			if len(storeUpgrades.Deleted) > 0 || len(storeUpgrades.Added) > 0 {
+				stup := &storetypes.StoreUpgrades{
+					Added:   storeUpgrades.Added,
+					Deleted: storeUpgrades.Deleted,
+				}
+				return ms.LoadLatestVersionAndUpgrade(stup)
+			}
+		}
+
+		// Otherwise load default store loader
+		return baseapp.DefaultStoreLoader(ms)
+	}
+}

runtime/v2/app.go

@@ -46,6 +46,8 @@ type App[T transaction.Tx] struct {
 	// GRPCMethodsToMessageMap maps gRPC method name to a function that decodes the request
 	// bytes into a gogoproto.Message, which then can be passed to appmanager.
 	GRPCMethodsToMessageMap map[string]func() gogoproto.Message
+
+	storeLoader StoreLoader
 }
 
 // Name returns the app name.
@@ -68,9 +70,14 @@ func (a *App[T]) DefaultGenesis() map[string]json.RawMessage {
 	return a.moduleManager.DefaultGenesis()
 }
 
+// SetStoreLoader sets the store loader.
+func (a *App[T]) SetStoreLoader(loader StoreLoader) {
+	a.storeLoader = loader
+}
+
 // LoadLatest loads the latest version.
 func (a *App[T]) LoadLatest() error {
-	return a.db.LoadLatestVersion()
+	return a.storeLoader(a.db)
 }
 
 // LoadHeight loads a particular height

runtime/v2/module.go

@@ -130,6 +130,7 @@ func ProvideAppBuilder[T transaction.Tx](
 		msgRouterBuilder:        msgRouterBuilder,
 		queryRouterBuilder:      stf.NewMsgRouterBuilder(), // TODO dedicated query router
 		GRPCMethodsToMessageMap: map[string]func() proto.Message{},
+		storeLoader:             DefaultStoreLoader,
 	}
 	appBuilder := &AppBuilder[T]{app: app}
 

runtime/v2/store.go

@@ -1,6 +1,8 @@
 package runtime
 
 import (
+	"fmt"
+
 	"cosmossdk.io/core/store"
 	"cosmossdk.io/server/v2/stf"
 	storev2 "cosmossdk.io/store/v2"
@@ -55,3 +57,34 @@ type Store interface {
 	// LastCommitID returns the latest commit ID
 	LastCommitID() (proof.CommitID, error)
 }
+
+// StoreLoader allows for custom loading of the store, this is useful when upgrading the store from a previous version
+type StoreLoader func(store Store) error
+
+// DefaultStoreLoader just calls LoadLatestVersion on the store
+func DefaultStoreLoader(store Store) error {
+	return store.LoadLatestVersion()
+}
+
+// UpgradeStoreLoader upgrades the store if the upgrade height matches the current version, it is used as a replacement
+// for the DefaultStoreLoader when there are store upgrades
+func UpgradeStoreLoader(upgradeHeight int64, storeUpgrades *store.StoreUpgrades) StoreLoader {
+	return func(store Store) error {
+		latestVersion, err := store.GetLatestVersion()
+		if err != nil {
+			return err
+		}
+
+		if uint64(upgradeHeight) == latestVersion+1 {
+			if len(storeUpgrades.Deleted) > 0 || len(storeUpgrades.Added) > 0 {
+				if upgrader, ok := store.(storev2.UpgradeableStore); ok {
+					return upgrader.LoadVersionAndUpgrade(latestVersion, storeUpgrades)
+				}
+
+				return fmt.Errorf("store does not support upgrades")
+			}
+		}
+
+		return DefaultStoreLoader(store)
+	}
+}

simapp/v2/upgrades.go

@@ -5,6 +5,7 @@ import (
 
 	"cosmossdk.io/core/appmodule"
 	"cosmossdk.io/core/store"
+	"cosmossdk.io/runtime/v2"
 	"cosmossdk.io/x/accounts"
 	bankv2types "cosmossdk.io/x/bank/v2/types"
 	epochstypes "cosmossdk.io/x/epochs/types"
@@ -44,8 +45,6 @@ func (app *SimApp[T]) RegisterUpgradeHandlers() {
 			Deleted: []string{"crisis"}, // The SDK discontinued the crisis module in v0.52.0
 		}
 
-		// configure store loader that checks if version == upgradeHeight and applies store upgrades
-		_ = storeUpgrades
-		// app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, &storeUpgrades))
+		app.SetStoreLoader(runtime.UpgradeStoreLoader(upgradeInfo.Height, &storeUpgrades))
 	}
 }

x/upgrade/README.md

@@ -94,6 +94,9 @@ expected upgrade. It eliminiates the chances for the new binary to execute `Stor
 times every time on restart. Also if there are multiple upgrades planned on same height, the `Name`
 will ensure these `StoreUpgrades` takes place only in planned upgrade handler.
 
+**Note:** The `StoreLoader` helper function for StoreUpgrades in v2 is not part of the `x/upgrade` module; 
+instead, you can find it in the runtime v2 module.
+
 ### Proposal
 
 Typically, a `Plan` is proposed and submitted through governance via a proposal

x/upgrade/types/storeloader.go

@@ -1,28 +1,6 @@
 package types
 
-import (
-	corestore "cosmossdk.io/core/store"
-	storetypes "cosmossdk.io/store/types"
+import "github.com/cosmos/cosmos-sdk/runtime"
 
-	"github.com/cosmos/cosmos-sdk/baseapp"
-)
-
-// UpgradeStoreLoader is used to prepare baseapp with a fixed StoreLoader
-// pattern. This is useful for custom upgrade loading logic.
-func UpgradeStoreLoader(upgradeHeight int64, storeUpgrades *corestore.StoreUpgrades) baseapp.StoreLoader {
-	return func(ms storetypes.CommitMultiStore) error {
-		if upgradeHeight == ms.LastCommitID().Version+1 {
-			// Check if the current commit version and upgrade height matches
-			if len(storeUpgrades.Deleted) > 0 || len(storeUpgrades.Added) > 0 {
-				stup := &storetypes.StoreUpgrades{
-					Added:   storeUpgrades.Added,
-					Deleted: storeUpgrades.Deleted,
-				}
-				return ms.LoadLatestVersionAndUpgrade(stup)
-			}
-		}
-
-		// Otherwise load default store loader
-		return baseapp.DefaultStoreLoader(ms)
-	}
-}
+// UpgradeStoreLoader moved to runtime package, keeping this for backwards compatibility
+var UpgradeStoreLoader = runtime.UpgradeStoreLoader