Allow Ctrl-C when querying balances (solana-labs#27314)

CriesofCarrots · Aug 23, 2022 · 62c5679 · 62c5679
1 parent 326f43d
commit 62c5679
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/tokens/src/commands.rs b/tokens/src/commands.rs
@@ -833,7 +833,11 @@ fn check_payer_balances(
     Ok(())
 }
 
-pub fn process_balances(client: &RpcClient, args: &BalancesArgs) -> Result<(), Error> {
+pub fn process_balances(
+    client: &RpcClient,
+    args: &BalancesArgs,
+    exit: Arc<AtomicBool>,
+) -> Result<(), Error> {
     let allocations: Vec<Allocation> =
         read_allocations(&args.input_csv, None, false, args.spl_token_args.is_some())?;
     let allocations = merge_allocations(&allocations);
@@ -855,6 +859,10 @@ pub fn process_balances(client: &RpcClient, args: &BalancesArgs) -> Result<(), E
     );
 
     for allocation in &allocations {
+        if exit.load(Ordering::SeqCst) {
+            return Err(Error::ExitSignal);
+        }
+
         if let Some(spl_token_args) = &args.spl_token_args {
             print_token_balances(client, allocation, spl_token_args)?;
         } else {

diff --git a/tokens/src/main.rs b/tokens/src/main.rs
@@ -44,7 +44,7 @@ fn main() -> Result<(), Box<dyn Error>> {
         }
         Command::Balances(mut args) => {
             spl_token::update_decimals(&client, &mut args.spl_token_args)?;
-            commands::process_balances(&client, &args)?;
+            commands::process_balances(&client, &args, exit)?;
         }
         Command::TransactionLog(args) => {
             commands::process_transaction_log(&args)?;