ledger-tool: Add additional modes for accounts subcommand (#34925)

- Add mode to output individual pubkeys
- Add mode to output program accounts

ledger-tool/src/main.rs

@@ -6,7 +6,9 @@ use {
         blockstore::*,
         ledger_path::*,
         ledger_utils::*,
-        output::{output_account, AccountsOutputConfig, AccountsOutputStreamer},
+        output::{
+            output_account, AccountsOutputConfig, AccountsOutputMode, AccountsOutputStreamer,
+        },
         program::*,
     },
     clap::{
@@ -1312,6 +1314,7 @@ fn main() {
                 .arg(&geyser_plugin_args)
                 .arg(&accounts_data_encoding_arg)
                 .arg(&use_snapshot_archives_at_startup)
+                .arg(&max_genesis_archive_unpacked_size_arg)
                 .arg(
                     Arg::with_name("include_sysvars")
                         .long("include-sysvars")
@@ -1333,7 +1336,27 @@ fn main() {
                         .takes_value(false)
                         .help("Do not print account data when printing account contents."),
                 )
-                .arg(&max_genesis_archive_unpacked_size_arg),
+                .arg(
+                    Arg::with_name("account")
+                        .long("account")
+                        .takes_value(true)
+                        .value_name("PUBKEY")
+                        .validator(is_pubkey)
+                        .multiple(true)
+                        .help(
+                            "Limit output to accounts corresponding to the specified pubkey(s), \
+                            may be specified multiple times",
+                        ),
+                )
+                .arg(
+                    Arg::with_name("program_accounts")
+                        .long("program-accounts")
+                        .takes_value(true)
+                        .value_name("PUBKEY")
+                        .validator(is_pubkey)
+                        .conflicts_with("account")
+                        .help("Limit output to accounts owned by the provided program pubkey"),
+                ),
         )
         .subcommand(
             SubCommand::with_name("capitalization")
@@ -2179,7 +2202,18 @@ fn main() {
                     let include_account_contents = !arg_matches.is_present("no_account_contents");
                     let include_account_data = !arg_matches.is_present("no_account_data");
                     let account_data_encoding = parse_encoding_format(arg_matches);
+                    let mode = if let Some(pubkeys) = pubkeys_of(arg_matches, "account") {
+                        info!("Scanning individual accounts: {pubkeys:?}");
+                        AccountsOutputMode::Individual(pubkeys)
+                    } else if let Some(pubkey) = pubkey_of(arg_matches, "program_accounts") {
+                        info!("Scanning program accounts for {pubkey}");
+                        AccountsOutputMode::Program(pubkey)
+                    } else {
+                        info!("Scanning all accounts");
+                        AccountsOutputMode::All
+                    };
                     let config = AccountsOutputConfig {
+                        mode,
                         include_sysvars,
                         include_account_contents,
                         include_account_data,

ledger-tool/src/output.rs

@@ -6,6 +6,7 @@ use {
         Deserialize, Serialize,
     },
     solana_account_decoder::{UiAccount, UiAccountData, UiAccountEncoding},
+    solana_accounts_db::accounts_index::ScanConfig,
     solana_cli_output::{
         display::writeln_transaction, CliAccount, CliAccountNewConfig, OutputFormat, QuietDisplay,
         VerboseDisplay,
@@ -572,7 +573,14 @@ pub struct AccountsOutputStreamer {
     output_format: OutputFormat,
 }
 
+pub enum AccountsOutputMode {
+    All,
+    Individual(Vec<Pubkey>),
+    Program(Pubkey),
+}
+
 pub struct AccountsOutputConfig {
+    pub mode: AccountsOutputMode,
     pub include_sysvars: bool,
     pub include_account_contents: bool,
     pub include_account_data: bool,
@@ -608,7 +616,10 @@ impl AccountsOutputStreamer {
                     .serialize_field("summary", &*self.total_accounts_stats.borrow())
                     .map_err(|err| format!("unable to serialize accounts summary: {err}"))?;
                 SerializeStruct::end(struct_serializer)
-                    .map_err(|err| format!("unable to end serialization: {err}"))
+                    .map_err(|err| format!("unable to end serialization: {err}"))?;
+                // The serializer doesn't give us a trailing newline so do it ourselves
+                println!();
+                Ok(())
             }
             _ => {
                 // The compiler needs a placeholder type to satisfy the generic
@@ -637,6 +648,33 @@ impl AccountsScanner {
             && (self.config.include_sysvars || !solana_sdk::sysvar::is_sysvar_id(pubkey))
     }
 
+    fn maybe_output_account<S>(
+        &self,
+        seq_serializer: &mut Option<S>,
+        pubkey: &Pubkey,
+        account: &AccountSharedData,
+        slot: Option<Slot>,
+        cli_account_new_config: &CliAccountNewConfig,
+    ) where
+        S: SerializeSeq,
+    {
+        if self.config.include_account_contents {
+            if let Some(serializer) = seq_serializer {
+                let cli_account =
+                    CliAccount::new_with_config(pubkey, account, cli_account_new_config);
+                serializer.serialize_element(&cli_account).unwrap();
+            } else {
+                output_account(
+                    pubkey,
+                    account,
+                    slot,
+                    self.config.include_account_data,
+                    self.config.account_data_encoding,
+                );
+            }
+        }
+    }
+
     pub fn output<S>(&self, seq_serializer: &mut Option<S>)
     where
         S: SerializeSeq,
@@ -654,26 +692,53 @@ impl AccountsScanner {
                 .filter(|(pubkey, account, _)| self.should_process_account(account, pubkey))
             {
                 total_accounts_stats.accumulate_account(pubkey, &account, rent_collector);
-
+                self.maybe_output_account(
-                if self.config.include_account_contents {
+                    seq_serializer,
-                    if let Some(serializer) = seq_serializer {
+                    pubkey,
-                        let cli_account =
+                    &account,
-                            CliAccount::new_with_config(pubkey, &account, &cli_account_new_config);
+                    Some(slot),
-                        serializer.serialize_element(&cli_account).unwrap();
+                    &cli_account_new_config,
-                    } else {
+                );
-                        output_account(
-                            pubkey,
-                            &account,
-                            Some(slot),
-                            self.config.include_account_data,
-                            self.config.account_data_encoding,
-                        );
-                    }
-                }
             }
         };
 
-        self.bank.scan_all_accounts(scan_func).unwrap();
+        match &self.config.mode {
+            AccountsOutputMode::All => {
+                self.bank.scan_all_accounts(scan_func).unwrap();
+            }
+            AccountsOutputMode::Individual(pubkeys) => pubkeys.iter().for_each(|pubkey| {
+                if let Some((account, slot)) = self
+                    .bank
+                    .get_account_modified_slot_with_fixed_root(pubkey)
+                    .filter(|(account, _)| self.should_process_account(account, pubkey))
+                {
+                    total_accounts_stats.accumulate_account(pubkey, &account, rent_collector);
+                    self.maybe_output_account(
+                        seq_serializer,
+                        pubkey,
+                        &account,
+                        Some(slot),
+                        &cli_account_new_config,
+                    );
+                }
+            }),
+            AccountsOutputMode::Program(program_pubkey) => self
+                .bank
+                .get_program_accounts(program_pubkey, &ScanConfig::default())
+                .unwrap()
+                .iter()
+                .filter(|(pubkey, account)| self.should_process_account(account, pubkey))
+                .for_each(|(pubkey, account)| {
+                    total_accounts_stats.accumulate_account(pubkey, account, rent_collector);
+                    self.maybe_output_account(
+                        seq_serializer,
+                        pubkey,
+                        account,
+                        None,
+                        &cli_account_new_config,
+                    );
+                }),
+        }
     }
 }
 

runtime/src/bank.rs

@@ -6039,10 +6039,18 @@ impl Bank {
     // pro: safer assertion can be enabled inside AccountsDb
     // con: panics!() if called from off-chain processing
     pub fn get_account_with_fixed_root(&self, pubkey: &Pubkey) -> Option<AccountSharedData> {
-        self.load_slow_with_fixed_root(&self.ancestors, pubkey)
+        self.get_account_modified_slot_with_fixed_root(pubkey)
             .map(|(acc, _slot)| acc)
     }
 
+    // See note above get_account_with_fixed_root() about when to prefer this function
+    pub fn get_account_modified_slot_with_fixed_root(
+        &self,
+        pubkey: &Pubkey,
+    ) -> Option<(AccountSharedData, Slot)> {
+        self.load_slow_with_fixed_root(&self.ancestors, pubkey)
+    }
+
     pub fn get_account_modified_slot(&self, pubkey: &Pubkey) -> Option<(AccountSharedData, Slot)> {
         self.load_slow(&self.ancestors, pubkey)
     }