From 6016da873f84a62f41a026e6d90e47ceac9b98d2 Mon Sep 17 00:00:00 2001 From: Ashok Menon Date: Sat, 2 Nov 2024 00:49:52 +0000 Subject: [PATCH] refactor(indexer-alt): graceful shutdown returns values ## Description Allow `graceful_shutdown` to accept tasks that return something other than `()`. ## Test plan CI --- crates/sui-indexer-alt/src/task.rs | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/crates/sui-indexer-alt/src/task.rs b/crates/sui-indexer-alt/src/task.rs index d027541a78310..f0c59183942f7 100644 --- a/crates/sui-indexer-alt/src/task.rs +++ b/crates/sui-indexer-alt/src/task.rs @@ -9,10 +9,12 @@ use tokio_util::sync::CancellationToken; /// Manages cleanly exiting the process, either because one of its constituent services has stopped /// or because an interrupt signal was sent to the process. -pub async fn graceful_shutdown( - services: impl IntoIterator>, +/// +/// Returns the exit values from all services that exited successfully. +pub async fn graceful_shutdown( + services: impl IntoIterator>, cancel: CancellationToken, -) { +) -> Vec { // If the service is naturalling winding down, we don't need to wait for an interrupt signal. // This channel is used to short-circuit the await in that case. let (cancel_ctrl_c_tx, cancel_ctrl_c_rx) = oneshot::channel(); @@ -24,20 +26,23 @@ pub async fn graceful_shutdown( _ = signal::ctrl_c() => cancel.cancel(), } - Ok(()) + None }; tokio::pin!(interrupt); let futures: Vec<_> = services .into_iter() - .map(Either::Left) + .map(|s| Either::Left(Box::pin(async move { s.await.ok() }))) .chain(iter::once(Either::Right(interrupt))) .collect(); // Wait for the first service to finish, or for an interrupt signal. - let (_, _, rest) = future::select_all(futures).await; + let (first, _, rest) = future::select_all(futures).await; let _ = cancel_ctrl_c_tx.send(()); // Wait for the remaining services to finish. - let _ = future::join_all(rest).await; + let mut results = vec![]; + results.extend(first); + results.extend(future::join_all(rest).await.into_iter().flatten()); + results }