Add database analyze-contract-sizes (#11249)

Add a command which looks at the current state and finds accounts with
the biggest contract sizes.

Example usage:

```
$ ./neard database analyze-contract-sizes --topn 10
...
Analyzed 33165 accounts
Accounts per shard:
s0.v3: 6443
s1.v3: 1
s2.v3: 7916
s3.v3: 3643
s4.v3: 11208
s5.v3: 3954

Top 10 accounts by contract size:
3.9 MB: nearpay-portals.near
3.8 MB: moodev.near
2.8 MB: v2_1_0.perp.spin-fi.near
2.8 MB: v2_0_2.perp.spin-fi.near
2.8 MB: v2.perp.spin-fi.near
2.7 MB: switchboard-v2.near
2.6 MB: nft.contented.near
2.5 MB: spot.spin-fi.near
2.5 MB: nftstaking.jumpfinance.near
2.5 MB: exchange.slowisfast.near
```

tools/database/Cargo.toml

@@ -21,6 +21,8 @@ rayon.workspace = true
 rocksdb.workspace = true
 strum.workspace = true
 tempfile.workspace = true
+bytesize.workspace = true
+zstd.workspace = true
 
 nearcore.workspace = true
 near-epoch-manager.workspace = true

tools/database/src/analyze_contract_sizes.rs

@@ -0,0 +1,129 @@
+use std::collections::BTreeMap;
+use std::path::PathBuf;
+use std::rc::Rc;
+use std::str::FromStr;
+
+use bytesize::ByteSize;
+use clap::Parser;
+use near_chain::{ChainStore, ChainStoreAccess};
+use near_chain_configs::GenesisValidationMode;
+use near_epoch_manager::EpochManager;
+use near_primitives::trie_key::col;
+use near_primitives::types::AccountId;
+use near_store::{ShardUId, Trie, TrieDBStorage};
+use nearcore::{load_config, open_storage};
+
+#[derive(Parser)]
+pub(crate) struct AnalyzeContractSizesCommand {
+    /// Show top N contracts by size.
+    #[arg(short, long, default_value_t = 50)]
+    topn: usize,
+
+    /// Compress contract code before calculating size.
+    #[arg(long, default_value_t = false)]
+    compressed: bool,
+}
+
+struct ContractSizeStats {
+    topn: usize,
+    top_accounts: BTreeMap<ByteSize, AccountId>,
+    total_accounts: usize,
+    shard_accounts: BTreeMap<ShardUId, usize>,
+}
+
+impl ContractSizeStats {
+    pub fn new(topn: usize) -> ContractSizeStats {
+        ContractSizeStats {
+            topn,
+            top_accounts: BTreeMap::new(),
+            total_accounts: 0,
+            shard_accounts: BTreeMap::new(),
+        }
+    }
+
+    pub fn add_info(&mut self, shard_uid: ShardUId, account_id: AccountId, contract_size: usize) {
+        self.total_accounts += 1;
+        *self.shard_accounts.entry(shard_uid).or_insert(0) += 1;
+
+        self.top_accounts.insert(ByteSize::b(contract_size as u64), account_id);
+        if self.top_accounts.len() > self.topn {
+            self.top_accounts.pop_first();
+        }
+    }
+
+    pub fn print_stats(&self, compressed: bool) {
+        println!("");
+        println!("Analyzed {} accounts", self.total_accounts);
+        println!("Accounts per shard:");
+        for (shard_uid, count) in self.shard_accounts.iter() {
+            println!("{}: {}", shard_uid, count);
+        }
+        println!("");
+        if compressed {
+            println!("Top {} accounts by compressed contract size:", self.topn);
+        } else {
+            println!("Top {} accounts by contract size:", self.topn);
+        }
+        for (size, account_id) in self.top_accounts.iter().rev() {
+            println!("{}: {}", size, account_id);
+        }
+    }
+}
+
+impl AnalyzeContractSizesCommand {
+    pub(crate) fn run(&self, home: &PathBuf) -> anyhow::Result<()> {
+        let mut near_config = load_config(home, GenesisValidationMode::Full).unwrap();
+        let node_storage = open_storage(&home, &mut near_config).unwrap();
+        let store = node_storage.get_split_store().unwrap_or_else(|| node_storage.get_hot_store());
+        let chain_store = Rc::new(ChainStore::new(
+            store.clone(),
+            near_config.genesis.config.genesis_height,
+            false,
+        ));
+
+        let head = chain_store.head().unwrap();
+        let epoch_manager =
+            EpochManager::new_from_genesis_config(store.clone(), &near_config.genesis.config)
+                .unwrap();
+        let shard_layout = epoch_manager.get_shard_layout(&head.epoch_id).unwrap();
+
+        let mut stats = ContractSizeStats::new(self.topn);
+        for shard_uid in shard_layout.shard_uids() {
+            println!("Analyzing chunk with uid: {}", shard_uid);
+
+            let chunk_extra =
+                chain_store.get_chunk_extra(&head.last_block_hash, &shard_uid).unwrap();
+
+            let state_root = chunk_extra.state_root();
+            let trie_storage = Rc::new(TrieDBStorage::new(store.clone(), shard_uid));
+            let trie = Trie::new(trie_storage, *state_root, None);
+
+            let mut iterator = trie.iter().unwrap();
+            iterator.seek_prefix(&[col::CONTRACT_CODE]).unwrap();
+
+            for (i, item) in iterator.enumerate() {
+                let (key, value) = item.unwrap();
+                if key.is_empty() || key[0] != col::CONTRACT_CODE {
+                    break;
+                }
+                let account_id_bytes = &key[1..];
+                let account_id_str = std::str::from_utf8(&account_id_bytes).unwrap();
+                let account_id = AccountId::from_str(account_id_str).unwrap();
+
+                let contract_size = if self.compressed {
+                    zstd::encode_all(value.as_slice(), 3).unwrap().len()
+                } else {
+                    value.len()
+                };
+
+                stats.add_info(shard_uid, account_id, contract_size);
+                if i % 1000 == 0 {
+                    println!("Processed {} contracts...", i);
+                }
+            }
+        }
+
+        stats.print_stats(self.compressed);
+        Ok(())
+    }
+}

tools/database/src/commands.rs

@@ -2,6 +2,7 @@ use crate::adjust_database::ChangeDbKindCommand;
 use crate::analyse_data_size_distribution::AnalyseDataSizeDistributionCommand;
 use crate::analyse_gas_usage::AnalyseGasUsageCommand;
 use crate::analyse_high_load::HighLoadStatsCommand;
+use crate::analyze_contract_sizes::AnalyzeContractSizesCommand;
 use crate::analyze_delayed_receipt::AnalyzeDelayedReceiptCommand;
 use crate::compact::RunCompactionCommand;
 use crate::corrupt::CorruptStateSnapshotCommand;
@@ -56,6 +57,8 @@ enum SubCommand {
     HighLoadStats(HighLoadStatsCommand),
     // Analyze congestion through delayed receipts
     AnalyzeDelayedReceipt(AnalyzeDelayedReceiptCommand),
+    /// Analyze size of contracts present in the current state
+    AnalyzeContractSizes(AnalyzeContractSizesCommand),
 }
 
 impl DatabaseCommand {
@@ -87,6 +90,7 @@ impl DatabaseCommand {
             SubCommand::WriteCryptoHash(cmd) => cmd.run(home),
             SubCommand::HighLoadStats(cmd) => cmd.run(home),
             SubCommand::AnalyzeDelayedReceipt(cmd) => cmd.run(home),
+            SubCommand::AnalyzeContractSizes(cmd) => cmd.run(home),
         }
     }
 }

tools/database/src/lib.rs

@@ -2,6 +2,7 @@ mod adjust_database;
 mod analyse_data_size_distribution;
 mod analyse_gas_usage;
 mod analyse_high_load;
+mod analyze_contract_sizes;
 mod analyze_delayed_receipt;
 mod block_iterators;
 pub mod commands;