feat(tcp-connector): enhance socket recycling and cancellation handling in accept logic

Author: test

aimdb-tcp-connector/src/embassy_transport.rs

@@ -144,6 +144,12 @@ where
 
 struct TcpSocketSlot {
     socket: RefCell<Option<TcpSocket<'static>>>,
+    // Single `Waker`: at most one accept ever waits on a given slot. Every shipped
+    // accept path holds one waiter per slot by construction — the `&mut self`
+    // `Listener` impl serializes accepts on the sole slot, and `TcpServer<N>` runs
+    // exactly one worker per slot. The pooled per-slot accept is test-only (see
+    // `accept_on`), and its single caller drives one accept per index. So this
+    // waker is never clobbered.
     waker: RefCell<Option<Waker>>,
 }
 
@@ -190,6 +196,47 @@ impl TcpSocketSlot {
     }
 }
 
+/// Owns a socket taken out of its slot for the duration of an `accept()`, and
+/// returns it to the slot if dropped before the accept succeeds. Without this,
+/// an accept future dropped mid-`accept` (a `select!` timeout or shutdown branch
+/// winning the race) would drop the socket instead of recycling it, leaving the
+/// slot permanently empty so every later accept on that index waits forever.
+/// [`SlotReturn::into_socket`] defuses it once the socket moves into a
+/// [`TcpConnection`] on the success path.
+struct SlotReturn<'a> {
+    slot: &'a Arc<TcpSocketSlot>,
+    socket: Option<TcpSocket<'static>>,
+}
+
+impl<'a> SlotReturn<'a> {
+    fn new(slot: &'a Arc<TcpSocketSlot>, socket: TcpSocket<'static>) -> Self {
+        Self {
+            slot,
+            socket: Some(socket),
+        }
+    }
+
+    fn socket_mut(&mut self) -> &mut TcpSocket<'static> {
+        self.socket
+            .as_mut()
+            .expect("socket present until into_socket")
+    }
+
+    /// Take the socket back, defusing the guard so its `Drop` becomes a no-op.
+    fn into_socket(mut self) -> TcpSocket<'static> {
+        self.socket.take().expect("socket taken exactly once")
+    }
+}
+
+impl Drop for SlotReturn<'_> {
+    fn drop(&mut self) {
+        if let Some(mut socket) = self.socket.take() {
+            socket.abort();
+            self.slot.put(socket);
+        }
+    }
+}
+
 /// A reusable Embassy TCP dialer backed by one caller-owned socket.
 ///
 /// Unlike moved-in UART peripherals, `embassy-net` TCP sockets can be reused
@@ -292,14 +339,17 @@ impl<const N: usize> TcpListener<N> {
         }
     }
 
-    /// Accept one connection on pooled socket `index`, recycling that socket
-    /// back into its slot when the returned connection drops.
+    /// Accept one connection on pooled socket `index`, recycling that socket back
+    /// into its slot when the returned connection drops. **Test-only** (gated on
+    /// `_test-embassy-loopback`): the shipped pooled path is `TcpServer<N>`, which
+    /// runs one worker per slot; this lets the transport-level loopback test drive
+    /// the same-port fan-out directly, without the session engine.
     ///
