diff --git a/crates/irtt-client/src/async_client.rs b/crates/irtt-client/src/async_client.rs index 8523746..ecee6c3 100644 --- a/crates/irtt-client/src/async_client.rs +++ b/crates/irtt-client/src/async_client.rs @@ -34,6 +34,13 @@ use crate::{ timing::ClientTimestamp, }; +/// Maximum amount of immediately available opening work performed by one poll. +/// +/// Opening may consume ignored UDP datagrams, roll over an expired attempt, or +/// encounter readiness false positives without ever awaiting. Bound that work +/// so one managed target cannot retain the current-thread runtime. +const OPEN_POLL_WORK_BUDGET: usize = 64; + #[cfg(test)] use crate::client::ProbeSendTimestamps; @@ -315,7 +322,7 @@ impl AsyncClient { unreachable!("connected async clients retain an open request"); } - loop { + for _ in 0..OPEN_POLL_WORK_BUDGET { if state.cleanup.is_some() { return self.poll_open_cleanup(state, cx); } @@ -430,6 +437,11 @@ impl AsyncClient { }; return Poll::Ready(Ok(self.commit_async_open(prepared))); } + + // We deliberately truncated known immediate opening work. Arrange for + // another poll even if no new socket-readiness edge occurs. + cx.waker().wake_by_ref(); + Poll::Pending } pub(crate) fn poll_send_probe( @@ -855,49 +867,48 @@ impl AsyncClient { state: &mut AsyncOpenState, cx: &mut Context<'_>, ) -> Poll> { - loop { - let Some(packet) = state - .cleanup - .as_ref() - .expect("cleanup polling requires a retained primary error") - .packet() - else { - return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); - }; + let Some(packet) = state + .cleanup + .as_ref() + .expect("cleanup polling requires a retained primary error") + .packet() + else { + return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); + }; - #[cfg(test)] - if self - .test_hooks - .pause_open_cleanup_before_send - .replace(false) - { - return Poll::Pending; - } + #[cfg(test)] + if self + .test_hooks + .pause_open_cleanup_before_send + .replace(false) + { + return Poll::Pending; + } - #[cfg(test)] - if self.test_hooks.fail_cleanup_send.replace(false) { - return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); - } - #[cfg(test)] - let send_result = self.test_hooks.try_send(&self.socket, packet); - #[cfg(not(test))] - let send_result = self.socket.try_send(packet); - match send_result { - Ok(_) => { - return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); - } - Err(error) if error.kind() == io::ErrorKind::WouldBlock => {} - Err(_) => { - return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); - } - } + #[cfg(test)] + if self.test_hooks.fail_cleanup_send.replace(false) { + return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); + } + #[cfg(test)] + let send_result = self.test_hooks.try_send(&self.socket, packet); + #[cfg(not(test))] + let send_result = self.socket.try_send(packet); + match send_result { + Ok(_) => return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())), + Err(error) if error.kind() == io::ErrorKind::WouldBlock => {} + Err(_) => return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())), + } - match poll_writable_until(&self.socket, state.deadline_timer(), cx) { - Poll::Pending => return Poll::Pending, - Poll::Ready(Ok(true)) => {} - Poll::Ready(Ok(false) | Err(_)) => { - return Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())); - } + match poll_writable_until(&self.socket, state.deadline_timer(), cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Ok(true)) => { + // Retain cleanup state and resume the send in a fresh poll. A + // writable readiness false positive must not spin this future. + cx.waker().wake_by_ref(); + Poll::Pending + } + Poll::Ready(Ok(false) | Err(_)) => { + Poll::Ready(Err(state.cleanup.take().unwrap().into_primary())) } } } diff --git a/crates/irtt-client/src/async_client/tests.rs b/crates/irtt-client/src/async_client/tests.rs index ec4df19..6447c5a 100644 --- a/crates/irtt-client/src/async_client/tests.rs +++ b/crates/irtt-client/src/async_client/tests.rs @@ -334,6 +334,56 @@ fn ignored_open_traffic_uses_one_request_and_deadline_per_attempt() { assert_eq!(server.finish().len(), 2); } +#[test] +fn opening_poll_budget_yields_then_public_open_completes() { + let (queued, queue_ready) = mpsc::channel(); + let server = start_server(move |socket, tx| { + let (open_packet, peer) = recv_packet(&socket, &tx); + let request = decode_open_request(&open_packet, None).unwrap(); + for _ in 0..=OPEN_POLL_WORK_BUDGET { + socket.send_to(&[0_u8], peer).unwrap(); + } + send_open_reply( + &socket, + peer, + request.params, + None, + flags::FLAG_OPEN | flags::FLAG_REPLY, + TOKEN, + ); + queued.send(()).unwrap(); + }); + + runtime().block_on(async { + let mut client = AsyncClient::connect(config(server.addr, None, 0)) + .await + .unwrap(); + client.socket.writable().await.unwrap(); + client.test_hooks.pause_open_before_readable.set(true); + let mut opening = Box::pin(client.open()); + + // Submit the one opening request, then let the server queue all noise + // before performing the explicit bounded poll below. + assert!(matches!( + poll_once(Pin::as_mut(&mut opening)), + Poll::Pending + )); + queue_ready.recv_timeout(Duration::from_secs(2)).unwrap(); + + // The valid reply is behind more ignored datagrams than one opening + // poll may consume, so this cannot complete without the budget yield. + assert!(matches!( + poll_once(Pin::as_mut(&mut opening)), + Poll::Pending + )); + assert!(matches!( + opening.await.unwrap(), + OpenOutcome::Started { .. } + )); + }); + assert_eq!(server.finish().len(), 1); +} + #[test] fn ignored_open_traffic_cannot_extend_the_absolute_attempt_deadline() { let server = start_server(move |socket, tx| {