Merge pull request #45 from Lochnair/fix/transactional-close

fix(client): make session closing transactional

crates/irtt-client/src/client.rs

@@ -55,6 +55,12 @@ struct PreparedClientOpenFailure {
 struct ClientTestHooks {
     fail_open_dscp: Cell<bool>,
     fail_open_timeout_restore: Cell<bool>,
+    fail_close_event_reserve: Cell<bool>,
+    fail_close_dscp_clear: Cell<bool>,
+    fail_close_send: Cell<bool>,
+    close_reported_len: Cell<Option<usize>>,
+    close_send_attempts: Cell<usize>,
+    close_sent_at: Cell<Option<ClientTimestamp>>,
     fail_cleanup_send: Cell<bool>,
     recv_buffer_len_override: Cell<Option<usize>>,
     fail_dscp_restore: Cell<bool>,
@@ -162,16 +168,52 @@ impl Client {
     /// Send a close request and emit a [`ClientEvent::SessionClosed`] event.
     ///
     /// The close event means the client has sent its close packet and stopped
-    /// tracking the session locally; it is not a server acknowledgement.
+    /// tracking the session locally; it is not a server acknowledgement. If the
+    /// send fails, the negotiated DSCP is restored best-effort and the session
+    /// and probe schedule remain open.
     pub fn close(&mut self) -> Result<Vec<ClientEvent>, ClientError> {
-        let socket = &self.socket;
+        let prepared = self.runtime.prepare_close()?;
-        let remote = self.remote;
+        let mut events = Vec::new();
-        let events = self.runtime.close_with(|packet| {
+        #[cfg(test)]
-            clear_dscp_on_socket(socket, remote)?;
+        if self.test_hooks.fail_close_event_reserve.replace(false) {
-            socket.send(packet)?;
+            events
-            Ok(())
+                .try_reserve(usize::MAX)
-        })?;
+                .map_err(|source| ClientError::AllocationFailed {
+                    operation: "close event result",
+                    source,
+                })?;
+        }
+        events
+            .try_reserve(1)
+            .map_err(|source| ClientError::AllocationFailed {
+                operation: "close event result",
+                source,
+            })?;
+        let previous_dscp = self.applied_dscp;
+        let expected_bytes = prepared.bytes.len();
+
+        self.clear_close_dscp()?;
+        let bytes = match self.send_close_datagram(prepared.bytes) {
+            Ok(bytes) => bytes,
+            Err(err) => {
+                self.restore_dscp_best_effort(previous_dscp);
+                return Err(ClientError::Socket(err));
+            }
+        };
+        #[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;
+
+        validate_datagram_length(expected_bytes, bytes)?;
+        events.push(event);
         Ok(events)
     }
 
@@ -244,8 +286,11 @@ impl Client {
             datagram.received_at,
             datagram.meta,
         )?;
-        if self.runtime.is_terminal() {
+        if self.runtime.is_peer_closed() {
             self.schedule = None;
+            if self.clear_close_dscp().is_ok() {
+                self.applied_dscp = None;
+            }
         }
         Ok(events)
     }
@@ -607,6 +652,18 @@ impl Client {
         };
     }
 
+    fn clear_close_dscp(&self) -> Result<(), ClientError> {
+        #[cfg(test)]
+        if self.test_hooks.fail_close_dscp_clear.replace(false) {
+            return Err(ClientError::SocketOption {
+                operation: "clear negotiated DSCP",
+                remote: self.remote,
+                source: io::Error::other("injected negotiated DSCP clear failure"),
+            });
+        }
+        clear_dscp_on_socket(&self.socket, self.remote)
+    }
+
     fn send_cleanup_close_best_effort(&self, packet: Option<&[u8]>) {
         if let Some(packet) = packet {
             #[cfg(test)]
@@ -640,6 +697,24 @@ impl Client {
         }
         self.socket.set_read_timeout(timeout)
     }
+
+    fn send_close_datagram(&self, packet: &[u8]) -> io::Result<usize> {
+        #[cfg(test)]
+        {
+            self.test_hooks
+                .close_send_attempts
+                .set(self.test_hooks.close_send_attempts.get() + 1);
+            if self.test_hooks.fail_close_send.replace(false) {
+                return Err(io::Error::other("injected close send failure"));
+            }
+        }
+        let bytes = self.socket.send(packet)?;
+        #[cfg(test)]
+        if let Some(reported) = self.test_hooks.close_reported_len.take() {
+            return Ok(reported);
+        }
+        Ok(bytes)
+    }
 }
 
 pub(crate) fn echo_sent_event(

crates/irtt-client/src/client/tests/close.rs

@@ -29,8 +29,23 @@ 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_eq!(events.len(), 1);
+    assert!(matches!(
+        events.as_slice(),
+        [ClientEvent::SessionClosed {
+            token: TOKEN,
+            at,
+            ..
+        }] if *at == sent_at
+    ));
+    assert!(!client.runtime.is_open());
+    assert!(client.schedule.is_none());
+    assert_eq!(client.applied_dscp, None);
     let packets: Vec<_> = server.rx.iter().take(2).collect();
     let close = &packets[1];
     assert_eq!(close[3], flags::FLAG_CLOSE);
@@ -38,6 +53,155 @@ fn close_sends_one_close_packet_with_negotiated_token() {
     server.join();
 }
 
+#[test]
+fn close_event_reservation_failure_precedes_dscp_clear_and_send() {
+    let params = default_params();
+    let server = start_fake_server(move |socket, tx| {
+        let (_, peer) = recv_request(&socket, &tx);
+        socket
+            .send_to(
+                &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params, None),
+                peer,
+            )
+            .unwrap();
+        let _ = recv_request(&socket, &tx);
+    });
+    let mut client = Client::connect(default_test_config(server.addr)).unwrap();
+    assert_open_started(client.open().unwrap());
+    client.test_hooks.fail_close_event_reserve.set(true);
+    client.test_hooks.fail_close_dscp_clear.set(true);
+
+    assert!(matches!(
+        client.close(),
+        Err(ClientError::AllocationFailed {
+            operation: "close event result",
+            ..
+        })
+    ));
+    assert_eq!(client.test_hooks.close_send_attempts.get(), 0);
+    assert!(client.test_hooks.fail_close_dscp_clear.get());
+    assert!(client.runtime.is_open());
+    assert!(client.schedule.is_some());
+    assert_eq!(client.applied_dscp, Some(0));
+
+    client.test_hooks.fail_close_dscp_clear.set(false);
+    client.close().unwrap();
+    assert_eq!(client.test_hooks.close_send_attempts.get(), 1);
+    server.join();
+}
+
+#[test]
+fn close_dscp_clear_failure_precedes_send_and_preserves_open_state() {
+    let params = default_params();
+    let server = start_fake_server(move |socket, tx| {
+        let (_, peer) = recv_request(&socket, &tx);
+        socket
+            .send_to(
+                &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params, None),
+                peer,
+            )
+            .unwrap();
+        let _ = recv_request(&socket, &tx);
+    });
+    let mut client = Client::connect(default_test_config(server.addr)).unwrap();
+    assert_open_started(client.open().unwrap());
+    client.test_hooks.fail_close_dscp_clear.set(true);
+
+    assert!(matches!(
+        client.close(),
+        Err(ClientError::SocketOption { .. })
+    ));
+    assert_eq!(client.test_hooks.close_send_attempts.get(), 0);
+    assert!(client.runtime.is_open());
+    assert!(client.schedule.is_some());
+    assert_eq!(client.applied_dscp, Some(0));
+
+    client.close().unwrap();
+    assert_eq!(client.test_hooks.close_send_attempts.get(), 1);
+    server.join();
+}
+
+#[test]
+fn close_send_failure_leaves_machine_and_schedule_open_for_retry() {
+    let params = default_params();
+    let server = start_fake_server(move |socket, tx| {
+        let (_, peer) = recv_request(&socket, &tx);
+        socket
+            .send_to(
+                &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params, None),
+                peer,
+            )
+            .unwrap();
+        let _ = recv_request(&socket, &tx);
+    });
+    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));
+    assert_eq!(client.test_hooks.close_send_attempts.get(), 1);
+
+    assert!(matches!(
+        client.close().unwrap().as_slice(),
+        [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();
+}
+
+#[test]
+fn short_successful_close_commits_before_length_mismatch() {
+    let params = default_params();
+    let server = start_fake_server(move |socket, tx| {
+        let (_, peer) = recv_request(&socket, &tx);
+        socket
+            .send_to(
+                &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params, None),
+                peer,
+            )
+            .unwrap();
+        let _ = recv_request(&socket, &tx);
+    });
+    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!(
+        client.close(),
+        Err(ClientError::DatagramLengthMismatch {
+            expected: error_expected,
+            actual,
+        }) if error_expected == expected && actual + 1 == expected
+    ));
+    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();
+}
+
 #[test]
 fn send_probe_fails_after_close() {
     let params = default_params();
@@ -132,6 +296,117 @@ fn close_flagged_echo_reply_emits_reply_then_closes_without_sending_close() {
     server.join();
 }
 
+#[test]
+fn peer_close_dscp_cleanup_failure_preserves_events_and_closed_state() {
+    let params = default_params();
+    let server = start_fake_server({
+        let params = params.clone();
+        move |socket, tx| {
+            let (_, peer) = recv_request(&socket, &tx);
+            socket
+                .send_to(
+                    &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params, None),
+                    peer,
+                )
+                .unwrap();
+            let (request, _) = recv_request(&socket, &tx);
+            let seq = u32::from_le_bytes(request[12..16].try_into().unwrap());
+            socket
+                .send_to(
+                    &echo_reply_packet_with_flags(
+                        TOKEN,
+                        seq,
+                        &params,
+                        &TimestampFields::default(),
+                        None,
+                        FLAG_REPLY | flags::FLAG_CLOSE,
+                    ),
+                    peer,
+                )
+                .unwrap();
+            socket
+                .set_read_timeout(Some(Duration::from_millis(250)))
+                .unwrap();
+            while recv_request_timeout(&socket, &tx).is_some() {}
+        }
+    });
+    let mut client = Client::connect(ClientConfig {
+        socket_config: crate::SocketConfig {
+            recv_timeout: Some(Duration::from_millis(200)),
+            ..Default::default()
+        },
+        ..default_test_config(server.addr)
+    })
+    .unwrap();
+    assert_open_started(client.open().unwrap());
+    client.send_probe().unwrap();
+    client.test_hooks.fail_close_dscp_clear.set(true);
+
+    assert!(matches!(
+        client.recv_once().unwrap().as_slice(),
+        [
+            ClientEvent::EchoReply { .. },
+            ClientEvent::SessionClosed { token: TOKEN, .. }
+        ]
+    ));
+    assert!(client.is_peer_closed());
+    assert!(client.schedule.is_none());
+    assert_eq!(client.applied_dscp, Some(0));
+    assert_eq!(client.test_hooks.close_send_attempts.get(), 0);
+    let open = server.rx.recv_timeout(Duration::from_millis(100)).unwrap();
+    let echo = server.rx.recv_timeout(Duration::from_millis(100)).unwrap();
+    assert_ne!(open[3] & FLAG_OPEN, 0);
+    assert_eq!(echo[3] & flags::FLAG_CLOSE, 0);
+    assert!(server.rx.recv_timeout(Duration::from_millis(400)).is_err());
+    server.join();
+}
+
+#[test]
+fn recv_available_stops_after_peer_close() {
+    let params = default_params();
+    let server = start_fake_server({
+        let params = params.clone();
+        move |socket, tx| {
+            let (_, peer) = recv_request(&socket, &tx);
+            socket
+                .send_to(
+                    &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params, None),
+                    peer,
+                )
+                .unwrap();
+            let (request, _) = recv_request(&socket, &tx);
+            let seq = u32::from_le_bytes(request[12..16].try_into().unwrap());
+            let close = echo_reply_packet_with_flags(
+                TOKEN,
+                seq,
+                &params,
+                &TimestampFields::default(),
+                None,
+                FLAG_REPLY | flags::FLAG_CLOSE,
+            );
+            socket.send_to(&close, peer).unwrap();
+            socket.send_to(&close, peer).unwrap();
+        }
+    });
+    let mut client = Client::connect(default_test_config(server.addr)).unwrap();
+    assert_open_started(client.open().unwrap());
+    client.send_probe().unwrap();
+
+    let events = client
+        .recv_available(RecvBudget { max_packets: 8 })
+        .unwrap();
+
+    assert!(matches!(
+        events.as_slice(),
+        [
+            ClientEvent::EchoReply { .. },
+            ClientEvent::SessionClosed { token: TOKEN, .. }
+        ]
+    ));
+    assert!(client.is_peer_closed());
+    server.join();
+}
+
 #[test]
 fn normal_echo_reply_does_not_close_session() {
     let params = default_params();

crates/irtt-client/src/client/tests/dscp.rs

@@ -73,3 +73,132 @@ fn normal_open_uses_negotiated_dscp_not_requested_dscp() {
     );
     server.join();
 }
