Merge pull request #98 from sei-protocol/yzang/speed-up-state-analysis

Optimization for state size tool

Author: yzang2019

tools/cmd/seidb/operations/state_size.go

@@ -24,6 +24,14 @@ func StateSizeCmd() *cobra.Command {
 	return cmd
 }
 
+type contractSizeEntry struct {
+	Address   string
+	KeySize   int64
+	ValueSize int64
+	TotalSize int64
+	KeyCount  int
+}
+
 func executeStateSize(cmd *cobra.Command, _ []string) {
 	module, _ := cmd.Flags().GetString("module")
 	dbDir, _ := cmd.Flags().GetString("db-dir")
@@ -42,6 +50,7 @@ func executeStateSize(cmd *cobra.Command, _ []string) {
 		panic(err)
 	}
 	defer db.Close()
+	fmt.Printf("Finished opening db, calculating state size for module: %s\n", module)
 	err = PrintStateSize(module, db)
 	if err != nil {
 		panic(err)
@@ -56,20 +65,25 @@ func PrintStateSize(module string, db *memiavl.DB) error {
 	} else {
 		modules = append(modules, module)
 	}
+	// Print top 100 contracts by total size
+	numToShow := 100
 
 	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)
 			keySizeByPrefix := map[string]int64{}
 			valueSizeByPrefix := map[string]int64{}
 			numKeysByPrefix := map[string]int64{}
+			contractSizes := make(map[string]*contractSizeEntry)
+			// Scan again to collect per-contract statistics
 			tree.ScanPostOrder(func(node memiavl.Node) bool {
 				if node.IsLeaf() {
 					totalNumKeys++
@@ -78,11 +92,28 @@ func PrintStateSize(module string, db *memiavl.DB) error {
 					totalKeySize += keySize
 					totalValueSize += valueSize
 					totalSize += keySize + valueSize
-					prefix := fmt.Sprintf("%X", node.Key())
-					prefix = prefix[:2]
+					prefixKey := fmt.Sprintf("%X", node.Key())
+					prefix := prefixKey[:2]
 					keySizeByPrefix[prefix] += int64(keySize)
 					valueSizeByPrefix[prefix] += int64(valueSize)
 					numKeysByPrefix[prefix]++
+
+					if module == "evm" && prefix == "03" {
+						// Extract contract address from key (assuming it follows after "03")
+						addr := prefixKey[2:42] // Adjust indices based on your key format
+						if _, exists := contractSizes[addr]; !exists {
+							contractSizes[addr] = &contractSizeEntry{Address: addr}
+						}
+						entry := contractSizes[addr]
+						entry.KeySize += int64(len(node.Key()))
+						entry.ValueSize += int64(len(node.Value()))
+						entry.TotalSize = entry.KeySize + entry.ValueSize
+						entry.KeyCount++
+					}
+
+					if totalNumKeys%1000000 == 0 {
+						fmt.Printf("Scanned %d keys for module %s\n", totalNumKeys, moduleName)
+					}
 				}
 				return true
 			})
@@ -94,66 +125,32 @@ func PrintStateSize(module string, db *memiavl.DB) error {
 			numKeysResult, _ := json.MarshalIndent(numKeysByPrefix, "", "  ")
 			fmt.Printf("Module %s prefix num of keys breakdown: %s \n", moduleName, numKeysResult)
 
-			// Print top 20 contracts by total size
-			numToShow := 20
-			if valueSizeByPrefix["03"] > 0 || keySizeByPrefix["03"] > 0 {
-				type contractSizeEntry struct {
-					Address   string
-					KeySize   int64
-					ValueSize int64
-					TotalSize int64
-					KeyCount  int
-				}
-
-				contractSizes := make(map[string]*contractSizeEntry)
-
-				// Scan again to collect per-contract statistics
-				tree.ScanPostOrder(func(node memiavl.Node) bool {
-					if node.IsLeaf() {
-						prefix := fmt.Sprintf("%X", node.Key())
-						if prefix[:2] == "03" {
-							// Extract contract address from key (assuming it follows after "03")
-							addr := prefix[2:42] // Adjust indices based on your key format
-							if _, exists := contractSizes[addr]; !exists {
-								contractSizes[addr] = &contractSizeEntry{Address: addr}
-							}
-							entry := contractSizes[addr]
-							entry.KeySize += int64(len(node.Key()))
-							entry.ValueSize += int64(len(node.Value()))
-							entry.TotalSize = entry.KeySize + entry.ValueSize
-							entry.KeyCount++
-						}
-					}
-					return true
-				})
-
-				// Convert map to slice
-				var sortedContracts []contractSizeEntry
-				for _, entry := range contractSizes {
-					sortedContracts = append(sortedContracts, *entry)
-				}
+			// Convert map to slice
+			var sortedContracts []contractSizeEntry
+			for _, entry := range contractSizes {
+				sortedContracts = append(sortedContracts, *entry)
+			}
 
-				// Sort by total size in descending order
-				sort.Slice(sortedContracts, func(i, j int) bool {
-					return sortedContracts[i].TotalSize > sortedContracts[j].TotalSize
-				})
+			// Sort by total size in descending order
+			sort.Slice(sortedContracts, func(i, j int) bool {
+				return sortedContracts[i].TotalSize > sortedContracts[j].TotalSize
+			})
 
-				fmt.Printf("\nDetailed breakdown for 0x03 prefix (top 20 contracts by total size):\n")
-				fmt.Printf("%-42s %15s %15s %15s %10s\n", "Contract Address", "Key Size", "Value Size", "Total Size", "Key Count")
-				fmt.Printf("%s\n", strings.Repeat("-", 100))
+			fmt.Printf("\nDetailed breakdown for 0x03 prefix (top 100 contracts by total size):\n")
+			fmt.Printf("%-42s %15s %15s %15s %10s\n", "Contract Address", "Key Size", "Value Size", "Total Size", "Key Count")
+			fmt.Printf("%s\n", strings.Repeat("-", 100))
 
-				if len(sortedContracts) < numToShow {
-					numToShow = len(sortedContracts)
-				}
-				for i := 0; i < numToShow; i++ {
-					contract := sortedContracts[i]
-					fmt.Printf("0x%-40s %15d %15d %15d %10d\n",
-						contract.Address,
-						contract.KeySize,
-						contract.ValueSize,
-						contract.TotalSize,
-						contract.KeyCount)
-				}
+			if len(sortedContracts) < numToShow {
+				numToShow = len(sortedContracts)
+			}
+			for i := 0; i < numToShow; i++ {
+				contract := sortedContracts[i]
+				fmt.Printf("0x%-40s %15d %15d %15d %10d\n",
+					contract.Address,
+					contract.KeySize,
+					contract.ValueSize,
+					contract.TotalSize,
+					contract.KeyCount)
 			}
 		}
 	}