Skip to content

Commit 7d6f864

Browse files
committed
fix: Missed wakeup in poll_trailers
poll_trailers returned Pending without storing a waker when a non-trailer event was at the front of pending_recv.
1 parent 8fea6b2 commit 7d6f864

2 files changed

Lines changed: 88 additions & 1 deletion

File tree

src/proto/streams/recv.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1230,7 +1230,7 @@ impl Recv {
12301230
Some(event) => {
12311231
// Frame is not trailers.. not ready to poll trailers yet.
12321232
stream.pending_recv.push_front(&mut self.buffer, event);
1233-
1233+
stream.recv_task = Some(cx.waker().clone());
12341234
Poll::Pending
12351235
}
12361236
None => self.schedule_recv(cx, stream),

tests/h2-tests/tests/trailers.rs

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
use futures::StreamExt;
22
use h2_support::prelude::*;
3+
use std::task::Poll;
4+
use tokio::sync::oneshot;
35

46
#[tokio::test]
57
async fn recv_trailers_only() {
@@ -110,6 +112,91 @@ fn recv_trailers_without_eos() {
110112
// This should be a protocol error?
111113
}
112114

115+
#[tokio::test]
116+
async fn poll_trailers_before_data_is_consumed() {
117+
h2_support::trace_init!();
118+
let (io, mut srv) = mock::new();
119+
let (frames_ready_tx, frames_ready_rx) = oneshot::channel();
120+
121+
let srv = async move {
122+
let settings = srv.assert_client_handshake().await;
123+
assert_default_settings!(settings);
124+
125+
// 2. Receive the request.
126+
srv.recv_frame(
127+
frames::headers(1)
128+
.request("GET", "https://example.com/")
129+
.eos(),
130+
)
131+
.await;
132+
133+
// 3. Send response HEADERS followed by DATA and trailers.
134+
srv.send_frame(frames::headers(1).response(200)).await;
135+
srv.send_frame(frames::data(1, "hello")).await;
136+
srv.send_frame(frames::headers(1).field("trailer-key", "trailer-val").eos())
137+
.await;
138+
139+
// 4. Ensure all preceding frames have been processed by the client.
140+
srv.ping_pong([1; 8]).await;
141+
frames_ready_tx.send(()).unwrap();
142+
};
143+
144+
let client = async move {
145+
let (mut client, conn) = client::handshake(io).await.expect("handshake");
146+
let conn = tokio::spawn(async move {
147+
conn.await.expect("client");
148+
});
149+
150+
// 1. Send the request and wait for response HEADERS.
151+
let resp = client.get("https://example.com/").await.expect("response");
152+
assert_eq!(resp.status(), StatusCode::OK);
153+
154+
let mut body = resp.into_body();
155+
frames_ready_rx.await.unwrap();
156+
let mut first_poll = true;
157+
158+
let trailers = tokio::time::timeout(
159+
Duration::from_secs(1),
160+
poll_fn(|cx| {
161+
if first_poll {
162+
// 5. Poll trailers while DATA is at the front of pending_recv.
163+
// This returns Pending and registers this future's waker.
164+
first_poll = false;
165+
assert!(
166+
matches!(body.poll_trailers(cx), Poll::Pending),
167+
"poll_trailers should be Pending when DATA is buffered"
168+
);
169+
170+
// 6. Consume the DATA frame. The next poll reaches the
171+
// queued trailers and wakes the waker registered in 5.
172+
match body.poll_data(cx) {
173+
Poll::Ready(Some(Ok(data))) => assert_eq!(data, "hello"),
174+
other => panic!("expected DATA, got {:?}", other),
175+
}
176+
assert!(matches!(body.poll_data(cx), Poll::Ready(None)));
177+
178+
Poll::Pending
179+
} else {
180+
// 7. This future must only be polled again after
181+
// poll_data's notify_recv wakes it.
182+
body.poll_trailers(cx)
183+
}
184+
})
185+
.wakened(),
186+
)
187+
.await
188+
.expect("poll_trailers was not woken")
189+
.expect("trailers result")
190+
.expect("should have trailers");
191+
192+
assert_eq!(trailers["trailer-key"], "trailer-val");
193+
194+
conn.await.unwrap();
195+
};
196+
197+
join(srv, client).await;
198+
}
199+
113200
#[tokio::test]
114201
async fn send_trailers_rejects_connection_specific_headers() {
115202
// RFC 9113 §8.2.2: endpoints MUST NOT *generate* an HTTP/2 message containing

0 commit comments

Comments
 (0)