+
+#[test]
+#[cfg(not(any(
+    target_os = "fuchsia",
+    target_os = "redox",
+    target_os = "solaris",
+    target_os = "illumos",
+    target_os = "haiku",
+)))]
+fn failed_close_send_restores_negotiated_dscp_and_keeps_session_open() {
+    let mut params = default_params();
+    params.dscp = 46;
+    let server = start_fake_server(move |socket, tx| {
+        let (_, peer) = recv_request(&socket, &tx);
+        socket
+            .send_to(
+                &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params, None),
+                peer,
+            )
+            .unwrap();
+        let _ = recv_request(&socket, &tx);
+    });
+    let mut config = default_test_config(server.addr);
+    config.dscp = 46;
+    let mut client = Client::connect(config).unwrap();
+    assert_open_started(client.open().unwrap());
+    client.test_hooks.fail_close_send.set(true);
+
+    let error = client.close().unwrap_err();
+
+    assert!(matches!(error, ClientError::Socket(_)));
+    assert!(client.runtime.is_open());
+    assert!(client.schedule.is_some());
+    assert_eq!(client.applied_dscp, Some(46));
+    assert_eq!(
+        socket_traffic_class(&client.socket, client.remote).unwrap() & 0xfc,
+        184
+    );
+    client.close().unwrap();
+    server.join();
+}
+
+#[test]
+fn failed_close_keeps_send_error_primary_when_dscp_restoration_also_fails() {
+    let params = default_params();
+    let server = start_fake_server(move |socket, tx| {
+        let (_, peer) = recv_request(&socket, &tx);
+        socket
+            .send_to(
+                &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params, None),
+                peer,
+            )
+            .unwrap();
+        let _ = recv_request(&socket, &tx);
+    });
+    let mut client = Client::connect(default_test_config(server.addr)).unwrap();
+    assert_open_started(client.open().unwrap());
+    client.test_hooks.fail_close_send.set(true);
+    client.test_hooks.fail_dscp_restore.set(true);
+
+    let error = client.close().unwrap_err();
+
+    assert!(matches!(error, ClientError::Socket(_)));
+    assert!(error.to_string().contains("injected close send failure"));
+    assert!(client.runtime.is_open());
+    assert!(client.schedule.is_some());
+    assert_eq!(client.applied_dscp, Some(0));
+    client.close().unwrap();
+    server.join();
+}
+
+#[test]
+#[cfg(not(any(
+    target_os = "fuchsia",
+    target_os = "redox",
+    target_os = "solaris",
+    target_os = "illumos",
+    target_os = "haiku",
+)))]
+fn authenticated_peer_close_clears_schedule_and_negotiated_dscp() {
+    let mut params = default_params();
+    params.dscp = 46;
+    let server = start_fake_server(move |socket, tx| {
+        let (_, peer) = recv_request(&socket, &tx);
+        socket
+            .send_to(
+                &open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params, None),
+                peer,
+            )
+            .unwrap();
+        let (request, _) = recv_request(&socket, &tx);
+        let seq = u32::from_le_bytes(request[12..16].try_into().unwrap());
+        socket
+            .send_to(
+                &echo_reply_packet_with_flags(
+                    TOKEN,
+                    seq,
+                    &params,
+                    &TimestampFields::default(),
+                    None,
+                    FLAG_REPLY | flags::FLAG_CLOSE,
+                ),
+                peer,
+            )
+            .unwrap();
+    });
+    let mut config = default_test_config(server.addr);
+    config.dscp = 46;
+    let mut client = Client::connect(config).unwrap();
+    assert_open_started(client.open().unwrap());
+    client.send_probe().unwrap();
+
+    let events = client.recv_once().unwrap();
+
+    assert!(matches!(
+        events.as_slice(),
+        [
+            ClientEvent::EchoReply { .. },
+            ClientEvent::SessionClosed { .. }
+        ]
+    ));
+    assert!(client.schedule.is_none());
+    assert_eq!(client.applied_dscp, None);
+    assert_eq!(
+        socket_traffic_class(&client.socket, client.remote).unwrap(),
+        0
+    );
+    server.join();
+}

