From 62c56795574429f21e7c5c814079d3a0711032d5 Mon Sep 17 00:00:00 2001
From: Brooks Prumo <brooks@solana.com>
Date: Tue, 23 Aug 2022 10:23:30 -0400
Subject: [PATCH] Allow Ctrl-C when querying balances (#27314)

---
 tokens/src/commands.rs | 10 +++++++++-
 tokens/src/main.rs     |  2 +-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/tokens/src/commands.rs b/tokens/src/commands.rs
index f7ce05aa4..744770486 100644
--- 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
index bbd8c0e43..1bf81b01f 100644
--- 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)?;