Add command line tool to dump chain state size

tools/cmd/seidb/main.go

@@ -23,7 +23,8 @@ func main() {
 		benchmark.DBReverseIterationCmd(),
 		operations.DumpDbCmd(),
 		operations.PruneCmd(),
-		operations.DumpIAVLCmd())
+		operations.DumpIAVLCmd(),
+		operations.StateSizeCmd())
 	if err := rootCmd.Execute(); err != nil {
 		fmt.Println(err)
 		os.Exit(1)

tools/cmd/seidb/operations/dump_db.go

@@ -18,7 +18,7 @@ func DumpDbCmd() *cobra.Command {
 	dumpDbCmd := &cobra.Command{
 		Use:   "dump-db",
 		Short: "For a given State Store DB, dump-db iterates over all keys and values for a specific store and writes them to a file",
-		Run:   dump,
+		Run:   executeDump,
 	}
 
 	dumpDbCmd.PersistentFlags().StringP("output-dir", "o", "", "Output Directory")
@@ -30,7 +30,7 @@ func DumpDbCmd() *cobra.Command {
 	return dumpDbCmd
 }
 
-func dump(cmd *cobra.Command, _ []string) {
+func executeDump(cmd *cobra.Command, _ []string) {
 	outputDir, _ := cmd.Flags().GetString("output-dir")
 	module, _ := cmd.Flags().GetString("module")
 	dbDir, _ := cmd.Flags().GetString("db-dir")

tools/cmd/seidb/operations/dump_iavl.go

@@ -9,15 +9,11 @@ import (
 	"github.com/spf13/cobra"
 )
 
-var AllModules = []string{
-	"dex", "wasm", "aclaccesscontrol", "oracle", "epoch", "mint", "acc", "bank", "crisis", "feegrant", "staking", "distribution", "slashing", "gov", "params", "ibc", "upgrade", "evidence", "transfer", "tokenfactory",
-}
-
 func DumpIAVLCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "dump-iavl",
 		Short: "Iterate and dump memIAVL data",
-		Run:   execute,
+		Run:   executeDumpIAVL,
 	}
 
 	cmd.PersistentFlags().StringP("db-dir", "d", "", "Database Directory")
@@ -27,7 +23,7 @@ func DumpIAVLCmd() *cobra.Command {
 	return cmd
 }
 
-func execute(cmd *cobra.Command, _ []string) {
+func executeDumpIAVL(cmd *cobra.Command, _ []string) {
 	module, _ := cmd.Flags().GetString("module")
 	dbDir, _ := cmd.Flags().GetString("db-dir")
 	outputDir, _ := cmd.Flags().GetString("output-dir")
@@ -41,24 +37,24 @@ func execute(cmd *cobra.Command, _ []string) {
 		panic("Must provide output dir")
 	}
 
-	err := DumpIAVLData(module, dbDir, outputDir, height)
-	if err != nil {
-		panic(err)
-	}
-}
-
-// DumpIAVLData print the raw keys and values for given module at given height for memIAVL tree
-func DumpIAVLData(module string, dbDir string, outputDir string, height int64) error {
 	opts := memiavl.Options{
 		Dir:             dbDir,
 		ZeroCopy:        true,
 		CreateIfMissing: false,
 	}
 	db, err := memiavl.OpenDB(logger.NewNopLogger(), height, opts)
 	if err != nil {
-		return err
+		panic(err)
 	}
 	defer db.Close()
+	err = DumpIAVLData(module, db, outputDir)
+	if err != nil {
+		panic(err)
+	}
+}
+
+// DumpIAVLData print the raw keys and values for given module at given height for memIAVL tree
+func DumpIAVLData(module string, db *memiavl.DB, outputDir string) error {
 	modules := []string{}
 	if module == "" {
 		modules = AllModules

tools/cmd/seidb/operations/module.go

@@ -0,0 +1,5 @@
+package operations
+
+var AllModules = []string{
+	"evm", "dex", "wasm", "aclaccesscontrol", "oracle", "epoch", "mint", "acc", "bank", "crisis", "feegrant", "staking", "distribution", "slashing", "gov", "params", "ibc", "upgrade", "evidence", "transfer", "tokenfactory",
+}

tools/cmd/seidb/operations/prune.go

@@ -13,7 +13,7 @@ func PruneCmd() *cobra.Command {
 	pruneDbCmd := &cobra.Command{
 		Use:   "prune",
 		Short: "Prune a db at a given height",
-		Run:   prune,
+		Run:   executePrune,
 	}
 
 	pruneDbCmd.PersistentFlags().StringP("db-dir", "d", "", "Database Directory")
@@ -23,7 +23,7 @@ func PruneCmd() *cobra.Command {
 	return pruneDbCmd
 }
 
-func prune(cmd *cobra.Command, _ []string) {
+func executePrune(cmd *cobra.Command, _ []string) {
 	dbDir, _ := cmd.Flags().GetString("db-dir")
 	dbBackend, _ := cmd.Flags().GetString("db-backend")
 	version, _ := cmd.Flags().GetInt64("version")

tools/cmd/seidb/operations/state_size.go

@@ -0,0 +1,87 @@
+package operations
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/sei-protocol/sei-db/common/logger"
+	"github.com/sei-protocol/sei-db/sc/memiavl"
+	"github.com/spf13/cobra"
+)
+
+func StateSizeCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "state-size",
+		Short: "Print analytical results for state size",
+		Run:   executeStateSize,
+	}
+
+	cmd.PersistentFlags().StringP("db-dir", "d", "", "Database Directory")
+	cmd.PersistentFlags().Int64("height", 0, "Block Height")
+	cmd.PersistentFlags().StringP("module", "m", "", "Module to export. Default to export all")
+	return cmd
+}
+
+func executeStateSize(cmd *cobra.Command, _ []string) {
+	module, _ := cmd.Flags().GetString("module")
+	dbDir, _ := cmd.Flags().GetString("db-dir")
+	height, _ := cmd.Flags().GetInt64("height")
+	if dbDir == "" {
+		panic("Must provide database dir")
+	}
+
+	opts := memiavl.Options{
+		Dir:             dbDir,
+		ZeroCopy:        true,
+		CreateIfMissing: false,
+	}
+	db, err := memiavl.OpenDB(logger.NewNopLogger(), height, opts)
+	if err != nil {
+		panic(err)
+	}
+	defer db.Close()
+	err = PrintStateSize(module, db)
+	if err != nil {
+		panic(err)
+	}
+}
+
+// PrintStateSize print the raw keys and values for given module at given height for memIAVL tree
+func PrintStateSize(module string, db *memiavl.DB) error {
+	modules := []string{}
+	if module == "" {
+		modules = AllModules
+	} else {
+		modules = append(modules, module)
+	}
+
+	for _, moduleName := range modules {
+		tree := db.TreeByName(moduleName)
+		totalNumKeys := 0
+		totalKeySize := 0
+		totalValueSize := 0
+		totalSize := 0
+		if tree == nil {
+			fmt.Printf("Tree does not exist for module %s \n", moduleName)
+		} else {
+			fmt.Printf("Calculating for module: %s \n", moduleName)
+			sizeByPrefix := map[string]int{}
+			tree.ScanPostOrder(func(node memiavl.Node) bool {
+				if node.IsLeaf() {
+					totalNumKeys++
+					totalKeySize += len(node.Key())
+					totalValueSize += len(node.Value())
+					totalSize += len(node.Key()) + len(node.Value())
+					prefix := fmt.Sprintf("%X", node.Key())
+					prefix = prefix[:2]
+					sizeByPrefix[prefix] += len(node.Value())
+				}
+				return true
+			})
+			fmt.Printf("Module %s total numKeys:%d, total keySize:%d, total valueSize:%d, totalSize: %d \n", moduleName, totalNumKeys, totalKeySize, totalValueSize, totalSize)
+			result, _ := json.MarshalIndent(sizeByPrefix, "", "  ")
+			fmt.Printf("Module %s prefix breakdown: %s \n", moduleName, result)
+		}
+	}
+	return nil
+}