Skip to content

Commit

Permalink
refactor(indexer-alt): graceful shutdown returns values
Browse files Browse the repository at this point in the history
## Description

Allow `graceful_shutdown` to accept tasks that return something other
than `()`.

## Test plan

CI
  • Loading branch information
amnn committed Nov 2, 2024
1 parent 48d3410 commit 6016da8
Showing 1 changed file with 12 additions and 7 deletions.
19 changes: 12 additions & 7 deletions crates/sui-indexer-alt/src/task.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<Item = JoinHandle<()>>,
///
/// Returns the exit values from all services that exited successfully.
pub async fn graceful_shutdown<T>(
services: impl IntoIterator<Item = JoinHandle<T>>,
cancel: CancellationToken,
) {
) -> Vec<T> {
// 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();
Expand All @@ -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
}

0 comments on commit 6016da8

Please sign in to comment.