Skip to content

Commit

Permalink
fix(conn/auto): hang on graceful shutdown when reading version
Browse files Browse the repository at this point in the history
  • Loading branch information
dswij authored and seanmonstar committed Feb 26, 2024
1 parent 8c7b13b commit 62ceda7
Showing 1 changed file with 13 additions and 2 deletions.
15 changes: 13 additions & 2 deletions src/server/conn/auto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,7 @@ where
buf: [MaybeUninit::uninit(); 24],
filled: 0,
version: Version::H2,
cancelled: false,
_pin: PhantomPinned,
}
}
Expand All @@ -185,12 +186,19 @@ pin_project! {
// the amount of `buf` thats been filled
filled: usize,
version: Version,
cancelled: bool,
// Make this future `!Unpin` for compatibility with async trait methods.
#[pin]
_pin: PhantomPinned,
}
}

impl<I> ReadVersion<I> {
pub fn cancel(self: Pin<&mut Self>) {
*self.project().cancelled = true;
}
}

impl<I> Future for ReadVersion<I>
where
I: Read + Unpin,
Expand All @@ -199,6 +207,9 @@ where

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
if *this.cancelled {
return Poll::Ready(Err(io::Error::new(io::ErrorKind::Interrupted, "Cancelled")));
}

let mut buf = ReadBuf::uninit(&mut *this.buf);
// SAFETY: `this.filled` tracks how many bytes have been read (and thus initialized) and
Expand Down Expand Up @@ -296,7 +307,7 @@ where
/// `Connection::poll` has resolved, this does nothing.
pub fn graceful_shutdown(self: Pin<&mut Self>) {
match self.project().state.project() {
ConnStateProj::ReadVersion { .. } => {}
ConnStateProj::ReadVersion { read_version, .. } => read_version.cancel(),
#[cfg(feature = "http1")]
ConnStateProj::H1 { conn } => conn.graceful_shutdown(),
#[cfg(feature = "http2")]
Expand Down Expand Up @@ -420,7 +431,7 @@ where
/// called after `UpgradeableConnection::poll` has resolved, this does nothing.
pub fn graceful_shutdown(self: Pin<&mut Self>) {
match self.project().state.project() {
UpgradeableConnStateProj::ReadVersion { .. } => {}
UpgradeableConnStateProj::ReadVersion { read_version, .. } => read_version.cancel(),
#[cfg(feature = "http1")]
UpgradeableConnStateProj::H1 { conn } => conn.graceful_shutdown(),
#[cfg(feature = "http2")]
Expand Down

0 comments on commit 62ceda7

Please sign in to comment.