fix(client): transact blocking session close

crates/irtt-client/src/client.rs

@@ -55,6 +55,11 @@ 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>,
     fail_cleanup_send: Cell<bool>,
     recv_buffer_len_override: Cell<Option<usize>>,
     fail_dscp_restore: Cell<bool>,
@@ -162,16 +167,44 @@ 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 remote = self.remote;
-        let events = self.runtime.close_with(|packet| {
-            clear_dscp_on_socket(socket, remote)?;
-            socket.send(packet)?;
-            Ok(())
-        })?;
+        let prepared = self.runtime.prepare_close()?;
+        let mut events = Vec::new();
+        #[cfg(test)]
+        if self.test_hooks.fail_close_event_reserve.replace(false) {
+            events
+                .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));
+            }
+        };
+        let event = self.runtime.commit_local_close(prepared.commit);
         self.schedule = None;
+        self.applied_dscp = None;
+
+        validate_datagram_length(expected_bytes, bytes)?;
+        events.push(event);
         Ok(events)
     }
 
@@ -244,8 +277,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 +643,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 +688,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

@@ -30,7 +30,13 @@ 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 events = client.close().unwrap();
-    assert_eq!(events.len(), 1);
+    assert!(matches!(
+        events.as_slice(),
+        [ClientEvent::SessionClosed { token: TOKEN, .. }]
+    ));
+    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 +44,138 @@ 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());
+    client.test_hooks.fail_close_send.set(true);
+
+    assert!(matches!(client.close(), Err(ClientError::Socket(_))));
+    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, .. }]
+    ));
+    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();
+    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!(matches!(client.close(), Err(ClientError::AlreadyClosed)));
+    server.join();
+}
+
 #[test]
 fn send_probe_fails_after_close() {
     let params = default_params();
@@ -132,6 +270,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();
+}