fix(iq): cancellation-safe IQ response waiters (unblocks try_join!) + named fan-out consts

src/client.rs

@@ -467,6 +467,10 @@ pub(crate) struct OfflineSyncMetrics {
     pub start_time: std::sync::Mutex<Option<wacore::time::Instant>>,
 }
 
+/// Map of pending IQ/ack response waiters, keyed by request id.
+pub(crate) type ResponseWaiterMap =
+    HashMap<String, futures::channel::oneshot::Sender<Arc<wacore_binary::OwnedNodeRef>>>;
+
 pub struct Client {
     pub(crate) runtime: Arc<dyn Runtime>,
     pub(crate) core: wacore::client::CoreClient,
@@ -520,9 +524,13 @@ pub struct Client {
     pub(crate) transport_factory: Arc<dyn crate::transport::TransportFactory>,
     pub(crate) noise_socket: Arc<Mutex<Option<Arc<NoiseSocket>>>>,
 
-    pub(crate) response_waiters: Arc<
-        Mutex<HashMap<String, futures::channel::oneshot::Sender<Arc<wacore_binary::OwnedNodeRef>>>>,
-    >,
+    /// Pending IQ/ack response waiters keyed by request id.
+    ///
+    /// A `std::sync::Mutex` (like the `node_waiters` sibling below): the critical
+    /// section is a trivial map op never held across an `.await`, and a sync lock
+    /// is what lets `ResponseWaiterGuard` remove a cancelled waiter from `Drop`
+    /// (an async lock couldn't). See `send_and_wait_iq`.
+    pub(crate) response_waiters: Arc<std::sync::Mutex<ResponseWaiterMap>>,
 
     /// Generic node waiters for waiting on specific stanzas by tag/attributes.
     /// Uses std::sync::Mutex (not tokio) since the critical section is trivial.

src/client/accessors.rs

@@ -140,7 +140,7 @@ impl Client {
             .await;
 
         // Each count read into a local so no two guards are ever held at once.
-        let response_waiters = self.response_waiters.lock().await.len();
+        let response_waiters = self.response_waiters_guard().len();
         let presence_subscriptions = self.presence_subscriptions.lock().await.len();
         let app_state_key_requests = self.app_state_key_requests.lock().await.len();
         let app_state_syncing = self.app_state_syncing.lock().await.len();
@@ -374,6 +374,19 @@ impl Client {
         rx
     }
 
+    /// Poison-recovering lock of the `response_waiters` map. Centralizes the
+    /// `unwrap_or_else(into_inner)` so no call site reaches for a bare
+    /// `.lock().unwrap()` that would panic if a holder ever panicked. The critical
+    /// section is a trivial map op, never held across an `.await`.
+    #[inline]
+    pub(crate) fn response_waiters_guard(
+        &self,
+    ) -> std::sync::MutexGuard<'_, crate::client::ResponseWaiterMap> {
+        self.response_waiters
+            .lock()
+            .unwrap_or_else(|p| p.into_inner())
+    }
+
     /// Check pending node waiters against an incoming node.
     /// Only called when `node_waiter_count > 0`.
     pub(crate) fn resolve_node_waiters(&self, node: &Arc<wacore_binary::OwnedNodeRef>) {

src/client/lifecycle.rs

@@ -140,7 +140,7 @@ impl Client {
             transport_factory,
             noise_socket: Arc::new(Mutex::new(None)),
 
-            response_waiters: Arc::new(Mutex::new(HashMap::new())),
+            response_waiters: Arc::new(std::sync::Mutex::new(HashMap::new())),
             node_waiters: std::sync::Mutex::new(Vec::new()),
             node_waiter_count: AtomicUsize::new(0),
             sent_node_waiters: std::sync::Mutex::new(Vec::new()),
@@ -876,12 +876,16 @@ impl Client {
         self.history_sync_idle_notifier.notify(usize::MAX);
         // Drain all pending IQ waiters so they fail fast with InternalChannelClosed
         // instead of hanging until the 75s timeout.
-        let mut waiters_map = self.response_waiters.lock().await;
-        let waiter_count = waiters_map.len();
-        // Replace with new map to release backing storage; old senders drop here,
-        // causing receivers to get RecvError → IqError::InternalChannelClosed
-        *waiters_map = HashMap::new();
-        drop(waiters_map);
+        // Scoped so the sync guard is dropped before the awaits below (a
+        // std::sync::MutexGuard held across an await would make this future !Send).
+        let waiter_count = {
+            let mut waiters_map = self.response_waiters_guard();
+            let count = waiters_map.len();
+            // Replace with new map to release backing storage; old senders drop
+            // here, causing receivers to get RecvError → InternalChannelClosed.
+            *waiters_map = HashMap::new();
+            count
+        };
         if waiter_count > 0 {
             debug!(
                 "Dropping {} orphaned IQ response waiter(s) on disconnect",

src/client/messaging.rs

@@ -231,14 +231,13 @@ impl Client {
 
     /// Register a oneshot waiter for a server ack by message ID.
     /// Returns the receiver — caller sends the node separately and awaits this in background.
-    pub(crate) async fn register_ack_waiter(
+    /// Sync: registration is just a `std::sync::Mutex` insert (no await).
+    pub(crate) fn register_ack_waiter(
         &self,
         message_id: &str,
     ) -> futures::channel::oneshot::Receiver<std::sync::Arc<wacore_binary::OwnedNodeRef>> {
         let (tx, rx) = futures::channel::oneshot::channel();
-        self.response_waiters
-            .lock()
-            .await
+        self.response_waiters_guard()
             .insert(message_id.to_string(), tx);
         rx
     }

src/client/node_io.rs

@@ -430,15 +430,12 @@ impl Client {
 
         if nr.tag.as_ref() == "iq"
             && let Some(id) = nr.get_attr("id").map(|v| v.as_str())
+            && let Some(waiter) = self.response_waiters_guard().remove(id.as_ref())
         {
-            // Single lock acquisition: try to remove the waiter directly.
-            let waiter = self.response_waiters.lock().await.remove(id.as_ref());
-            if let Some(waiter) = waiter {
-                if waiter.send(Arc::clone(&node)).is_err() {
-                    warn!(target: "Client/IQ", "Failed to send IQ response to waiter. Receiver was likely dropped.");
-                }
-                return;
+            if waiter.send(Arc::clone(&node)).is_err() {
+                warn!(target: "Client/IQ", "Failed to send IQ response to waiter. Receiver was likely dropped.");
             }
+            return;
         }
 
         // Dispatch to appropriate handler using the router
@@ -1114,7 +1111,7 @@ impl Client {
         }
 
         if let Some(id) = node.get_attr("id").map(|v| v.as_str())
-            && let Some(waiter) = self.response_waiters.lock().await.remove(id.as_ref())
+            && let Some(waiter) = self.response_waiters_guard().remove(id.as_ref())
         {
             // ACK responses are infrequent; re-encode into OwnedNodeRef for the channel.
             // marshal_ref prepends a leading 0x00 format byte; OwnedNodeRef::new expects raw

src/client/sessions.rs

@@ -298,6 +298,7 @@ impl Client {
         // otherwise serialize the per-device DB reads (warm hits serialize on the
         // cache mutex anyway). Order is irrelevant — misses are chunked for the fetch.
         use futures::StreamExt;
+        const SESSION_PROBE_CONCURRENCY: usize = 16;
         let backend = device_snapshot.backend.clone();
         let jids_needing_sessions: Vec<Jid> = futures::stream::iter(jids)
             .map(|jid| {
@@ -315,7 +316,7 @@ impl Client {
                     }
                 }
             })
-            .buffer_unordered(16)
+            .buffer_unordered(SESSION_PROBE_CONCURRENCY)
             .filter_map(|needed| async move { needed })
             .collect()
             .await;

src/client/tests.rs

@@ -130,13 +130,9 @@ async fn test_ack_waiter_resolves() {
     // 1. Insert a waiter for a specific ID
     let test_id = "ack-test-123".to_string();
     let (tx, rx) = oneshot::channel();
-    client
-        .response_waiters
-        .lock()
-        .await
-        .insert(test_id.clone(), tx);
+    client.response_waiters_guard().insert(test_id.clone(), tx);
     assert!(
-        client.response_waiters.lock().await.contains_key(&test_id),
+        client.response_waiters_guard().contains_key(&test_id),
         "Waiter should be inserted before handling ack"
     );
 
@@ -170,7 +166,7 @@ async fn test_ack_waiter_resolves() {
 
     // 5. Verify the waiter was removed
     assert!(
-        !client.response_waiters.lock().await.contains_key(&test_id),
+        !client.response_waiters_guard().contains_key(&test_id),
         "Waiter should be removed after handling"
     );
 

src/features/contacts.rs

@@ -175,14 +175,14 @@ impl<'a> Contacts<'a> {
                     .await
             }
         };
-        // join!, NOT try_join!: a fail-fast try_join! would drop the sibling IQ
-        // future the instant one errored, leaking its `response_waiters` entry —
-        // send_and_wait_iq only removes the waiter on send-failure/timeout/shutdown,
-        // not on cancellation-via-drop, and a lingering waiter suppresses
-        // keepalives. Awaiting both lets each clean up its own waiter.
-        let (pn_results, lid_results) = futures::join!(pn_fut, lid_fut);
-        let mut results = pn_results?;
-        results.extend(lid_results?);
+        // try_join! fails fast: it returns the instant either query errors and
+        // drops the sibling in-flight future. That's now safe — `send_and_wait_iq`
+        // registers a `ResponseWaiterGuard` that removes the waiter on drop, so a
+        // cancelled sibling can't leak its `response_waiters` entry (which would
+        // otherwise suppress keepalives). The old sequential code also failed on
+        // the first error, so fail-fast matches the original latency profile.
+        let (mut results, lid_results) = futures::try_join!(pn_fut, lid_fut)?;
+        results.extend(lid_results);
 
         self.persist_lid_mappings(results.iter().map(forward_lid_pair))
             .await;

src/features/groups.rs

@@ -373,6 +373,7 @@ impl<'a> Groups<'a> {
         // Cache hits are in-memory, but a cold cache falls back to the DB and a
         // large group would otherwise serialize those lookups — bounded fan-out.
         use futures::StreamExt;
+        const LID_PN_RESOLVE_CONCURRENCY: usize = 16;
         let resolved: Vec<(usize, Jid)> = futures::stream::iter(pending)
             .map(|(i, jid)| async move {
                 let pn = self
@@ -384,7 +385,7 @@ impl<'a> Groups<'a> {
                     .map(|e| Jid::pn(&*e.phone_number));
                 (i, pn)
             })
-            .buffer_unordered(16)
+            .buffer_unordered(LID_PN_RESOLVE_CONCURRENCY)
             .filter_map(|(i, pn)| async move { pn.map(|pn| (i, pn)) })
             .collect()
             .await;

src/keepalive.rs

@@ -77,7 +77,7 @@ impl Client {
 
         // WA Web: skip ping if there are pending IQs
         // (`activePing || ackHandlers.length || pendingIqs.size`)
-        let has_pending = !self.response_waiters.lock().await.is_empty();
+        let has_pending = !self.response_waiters_guard().is_empty();
         if has_pending {
             debug!(target: "Client/Keepalive", "Skipping ping: IQ responses pending");
             return KeepaliveResult::Ok;

src/prekeys.rs

@@ -181,14 +181,15 @@ impl Client {
         // Fan out the per-companion identity loads (independent cache/DB reads) —
         // the keyless-companion set can be large on a cold group send. Owned Vec so
         // the stream doesn't borrow `jids` through buffer_unordered (Send bound).
+        const COMPANION_IDENTITY_LOAD_CONCURRENCY: usize = 16;
         let companions: Vec<Jid> = jids.iter().filter(|j| j.device != 0).cloned().collect();
         futures::stream::iter(companions)
             .map(|jid| async move {
                 self.load_account_identity(&jid)
                     .await
                     .map(|id| (jid.normalize_for_prekey_bundle(), id))
             })
-            .buffer_unordered(16)
+            .buffer_unordered(COMPANION_IDENTITY_LOAD_CONCURRENCY)
             .filter_map(|entry| async move { entry })
             .collect()
             .await

src/request.rs

@@ -22,6 +22,29 @@ type IqSendFuture<'a> =
 type IqSendFuture<'a> =
     std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), ClientError>> + 'a>>;
 
+/// Removes a pending `response_waiters` entry when dropped.
+///
+/// `send_and_wait_iq` can be cancelled mid-await — e.g. the losing side of a
+/// `futures::try_join!` is dropped the instant its sibling errors. Without this
+/// guard the registered waiter would linger in the map: the explicit cleanups
+/// only fired on the send-fail / timeout / shutdown paths, never on
+/// cancellation-via-drop, and a lingering waiter suppresses keepalives for the
+/// life of the connection. Dropping the guard removes the entry on every exit
+/// path; on success `resolve_waiters` already removed it, so it's a no-op.
+struct ResponseWaiterGuard {
+    waiters: Arc<std::sync::Mutex<crate::client::ResponseWaiterMap>>,
+    req_id: String,
+}
+
+impl Drop for ResponseWaiterGuard {
+    fn drop(&mut self) {
+        self.waiters
+            .lock()
+            .unwrap_or_else(|p| p.into_inner())
+            .remove(&self.req_id);
+    }
+}
+
 #[derive(Debug, Error)]
 #[non_exhaustive]
 pub enum IqError {
@@ -279,23 +302,38 @@ impl Client {
         }
 
         let (tx, rx) = futures::channel::oneshot::channel();
-        self.response_waiters
-            .lock()
-            .await
-            .insert(req_id.clone(), tx);
+        {
+            let mut waiters = self.response_waiters_guard();
+            // req_ids come from the monotonic generate_request_id(), so a given
+            // id is never in flight twice — the invariant that makes the guard's
+            // remove-by-id unambiguous (an overwrite could otherwise let an older
+            // guard evict a newer waiter). Assert it so a future caller passing a
+            // duplicate id is caught in tests instead of silently.
+            debug_assert!(
+                !waiters.contains_key(&req_id),
+                "duplicate in-flight IQ request id: {req_id}"
+            );
+            waiters.insert(req_id.clone(), tx);
+        }
+        // RAII cleanup covers every exit below — including this future being
+        // dropped mid-await (cancellation), which the explicit paths can't
+        // catch. So the send-fail / timeout / shutdown arms no longer remove
+        // the waiter by hand; the guard does it on drop.
+        let _waiter_guard = ResponseWaiterGuard {
+            waiters: self.response_waiters.clone(),
+            req_id,
+        };
 
         // Per-connection: pending IQ requests are bound to the current socket;
         // a reconnect aborts them (sender retries on the new connection).
         let shutdown = wacore::runtime::wait_for_shutdown(&self.connection_shutdown_signal());
 
         if !self.is_running.load(Ordering::Acquire) {
-            self.response_waiters.lock().await.remove(&req_id);
             wacore::telemetry::iq("error");
             return Err(IqError::NotConnected);
         }
 
         if let Err(e) = send_fn.await {
-            self.response_waiters.lock().await.remove(&req_id);
             wacore::telemetry::iq("error");
             return match e {
                 ClientError::Socket(s_err) => Err(IqError::Socket(s_err)),
@@ -317,16 +355,10 @@ impl Client {
                         Err(e) => Err(e.into()),
                     },
                     Ok(Err(_)) => Err(IqError::InternalChannelClosed),
-                    Err(_) => {
-                        self.response_waiters.lock().await.remove(&req_id);
-                        Err(IqError::Timeout)
-                    }
+                    Err(_) => Err(IqError::Timeout),
                 }
             }
-            _ = shutdown.fuse() => {
-                self.response_waiters.lock().await.remove(&req_id);
-                Err(IqError::NotConnected)
-            }
+            _ = shutdown.fuse() => Err(IqError::NotConnected),
         };
         wacore::telemetry::iq(match &result {
             Ok(_) => "ok",
@@ -339,7 +371,9 @@ impl Client {
 
 #[cfg(test)]
 mod tests {
-    use super::IqError;
+    use super::{IqError, ResponseWaiterGuard};
+    use std::collections::HashMap;
+    use std::sync::{Arc, Mutex};
 
     #[test]
     fn converts_unexpected_response_type() {
@@ -352,4 +386,43 @@ mod tests {
             other => panic!("expected UnexpectedResponseType, got {other:?}"),
         }
     }
+
+    // Cancellation cleanup: dropping a `send_and_wait_iq` future mid-await (e.g.
+    // the loser of a `try_join!`) must remove its still-pending waiter, or a
+    // leaked entry suppresses keepalives for the life of the connection.
+    #[test]
+    fn waiter_guard_removes_pending_entry_on_drop() {
+        let waiters: Arc<Mutex<crate::client::ResponseWaiterMap>> =
+            Arc::new(Mutex::new(HashMap::new()));
+        let (tx, _rx) = futures::channel::oneshot::channel();
+        waiters.lock().unwrap().insert("req-1".to_string(), tx);
+        assert!(waiters.lock().unwrap().contains_key("req-1"));
+
+        {
+            let _guard = ResponseWaiterGuard {
+                waiters: waiters.clone(),
+                req_id: "req-1".to_string(),
+            };
+        }
+        assert!(
+            !waiters.lock().unwrap().contains_key("req-1"),
+            "dropping the guard must remove the pending waiter"
+        );
+    }
+
+    // On the success path the resolver already removed the entry before the
+    // guard drops, so the guard's removal must be a harmless no-op.
+    #[test]
+    fn waiter_guard_drop_is_noop_when_already_resolved() {
+        let waiters: Arc<Mutex<crate::client::ResponseWaiterMap>> =
+            Arc::new(Mutex::new(HashMap::new()));
+        // Map empty = resolver already delivered + removed this request's waiter.
+        {
+            let _guard = ResponseWaiterGuard {
+                waiters: waiters.clone(),
+                req_id: "req-1".to_string(),
+            };
+        }
+        assert!(waiters.lock().unwrap().is_empty());
+    }
 }