Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(http2/client): avoid double-polling a Select future #3290

Merged
merged 3 commits into from
Aug 21, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 49 additions & 28 deletions src/proto/h2/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use crate::rt::{Read, Write};
use bytes::Bytes;
use futures_channel::mpsc::{Receiver, Sender};
use futures_channel::{mpsc, oneshot};
use futures_util::future::{self, Either, FutureExt as _, Select};
use futures_util::future::{Either, FusedFuture, FutureExt as _};
use futures_util::stream::{StreamExt as _, StreamFuture};
use h2::client::{Builder, Connection, SendRequest};
use h2::SendStream;
Expand Down Expand Up @@ -143,7 +143,10 @@ where
} else {
(Either::Right(conn), ping::disabled())
};
let conn: ConnMapErr<T, B> = ConnMapErr { conn };
let conn: ConnMapErr<T, B> = ConnMapErr {
conn,
is_terminated: false,
};

exec.execute_h2_future(H2ClientFuture::Task {
task: ConnTask::new(conn, conn_drop_rx, cancel_tx),
Expand Down Expand Up @@ -218,6 +221,8 @@ pin_project! {
{
#[pin]
conn: Either<Conn<T, B>, Connection<Compat<T>, SendBuf<<B as Body>::Data>>>,
#[pin]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No need to pin a bool.

is_terminated: bool,
}
}

Expand All @@ -229,10 +234,26 @@ where
type Output = Result<(), ()>;

fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
self.project()
.conn
.poll(cx)
.map_err(|e| debug!("connection error: {}", e))
let mut this = self.project();

if *this.is_terminated {
return Poll::Pending;
}
let polled = this.conn.poll(cx);
if polled.is_ready() {
*this.is_terminated = true;
}
polled.map_err(|e| debug!("connection error: {}", e))
}
}

impl<T, B> FusedFuture for ConnMapErr<T, B>
where
B: Body,
T: Read + Write + Unpin,
{
fn is_terminated(&self) -> bool {
self.is_terminated
}
}

Expand All @@ -245,10 +266,11 @@ pin_project! {
T: Unpin,
{
#[pin]
select: Select<ConnMapErr<T, B>, StreamFuture<Receiver<Never>>>,
seanmonstar marked this conversation as resolved.
Show resolved Hide resolved
drop_rx: StreamFuture<Receiver<Never>>,
#[pin]
cancel_tx: Option<oneshot::Sender<Never>>,
conn: Option<ConnMapErr<T, B>>,
#[pin]
conn: ConnMapErr<T, B>,
}
}

Expand All @@ -263,9 +285,9 @@ where
cancel_tx: oneshot::Sender<Never>,
) -> Self {
Self {
select: future::select(conn, drop_rx),
drop_rx,
cancel_tx: Some(cancel_tx),
conn: None,
conn,
}
}
}
Expand All @@ -280,25 +302,24 @@ where
fn poll(self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
let mut this = self.project();

if let Some(conn) = this.conn {
conn.poll_unpin(cx).map(|_| ())
} else {
match ready!(this.select.poll_unpin(cx)) {
Either::Left((_, _)) => {
// ok or err, the `conn` has finished
return Poll::Ready(());
}
Either::Right((_, b)) => {
// mpsc has been dropped, hopefully polling
// the connection some more should start shutdown
// and then close
trace!("send_request dropped, starting conn shutdown");
drop(this.cancel_tx.take().expect("Future polled twice"));
this.conn = &mut Some(b);
return Poll::Pending;
}
}
if !this.conn.is_terminated() {
if let Poll::Ready(_) = this.conn.poll_unpin(cx) {
// ok or err, the `conn` has finished.
return Poll::Ready(());
};
}

if !this.drop_rx.is_terminated() {
if let Poll::Ready(_) = this.drop_rx.poll_unpin(cx) {
// mpsc has been dropped, hopefully polling
// the connection some more should start shutdown
// and then close.
trace!("send_request dropped, starting conn shutdown");
drop(this.cancel_tx.take().expect("ConnTask Future polled twice"));
}
};

Poll::Pending
}
}

Expand Down