From 260a680ad6c7b6e4737fc77e938353fe33735eb8 Mon Sep 17 00:00:00 2001
From: Andre da Silva <2917611+ndr-ds@users.noreply.github.com>
Date: Fri, 14 Feb 2025 00:05:25 -0300
Subject: [PATCH] Gracefully shutdown benchmark (#3326)

## Motivation

Right now if you interrupt the benchmark, it might leave you in a state where the local node is inconsistent with the validators, the wallet hasn't been saved... etc

## Proposal

Gracefully shutdown the benchmark on interruption signals, making sure we never exit in the middle of sending a block proposal, and only saving the wallet on program exit.

## Test Plan

Ran the benchmark locally, did Ctrl+C and saw it gracefully exit.

## Release Plan

- Nothing to do / These changes follow the usual release cycle.
---
 Cargo.lock                          |  1 +
 linera-client/Cargo.toml            |  3 ++-
 linera-client/src/client_context.rs | 12 +++++++++++-
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 2767d920c767..c84bdfb371fa 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4478,6 +4478,7 @@ dependencies = [
  "thiserror-context",
  "tokio",
  "tokio-stream",
+ "tokio-util",
  "tracing",
  "trait-variant",
  "wasm-bindgen-futures",
diff --git a/linera-client/Cargo.toml b/linera-client/Cargo.toml
index d08c08753dee..8576cd030e85 100644
--- a/linera-client/Cargo.toml
+++ b/linera-client/Cargo.toml
@@ -13,7 +13,7 @@ version.workspace = true
 
 [features]
 test = ["linera-views/test", "linera-execution/test"]
-benchmark = ["linera-base/test", "dep:linera-sdk"]
+benchmark = ["linera-base/test", "dep:linera-sdk", "dep:tokio-util"]
 wasmer = [
     "linera-core/wasmer",
     "linera-execution/wasmer",
@@ -75,6 +75,7 @@ thiserror.workspace = true
 thiserror-context.workspace = true
 tokio.workspace = true
 tokio-stream.workspace = true
+tokio-util = { workspace = true, optional = true }
 tracing.workspace = true
 trait-variant.workspace = true
 
diff --git a/linera-client/src/client_context.rs b/linera-client/src/client_context.rs
index b8d198a2b403..ed0a28a2223e 100644
--- a/linera-client/src/client_context.rs
+++ b/linera-client/src/client_context.rs
@@ -43,6 +43,7 @@ use {
         data_types::Amount,
         hashed::Hashed,
         identifiers::{AccountOwner, ApplicationId, Owner},
+        listen_for_shutdown_signals,
     },
     linera_chain::{
         data_types::{BlockProposal, ExecutedBlock, ProposedBlock, SignatureAggregator, Vote},
@@ -57,6 +58,7 @@ use {
     linera_sdk::abis::fungible,
     std::{collections::HashMap, iter},
     tokio::task,
+    tokio_util::sync::CancellationToken,
     tracing::{error, trace, warn},
 };
 
@@ -571,9 +573,17 @@ where
         transactions_per_block: usize,
         epoch: Epoch,
     ) -> Result<(), Error> {
+        let shutdown_notifier = CancellationToken::new();
+        tokio::spawn(listen_for_shutdown_signals(shutdown_notifier.clone()));
+
         let mut num_sent_proposals = 0;
         let mut start = Instant::now();
         for (chain_id, operations, key_pair) in blocks_infos_iter {
+            if shutdown_notifier.is_cancelled() {
+                info!("Shutdown signal received, stopping benchmark");
+                self.save_wallet().await?;
+                return Ok(());
+            }
             let chain = self.wallet.get(*chain_id).expect("should have chain");
             let block = ProposedBlock {
                 epoch,
@@ -664,6 +674,7 @@ where
         }
 
         if bps.is_none() {
+            self.save_wallet().await?;
             let elapsed = start.elapsed();
             let bps = num_sent_proposals as f64 / elapsed.as_secs_f64();
             info!(
@@ -1046,7 +1057,6 @@ where
             chain.block_hash = info.block_hash;
             chain.next_block_height = info.next_block_height;
         }
-        self.save_wallet().await.unwrap();
     }
 
     /// Creates a fungible token transfer operation.