From ff737730f4c6bbc6c9840163c5f792c00833a771 Mon Sep 17 00:00:00 2001 From: Jiahao XU Date: Wed, 30 Nov 2022 14:05:52 +1100 Subject: [PATCH] Optimize use of `tokio::select!`: Use biased selection (#580) as there is no need to randomize the first one to be polled. For `cancel_on_user_sig_term` and `StreamReadable::fill_buf`, the cancellation future should always to be polled first so that user would feel responsive. Signed-off-by: Jiahao XU --- .../binstalk-downloader/src/download/stream_readable.rs | 4 +++- crates/binstalk/src/helpers/signal.rs | 8 +++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/crates/binstalk-downloader/src/download/stream_readable.rs b/crates/binstalk-downloader/src/download/stream_readable.rs index af6e9c67..2db0666a 100644 --- a/crates/binstalk-downloader/src/download/stream_readable.rs +++ b/crates/binstalk-downloader/src/download/stream_readable.rs @@ -122,10 +122,12 @@ where let option = self.handle.block_on(async { if let Some(cancellation_future) = self.cancellation_future.as_mut() { tokio::select! { - res = next_stream(&mut self.stream) => res, + biased; + res = cancellation_future => { Err(res.err().unwrap_or_else(|| io::Error::from(DownloadError::UserAbort))) }, + res = next_stream(&mut self.stream) => res, } } else { next_stream(&mut self.stream).await diff --git a/crates/binstalk/src/helpers/signal.rs b/crates/binstalk/src/helpers/signal.rs index 30eb57bf..5b6b7465 100644 --- a/crates/binstalk/src/helpers/signal.rs +++ b/crates/binstalk/src/helpers/signal.rs @@ -20,11 +20,13 @@ pub async fn cancel_on_user_sig_term( ignore_signals()?; tokio::select! { - res = handle => res, + biased; + res = wait_on_cancellation_signal() => { res.map_err(BinstallError::Io) .and(Err(BinstallError::UserAbort)) } + res = handle => res, } } @@ -59,6 +61,8 @@ async fn wait_on_cancellation_signal_inner() -> Result<(), io::Error> { } tokio::select! { + biased; + res = signal::ctrl_c() => res, res = inner() => res, } @@ -72,6 +76,8 @@ mod unix { /// Same as [`wait_on_cancellation_signal`] but is only available on unix. pub async fn wait_on_cancellation_signal_unix() -> Result<(), io::Error> { tokio::select! { + biased; + res = wait_for_signal_unix(SignalKind::interrupt()) => res, res = wait_for_signal_unix(SignalKind::hangup()) => res, res = wait_for_signal_unix(SignalKind::terminate()) => res,