From 3948e3dec97f5029014e49d843750dc2c57c6a09 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Tue, 4 Aug 2026 13:44:33 +0700 Subject: [PATCH] fix(client): timestamp close events after send --- crates/irtt-client/src/client.rs | 11 +++- crates/irtt-client/src/client/tests/close.rs | 30 ++++++++++- crates/irtt-client/src/managed/group.rs | 57 +++++++++++++++----- crates/irtt-client/src/session/machine.rs | 45 +++++++++++----- 4 files changed, 114 insertions(+), 29 deletions(-) diff --git a/crates/irtt-client/src/client.rs b/crates/irtt-client/src/client.rs index 1d1c85f..c9e9a99 100644 --- a/crates/irtt-client/src/client.rs +++ b/crates/irtt-client/src/client.rs @@ -60,6 +60,7 @@ struct ClientTestHooks { fail_close_send: Cell, close_reported_len: Cell>, close_send_attempts: Cell, + close_sent_at: Cell>, fail_cleanup_send: Cell, recv_buffer_len_override: Cell>, fail_dscp_restore: Cell, @@ -199,7 +200,15 @@ impl Client { return Err(ClientError::Socket(err)); } }; - let event = self.runtime.commit_local_close(prepared.commit); + #[cfg(not(test))] + let sent_at = ClientTimestamp::now(); + #[cfg(test)] + let sent_at = self + .test_hooks + .close_sent_at + .take() + .unwrap_or_else(ClientTimestamp::now); + let event = self.runtime.commit_local_close(prepared.commit, sent_at); self.schedule = None; self.applied_dscp = None; diff --git a/crates/irtt-client/src/client/tests/close.rs b/crates/irtt-client/src/client/tests/close.rs index 656692f..79813cb 100644 --- a/crates/irtt-client/src/client/tests/close.rs +++ b/crates/irtt-client/src/client/tests/close.rs @@ -29,10 +29,19 @@ fn close_sends_one_close_packet_with_negotiated_token() { }); let mut client = Client::connect(default_test_config(server.addr)).unwrap(); assert_open_started(client.open().unwrap()); + let sent_at = ClientTimestamp { + wall: SystemTime::UNIX_EPOCH + Duration::from_secs(1_234), + mono: Instant::now() + Duration::from_secs(5), + }; + client.test_hooks.close_sent_at.set(Some(sent_at)); let events = client.close().unwrap(); assert!(matches!( events.as_slice(), - [ClientEvent::SessionClosed { token: TOKEN, .. }] + [ClientEvent::SessionClosed { + token: TOKEN, + at, + .. + }] if *at == sent_at )); assert!(!client.runtime.is_open()); assert!(client.schedule.is_none()); @@ -127,9 +136,15 @@ fn close_send_failure_leaves_machine_and_schedule_open_for_retry() { }); let mut client = Client::connect(default_test_config(server.addr)).unwrap(); assert_open_started(client.open().unwrap()); + let sent_at = ClientTimestamp { + wall: SystemTime::UNIX_EPOCH + Duration::from_secs(2_345), + mono: Instant::now() + Duration::from_secs(6), + }; + client.test_hooks.close_sent_at.set(Some(sent_at)); client.test_hooks.fail_close_send.set(true); assert!(matches!(client.close(), Err(ClientError::Socket(_)))); + assert_eq!(client.test_hooks.close_sent_at.get(), Some(sent_at)); assert!(client.runtime.is_open()); assert!(client.schedule.is_some()); assert_eq!(client.applied_dscp, Some(0)); @@ -137,8 +152,13 @@ fn close_send_failure_leaves_machine_and_schedule_open_for_retry() { assert!(matches!( client.close().unwrap().as_slice(), - [ClientEvent::SessionClosed { token: TOKEN, .. }] + [ClientEvent::SessionClosed { + token: TOKEN, + at, + .. + }] if *at == sent_at )); + assert_eq!(client.test_hooks.close_sent_at.get(), None); assert_eq!(client.test_hooks.close_send_attempts.get(), 2); assert!(matches!(client.close(), Err(ClientError::AlreadyClosed))); server.join(); @@ -160,6 +180,11 @@ fn short_successful_close_commits_before_length_mismatch() { let mut client = Client::connect(default_test_config(server.addr)).unwrap(); assert_open_started(client.open().unwrap()); let expected = client.runtime.prepare_close().unwrap().bytes.len(); + let sent_at = ClientTimestamp { + wall: SystemTime::UNIX_EPOCH + Duration::from_secs(3_456), + mono: Instant::now() + Duration::from_secs(7), + }; + client.test_hooks.close_sent_at.set(Some(sent_at)); client.test_hooks.close_reported_len.set(Some(expected - 1)); assert!(matches!( @@ -172,6 +197,7 @@ fn short_successful_close_commits_before_length_mismatch() { assert!(!client.runtime.is_open()); assert!(client.schedule.is_none()); assert_eq!(client.applied_dscp, None); + assert_eq!(client.test_hooks.close_sent_at.get(), None); assert!(matches!(client.close(), Err(ClientError::AlreadyClosed))); server.join(); } diff --git a/crates/irtt-client/src/managed/group.rs b/crates/irtt-client/src/managed/group.rs index 9d1145e..0e74ab9 100644 --- a/crates/irtt-client/src/managed/group.rs +++ b/crates/irtt-client/src/managed/group.rs @@ -733,6 +733,8 @@ struct TargetState { close_send_attempts: usize, #[cfg(test)] close_event_reserve_error: bool, + #[cfg(test)] + close_sent_at: Option, } type ScheduledTarget = (Arc>, Instant); @@ -855,6 +857,8 @@ impl TargetState { close_send_attempts: 0, #[cfg(test)] close_event_reserve_error: false, + #[cfg(test)] + close_sent_at: None, }) } @@ -1872,7 +1876,14 @@ fn close_locked_target( return; } }; - let event = target.runtime.commit_local_close(prepared.commit); + #[cfg(not(test))] + let sent_at = ClientTimestamp::now(); + #[cfg(test)] + let sent_at = target + .close_sent_at + .take() + .unwrap_or_else(ClientTimestamp::now); + let event = target.runtime.commit_local_close(prepared.commit, sent_at); target.schedule = None; #[cfg(test)] let bytes = reported_bytes.unwrap_or(bytes); @@ -3046,8 +3057,13 @@ mod tests { let fixture = active_target_fixture_at(peer.local_addr().unwrap()); let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); + let sent_at = ClientTimestamp { + wall: std::time::SystemTime::UNIX_EPOCH + Duration::from_secs(4_567), + mono: Instant::now() + Duration::from_secs(8), + }; { let mut target = fixture.target.lock().expect("target mutex poisoned"); + target.close_sent_at = Some(sent_at); close_locked_target( &socket, &fixture.hub, @@ -3069,19 +3085,20 @@ mod tests { assert_eq!(u64::from_le_bytes(packet[4..12].try_into().unwrap()), TOKEN); assert_eq!(len, 12); let events = drain_available_group_events(&fixture.subscription); - assert_eq!( - events - .iter() - .filter(|event| matches!( - event, - ManagedGroupEvent::Client(TargetEvent { - event: ClientEvent::SessionClosed { token: TOKEN, .. }, - .. - }) - )) - .count(), - 1 - ); + let closes: Vec<_> = events + .iter() + .filter_map(|event| match event { + ManagedGroupEvent::Client(TargetEvent { + event: + ClientEvent::SessionClosed { + token: TOKEN, at, .. + }, + .. + }) => Some(*at), + _ => None, + }) + .collect(); + assert_eq!(closes, [sent_at]); } #[test] @@ -3127,8 +3144,13 @@ mod tests { let fixture = active_target_fixture_at(peer.local_addr().unwrap()); let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); + let sent_at = ClientTimestamp { + wall: std::time::SystemTime::UNIX_EPOCH + Duration::from_secs(5_678), + mono: Instant::now() + Duration::from_secs(9), + }; { let mut target = fixture.target.lock().expect("target mutex poisoned"); + target.close_sent_at = Some(sent_at); target.close_send_error = true; close_locked_target( &socket, @@ -3139,6 +3161,7 @@ mod tests { assert!(target.runtime.is_open()); assert!(target.schedule.is_some()); assert_eq!(target.close_send_attempts, 1); + assert_eq!(target.close_sent_at, Some(sent_at)); assert!(matches!( target.final_reason, Some(ManagedTargetEndReason::Failed(_)) @@ -3169,9 +3192,14 @@ mod tests { let fixture = active_target_fixture_at(peer.local_addr().unwrap()); let socket = UdpSocket::bind("127.0.0.1:0").unwrap(); + let sent_at = ClientTimestamp { + wall: std::time::SystemTime::UNIX_EPOCH + Duration::from_secs(6_789), + mono: Instant::now() + Duration::from_secs(10), + }; { let mut target = fixture.target.lock().expect("target mutex poisoned"); let expected = target.runtime.prepare_close().unwrap().bytes.len(); + target.close_sent_at = Some(sent_at); target.close_reported_len = Some(expected - 1); close_locked_target( &socket, @@ -3182,6 +3210,7 @@ mod tests { assert!(!target.runtime.is_open()); assert!(target.schedule.is_none()); assert_eq!(target.close_send_attempts, 1); + assert_eq!(target.close_sent_at, None); let failure = target .final_reason .as_ref() diff --git a/crates/irtt-client/src/session/machine.rs b/crates/irtt-client/src/session/machine.rs index 60a2f86..5072ced 100644 --- a/crates/irtt-client/src/session/machine.rs +++ b/crates/irtt-client/src/session/machine.rs @@ -115,7 +115,7 @@ pub(crate) struct PreparedClose<'a> { #[derive(Debug)] pub(crate) struct CloseCommit { packets_sent: u64, - event: ClientEvent, + token: u64, } #[derive(Debug)] @@ -425,16 +425,16 @@ impl SessionMachine { bytes: &session.local_close_packet, commit: CloseCommit { packets_sent: session.packets_sent, - event: ClientEvent::SessionClosed { - remote: self.remote, - token: session.token, - at: ClientTimestamp::now(), - }, + token: session.token, }, }) } - pub(crate) fn commit_local_close(&mut self, commit: CloseCommit) -> ClientEvent { + pub(crate) fn commit_local_close( + &mut self, + commit: CloseCommit, + sent_at: ClientTimestamp, + ) -> ClientEvent { debug_assert!( matches!(self.state, MachineState::Open(_)), "local close commits only from open state" @@ -443,7 +443,11 @@ impl SessionMachine { source: CloseSource::Local, packets_sent: commit.packets_sent, }; - commit.event + ClientEvent::SessionClosed { + remote: self.remote, + token: commit.token, + at: sent_at, + } } pub(crate) fn pending_is_empty(&self) -> bool { @@ -1370,9 +1374,8 @@ mod tests { } #[test] - fn prepared_close_is_stable_inert_and_commits_preserved_state() { + fn prepare_close_is_stable_inert_and_does_not_determine_event_timestamp() { let mut machine = open_machine(4, Duration::from_secs(1)); - let remote = machine.remote; active_mut(&mut machine).packets_sent = 7; let first_bytes = { @@ -1387,6 +1390,8 @@ mod tests { ) .unwrap() ); + assert_eq!(prepared.commit.packets_sent, 7); + assert_eq!(prepared.commit.token, 0x0102_0304_0506_0708); prepared.bytes.to_vec() }; assert!(machine.is_open()); @@ -1394,15 +1399,31 @@ mod tests { let prepared = machine.prepare_close().unwrap(); assert_eq!(prepared.bytes, first_bytes); - let event = machine.commit_local_close(prepared.commit); + assert_eq!(prepared.commit.packets_sent, 7); + assert_eq!(prepared.commit.token, 0x0102_0304_0506_0708); + assert!(machine.is_open()); + } + + #[test] + fn local_close_commit_uses_supplied_exact_timestamp() { + let mut machine = open_machine(4, Duration::from_secs(1)); + let remote = machine.remote; + active_mut(&mut machine).packets_sent = 7; + let prepared = machine.prepare_close().unwrap(); + let sent_at = ClientTimestamp { + wall: UNIX_EPOCH + Duration::from_secs(1_234), + mono: Instant::now() + Duration::from_secs(5), + }; + let event = machine.commit_local_close(prepared.commit, sent_at); assert!(matches!( event, ClientEvent::SessionClosed { remote: event_remote, token: 0x0102_0304_0506_0708, - .. + at, } if event_remote == remote + && at == sent_at )); assert!(matches!( machine.state,