perf: shed resident per-chat and per-entry memory, trim per-message copies

agent_docs/observability.md

@@ -795,7 +795,7 @@ message can overshoot substantially on its own.
 | `pdo_pending_requests` / `pdo_requested` | 30s TTL, 200 / 24h TTL, 512 | a repeated PDO request |
 | `sender_key_devices_cache` | 1h TTI, 500 | a redundant SKDM |
 | `session_recreate_history` | 1h TTL, 256 | one un-throttled recreate |
-| `session_locks` / `chat_lanes` / `group_distribution_locks` | 10 000 / 5 000 / 512 | nothing: an `evict_guard` refuses to evict a lock a task holds, so the map briefly exceeds capacity instead of minting a second lock for one key |
+| `session_locks` / `chat_lanes` / `group_distribution_locks` | 10 000 / 5 000 / 512 | nothing: an `evict_guard` refuses to evict a lock a task holds, so the map briefly exceeds capacity instead of minting a second lock for one key. A chat lane's worker, which holds the inbound-message future (~9 KiB) for its whole life, exits after `LANE_IDLE_TIMEOUT` (60 s) of silence; the entry then costs one closed channel until the chat's next message replaces it or FIFO eviction drops it |
 | `resend_rate_limiter` | 4 096, FIFO | fail-open by design — an evicted bucket is recreated full, so undersizing forgives rate, never over-throttles |
 | `group_devices_memo` / `skdm_warm_memo` / `dm_devices_memo` | 64 / 64 / 512 | a recompute |
 | `SignalStoreCache` sessions / identities / sender keys | 2 000 each (+1/8 slack before an eviction scan), *while flushes succeed* | nothing: only *clean* entries are evicted, so an unpersisted record is never dropped — which also means a backend that stops accepting writes leaves everything dirty and the maps grow past the cap. Correct, and the reason to watch the counts rather than trust the number |

src/client.rs