crates/irtt-client/src/managed/group.rs

@@ -725,6 +725,16 @@ struct TargetState {
     probe_send_error: bool,
     #[cfg(test)]
     probe_send_timestamps: Option<ProbeSendTimestamps>,
+    #[cfg(test)]
+    close_reported_len: Option<usize>,
+    #[cfg(test)]
+    close_send_error: bool,
+    #[cfg(test)]
+    close_send_attempts: usize,
+    #[cfg(test)]
+    close_event_reserve_error: bool,
+    #[cfg(test)]
+    close_sent_at: Option<ClientTimestamp>,
 }
 
 type ScheduledTarget = (Arc<Mutex<TargetState>>, Instant);
@@ -839,6 +849,16 @@ impl TargetState {
             probe_send_error: false,
             #[cfg(test)]
             probe_send_timestamps: None,
+            #[cfg(test)]
+            close_reported_len: None,
+            #[cfg(test)]
+            close_send_error: false,
+            #[cfg(test)]
+            close_send_attempts: 0,
+            #[cfg(test)]
+            close_event_reserve_error: false,
+            #[cfg(test)]
+            close_sent_at: None,
         })
     }
 
@@ -1795,21 +1815,86 @@ fn close_locked_target(
 ) {
     if target.runtime.is_open() && !target.runtime.is_peer_closed() {
         let remote = target.remote;
-        match target.runtime.close_with(|packet| {
+        let prepared = match target.runtime.prepare_close() {
-            socket.send_to(packet, remote)?;
+            Ok(prepared) => prepared,
-            Ok(())
-        }) {
-            Ok(events) => {
-                publish_events(hub, target, events);
-                target.schedule = None;
-            }
             Err(err) => {
                 target.mark_finished(ManagedTargetEndReason::Failed(
                     ManagedTargetFailure::runtime(&err),
                 ));
                 return;
             }
+        };
+        let mut events = Vec::new();
+        #[cfg(test)]
+        if std::mem::take(&mut target.close_event_reserve_error) {
+            if let Err(source) = events.try_reserve(usize::MAX) {
+                let err = ClientError::AllocationFailed {
+                    operation: "managed close event result",
+                    source,
+                };
+                target.mark_finished(ManagedTargetEndReason::Failed(
+                    ManagedTargetFailure::runtime(&err),
+                ));
+                return;
+            }
         }
+        if let Err(source) = events.try_reserve(1) {
+            let err = ClientError::AllocationFailed {
+                operation: "managed close event result",
+                source,
+            };
+            target.mark_finished(ManagedTargetEndReason::Failed(
+                ManagedTargetFailure::runtime(&err),
+            ));
+            return;
+        }
+        let expected_bytes = prepared.bytes.len();
+        #[cfg(test)]
+        let reported_bytes = target.close_reported_len.take();
+        #[cfg(test)]
+        let fail_send = std::mem::take(&mut target.close_send_error);
+        #[cfg(test)]
+        {
+            target.close_send_attempts += 1;
+            if fail_send {
+                let err = ClientError::Socket(std::io::Error::other(
+                    "injected managed close send failure",
+                ));
+                target.mark_finished(ManagedTargetEndReason::Failed(
+                    ManagedTargetFailure::runtime(&err),
+                ));
+                return;
+            }
+        }
+        let bytes = match socket.send_to(prepared.bytes, remote) {
+            Ok(bytes) => bytes,
+            Err(err) => {
+                let err = ClientError::Socket(err);
+                target.mark_finished(ManagedTargetEndReason::Failed(
+                    ManagedTargetFailure::runtime(&err),
+                ));
+                return;
+            }
+        };
+        #[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);
+        if let Err(err) = validate_datagram_length(expected_bytes, bytes) {
+            target.mark_finished(ManagedTargetEndReason::Failed(
+                ManagedTargetFailure::runtime(&err),
+            ));
+            return;
+        }
+        events.push(event);
+        publish_events(hub, target, events);
     }
     target.mark_finished(reason);
 }
@@ -2085,7 +2170,8 @@ mod tests {
         compute_hmac_in_place, echo_packet_len,
         flags::{self, FLAG_OPEN, FLAG_REPLY},
         layout::PacketLayout,
-        Clock, Params, ReceivedStats, StampAt, TimestampFields, HMAC_SIZE, MAGIC, PROTOCOL_VERSION,
+        Clock, OpenReply, Params, ReceivedStats, StampAt, TimestampFields, HMAC_SIZE, MAGIC,
+        PROTOCOL_VERSION,
     };
     use std::sync::mpsc;
 
@@ -2172,10 +2258,13 @@ mod tests {
     }
 
     fn active_target_fixture() -> ActiveTargetFixture {
+        active_target_fixture_at("127.0.0.1:2112".parse().unwrap())
+    }
+
+    fn active_target_fixture_at(remote: SocketAddr) -> ActiveTargetFixture {
         let interval = Duration::from_millis(10);
         let params = test_params(None, interval);
         let client_config = group_config(None, interval, ManagedGroupPacing::Staggered).client;
-        let remote: SocketAddr = "127.0.0.1:2112".parse().unwrap();
         let hub = EventHub::new();
         let subscription = hub
             .subscribe(SubscriberConfig {
@@ -2894,6 +2983,7 @@ mod tests {
                 ReceiveMeta::default(),
             ));
             assert!(matches!(target.status, TargetStatus::Finished));
+            assert!(target.schedule.is_none());
             assert_eq!(
                 target.final_reason,
                 Some(ManagedTargetEndReason::PeerClosed)
@@ -2959,6 +3049,290 @@ mod tests {
         assert_eq!(outcome.targets[0].replies_received, 1);
     }
 
+    #[test]
+    fn managed_close_success_commits_clears_schedule_and_publishes_once() {
+        let peer = UdpSocket::bind("127.0.0.1:0").unwrap();
+        peer.set_read_timeout(Some(Duration::from_millis(200)))
+            .unwrap();
+        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,
+                &mut target,
+                ManagedTargetEndReason::TestComplete,
+            );
+            assert!(!target.runtime.is_open());
+            assert!(target.schedule.is_none());
+            assert_eq!(target.close_send_attempts, 1);
+            assert_eq!(
+                target.final_reason,
+                Some(ManagedTargetEndReason::TestComplete)
+            );
+        }
+
+        let mut packet = [0_u8; 512];
+        let (len, _) = peer.recv_from(&mut packet).unwrap();
+        assert_eq!(packet[3], flags::FLAG_CLOSE);
+        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);
+        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]
+    fn managed_removal_and_cancellation_each_send_one_close() {
+        for reason in [
+            ManagedTargetEndReason::Removed,
+            ManagedTargetEndReason::Cancelled,
+        ] {
+            let peer = UdpSocket::bind("127.0.0.1:0").unwrap();
+            peer.set_read_timeout(Some(Duration::from_millis(200)))
+                .unwrap();
+            let fixture = active_target_fixture_at(peer.local_addr().unwrap());
+            let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
+
+            {
+                let mut target = fixture.target.lock().expect("target mutex poisoned");
+                close_locked_target(&socket, &fixture.hub, &mut target, reason.clone());
+                assert_eq!(target.final_reason, Some(reason));
+                assert_eq!(target.close_send_attempts, 1);
+                assert!(!target.runtime.is_open());
+                assert!(target.schedule.is_none());
+            }
+
+            let mut packet = [0_u8; 512];
+            let (len, _) = peer.recv_from(&mut packet).unwrap();
+            assert_eq!(packet[3], flags::FLAG_CLOSE);
+            assert_eq!(len, 12);
+            assert!(matches!(
+                peer.recv_from(&mut packet),
+                Err(err)
+                    if matches!(
+                        err.kind(),
+                        std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut
+                    )
+            ));
+        }
+    }
+
+    #[test]
+    fn managed_close_send_failure_preserves_machine_and_requested_reason_is_not_recorded() {
+        let peer = UdpSocket::bind("127.0.0.1:0").unwrap();
+        peer.set_nonblocking(true).unwrap();
+        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,
+                &fixture.hub,
+                &mut target,
+                ManagedTargetEndReason::Removed,
+            );
+            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(_))
+            ));
+        }
+
+        let mut packet = [0_u8; 512];
+        assert!(matches!(
+            peer.recv_from(&mut packet),
+            Err(err) if err.kind() == std::io::ErrorKind::WouldBlock
+        ));
+        assert!(!drain_available_group_events(&fixture.subscription)
+            .iter()
+            .any(|event| matches!(
+                event,
+                ManagedGroupEvent::Client(TargetEvent {
+                    event: ClientEvent::SessionClosed { .. },
+                    ..
+                })
+            )));
+    }
+
+    #[test]
+    fn managed_short_close_commits_and_clears_schedule_before_failure() {
+        let peer = UdpSocket::bind("127.0.0.1:0").unwrap();
+        peer.set_read_timeout(Some(Duration::from_millis(200)))
+            .unwrap();
+        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,
+                &fixture.hub,
+                &mut target,
+                ManagedTargetEndReason::Cancelled,
+            );
+            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()
+                .and_then(ManagedTargetEndReason::failure)
+                .expect("short close must use managed failure policy");
+            assert_eq!(failure.kind, ManagedTargetFailureKind::RuntimeProtocol);
+            assert!(failure.message.contains("UDP accepted"));
+        }
+
+        let mut packet = [0_u8; 512];
+        peer.recv_from(&mut packet).unwrap();
+        assert!(!drain_available_group_events(&fixture.subscription)
+            .iter()
+            .any(|event| matches!(
+                event,
+                ManagedGroupEvent::Client(TargetEvent {
+                    event: ClientEvent::SessionClosed { .. },
+                    ..
+                })
+            )));
+    }
+
+    #[test]
+    fn managed_close_reservation_failure_precedes_send() {
+        let peer = UdpSocket::bind("127.0.0.1:0").unwrap();
+        peer.set_nonblocking(true).unwrap();
+        let fixture = active_target_fixture_at(peer.local_addr().unwrap());
+        let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
+
+        {
+            let mut target = fixture.target.lock().expect("target mutex poisoned");
+            target.close_event_reserve_error = true;
+            close_locked_target(
+                &socket,
+                &fixture.hub,
+                &mut target,
+                ManagedTargetEndReason::TestComplete,
+            );
+            assert!(target.runtime.is_open());
+            assert!(target.schedule.is_some());
+            assert_eq!(target.close_send_attempts, 0);
+            assert!(matches!(
+                target.final_reason,
+                Some(ManagedTargetEndReason::Failed(_))
+            ));
+        }
+
+        let mut packet = [0_u8; 512];
+        assert!(matches!(
+            peer.recv_from(&mut packet),
+            Err(err) if err.kind() == std::io::ErrorKind::WouldBlock
+        ));
+    }
+
+    #[test]
+    fn managed_peer_close_affects_only_the_matching_target() {
+        let peer = active_target_fixture_at("127.0.0.1:2112".parse().unwrap());
+        let other = active_target_fixture_at("127.0.0.1:2113".parse().unwrap());
+
+        {
+            let mut target = peer.target.lock().expect("target mutex poisoned");
+            assert!(process_active_target_packet(
+                &peer.hub,
+                &mut target,
+                &peer.close_reply,
+                ClientTimestamp::now(),
+                ReceiveMeta::default(),
+            ));
+            assert!(target.runtime.is_peer_closed());
+            assert!(target.schedule.is_none());
+            assert_eq!(
+                target.final_reason,
+                Some(ManagedTargetEndReason::PeerClosed)
+            );
+        }
+
+        let target = other.target.lock().expect("target mutex poisoned");
+        assert!(target.runtime.is_open());
+        assert!(matches!(target.status, TargetStatus::Active));
+        assert!(target.schedule.is_some());
+        assert_eq!(target.final_reason, None);
+    }
+
+    #[test]
+    fn managed_non_active_states_do_not_attempt_local_close() {
+        let remote: SocketAddr = "127.0.0.1:2112".parse().unwrap();
+        let mut normal_config = group_config(
+            None,
+            Duration::from_millis(10),
+            ManagedGroupPacing::Staggered,
+        )
+        .client;
+        let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
+        let hub = EventHub::new();
+        let mut opening =
+            TargetState::new(&normal_config, target("opening", remote), 0, Instant::now()).unwrap();
+        close_locked_target(&socket, &hub, &mut opening, ManagedTargetEndReason::Removed);
+        assert_eq!(opening.close_send_attempts, 0);
+
+        normal_config.run_mode = crate::RunMode::NoTest;
+        let no_test_params = params_from_config(&normal_config).unwrap();
+        let mut no_test =
+            TargetState::new(&normal_config, target("no-test", remote), 0, Instant::now()).unwrap();
+        let reply = OpenReply {
+            flags: FLAG_OPEN | FLAG_REPLY | flags::FLAG_CLOSE,
+            token: 0,
+            params: no_test_params,
+        };
+        let prepared = no_test
+            .runtime
+            .prepare_open_acceptance(reply, ClientTimestamp::now())
+            .unwrap();
+        no_test.runtime.commit_open(prepared);
+        close_locked_target(
+            &socket,
+            &hub,
+            &mut no_test,
+            ManagedTargetEndReason::Cancelled,
+        );
+        assert_eq!(no_test.close_send_attempts, 0);
+    }
+
     #[test]
     fn authoritative_status_counts_only_peer_closed_outcomes() {
         let remote: SocketAddr = "127.0.0.1:2112".parse().unwrap();

crates/irtt-client/src/session/machine.rs

@@ -53,6 +53,7 @@ enum CloseSource {
 struct ActiveSession {
     token: u64,
     negotiated: NegotiatedParams,
+    local_close_packet: Box<[u8]>,
     next_wire_seq: u32,
     highest_received_seq: Option<u32>,
     packets_sent: u64,
@@ -76,7 +77,6 @@ pub(crate) enum OpenDatagramDisposition {
 pub(crate) struct PreparedOpenAcceptance {
     next_state: MachineState,
     outcome: OpenOutcome,
-    cleanup_close: Option<Box<[u8]>>,
 }
 
 impl PreparedOpenAcceptance {
@@ -91,7 +91,13 @@ impl PreparedOpenAcceptance {
     }
 
     pub(crate) fn cleanup_close_packet(&self) -> Option<&[u8]> {
-        self.cleanup_close.as_deref()
+        match &self.next_state {
+            MachineState::Open(session) => Some(&session.local_close_packet),
+            MachineState::NoTestCompleted => None,
+            MachineState::Connected | MachineState::Closed { .. } => {
+                unreachable!("open acceptance only prepares open or no-test state")
+            }
+        }
     }
 
     #[cfg(test)]
@@ -100,6 +106,18 @@ impl PreparedOpenAcceptance {
     }
 }
 
+#[derive(Debug)]
+pub(crate) struct PreparedClose<'a> {
+    pub(crate) bytes: &'a [u8],
+    pub(crate) commit: CloseCommit,
+}
+
+#[derive(Debug)]
+pub(crate) struct CloseCommit {
+    packets_sent: u64,
+    token: u64,
+}
+
 #[derive(Debug)]
 pub(crate) struct OpenAcceptanceFailure {
     pub(crate) primary: ClientError,
@@ -401,28 +419,35 @@ impl SessionMachine {
         Ok(events)
     }
 
-    pub(crate) fn close_with<F>(&mut self, send: F) -> Result<Vec<ClientEvent>, ClientError>
+    pub(crate) fn prepare_close(&self) -> Result<PreparedClose<'_>, ClientError> {
-    where
+        let session = self.open_session()?;
-        F: FnOnce(&[u8]) -> Result<(), ClientError>,
+        Ok(PreparedClose {
-    {
+            bytes: &session.local_close_packet,
-        let token = match &self.state {
+            commit: CloseCommit {
-            MachineState::Open(session) => session.token,
+                packets_sent: session.packets_sent,
-            MachineState::Closed { .. } => return Err(ClientError::AlreadyClosed),
+                token: session.token,
-            MachineState::Connected | MachineState::NoTestCompleted => {
+            },
-                return Err(ClientError::NotOpen)
+        })
-            }
+    }
+
+    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"
+        );
+        self.state = MachineState::Closed {
+            source: CloseSource::Local,
+            packets_sent: commit.packets_sent,
         };
-
+        ClientEvent::SessionClosed {
-        let packet =
-            encode_close_request(&CloseRequest { token }, self.config.hmac_key.as_deref())?;
-        send(&packet)?;
-        self.transition_to_closed(CloseSource::Local);
-
-        Ok(vec![ClientEvent::SessionClosed {
             remote: self.remote,
-            token,
+            token: commit.token,
-            at: ClientTimestamp::now(),
+            at: sent_at,
-        }])
+        }
     }
 
     pub(crate) fn pending_is_empty(&self) -> bool {
@@ -488,9 +513,12 @@ impl SessionMachine {
             Ok(negotiated) => negotiated,
             Err(primary) => return Err(OpenAcceptanceFailure::new(primary, cleanup_close)),
         };
+        let local_close_packet =
+            cleanup_close.expect("normal non-zero-token replies prepare a cleanup close");
         let next_state = MachineState::Open(Box::new(ActiveSession {
             token,
             negotiated: negotiated.clone(),
+            local_close_packet,
             next_wire_seq: 0,
             highest_received_seq: None,
             packets_sent: 0,
@@ -515,7 +543,6 @@ impl SessionMachine {
         Ok(PreparedOpenAcceptance {
             next_state,
             outcome,
-            cleanup_close,
         })
     }
 
@@ -543,7 +570,6 @@ impl SessionMachine {
         Ok(PreparedOpenAcceptance {
             next_state: MachineState::NoTestCompleted,
             outcome,
-            cleanup_close: None,
         })
     }
 
@@ -981,6 +1007,14 @@ mod tests {
         machine.state = MachineState::Open(Box::new(ActiveSession {
             token: 0x0102_0304_0506_0708,
             negotiated,
+            local_close_packet: encode_close_request(
+                &CloseRequest {
+                    token: 0x0102_0304_0506_0708,
+                },
+                None,
+            )
+            .unwrap()
+            .into_boxed_slice(),
             next_wire_seq: 0,
             highest_received_seq: None,
             packets_sent: 0,
@@ -1268,6 +1302,7 @@ mod tests {
             )
             .unwrap();
 
+        let expected_close = prepared.cleanup_close_packet().unwrap().to_vec();
         let outcome = machine.commit_open(prepared);
 
         assert!(matches!(
@@ -1278,6 +1313,7 @@ mod tests {
             } if outcome_token == token
         ));
         assert!(machine.is_open());
+        assert_eq!(active(&machine).local_close_packet.as_ref(), expected_close);
         assert!(matches!(
             machine.prepare_open_request(),
             Err(ClientError::AlreadyOpen)
@@ -1337,6 +1373,146 @@ mod tests {
         assert!(matches!(machine.state, MachineState::Connected));
     }
 
+    #[test]
+    fn prepare_close_is_stable_inert_and_does_not_determine_event_timestamp() {
+        let mut machine = open_machine(4, Duration::from_secs(1));
+        active_mut(&mut machine).packets_sent = 7;
+
+        let first_bytes = {
+            let prepared = machine.prepare_close().unwrap();
+            assert_eq!(
+                prepared.bytes,
+                encode_close_request(
+                    &CloseRequest {
+                        token: 0x0102_0304_0506_0708,
+                    },
+                    None,
+                )
+                .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());
+        assert_eq!(active(&machine).packets_sent, 7);
+
+        let prepared = machine.prepare_close().unwrap();
+        assert_eq!(prepared.bytes, first_bytes);
+        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,
+            MachineState::Closed {
+                source: CloseSource::Local,
+                packets_sent: 7,
+            }
+        ));
+        assert!(matches!(
+            machine.prepare_close(),
+            Err(ClientError::AlreadyClosed)
+        ));
+    }
+
+    #[test]
+    fn prepared_close_reuses_hmac_packet_built_during_open_acceptance() {
+        let key = b"close-key".to_vec();
+        let mut machine = connected_machine(ClientConfig {
+            hmac_key: Some(key.clone()),
+            ..ClientConfig::default()
+        });
+        let token = 0x1020_3040_5060_7080;
+        let prepared_open = machine
+            .prepare_open_acceptance(
+                normal_open_reply(&machine, token),
+                timestamp(Instant::now()),
+            )
+            .unwrap();
+        let opening_close = prepared_open.cleanup_close_packet().unwrap().to_vec();
+        machine.commit_open(prepared_open);
+
+        let prepared_close = machine.prepare_close().unwrap();
+        assert_eq!(prepared_close.bytes, opening_close);
+        assert_eq!(
+            prepared_close.bytes[3],
+            flags::FLAG_CLOSE | flags::FLAG_HMAC
+        );
+        irtt_proto::verify_hmac(&key, prepared_close.bytes, 4).unwrap();
+        assert_eq!(
+            u64::from_le_bytes(
+                prepared_close.bytes[4 + irtt_proto::HMAC_SIZE..12 + irtt_proto::HMAC_SIZE]
+                    .try_into()
+                    .unwrap()
+            ),
+            token
+        );
+    }
+
+    #[test]
+    fn close_preparation_preserves_non_open_errors_and_peer_source() {
+        let connected = connected_machine(ClientConfig::default());
+        assert!(matches!(
+            connected.prepare_close(),
+            Err(ClientError::NotOpen)
+        ));
+
+        let mut no_test = connected_machine(ClientConfig {
+            run_mode: RunMode::NoTest,
+            ..ClientConfig::default()
+        });
+        let reply = OpenReply {
+            flags: flags::FLAG_OPEN | flags::FLAG_REPLY | flags::FLAG_CLOSE,
+            token: 0,
+            params: no_test.requested.clone(),
+        };
+        let prepared = no_test
+            .prepare_open_acceptance(reply, timestamp(Instant::now()))
+            .unwrap();
+        no_test.commit_open(prepared);
+        assert!(matches!(
+            no_test.prepare_close(),
+            Err(ClientError::AlreadyCompleted)
+        ));
+
+        let mut peer_closed = open_machine(4, Duration::from_secs(1));
+        peer_closed.transition_to_closed(CloseSource::Peer);
+        assert!(matches!(
+            peer_closed.state,
+            MachineState::Closed {
+                source: CloseSource::Peer,
+                ..
+            }
+        ));
+        assert!(matches!(
+            peer_closed.prepare_close(),
+            Err(ClientError::AlreadyClosed)
+        ));
+    }
+
     #[test]
     fn dropping_prepared_probe_changes_nothing() {
         let machine = open_machine(4, Duration::from_secs(1));