diff --git a/crates/irtt-cli/src/shared/client/session.rs b/crates/irtt-cli/src/shared/client/session.rs index 116ee37..fdc47e9 100644 --- a/crates/irtt-cli/src/shared/client/session.rs +++ b/crates/irtt-cli/src/shared/client/session.rs @@ -40,12 +40,16 @@ impl ClientSession { events.extend(self.client.send_probe()?); } events.extend(self.client.recv_available(RECV_BUDGET)?); + if self.client.is_peer_closed() { + return Ok(events); + } events.extend(self.client.poll_timeouts()?); Ok(events) } pub fn should_continue(&self, shutdown_requested: &AtomicBool) -> bool { !is_shutdown_requested(shutdown_requested) + && !self.client.is_peer_closed() && (self.continuous || self.client.next_send_deadline().is_some()) } @@ -59,12 +63,20 @@ impl ClientSession { { let deadline = Instant::now() + final_drain_duration(self.client.probe_timeout()); loop { + if self.client.is_peer_closed() { + break; + } + let mut received = false; let events = self.client.recv_available(RECV_BUDGET)?; received |= !events.is_empty(); on_events(&events); + if self.client.is_peer_closed() { + break; + } + let events = self.client.poll_timeouts()?; received |= !events.is_empty(); on_events(&events); @@ -81,10 +93,16 @@ impl ClientSession { } pub fn poll_timeouts(&mut self) -> Result, irtt_client::ClientError> { + if self.client.is_peer_closed() { + return Ok(vec![]); + } self.client.poll_timeouts() } pub fn close(&mut self) -> Result, irtt_client::ClientError> { + if self.client.is_peer_closed() { + return Ok(vec![]); + } self.client.close() } diff --git a/crates/irtt-client/src/client.rs b/crates/irtt-client/src/client.rs index 33e4bd0..c6a0f94 100644 --- a/crates/irtt-client/src/client.rs +++ b/crates/irtt-client/src/client.rs @@ -359,6 +359,9 @@ impl Client { break; } all_events.extend(events); + if self.is_peer_closed() { + break; + } } Ok(all_events) } diff --git a/crates/irtt-client/src/managed/runner.rs b/crates/irtt-client/src/managed/runner.rs index 97b0d48..bf80863 100644 --- a/crates/irtt-client/src/managed/runner.rs +++ b/crates/irtt-client/src/managed/runner.rs @@ -224,7 +224,9 @@ fn run_client( &mut counters, client.recv_available(MANAGED_RECV_BUDGET)?, ); - publish_events(&hub, &mut counters, client.poll_timeouts()?); + if !client.is_peer_closed() { + publish_events(&hub, &mut counters, client.poll_timeouts()?); + } break; } @@ -243,6 +245,9 @@ fn run_client( &mut counters, client.recv_available(MANAGED_RECV_BUDGET)?, ); + if client.is_peer_closed() { + break; + } publish_events(&hub, &mut counters, client.poll_timeouts()?); if client.is_run_complete() { @@ -257,8 +262,10 @@ fn run_client( } let packets_sent = client.packets_sent(); - let close_events = client.close()?; - publish_events(&hub, &mut counters, close_events); + if !client.is_peer_closed() { + let close_events = client.close()?; + publish_events(&hub, &mut counters, close_events); + } Ok(SessionOutcome { end_reason: if cancelled { @@ -295,12 +302,20 @@ fn drain_final_late_replies( let deadline = Instant::now() + MANAGED_FINAL_DRAIN; while Instant::now() < deadline && client.has_timed_out_metadata() { + if client.is_peer_closed() { + break; + } + let mut published = false; let events = client.recv_available(MANAGED_RECV_BUDGET)?; published |= !events.is_empty(); publish_events(hub, counters, events); + if client.is_peer_closed() { + break; + } + let events = client.poll_timeouts()?; published |= !events.is_empty(); publish_events(hub, counters, events);