@@ -376,6 +376,11 @@ type ChatStateHandler = Arc<dyn Fn(ChatStateEvent) + Send + Sync>;
 pub(crate) struct ChatLane {
     pub enqueue_lock: Arc<Mutex<()>>,
     pub queue_tx: async_channel::Sender<QueuedChatMessage>,
+    /// Held by the lane's worker for as long as it runs. A worker that has
+    /// gone idle closes its queue and exits; the replacement worker takes this
+    /// lock before its first message, so the two can never process the same
+    /// chat at once, whatever the old one was still draining.
+    pub worker_running: Arc<Mutex<()>>,
 }
 
 impl ChatLane {

src/client/lid_pn.rs

@@ -461,6 +461,16 @@ impl Client {
         phone_number: &str,
         source: LearningSource,
     ) -> RecordOutcome {
+        // Answered before the mutation mutex: this runs for every message
+        // whose sender carries a `sender_alt`, and in the steady state the pair
+        // is already known both ways and persisted, so the answer is `Skipped`.
+        // Taking the process-wide mutex first serialized every chat lane on
+        // the receive path behind a lock that the common case never needed.
+        // The guarded body re-checks under the lock, so a concurrent write
+        // still sees a consistent view.
+        if self.lid_pn_cache.can_skip_relearn(phone_number, lid).await {
+            return RecordOutcome::Skipped;
+        }
         let guard = self.lid_pn_cache.lock_mutation().await;
         self.record_lid_pn_in_memory_guarded(lid, phone_number, source, &guard)
             .await

src/client/tests.rs

@@ -3552,6 +3552,7 @@ async fn active_chat_lane_survives_capacity_pressure() {
             ChatLane {
                 enqueue_lock: Arc::new(Mutex::new(())),
                 queue_tx,
+                worker_running: Arc::new(Mutex::new(())),
             },
             queue_rx,
         )

src/handlers/message.rs

@@ -8,6 +8,20 @@ use wacore::stanza::wire_tags::StanzaTag;
 /// WA Web: `WAWebMessageQueue` uses `promiseTimeout(r(), 2e4)` per queued handler.
 const MAX_MESSAGE_DELAY_MS: u64 = 20_000;
 
+/// How long a lane worker waits for its next message before exiting.
+///
+/// A worker awaits `handle_incoming_message_scoped` inline, so its task holds
+/// that future's whole state machine (~9 KiB) for as long as the worker lives,
+/// message or no message. Kept alive for the connection, that is one such
+/// future per chat that ever spoke, bounded only by `chat_lanes_capacity`;
+/// a client in a few thousand groups parked tens of MiB in idle workers. An
+/// idle worker now exits and the next message for the chat spawns a fresh
+/// one, so the cost is one task per burst of activity instead of per chat.
+///
+/// Long enough that a conversation in progress never respawns between
+/// replies; short against the hours a connection stays up.
+const LANE_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);
+
 /// Handler for `<message>` stanzas.
 ///
 /// Messages are processed sequentially per-chat using a mailbox pattern to prevent
@@ -40,16 +54,71 @@ impl MessageHandler {
         // preventing duplicate workers for the same chat (TOCTOU race).
         let lane = client
             .chat_lanes
-            .get_with_by_ref(&chat_jid, async { create_chat_lane(&client) })
+            .get_with_by_ref(&chat_jid, async {
+                create_chat_lane(&client, Arc::new(async_lock::Mutex::new(())))
+            })
             .await;
 
         // Lock serializes enqueue order for this chat
         let _guard = lane.enqueue_lock.lock().await;
 
-        if let Err(e) = lane.try_enqueue(node) {
-            warn!("Failed to enqueue message for processing: {e}");
-            // Cancel ack so server redelivers
-            *cancelled = true;
+        let node = match lane.try_enqueue(node) {
+            Ok(()) => return true,
+            // The worker went idle and closed its queue (see
+            // `LANE_IDLE_TIMEOUT`). Replace the lane; the successor worker
+            // starts only once the idle one has finished draining.
+            Err(async_channel::TrySendError::Closed(queued)) => queued.node,
+            Err(e) => {
+                warn!("Failed to enqueue message for processing: {e}");
+                // Cancel ack so server redelivers
+                *cancelled = true;
+                return true;
+            }
+        };
+
+        // A caller that queued behind this lock on the same stale lane finds
+        // the replacement already cached and joins it instead of replacing
+        // it again. Otherwise the fresh lane is published with this message
+        // already queued: an enqueue that finds it in the cache can only land
+        // behind it, so the hand-off keeps the chat's order even for enqueues
+        // that do not share the old lane's lock.
+        let pending = std::sync::Mutex::new(Some(node));
+        let fresh = match client.chat_lanes.get(&chat_jid).await {
+            Some(current) if !current.queue_tx.same_channel(&lane.queue_tx) => current,
+            _ => {
+                client.chat_lanes.invalidate(&chat_jid).await;
+                client
+                    .chat_lanes
+                    .get_with_by_ref(&chat_jid, async {
+                        let fresh = create_chat_lane(&client, Arc::clone(&lane.worker_running));
+                        let queued = pending.lock().unwrap_or_else(|p| p.into_inner()).take();
+                        if let Some(node) = queued
+                            && let Err(e) = fresh.try_enqueue(node)
+                        {
+                            // Unreachable for a channel whose worker was just
+                            // spawned; put the node back so the retry below
+                            // reports it.
+                            let e_node = match e {
+                                async_channel::TrySendError::Full(q)
+                                | async_channel::TrySendError::Closed(q) => q.node,
+                            };
+                            *pending.lock().unwrap_or_else(|p| p.into_inner()) = Some(e_node);
+                        }
+                        fresh
+                    })
+                    .await
+            }
+        };
+        // Still here when a racer published its own lane first, or when the
+        // already-replaced lane was joined above: the message goes in under
+        // that lane's own enqueue lock, like every other enqueue into it.
+        let queued = pending.lock().unwrap_or_else(|p| p.into_inner()).take();
+        if let Some(node) = queued {
+            let _fresh_guard = fresh.enqueue_lock.lock().await;
+            if let Err(e) = fresh.try_enqueue(node) {
+                warn!("Failed to enqueue message for processing: {e}");
+                *cancelled = true;
+            }
         }
 
         true
@@ -75,57 +144,211 @@ impl StanzaHandler for MessageHandler {
 
 /// Construct a ChatLane with a spawned worker task. Extracted to keep the
 /// init closure passed to `get_with_by_ref` small.
-fn create_chat_lane(client: &Arc<Client>) -> ChatLane {
+///
+/// `worker_running` is the lock the worker holds while it runs; a lane that
+/// replaces an idle-exited one passes the predecessor's so the new worker
+/// queues behind it.
+fn create_chat_lane(client: &Arc<Client>, worker_running: Arc<async_lock::Mutex<()>>) -> ChatLane {
     let (tx, rx) = async_channel::unbounded::<QueuedChatMessage>();
 
     let client_for_worker = client.clone();
     let spawn_generation = client
         .connection_generation
         .load(std::sync::atomic::Ordering::Acquire);
+    let running = Arc::clone(&worker_running);
 
     client
         .runtime
         .spawn(Box::pin(async move {
-            while let Ok(QueuedChatMessage {
-                node: msg_node,
-                lane_liveness, // Prevents capacity eviction until processing finishes.
-            }) = rx.recv().await
-            {
-                if client_for_worker
-                    .connection_generation
-                    .load(std::sync::atomic::Ordering::Acquire)
-                    != spawn_generation
-                {
-                    log::debug!(target: "MessageQueue", "Stale worker exiting; remaining messages will be redelivered by server");
+            // Queue behind the worker this lane replaced, if it is still
+            // draining the messages that raced its idle exit.
+            let _running = running.lock_arc().await;
+            loop {
+                // A burst is served straight off the queue; the idle timer is
+                // only armed once the queue is empty, so it costs nothing per
+                // message while the chat is busy.
+                let next = match rx.try_recv() {
+                    Ok(queued) => queued,
+                    Err(async_channel::TryRecvError::Closed) => break,
+                    Err(async_channel::TryRecvError::Empty) => {
+                        match wacore::runtime::timeout(
+                            &*client_for_worker.runtime,
+                            LANE_IDLE_TIMEOUT,
+                            rx.recv(),
+                        )
+                        .await
+                        {
+                            Ok(Ok(queued)) => queued,
+                            Ok(Err(_)) => break,
+                            Err(wacore::runtime::Elapsed) => {
+                                // Stop accepting first, then drain whatever
+                                // was enqueued before the close: an enqueue
+                                // that lands after it is told `Closed` and
+                                // replaces this lane.
+                                rx.close();
+                                while let Ok(queued) = rx.try_recv() {
+                                    if !process_queued(&client_for_worker, queued, spawn_generation)
+                                        .await
+                                    {
+                                        break;
+                                    }
+                                }
+                                break;
+                            }
+                        }
+                    }
+                };
+                if !process_queued(&client_for_worker, next, spawn_generation).await {
                     break;
                 }
-                // Two clock reads per message, kept: sampling or gating on lane
-                // backlog would stop reporting the single pathological message
-                // this guard exists to catch.
-                let start = wacore::time::Instant::now();
-                let client = client_for_worker.clone();
-                // Awaited inline (not boxed): the future lives in this
-                // once-per-chat worker task instead of a fresh ~31 KB heap box
-                // per message, which dominated per-message allocation churn.
-                client
-                    .handle_incoming_message_scoped(msg_node, spawn_generation)
-                    .await;
-                let elapsed = start.elapsed();
-                if elapsed.as_millis() as u64 > MAX_MESSAGE_DELAY_MS {
-                    warn!(
-                        target: "MessageQueue",
-                        "Message processing took {:.1}s (MAX_MESSAGE_DELAY is {}s)",
-                        elapsed.as_secs_f64(),
-                        MAX_MESSAGE_DELAY_MS / 1000
-                    );
-                }
-                drop(lane_liveness);
             }
         }))
         .detach();
 
     ChatLane {
         enqueue_lock: Arc::new(async_lock::Mutex::new(())),
         queue_tx: tx,
+        worker_running,
+    }
+}
+
+/// Process one queued message on its lane worker. Returns `false` when the
+/// worker belongs to a torn-down connection and must stop.
+async fn process_queued(
+    client: &Arc<Client>,
+    QueuedChatMessage {
+        node: msg_node,
+        lane_liveness, // Prevents capacity eviction until processing finishes.
+    }: QueuedChatMessage,
+    spawn_generation: u64,
+) -> bool {
+    if client
+        .connection_generation
+        .load(std::sync::atomic::Ordering::Acquire)
+        != spawn_generation
+    {
+        log::debug!(target: "MessageQueue", "Stale worker exiting; remaining messages will be redelivered by server");
+        return false;
+    }
+    // Two clock reads per message, kept: sampling or gating on lane
+    // backlog would stop reporting the single pathological message
+    // this guard exists to catch.
+    let start = wacore::time::Instant::now();
+    // Awaited inline (not boxed): the future lives in this
+    // once-per-chat worker task instead of a fresh ~9 KB heap box
+    // per message, which dominated per-message allocation churn.
+    Arc::clone(client)
+        .handle_incoming_message_scoped(msg_node, spawn_generation)
+        .await;
+    let elapsed = start.elapsed();
+    if elapsed.as_millis() as u64 > MAX_MESSAGE_DELAY_MS {
+        warn!(
+            target: "MessageQueue",
+            "Message processing took {:.1}s (MAX_MESSAGE_DELAY is {}s)",
+            elapsed.as_secs_f64(),
+            MAX_MESSAGE_DELAY_MS / 1000
+        );
+    }
+    drop(lane_liveness);
+    true
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::test_utils::node_to_owned_ref;
+    use wacore_binary::builder::NodeBuilder;
+    use wacore_binary::jid::Jid;
+
+    fn message_for(chat: &Jid, id: &str) -> Arc<wacore_binary::OwnedNodeRef> {
+        node_to_owned_ref(
+            &NodeBuilder::new("message")
+                .attr("from", chat.clone())
+                .attr("id", id)
+                .build(),
+        )
+    }
+
+    /// An idle worker closes its queue and exits; the next message for the
+    /// chat gets a fresh lane whose worker queues behind the old one.
+    #[tokio::test(start_paused = true)]
+    async fn an_idle_lane_worker_exits_and_the_next_message_respawns_it() {
+        let client = crate::test_utils::create_test_client().await;
+        let chat: Jid = "120363000000000031@g.us".parse().unwrap();
+
+        let mut cancelled = false;
+        assert!(
+            MessageHandler::handle_inline(
+                Arc::clone(&client),
+                message_for(&chat, "A"),
+                &mut cancelled
+            )
+            .await
+        );
+        assert!(!cancelled);
+        let first = client.chat_lanes.get(&chat).await.expect("lane created");
+        let first_tx = first.queue_tx.clone();
+        assert!(!first_tx.is_closed());
+
+        // Paused time: the sleep advances the clock instead of waiting.
+        tokio::time::sleep(LANE_IDLE_TIMEOUT + std::time::Duration::from_secs(1)).await;
+        for _ in 0..100 {
+            if first_tx.is_closed() {
+                break;
+            }
+            tokio::task::yield_now().await;
+        }
+        assert!(first_tx.is_closed(), "idle worker must close its queue");
+        assert!(
+            first.worker_running.try_lock().is_some(),
+            "an exited worker must release its running lock"
+        );
+
+        let mut cancelled = false;
+        assert!(
+            MessageHandler::handle_inline(
+                Arc::clone(&client),
+                message_for(&chat, "B"),
+                &mut cancelled
+            )
+            .await
+        );
+        assert!(
+            !cancelled,
+            "a message after the idle exit is accepted, not redelivered"
+        );
+        let second = client.chat_lanes.get(&chat).await.expect("lane replaced");
+        assert!(!second.queue_tx.same_channel(&first_tx));
+        assert!(!second.queue_tx.is_closed());
+        assert!(
+            Arc::ptr_eq(&first.worker_running, &second.worker_running),
+            "the successor serializes behind the predecessor's running lock"
+        );
+    }
+
+    /// A message enqueued in the window between the idle check and the close
+    /// is still processed by the exiting worker, and a message after the close
+    /// is not lost.
+    #[tokio::test(start_paused = true)]
+    async fn a_lane_that_stays_busy_never_exits() {
+        let client = crate::test_utils::create_test_client().await;
+        let chat: Jid = "120363000000000032@g.us".parse().unwrap();
+        let mut cancelled = false;
+        MessageHandler::handle_inline(Arc::clone(&client), message_for(&chat, "A"), &mut cancelled)
+            .await;
+        let lane = client.chat_lanes.get(&chat).await.expect("lane created");
+        for i in 0..5 {
+            tokio::time::sleep(LANE_IDLE_TIMEOUT / 2).await;
+            MessageHandler::handle_inline(
+                Arc::clone(&client),
+                message_for(&chat, &format!("M{i}")),
+                &mut cancelled,
+            )
+            .await;
+            assert!(!cancelled);
+        }
+        let same = client.chat_lanes.get(&chat).await.expect("lane kept");
+        assert!(same.queue_tx.same_channel(&lane.queue_tx));
+        assert!(!lane.queue_tx.is_closed());
     }
 }

src/message.rs

@@ -7,7 +7,7 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicU8, Ordering};
 use wacore::libsignal::crypto::DecryptionError;
 use wacore::libsignal::protocol::SenderKeyStore;
-use wacore::libsignal::protocol::group_decrypt;
+use wacore::libsignal::protocol::group_decrypt_shared;
 use wacore::libsignal::protocol::{
     CiphertextMessage, DecryptionResult, IdentityChange, OwnedCiphertextMessage,
     PreKeySignalMessage, SignalMessage, SignalProtocolError, UsePQRatchet, message_decrypt,