-    /// The [`Listener`] impl (`N = 1`) and the `TcpServer` workers each accept a
-    /// single slot at a time; this exposes the same per-slot accept so a caller
-    /// can keep several pooled sockets in `accept()` on one endpoint at once
-    /// (the concurrent-listener property `with_buffers` exists for). Panics if
-    /// `index >= N`.
+    /// The slot stores a single `Waker`, so this takes `&self` on the contract
+    /// that the caller drives **at most one accept per `index`** (the test uses
+    /// distinct indices). Panics if `index >= N`.
+    #[cfg(feature = "_test-embassy-loopback")]
+    #[doc(hidden)]
     pub fn accept_on(
         &self,
         index: usize,
@@ -333,26 +383,42 @@ impl<const N: usize> TcpListener<N> {
 
 impl Listener for TcpListener<1> {
     fn accept(&mut self) -> BoxFut<'_, TransportResult<Box<dyn Connection>>> {
-        Box::pin(self.accept_on(0))
+        // `&mut self` already serializes accepts on the sole slot, so the
+        // one-waiter-per-slot invariant holds here.
+        let slot = self.slots[0].clone();
+        let local_endpoint = self.local_endpoint;
+        Box::pin(SendFutureWrapper(async move {
+            accept_on_slot(&slot, local_endpoint).await
+        }))
     }
 }
 
 /// Take the socket from `slot`, accept one inbound connection on it, and hand
-/// back a recyclable [`TcpConnection`]. Shared by [`TcpListener::accept_on`] and
-/// the per-slot server workers so all pooled accept paths behave identically.
+/// back a recyclable [`TcpConnection`]. Shared by the [`Listener`] impl, the
+/// per-slot server workers, and the test-only `accept_on` so all pooled accept
+/// paths behave identically.
 async fn accept_on_slot(
     slot: &Arc<TcpSocketSlot>,
     local_endpoint: IpListenEndpoint,
 ) -> TransportResult<Box<dyn Connection>> {
-    let mut socket = poll_fn(|cx| slot.poll_take(cx)).await;
-    socket.abort();
-    match socket.accept(local_endpoint).await {
+    let socket = poll_fn(|cx| slot.poll_take(cx)).await;
+    // The socket lives in this guard until it either moves into a `TcpConnection`
+    // (success) or is returned to the slot. If the whole future is dropped while
+    // `accept()` is still pending, the guard's `Drop` recycles the socket so the
+    // slot is never left permanently empty.
+    let mut guard = SlotReturn::new(slot, socket);
+    guard.socket_mut().abort();
+    // Bind the result before matching so the `accept()` future's borrow of
+    // `guard` ends here, freeing `guard` for `into_socket` / `drop` below.
+    let accepted = guard.socket_mut().accept(local_endpoint).await;
+    match accepted {
         Ok(()) => {
+            let socket = guard.into_socket();
             Ok(Box::new(TcpConnection::reusable(socket, slot.clone())) as Box<dyn Connection>)
         }
         Err(_) => {
-            socket.abort();
-            slot.put(socket);
+            // Dropping `guard` aborts the socket and returns it to the slot.
+            drop(guard);
             // `accept()` can fail synchronously (e.g. port-0 `InvalidPort`);
             // without this await the caller re-enters `accept()` immediately with
             // no yield point, starving the executor. Yield so a misconfig

aimdb-tcp-connector/tests/embassy_loopback.rs

@@ -8,9 +8,12 @@
 //!
 //! - recycle: accept -> exchange -> drop -> re-accept (`recycle_then_reaccept`);
 //! - concurrency: one pooled `N = 2` listener keeps both sockets in `accept()`
-//!   on a single port while two clients dial it (`two_concurrent_sessions`);
+//!   on a single port (via the test-only `accept_on`, one accept per index) while
+//!   two clients dial it (`two_concurrent_sessions`);
 //! - redial: after a failed connect and after a dropped link
-//!   (`dialer_redials_after_failure_and_drop`).
+//!   (`dialer_redials_after_failure_and_drop`);
+//! - cancellation: a cancelled accept returns its socket to the slot
+//!   (`cancelled_accept_recycles_socket`).
 #![cfg(feature = "_test-embassy-loopback")]
 
 extern crate alloc;
@@ -258,7 +261,7 @@ fn recycle_then_reaccept() {
 /// Each client lands on a distinct pooled slot; a broken pool (only one socket
 /// accepting, or both racing to the same slot) would hang the second session and
 /// trip the watchdog. (Wiring the pool into the AimX session engine via
-/// `TcpServer<N>` is covered elsewhere; this stays at the transport layer.)
+/// `TcpServer<N>` is intentionally outside this transport-level smoke test.)
 #[test]
 fn two_concurrent_sessions() {
     drive(|server_stack, client_stack| async move {
@@ -268,8 +271,9 @@ fn two_concurrent_sessions() {
         let dialer_a = TcpDialer::new(client_stack, endpoint(7001), buf(), buf());
         let dialer_b = TcpDialer::new(client_stack, endpoint(7001), buf(), buf());
 
-        // Both pooled sockets accept on 7001 while both clients dial it; each
-        // client establishes against a separate slot.
+        // Drive the pooled sockets directly via the test-only `accept_on`, one
+        // accept per index (its single-caller-per-index contract): both sockets
+        // accept on 7001 while both clients dial it, each landing on its own slot.
         let (a_srv, b_srv, a_cli, b_cli) = futures::join!(
             listener.accept_on(0),
             listener.accept_on(1),
@@ -323,3 +327,42 @@ fn dialer_redials_after_failure_and_drop() {
         exchange(server.as_mut(), client.as_mut(), b"again").await;
     });
 }
+
+/// A cancelled accept — its future dropped mid-`accept()`, as a `select!` timeout
+/// or shutdown branch would drop it — must return the pooled socket to its slot.
+/// Without the drop guard the socket is dropped instead of recycled, the slot
+/// stays empty, and the follow-up accept below would hang until the watchdog.
+#[test]
+fn cancelled_accept_recycles_socket() {
+    use futures::future::{ready, select, Either};
+    use futures::pin_mut;
+
+    drive(|server_stack, client_stack| async move {
+        let mut listener = TcpListener::new(server_stack, 7005u16, buf(), buf());
+        let dialer = TcpDialer::new(client_stack, endpoint(7005), buf(), buf());
+
+        // No client dials 7005, so this accept takes the pooled socket and then
+        // parks in `TcpSocket::accept`. Cancel it by letting a ready future win a
+        // `select`, dropping the accept future while its accept is still pending.
+        {
+            let accept = listener.accept();
+            pin_mut!(accept);
+            match select(accept, ready(())).await {
+                Either::Right(_) => {} // `ready` won -> the accept future is cancelled
+                Either::Left(_) => panic!("accept resolved with no client dialing"),
+            }
+        } // the parked accept future drops here, recycling its socket into the slot
+
+        // The recycled socket must be back in its slot: a real accept + dial now
+        // succeeds. A leaked slot would hang this accept and trip the watchdog.
+        let (accepted, connected) = futures::join!(listener.accept(), dialer.connect());
+        let mut server = accepted.expect("accept after a cancelled accept (socket recycled)");
+        let mut client = connected.expect("connect after a cancelled accept");
+        exchange(server.as_mut(), client.as_mut(), b"post-cancel").await;
+    });
+}
+
+// Note: there is no shipped public multi-slot accept to misuse — the pooled path
+// is `TcpServer<N>` (one worker per slot) and the `accept_on` used above is a
+// test-only, single-caller-per-index hook. So the one-waiter-per-slot invariant
+// is upheld by construction and there is nothing to assert at runtime.