Fix missing module for migration and add command to export IAVL (#1955)

* Archive Generate Stats

* leaf nodes

* Fix logs

* Fix logs

* Add tool command

* Add tool command

* Rename

* Refactor the new command

* Fix missing module

* Fix lint

---------

Co-authored-by: kbhat1 <[email protected]>

cmd/seid/cmd/root.go

@@ -312,7 +312,7 @@ func newApp(
 			homeDir := cast.ToString(appOpts.Get(flags.FlagHome))
 			stateStore := app.GetStateStore()
 			migrationHeight := cast.ToInt64(appOpts.Get("migrate-height"))
-			migrator := ss.NewMigrator(homeDir, db, stateStore)
+			migrator := ss.NewMigrator(db, stateStore)
 			if err := migrator.Migrate(migrationHeight, homeDir); err != nil {
 				panic(err)
 			}

tools/cmd.go

@@ -1,9 +1,10 @@
 package tools
 
 import (
+	"github.com/spf13/cobra"
+
 	migration "github.com/sei-protocol/sei-chain/tools/migration/cmd"
 	scanner "github.com/sei-protocol/sei-chain/tools/tx-scanner/cmd"
-	"github.com/spf13/cobra"
 )
 
 func ToolCmd() *cobra.Command {
@@ -14,5 +15,6 @@ func ToolCmd() *cobra.Command {
 	toolsCmd.AddCommand(scanner.ScanCmd())
 	toolsCmd.AddCommand(migration.MigrateCmd())
 	toolsCmd.AddCommand(migration.VerifyMigrationCmd())
+	toolsCmd.AddCommand(migration.GenerateStats())
 	return toolsCmd
 }

tools/migration/cmd/cmd.go

@@ -1,12 +1,16 @@
 package cmd
 
 import (
+	"bytes"
 	"fmt"
 	"path/filepath"
+	"time"
 
 	"github.com/cosmos/cosmos-sdk/store/rootmulti"
+	"github.com/cosmos/iavl"
 	"github.com/sei-protocol/sei-chain/tools/migration/sc"
 	"github.com/sei-protocol/sei-chain/tools/migration/ss"
+	"github.com/sei-protocol/sei-chain/tools/migration/utils"
 	"github.com/sei-protocol/sei-db/config"
 	sstypes "github.com/sei-protocol/sei-db/ss"
 	"github.com/spf13/cobra"
@@ -91,6 +95,76 @@ func verifySS(version int64, homeDir string, db dbm.DB) error {
 		return err
 	}
 
-	migrator := ss.NewMigrator(homeDir, db, stateStore)
+	migrator := ss.NewMigrator(db, stateStore)
 	return migrator.Verify(version)
 }
+
+func GenerateStats() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "iavl-stats",
+		Short: "A tool to generate archive node iavl stats like number of keys and size per module.",
+		Run:   generateIavlStats,
+	}
+	cmd.PersistentFlags().String("home-dir", "/root/.sei", "Sei home directory")
+	return cmd
+}
+
+func generateIavlStats(cmd *cobra.Command, _ []string) {
+	homeDir, _ := cmd.Flags().GetString("home-dir")
+	dataDir := filepath.Join(homeDir, "data")
+	db, err := dbm.NewGoLevelDB("application", dataDir)
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Println("Aggregating iavl module stats...")
+	for _, module := range utils.Modules {
+
+		startTimeModule := time.Now()
+		fmt.Printf("Aggregating stats for module %s...\n", module)
+
+		// Prepare the prefixed DB for the module
+		prefixDB := dbm.NewPrefixDB(db, []byte(utils.BuildRawPrefix(module)))
+
+		// Get an iterator for the prefixed DB
+		itr, err := prefixDB.Iterator(nil, nil)
+		if err != nil {
+			panic(fmt.Errorf("failed to create iterator: %w", err))
+		}
+		defer itr.Close()
+
+		// Aggregated stats for the current module
+		var totalKeys int
+		var totalKeySize int
+		var totalValueSize int
+
+		for ; itr.Valid(); itr.Next() {
+			value := bytes.Clone(itr.Value())
+			node, err := iavl.MakeNode(value)
+			if err != nil {
+				panic(fmt.Errorf("failed to make node: %w", err))
+			}
+
+			// Only export leaf nodes
+			if node.GetHeight() == 0 {
+				totalKeys++
+				totalKeySize += len(itr.Key())
+				totalValueSize += len(itr.Value())
+				if totalKeys%1000000 == 0 {
+					fmt.Printf("Module %s num of keys %d, key size %d, value size %d\n",
+						module, totalKeys, totalKeySize, totalValueSize)
+				}
+			}
+		}
+
+		if err := itr.Error(); err != nil {
+			panic(fmt.Errorf("iterator error for module %s: %w", module, err))
+		}
+
+		moduleDuration := time.Since(startTimeModule)
+		fmt.Printf("Module %s stats finished, total keys %d, total key size %d, total value size %d, time taken %s\n",
+			module, totalKeys, totalKeySize, totalValueSize, moduleDuration)
+	}
+
+	fmt.Println("SeiDB Archive Migration: Aggregation completed.")
+}

tools/migration/sc/migrator.go

@@ -15,30 +15,12 @@ import (
 	rootmulti2 "github.com/cosmos/cosmos-sdk/storev2/rootmulti"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
-	acltypes "github.com/cosmos/cosmos-sdk/x/accesscontrol/types"
 	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
-	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
-	authzkeeper "github.com/cosmos/cosmos-sdk/x/authz/keeper"
-	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
-	capabilitytypes "github.com/cosmos/cosmos-sdk/x/capability/types"
-	distrtypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
-	evidencetypes "github.com/cosmos/cosmos-sdk/x/evidence/types"
-	"github.com/cosmos/cosmos-sdk/x/feegrant"
-	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
 	paramskeeper "github.com/cosmos/cosmos-sdk/x/params/keeper"
 	paramstypes "github.com/cosmos/cosmos-sdk/x/params/types"
-	slashingtypes "github.com/cosmos/cosmos-sdk/x/slashing/types"
 	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
-	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
-	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
-	ibctransfertypes "github.com/cosmos/ibc-go/v3/modules/apps/transfer/types"
-	ibchost "github.com/cosmos/ibc-go/v3/modules/core/24-host"
 	"github.com/sei-protocol/sei-chain/app/params"
-	epochmoduletypes "github.com/sei-protocol/sei-chain/x/epoch/types"
-	evmtypes "github.com/sei-protocol/sei-chain/x/evm/types"
-	minttypes "github.com/sei-protocol/sei-chain/x/mint/types"
-	oracletypes "github.com/sei-protocol/sei-chain/x/oracle/types"
-	tokenfactorytypes "github.com/sei-protocol/sei-chain/x/tokenfactory/types"
+	"github.com/sei-protocol/sei-chain/tools/migration/utils"
 	"github.com/sei-protocol/sei-db/config"
 	"github.com/tendermint/tendermint/libs/log"
 	dbm "github.com/tendermint/tm-db"
@@ -51,20 +33,12 @@ type Migrator struct {
 	storeV2 store.CommitMultiStore
 }
 
-var Keys = sdk.NewKVStoreKeys(
-	acltypes.StoreKey, authtypes.StoreKey, authzkeeper.StoreKey, banktypes.StoreKey, stakingtypes.StoreKey,
-	minttypes.StoreKey, distrtypes.StoreKey, slashingtypes.StoreKey,
-	govtypes.StoreKey, paramstypes.StoreKey, ibchost.StoreKey, upgradetypes.StoreKey, feegrant.StoreKey,
-	evidencetypes.StoreKey, ibctransfertypes.StoreKey, capabilitytypes.StoreKey, oracletypes.StoreKey,
-	evmtypes.StoreKey, wasm.StoreKey, epochmoduletypes.StoreKey, tokenfactorytypes.StoreKey,
-)
-
 func NewMigrator(homeDir string, db dbm.DB) *Migrator {
 	logger := log.NewTMLogger(log.NewSyncWriter(os.Stdout))
 
 	// Creating CMS for store V1
 	cmsV1 := rootmulti.NewStore(db, logger)
-	for _, key := range Keys {
+	for _, key := range utils.ModuleKeys {
 		cmsV1.MountStoreWithDB(key, sdk.StoreTypeIAVL, nil)
 	}
 	err := cmsV1.LoadLatestVersion()
@@ -79,7 +53,7 @@ func NewMigrator(homeDir string, db dbm.DB) *Migrator {
 	ssConfig.Enable = true
 	ssConfig.KeepRecent = 0
 	cmsV2 := rootmulti2.NewStore(homeDir, logger, scConfig, ssConfig, true)
-	for _, key := range Keys {
+	for _, key := range utils.ModuleKeys {
 		cmsV2.MountStoreWithDB(key, sdk.StoreTypeIAVL, db)
 	}
 	err = cmsV2.LoadLatestVersion()
@@ -174,7 +148,7 @@ func CreateWasmSnapshotter(cms sdk.MultiStore, homeDir string) *keeper.WasmSnaps
 	pk := paramskeeper.NewKeeper(encodingConfig.Marshaler, encodingConfig.Amino, keyParams, tkeyParams)
 	wasmKeeper := keeper.NewKeeper(
 		encodingConfig.Marshaler,
-		Keys[wasm.StoreKey],
+		utils.ModuleKeys[wasm.StoreKey],
 		paramskeeper.Keeper{},
 		pk.Subspace("wasm"),
 		authkeeper.AccountKeeper{},

tools/migration/ss/migrator.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/armon/go-metrics"
 	"github.com/cosmos/iavl"
+	"github.com/sei-protocol/sei-chain/tools/migration/utils"
 	"github.com/sei-protocol/sei-db/ss/types"
 	dbm "github.com/tendermint/tm-db"
 )
@@ -21,11 +22,7 @@ const (
 	DefaultCacheSize int = 10000
 )
 
-var modules = []string{
-	"wasm", "aclaccesscontrol", "oracle", "epoch", "mint", "acc", "bank", "feegrant", "staking", "distribution", "slashing", "gov", "params", "ibc", "upgrade", "evidence", "transfer", "tokenfactory",
-}
-
-func NewMigrator(homeDir string, db dbm.DB, stateStore types.StateStore) *Migrator {
+func NewMigrator(db dbm.DB, stateStore types.StateStore) *Migrator {
 	return &Migrator{
 		iavlDB:     db,
 		stateStore: stateStore,
@@ -78,8 +75,8 @@ func (m *Migrator) Migrate(version int64, homeDir string) error {
 
 func (m *Migrator) Verify(version int64) error {
 	var verifyErr error
-	for _, module := range modules {
-		tree, err := ReadTree(m.iavlDB, version, []byte(buildTreePrefix(module)))
+	for _, module := range utils.Modules {
+		tree, err := ReadTree(m.iavlDB, version, []byte(utils.BuildTreePrefix(module)))
 		if err != nil {
 			fmt.Printf("Error reading tree %s: %s\n", module, err.Error())
 			return err
@@ -125,7 +122,7 @@ func ExportLeafNodesFromKey(db dbm.DB, ch chan<- types.RawSnapshotNode, startKey
 	var batchLeafNodeCount int
 	startModuleFound := startModule == "" // true if no start module specified
 
-	for _, module := range modules {
+	for _, module := range utils.Modules {
 		if !startModuleFound {
 			if module == startModule {
 				startModuleFound = true
@@ -136,12 +133,12 @@ func ExportLeafNodesFromKey(db dbm.DB, ch chan<- types.RawSnapshotNode, startKey
 		startTimeModule := time.Now() // Measure time for each module
 		fmt.Printf("SeiDB Archive Migration: Iterating through %s module...\n", module)
 
-		prefixDB := dbm.NewPrefixDB(db, []byte(buildRawPrefix(module)))
+		prefixDB := dbm.NewPrefixDB(db, []byte(utils.BuildRawPrefix(module)))
 		var itr dbm.Iterator
 		var err error
 
 		// If there is a starting key, seek to it, otherwise start from the beginning
-		if startKey != nil && bytes.HasPrefix(startKey, []byte(buildRawPrefix(module))) {
+		if startKey != nil && bytes.HasPrefix(startKey, []byte(utils.BuildRawPrefix(module))) {
 			itr, err = prefixDB.Iterator(startKey, nil) // Start from the latest key
 		} else {
 			itr, err = prefixDB.Iterator(nil, nil) // Start from the beginning
@@ -204,14 +201,6 @@ func ExportLeafNodesFromKey(db dbm.DB, ch chan<- types.RawSnapshotNode, startKey
 	return nil
 }
 
-func buildRawPrefix(moduleName string) string {
-	return fmt.Sprintf("s/k:%s/n", moduleName)
-}
-
-func buildTreePrefix(moduleName string) string {
-	return fmt.Sprintf("s/k:%s/", moduleName)
-}
-
 func ReadTree(db dbm.DB, version int64, prefix []byte) (*iavl.MutableTree, error) {
 	// TODO: Verify if we need a prefix here (or can just iterate through all modules)
 	if len(prefix) != 0 {

tools/migration/utils/helper.go

@@ -0,0 +1,67 @@
+package utils
+
+import (
+	"fmt"
+
+	"github.com/CosmWasm/wasmd/x/wasm"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	acltypes "github.com/cosmos/cosmos-sdk/x/accesscontrol/types"
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	authzkeeper "github.com/cosmos/cosmos-sdk/x/authz/keeper"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+	capabilitytypes "github.com/cosmos/cosmos-sdk/x/capability/types"
+	distrtypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
+	evidencetypes "github.com/cosmos/cosmos-sdk/x/evidence/types"
+	"github.com/cosmos/cosmos-sdk/x/feegrant"
+	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
+	paramstypes "github.com/cosmos/cosmos-sdk/x/params/types"
+	slashingtypes "github.com/cosmos/cosmos-sdk/x/slashing/types"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
+	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
+	ibctransfertypes "github.com/cosmos/ibc-go/v3/modules/apps/transfer/types"
+	ibchost "github.com/cosmos/ibc-go/v3/modules/core/24-host"
+	epochmoduletypes "github.com/sei-protocol/sei-chain/x/epoch/types"
+	evmtypes "github.com/sei-protocol/sei-chain/x/evm/types"
+	minttypes "github.com/sei-protocol/sei-chain/x/mint/types"
+	oracletypes "github.com/sei-protocol/sei-chain/x/oracle/types"
+	tokenfactorytypes "github.com/sei-protocol/sei-chain/x/tokenfactory/types"
+)
+
+var ModuleKeys = sdk.NewKVStoreKeys(
+	acltypes.StoreKey, authtypes.StoreKey, authzkeeper.StoreKey, banktypes.StoreKey, stakingtypes.StoreKey,
+	minttypes.StoreKey, distrtypes.StoreKey, slashingtypes.StoreKey,
+	govtypes.StoreKey, paramstypes.StoreKey, ibchost.StoreKey, upgradetypes.StoreKey, feegrant.StoreKey,
+	evidencetypes.StoreKey, ibctransfertypes.StoreKey, capabilitytypes.StoreKey, oracletypes.StoreKey,
+	evmtypes.StoreKey, wasm.StoreKey, epochmoduletypes.StoreKey, tokenfactorytypes.StoreKey,
+)
+
+var Modules = []string{
+	"acc",
+	"aclaccesscontrol",
+	"authz",
+	"bank",
+	"capability",
+	"distribution",
+	"epoch",
+	"evidence",
+	"evm",
+	"feegrant",
+	"gov",
+	"ibc",
+	"mint",
+	"oracle",
+	"params",
+	"slashing",
+	"staking",
+	"tokenfactory",
+	"transfer",
+	"upgrade",
+	"wasm"}
+
+func BuildRawPrefix(moduleName string) string {
+	return fmt.Sprintf("s/k:%s/n", moduleName)
+}
+
+func BuildTreePrefix(moduleName string) string {
+	return fmt.Sprintf("s/k:%s/", moduleName)
+}