diff --git a/Cargo.lock b/Cargo.lock index 9263cd2c55..dbe825f14e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3713,16 +3713,16 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.10" +version = "0.7.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5419f34732d9eb6ee4c3578b7989078579b7f039cbbb9ca2c4da015749371e15" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", - "tracing", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index e35ea56ae3..7dff36ab5d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -110,7 +110,7 @@ tokio = { version = "1", features = [ "macros", ] } tokio-serde = "0.9" -tokio-util = { version = "0.7", features = ["codec", "io"] } +tokio-util = { version = "0.7.18", features = ["codec", "io", "rt"] } toml = "0.9" tower-service = "0.3" typed-path = "0.12.0" diff --git a/src/mock_command.rs b/src/mock_command.rs index d7ffa03840..911fe7e13f 100644 --- a/src/mock_command.rs +++ b/src/mock_command.rs @@ -256,6 +256,7 @@ impl RunCommand for AsyncCommand { let token = self.jobserver.acquire().await?; let mut inner = tokio::process::Command::from(inner); let child = inner + .kill_on_drop(true) .spawn() .with_context(|| format!("failed to spawn {:?}", inner))?; diff --git a/src/server.rs b/src/server.rs index 143ea6375c..0ea4db10b3 100644 --- a/src/server.rs +++ b/src/server.rs @@ -37,7 +37,7 @@ use fs::metadata; use fs_err as fs; use futures::channel::mpsc; use futures::future::FutureExt; -use futures::{Sink, SinkExt, Stream, StreamExt, TryFutureExt, future, stream}; +use futures::{Sink, SinkExt, Stream, StreamExt, TryFutureExt, future}; use number_prefix::NumberPrefix; use serde::{Deserialize, Serialize}; use std::cell::Cell; @@ -989,7 +989,6 @@ where } use futures::TryStreamExt; -use futures::future::Either; impl SccacheService where @@ -1072,7 +1071,7 @@ where } } - fn bind(self, socket: T) -> impl Future> + Send + Sized + 'static + async fn bind(self, socket: T) -> Result<()> where T: AsyncRead + AsyncWrite + Unpin + Send + 'static, { @@ -1086,36 +1085,61 @@ where } let io = builder.new_framed(socket); - let (sink, stream) = SccacheTransport { + let (sink, mut stream) = SccacheTransport { inner: Framed::new(io.sink_err_into().err_into(), BincodeCodec), } .split(); - let sink = sink.sink_err_into::(); + let mut sink = sink.sink_err_into::(); + + let (reqs_tx, mut reqs_rx) = tokio::sync::mpsc::unbounded_channel(); let me = Arc::new(self); - stream - .err_into::() - .and_then(move |input| me.clone().call(input)) - .and_then(move |response| async move { - let fut = match response { - Message::WithoutBody(message) => { - let stream = stream::once(async move { Ok(Frame::Message { message }) }); - Either::Left(stream) + + let _handle = util::spawn(async move { + while let Some(req) = reqs_rx.recv().await { + match req { + Ok(req) => { + let res = match util::spawn(me.clone().call(req)).await? { + Ok(res) => res, + Err(err) => { + return Err(err); + } + }; + match res { + Message::WithoutBody(message) => { + sink.send(Frame::Message { message }).await?; + } + Message::WithBody(message, body) => { + sink.send(Frame::Message { message }).await?; + sink.send(Frame::Body { + chunk: Some(util::spawn(body).await??), + }) + .await?; + sink.send(Frame::Body { chunk: None }).await?; + } + } } - Message::WithBody(message, body) => { - let stream = stream::once(async move { Ok(Frame::Message { message }) }) - .chain( - body.into_stream() - .map_ok(|chunk| Frame::Body { chunk: Some(chunk) }), - ) - .chain(stream::once(async move { Ok(Frame::Body { chunk: None }) })); - Either::Right(stream) + Err(err) => { + return Err(err); } - }; - Ok(Box::pin(fut)) - }) - .try_flatten() - .forward(sink) + } + } + + Ok(()) + }); + + while let Some(req) = stream.next().await { + match req { + Ok(req) => { + reqs_tx.send(Ok(req))?; + } + Err(err) => { + return Err(err); + } + } + } + + Ok(()) } /// Get dist status. @@ -1445,8 +1469,12 @@ where let me = self.clone(); - self.rt - .spawn(async move { + // This redundant async block exists to reduce whitespace-only + // changes when comparing this diff with upstream/main. + // TODO: remove this before merging + #[allow(clippy::redundant_async_block)] + util::spawn_on(&self.rt, async move { + async move { let result = match me.dist_client.get_client().await { Ok(client) => std::panic::AssertUnwindSafe(hasher.get_cached_or_compile( &me, @@ -1651,9 +1679,11 @@ where } Ok(res) - }) - .map_err(anyhow::Error::new) - .await? + } + .await + }) + .map_err(anyhow::Error::new) + .await? } } diff --git a/src/util.rs b/src/util.rs index bc64c3d1b7..3a0e847b3b 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1020,6 +1020,25 @@ fn unhex(b: u8) -> std::io::Result { } } +pub fn spawn(future: F) -> tokio_util::task::AbortOnDropHandle +where + F: Future + Send + 'static, + F::Output: Send + 'static, +{ + tokio_util::task::AbortOnDropHandle::new(tokio::spawn(future)) +} + +pub fn spawn_on( + handle: &tokio::runtime::Handle, + future: F, +) -> tokio_util::task::AbortOnDropHandle +where + F: Future + Send + 'static, + F::Output: Send + 'static, +{ + tokio_util::task::AbortOnDropHandle::new(handle.spawn(future)) +} + /// A reverse version of std::ascii::escape_default pub fn ascii_unescape_default(s: &[u8]) -> std::io::Result> { let mut out = Vec::with_capacity(s.len() + 4);