diff --git a/.gitignore b/.gitignore index 5440f5dd0..a61542e68 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,9 @@ dhat-heap*.json # proptest persists failure seeds next to the test on failure; these are # environment-local replay caches, not source. *.proptest-regressions + +# Diagnostic store snapshots tests write into the repo root +# (PersistenceManager::create_snapshot). Generated artifacts, never sources. +# Two spellings: the plain name and the "file:" form the sqlite URI produces. +memdb_*.snapshot-* +file:memdb_*.snapshot-* diff --git a/src/client.rs b/src/client.rs index f99d1bb54..d8a20690f 100644 --- a/src/client.rs +++ b/src/client.rs @@ -288,6 +288,12 @@ pub struct MemoryReport { pub group_distribution_lock_eviction_blocks: u64, pub resend_rate_limiter_chats: u64, // -- Unbounded collections -- + /// Deferred acks queued for the transport-ack worker. Unbounded, and each + /// entry retains the full inbound node plus a flush guard, so a stalled + /// transport shows up here as a growing backlog. + pub transport_ack_queue: usize, + /// Delivery receipts queued for their worker, same shape as above. + pub delivery_receipt_queue: usize, pub response_waiters: usize, pub node_waiters: usize, pub pending_retries: usize, @@ -393,6 +399,12 @@ impl std::fmt::Display for MemoryReport { self.resend_rate_limiter_chats )?; writeln!(f, "--- Unbounded collections ---")?; + writeln!(f, " transport_ack_queue: {}", self.transport_ack_queue)?; + writeln!( + f, + " delivery_receipt_queue: {}", + self.delivery_receipt_queue + )?; writeln!(f, " response_waiters: {}", self.response_waiters)?; writeln!(f, " node_waiters: {}", self.node_waiters)?; writeln!(f, " pending_retries: {}", self.pending_retries)?; @@ -709,9 +721,36 @@ pub(crate) struct OfflineSyncMetrics { type ResponseWaiterSender = futures::channel::oneshot::Sender>; +/// What a pending ack/IQ entry is waiting to do once the response arrives. +/// +/// A phash check used to be an `Iq` waiter plus a spawned task holding the +/// receiver and a ten second timer, which is a task, a channel and a timer per +/// outgoing message for a comparison that almost always succeeds. Carrying the +/// expected value in the map instead lets the read loop compare it inline and +/// spawn only on the rare mismatch. +pub(crate) enum ResponseWaiter { + /// Classic request/response: hand the node to whoever is awaiting it. + Iq(ResponseWaiterSender), + /// Compare the server's `phash` against ours; act only if they differ. + Phash(PhashWaiter), +} + +pub(crate) struct PhashWaiter { + pub(crate) expected: wacore_binary::CompactString, + pub(crate) jid: Jid, + pub(crate) invalidate_group_cache: bool, + /// Sweep epoch this waiter was registered in. Expiry is counted in sweeps + /// rather than seconds: a wall deadline is subject to clock jumps (see + /// wacore::time) and would have to be derived from an instant sampled well + /// before registration, while reading a fresh clock here is what the send + /// clock budget forbids. Surviving one full sweep is the trigger, so the + /// window is one keepalive tick (15 to 30 s) instead of the old fixed 10 s. + pub(crate) registered_epoch: u64, +} + struct ResponseWaiterEntry { generation: NonZeroU64, - sender: ResponseWaiterSender, + waiter: ResponseWaiter, } /// Map of pending IQ/ack response waiters, keyed by request id. @@ -722,6 +761,9 @@ struct ResponseWaiterEntry { pub(crate) struct ResponseWaiterMap { entries: HashMap, last_generation: u64, + /// Advanced once per sweep. Registration reads it under the lock it already + /// takes, so a waiter records its age without touching a clock. + sweep_epoch: u64, } impl ResponseWaiterMap { @@ -737,14 +779,14 @@ impl ResponseWaiterMap { pub(crate) fn try_insert_guarded( &mut self, request_id: String, - sender: ResponseWaiterSender, + waiter: ResponseWaiter, ) -> Option { use std::collections::hash_map::Entry; let generation = self.next_generation(); match self.entries.entry(request_id) { Entry::Vacant(entry) => { - entry.insert(ResponseWaiterEntry { generation, sender }); + entry.insert(ResponseWaiterEntry { generation, waiter }); Some(generation) } Entry::Occupied(_) => None, @@ -754,16 +796,37 @@ impl ResponseWaiterMap { pub(crate) fn insert( &mut self, request_id: String, - sender: ResponseWaiterSender, - ) -> Option { + waiter: ResponseWaiter, + ) -> Option { let generation = self.next_generation(); self.entries - .insert(request_id, ResponseWaiterEntry { generation, sender }) - .map(|entry| entry.sender) + .insert(request_id, ResponseWaiterEntry { generation, waiter }) + .map(|entry| entry.waiter) } - pub(crate) fn remove(&mut self, request_id: &str) -> Option { - self.entries.remove(request_id).map(|entry| entry.sender) + pub(crate) fn remove(&mut self, request_id: &str) -> Option { + self.entries.remove(request_id).map(|entry| entry.waiter) + } + + /// The epoch a waiter registered now belongs to. + pub(crate) fn current_epoch(&self) -> u64 { + self.sweep_epoch + } + + /// Drop phash waiters that lived through a whole sweep without their ack. + /// + /// Runs on the keepalive tick, before the recent-activity early return: a + /// connection with steady inbound traffic skips the ping entirely, and + /// sweeping only inside the ping would let lost acks accumulate for as long + /// as traffic keeps flowing. The map is also what makes keepalive treat the + /// connection as "IQs pending", so a stranded waiter silences pings. + pub(crate) fn drop_expired_phash(&mut self) { + let epoch = self.sweep_epoch; + self.entries.retain(|_, entry| match &entry.waiter { + ResponseWaiter::Phash(waiter) => waiter.registered_epoch >= epoch, + ResponseWaiter::Iq(_) => true, + }); + self.sweep_epoch = self.sweep_epoch.wrapping_add(1); } pub(crate) fn remove_guarded(&mut self, request_id: &str, cleanup_generation: NonZeroU64) { @@ -1058,6 +1121,16 @@ pub struct Client { crate::flush_scope::FlushGuard, )>, >, + /// Feed of the persistent transport-ack worker, mirroring + /// [`Self::delivery_receipt_queue`]. Deferred acks used to be one spawned + /// task each; the queue also gives them FIFO order, which the spawns did + /// not guarantee. + pub(crate) transport_ack_queue: std::sync::OnceLock< + async_channel::Sender<( + Arc, + crate::flush_scope::FlushGuard, + )>, + >, /// Contacts with active presence subscriptions that must be re-subscribed on reconnect. pub(crate) presence_subscriptions: Arc>>, /// Metrics for granular offline sync logging diff --git a/src/client/accessors.rs b/src/client/accessors.rs index 02f6427dd..0a7bab1f5 100644 --- a/src/client/accessors.rs +++ b/src/client/accessors.rs @@ -257,6 +257,8 @@ impl Client { group_distribution_lock_evictions: group_distribution_locks.evictions, group_distribution_lock_eviction_blocks: group_distribution_locks.eviction_blocks, resend_rate_limiter_chats: self.resend_rate_limiter.entry_count(), + transport_ack_queue: self.transport_ack_queue.get().map_or(0, |tx| tx.len()), + delivery_receipt_queue: self.delivery_receipt_queue.get().map_or(0, |tx| tx.len()), response_waiters, node_waiters: self.node_waiter_count.load(Ordering::Relaxed), pending_retries: pending_retries_count, diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 250a73d3a..8474de48c 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -345,6 +345,7 @@ impl Client { history_sync_activity: Arc::new(crate::sync_task::HistorySyncActivity::new()), outbound_flush: Arc::new(crate::flush_scope::FlushScope::new()), delivery_receipt_queue: std::sync::OnceLock::new(), + transport_ack_queue: std::sync::OnceLock::new(), presence_subscriptions: Arc::new(Mutex::new(HashSet::new())), socket_ready_notifier: Arc::new(event_listener::Event::new()), is_ready: Arc::new(AtomicBool::new(false)), diff --git a/src/client/messaging.rs b/src/client/messaging.rs index d75ca7863..f442e63f3 100644 --- a/src/client/messaging.rs +++ b/src/client/messaging.rs @@ -269,16 +269,51 @@ 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. /// Sync: registration is just a `std::sync::Mutex` insert (no await). + /// Register a waiter that receives the ack node itself. + /// + /// Used where the caller needs the response: the VoIP offer reads the relay + /// out of its ack. A phash check does not, which is why that path uses + /// [`Self::register_phash_waiter`] and pays no channel per message. Gated on + /// the only consumer's feature, or it is dead code in a default build. + #[cfg(feature = "voip-runtime")] pub(crate) fn register_ack_waiter( &self, message_id: &str, ) -> futures::channel::oneshot::Receiver> { let (tx, rx) = futures::channel::oneshot::channel(); self.response_waiters_guard() - .insert(message_id.to_string(), tx); + .insert(message_id.to_string(), ResponseWaiter::Iq(tx)); rx } + /// Register the phash the server is expected to echo for this send. + /// + /// Nothing awaits the result: the read loop compares inline when the ack + /// lands and only acts on a mismatch, so a send costs a map entry instead of + /// a task, a oneshot and a timer. + pub(crate) fn register_phash_waiter( + &self, + message_id: &str, + expected: wacore_binary::CompactString, + jid: Jid, + invalidate_group_cache: bool, + ) { + let mut waiters = self.response_waiters_guard(); + // Stamped with the sweep epoch under the lock the insert already holds: + // a deadline derived from the instant the send started would already be + // stale here when preparation is slow, and a wall clock can jump. + let registered_epoch = waiters.current_epoch(); + waiters.insert( + message_id.to_string(), + ResponseWaiter::Phash(PhashWaiter { + expected, + jid, + invalidate_group_cache, + registered_epoch, + }), + ); + } + /// Creates a normalized ChatMessageId by resolving PN to LID JIDs. pub(crate) async fn make_chat_message_id(&self, chat: &Jid, id: &str) -> ChatMessageId { // Resolve chat JID to LID if possible diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 13e42bca8..f61c6bc19 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -1,6 +1,7 @@ //! Inbound node I/O: read loop, frame decryption, node routing, acks and stream errors. use super::*; +use crate::client::{PhashWaiter, ResponseWaiter}; use wacore::net::DisconnectReason; /// Non-error exits of [`Client::read_messages_loop`] — `ServerRecycle` keeps the @@ -478,8 +479,17 @@ impl Client { && let Some(id) = nr.get_attr("id").map(|v| v.as_str()) && let Some(waiter) = self.response_waiters_guard().remove(id.as_ref()) { - if waiter.send(Arc::clone(&node)).is_err() { - warn!(target: "Client/IQ", "Failed to send IQ response to waiter. Receiver was likely dropped."); + // An IQ id never carries a phash waiter (those are registered under + // message ids), so a mismatch here means the id space collided. + match waiter { + ResponseWaiter::Iq(sender) => { + if sender.send(Arc::clone(&node)).is_err() { + warn!(target: "Client/IQ", "Failed to send IQ response to waiter. Receiver was likely dropped."); + } + } + ResponseWaiter::Phash(_) => { + warn!(target: "Client/IQ", "IQ id collided with a pending phash waiter; dropping the phash check"); + } } return; } @@ -614,9 +624,12 @@ impl Client { } } - /// Possibly send a deferred ack: either immediately or via spawned task. - /// Handlers can cancel by setting `cancelled` to true. - /// Uses Arc to avoid cloning when spawning the async task. + /// Possibly send a deferred ack: either immediately or through the ack + /// worker. Handlers can cancel by setting `cancelled` to true. + /// Uses Arc so queueing does not clone the node. + /// + /// The deferred path feeds one persistent worker rather than spawning a + /// task per ack, which also makes acks leave in arrival order. async fn maybe_deferred_ack(self: &Arc, node: Arc) { if self.synchronous_ack { if let Err(e) = self.send_ack_for(node.get()).await @@ -624,18 +637,50 @@ impl Client { { warn!("Failed to send ack: {e:?}"); } - } else { - let this = self.clone(); - self.runtime - .spawn(Box::pin(async move { - if let Err(e) = this.send_ack_for(node.get()).await + return; + } + // A closed scope means disconnect is already running; the spawned task + // it replaces would have failed on an unavailable transport anyway. + let Some(guard) = self.outbound_flush.try_track() else { + return; + }; + let tx = self + .transport_ack_queue + .get_or_init(|| self.start_transport_ack_worker()); + // Only fails once the worker is gone (client teardown). + let _ = tx.try_send((node, guard)); + } + + /// Worker shared by every deferred ack. Holds a `Weak`, so a dropped + /// `Client` closes the channel and ends the task instead of keeping the + /// client alive. + fn start_transport_ack_worker( + self: &Arc, + ) -> async_channel::Sender<( + Arc, + crate::flush_scope::FlushGuard, + )> { + let (tx, rx) = async_channel::unbounded::<( + Arc, + crate::flush_scope::FlushGuard, + )>(); + let client = Arc::downgrade(self); + self.runtime + .spawn(Box::pin(async move { + while let Ok((node, guard)) = rx.recv().await { + let Some(client) = client.upgrade() else { + break; + }; + if let Err(e) = client.send_ack_for(node.get()).await && !e.is_transport_unavailable() { warn!("Failed to send ack: {e:?}"); } - })) - .detach(); - } + drop(guard); + } + })) + .detach(); + tx } #[inline] @@ -1269,12 +1314,20 @@ impl Client { /// Ack entry point for callers that already share the node: the waiter /// receives an `Arc` clone instead of a ~1 KB re-encode + re-parse. - pub(crate) fn handle_ack_response_arc(&self, node: &Arc) -> bool { + pub(crate) fn handle_ack_response_arc( + self: &Arc, + node: &Arc, + ) -> bool { let Some(waiter) = self.take_ack_waiter(node.get()) else { return false; }; - if let Err(rejected) = waiter.send(Arc::clone(node)) { - Self::warn_ack_waiter_dropped(&rejected); + match waiter { + ResponseWaiter::Iq(sender) => { + if let Err(rejected) = sender.send(Arc::clone(node)) { + Self::warn_ack_waiter_dropped(&rejected); + } + } + ResponseWaiter::Phash(waiter) => self.check_phash_against_ack(node.get(), waiter), } true } @@ -1282,16 +1335,54 @@ impl Client { /// Ack entry point for the read-loop fast path, which owns the node: the /// `Arc` is built from the existing allocation, and only when a waiter is /// actually waiting. - pub(crate) fn handle_ack_response_owned(&self, node: wacore_binary::OwnedNodeRef) -> bool { + pub(crate) fn handle_ack_response_owned( + self: &Arc, + node: wacore_binary::OwnedNodeRef, + ) -> bool { let Some(waiter) = self.take_ack_waiter(node.get()) else { return false; }; - if let Err(rejected) = waiter.send(Arc::new(node)) { - Self::warn_ack_waiter_dropped(&rejected); + match waiter { + ResponseWaiter::Iq(sender) => { + if let Err(rejected) = sender.send(Arc::new(node)) { + Self::warn_ack_waiter_dropped(&rejected); + } + } + ResponseWaiter::Phash(waiter) => self.check_phash_against_ack(node.get(), waiter), } true } + /// Inline half of the phash check. The comparison is a string equality on + /// the read loop; only a disagreement pays for a task, and that path + /// re-reads caches and can force a sender-key redistribution. + fn check_phash_against_ack( + self: &Arc, + node: &wacore_binary::NodeRef<'_>, + waiter: PhashWaiter, + ) { + let Some(server) = node.get_attr("phash") else { + return; + }; + if server.as_str() == waiter.expected { + return; + } + let client = Arc::clone(self); + let server = server.as_str().to_string(); + self.runtime + .spawn(Box::pin(async move { + client + .handle_phash_mismatch( + &waiter.jid, + &waiter.expected, + &server, + waiter.invalidate_group_cache, + ) + .await; + })) + .detach(); + } + fn warn_ack_waiter_dropped(rejected: &Arc) { warn!( target: "Client/Ack", @@ -1306,10 +1397,7 @@ impl Client { feature = "tracing", tracing::instrument(name = "wa.conn.ack_response", level = "debug", skip_all) )] - fn take_ack_waiter( - &self, - node: &wacore_binary::NodeRef<'_>, - ) -> Option>> { + fn take_ack_waiter(&self, node: &wacore_binary::NodeRef<'_>) -> Option { let ack_id = node.get_attr("id"); let ack_error = node.get_attr("error"); diff --git a/src/client/tests.rs b/src/client/tests.rs index fb6a75d46..5a37dd835 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -143,7 +143,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_guard().insert(test_id.clone(), tx); + client + .response_waiters_guard() + .insert(test_id.clone(), ResponseWaiter::Iq(tx)); assert!( client.response_waiters_guard().contains_key(&test_id), "Waiter should be inserted before handling ack" @@ -249,7 +251,7 @@ async fn ack_arc_delivery_shares_allocation() { let (tx, rx) = oneshot::channel(); client .response_waiters_guard() - .insert(test_id.to_string(), tx); + .insert(test_id.to_string(), ResponseWaiter::Iq(tx)); let node = Arc::new(owned_ack_node(test_id)); assert!(client.handle_ack_response_arc(&node)); @@ -277,7 +279,7 @@ async fn ack_owned_delivery_resolves_waiter() { let (tx, rx) = oneshot::channel(); client .response_waiters_guard() - .insert(test_id.to_string(), tx); + .insert(test_id.to_string(), ResponseWaiter::Iq(tx)); assert!(client.handle_ack_response_owned(owned_ack_node(test_id))); let received = tokio::time::timeout(Duration::from_secs(1), rx) @@ -370,7 +372,7 @@ async fn test_ack_dispatches_server_ack_event() { let (tx, rx) = oneshot::channel(); client .response_waiters_guard() - .insert("ack-evt-3".to_string(), tx); + .insert("ack-evt-3".to_string(), ResponseWaiter::Iq(tx)); let waited_ack = NodeBuilder::new("ack") .attr("id", "ack-evt-3") .attr("class", "message") @@ -4109,3 +4111,49 @@ async fn offline_preview_defaults_absent_counts_to_zero() { assert_eq!(preview.calls, 0); assert_eq!(preview.statuses, 0); } + +/// A phash waiter is resolved by an ack that may never arrive, and nothing +/// polls it. The sweep has to drop the stale one, or a non-empty map reads as +/// "IQ pending" and silences pings for the life of the connection. +#[test] +fn phash_waiter_sweep_drops_only_entries_that_lived_through_a_sweep() { + use crate::client::{PhashWaiter, ResponseWaiter, ResponseWaiterMap}; + use futures::channel::oneshot; + + let mut map = ResponseWaiterMap::default(); + let waiter = |registered_epoch: u64| { + ResponseWaiter::Phash(PhashWaiter { + expected: wacore_binary::CompactString::from("hash"), + jid: "13135550100@s.whatsapp.net".parse().expect("valid jid"), + invalidate_group_cache: false, + registered_epoch, + }) + }; + + let epoch = map.current_epoch(); + map.insert("first".to_string(), waiter(epoch)); + let (iq_tx, _iq_rx) = oneshot::channel(); + map.insert("iq".to_string(), ResponseWaiter::Iq(iq_tx)); + + // One sweep is not enough: the waiter registered in the current epoch is + // still within its window, so an ack in flight is not discarded early. + map.drop_expired_phash(); + assert!( + map.remove("first").is_some(), + "a waiter must survive the sweep of the epoch it registered in" + ); + + // Registered before a sweep, then swept again: now it is stale. + let epoch = map.current_epoch(); + map.insert("stale".to_string(), waiter(epoch)); + map.drop_expired_phash(); + map.drop_expired_phash(); + assert!( + map.remove("stale").is_none(), + "a waiter that lived through a full sweep must be dropped" + ); + assert!( + map.remove("iq").is_some(), + "the sweep must never touch IQ waiters, which have their own cleanup" + ); +} diff --git a/src/keepalive.rs b/src/keepalive.rs index 24c722358..a8c5486a2 100644 --- a/src/keepalive.rs +++ b/src/keepalive.rs @@ -162,6 +162,14 @@ impl Client { self.spawn_retention_cleanup(sent_msg_ttl); } + // Same reason as the retention sweep above: driven by the + // tick rather than by send_keepalive, because a connection + // with steady inbound traffic takes the early return below + // and would never sweep. A phash waiter whose ack was lost + // has nothing else to remove it, and it reads as an + // outstanding IQ. + self.response_waiters_guard().drop_expired_phash(); + let last_recv = self.stats.last_data_received_ms(); // WA Web: maybeScheduleHealthCheck — only send ping when idle. diff --git a/src/msg_secret_buffer.rs b/src/msg_secret_buffer.rs index def073a85..204507cf6 100644 --- a/src/msg_secret_buffer.rs +++ b/src/msg_secret_buffer.rs @@ -152,7 +152,15 @@ pub(crate) struct MsgSecretWriteBuffer { /// later writer always carries values at least as new as any earlier /// in-flight write, and a stale upsert can never land after a fresh one. write_lock: async_lock::Mutex<()>, - drain_in_flight: AtomicBool, + /// Wake signal for the drain worker. Capacity one: the token means "there + /// is work", so a burst collapses into a single wakeup and a full channel + /// is success, not backpressure. + wake_tx: async_channel::Sender<()>, + wake_rx: async_channel::Receiver<()>, + /// Guards the one-time worker start. The worker is not permanent by + /// ownership: it holds a Weak, and the Sender above lives here, so dropping + /// the buffer closes the channel and ends the worker. + worker_started: AtomicBool, backend: Arc, runtime: Arc, /// Batches written so far; test observability for the coalescing claim. @@ -179,13 +187,16 @@ impl MsgSecretWriteBuffer { pending_limit > 0, "pending message-secret limit must be positive" ); + let (wake_tx, wake_rx) = async_channel::bounded(1); Arc::new(Self { pending: Mutex::new(HashMap::with_hasher(RandomState::new())), pending_limit, capacity_available: event_listener::Event::new(), sealed: AtomicBool::new(false), write_lock: async_lock::Mutex::new(()), - drain_in_flight: AtomicBool::new(false), + wake_tx, + wake_rx, + worker_started: AtomicBool::new(false), backend, runtime, flushed_batches: AtomicU64::new(0), @@ -300,6 +311,9 @@ impl MsgSecretWriteBuffer { /// `disconnect()` right before its final flush. pub(crate) fn seal(&self) { self.sealed.store(true, Ordering::Release); + // From here every queue writes inline, so the worker has nothing left + // to do; closing releases it without waiting for a wakeup. + self.wake_tx.close(); } /// Buffered-first read. Returns `(secret, message_ts)` like @@ -315,38 +329,42 @@ impl MsgSecretWriteBuffer { .map(|e| (e.secret.to_vec(), e.message_ts)) } + /// Signal the drain worker, starting it on first use. + /// + /// This used to spawn a task that died as soon as the map emptied, which + /// under a steady stream is one task per message: a capture arrives after + /// the previous one has already been written. A single worker woken by a + /// channel keeps the same write-behind semantics without that per-message + /// task. fn schedule_drain(self: &Arc) { - if self.drain_in_flight.swap(true, Ordering::AcqRel) { + self.start_worker(); + // Full means a wakeup is already pending, which is exactly what this + // call wanted; closed means the worker is gone with the buffer. + let _ = self.wake_tx.try_send(()); + } + + fn start_worker(self: &Arc) { + if self.worker_started.swap(true, Ordering::AcqRel) { return; } - let buffer = Arc::clone(self); + // Weak so a live worker never keeps the buffer (and through it the + // backend) alive. The Sender lives in the buffer, so a dropped buffer + // closes the channel and the worker returns; there is no permanent task + // to abort at teardown. + let weak = Arc::downgrade(self); + let wake_rx = self.wake_rx.clone(); self.runtime .spawn(Box::pin(async move { - buffer.drain_loop().await; + while wake_rx.recv().await.is_ok() { + let Some(buffer) = weak.upgrade() else { + break; + }; + buffer.flush().await; + } })) .detach(); } - async fn drain_loop(self: Arc) { - loop { - if self.flush_pending_once().await { - continue; - } - self.drain_in_flight.store(false, Ordering::Release); - // An insert may have raced the flag clear; reclaim the drain - // only if work exists and nobody else took it. - let has_work = !self - .pending - .lock() - .unwrap_or_else(|p| p.into_inner()) - .is_empty(); - if has_work && !self.drain_in_flight.swap(true, Ordering::AcqRel) { - continue; - } - return; - } - } - /// Write one snapshot of the pending map. Returns whether anything was /// pending. Idempotent against a concurrent drain: the upsert repeats /// harmlessly and [`Self::finish_batch`] only removes what was written. @@ -721,6 +739,37 @@ mod tests { } } + /// Captures spread over separate scheduler turns must all be written. The + /// drain used to be a task per insert, so nothing depended on the worker + /// outliving one batch; now a single worker serves every later capture, and + /// a worker that exited early would strand these silently. + #[tokio::test] + async fn captures_across_turns_all_reach_the_backend() { + let buf = buffer().await; + for i in 0..8u8 { + buf.queue(vec![entry( + "g@g.us", + "a@s.whatsapp.net", + &format!("T{i}"), + i, + )]) + .await; + // Let the worker drain and park before the next capture, which is + // the steady-state shape this path sees under load. + buf.wait_flushed().await; + tokio::task::yield_now().await; + } + + for i in 0..8u8 { + let stored = buf + .backend + .get_msg_secret("g@g.us", "a@s.whatsapp.net", &format!("T{i}")) + .await + .expect("backend read"); + assert_eq!(stored.as_deref(), Some(&[i; 32][..]), "entry T{i}"); + } + } + /// A secret refreshed for the same key while its predecessor is being /// written (edit recapture) must survive the predecessor's post-flush /// removal and reach the backend on the next iteration. diff --git a/src/request.rs b/src/request.rs index 7538ccdd0..6b9e7fe20 100644 --- a/src/request.rs +++ b/src/request.rs @@ -1,5 +1,6 @@ use crate::client::Client; use crate::client::ClientError; +use crate::client::ResponseWaiter; use crate::socket::error::{EncryptSendError, SocketError}; use futures::FutureExt; use std::num::NonZeroU64; @@ -190,8 +191,9 @@ impl Client { /// second (see [`RequestUtils::generate_message_id_at`]). pub(crate) fn generate_message_id_at(&self, unix_secs: u64) -> String { let device_snapshot = self.persistence_manager.get_device_snapshot(); - self.get_request_utils() - .generate_message_id_at(device_snapshot.pn.as_ref(), unix_secs) + // Associated function on purpose: building a RequestUtils here cloned + // the unique id per message, and the derivation never reads it. + RequestUtils::message_id_at(device_snapshot.pn.as_ref(), unix_secs) } fn get_request_utils(&self) -> RequestUtils { @@ -420,7 +422,9 @@ impl Client { // Explicit IDs are accepted by both InfoQuery and send_iq_node. Never // overwrite an older waiter. The per-registration generation also // prevents an older guard from removing a later reuse of this ID. - let Some(cleanup_generation) = waiters.try_insert_guarded(req_id.clone(), tx) else { + let Some(cleanup_generation) = + waiters.try_insert_guarded(req_id.clone(), ResponseWaiter::Iq(tx)) + else { wacore::telemetry::iq("error"); return Err(IqError::DuplicateRequestId(req_id)); }; @@ -484,7 +488,7 @@ impl Client { #[cfg(test)] mod tests { use super::{IQ_ID_ATTR, IQ_TAG, IqError, ResponseWaiterGuard}; - use crate::client::ResponseWaiterMap; + use crate::client::{ResponseWaiter, ResponseWaiterMap}; use std::sync::atomic::Ordering; use std::sync::{Arc, Mutex}; use wacore_binary::builder::NodeBuilder; @@ -507,7 +511,7 @@ mod tests { let (tx, _rx) = futures::channel::oneshot::channel(); client .response_waiters_guard() - .insert(request_id.to_owned(), tx); + .insert(request_id.to_owned(), ResponseWaiter::Iq(tx)); let error = client .send_iq_node( @@ -548,7 +552,7 @@ mod tests { let cleanup_generation = waiters .lock() .unwrap() - .try_insert_guarded("req-1".to_string(), tx) + .try_insert_guarded("req-1".to_string(), ResponseWaiter::Iq(tx)) .expect("unique request ID"); assert!(waiters.lock().unwrap().contains_key("req-1")); @@ -575,7 +579,7 @@ mod tests { let cleanup_generation = waiters .lock() .unwrap() - .try_insert_guarded("req-1".to_string(), tx) + .try_insert_guarded("req-1".to_string(), ResponseWaiter::Iq(tx)) .expect("unique request ID"); // Map empty = resolver already delivered + removed this request's waiter. waiters.lock().unwrap().remove("req-1"); @@ -596,7 +600,7 @@ mod tests { let old_generation = waiters .lock() .unwrap() - .try_insert_guarded("reused-id".to_string(), old_tx) + .try_insert_guarded("reused-id".to_string(), ResponseWaiter::Iq(old_tx)) .expect("initial request ID"); let old_guard = ResponseWaiterGuard { waiters: waiters.clone(), @@ -611,7 +615,7 @@ mod tests { waiters .lock() .unwrap() - .try_insert_guarded("reused-id".to_string(), new_tx) + .try_insert_guarded("reused-id".to_string(), ResponseWaiter::Iq(new_tx)) .expect("reused request ID"); drop(old_guard); @@ -628,7 +632,7 @@ mod tests { let old_generation = waiters .lock() .unwrap() - .try_insert_guarded("reused-id".to_string(), old_tx) + .try_insert_guarded("reused-id".to_string(), ResponseWaiter::Iq(old_tx)) .expect("initial request ID"); let old_guard = ResponseWaiterGuard { waiters: waiters.clone(), @@ -643,7 +647,7 @@ mod tests { waiters .lock() .unwrap() - .try_insert_guarded("reused-id".to_string(), new_tx) + .try_insert_guarded("reused-id".to_string(), ResponseWaiter::Iq(new_tx)) .expect("reused request ID"); drop(old_guard); diff --git a/src/send/mod.rs b/src/send/mod.rs index fc822222c..198cb2217 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -1226,16 +1226,13 @@ impl Client { // (same rule as the DM/group send path); a failure aborts the send. self.persist_signal_state_pre_wire().await?; - let ack = if let Some(phash) = stanza + let ack = stanza .attrs() .optional_string("phash") - .map(|s| wacore_binary::CompactString::from(s.as_ref())) - { - let rx = self.register_ack_waiter(&request_id); - Some((rx, phash)) - } else { - None - }; + .map(|s| wacore_binary::CompactString::from(s.as_ref())); + if let Some(phash) = ack.clone() { + self.register_phash_waiter(&request_id, phash, to.clone(), true); + } if let Err(e) = self.send_node(stanza).await { if ack.is_some() { @@ -1244,10 +1241,6 @@ impl Client { return Err(e.into()); } - if let Some((rx, phash)) = ack { - self.spawn_phash_validation(rx, phash, to.clone(), true, request_id.clone()); - } - self.update_sender_key_devices(&to_str, &prepared.skdm_devices) .await; drop(distribution_guard); @@ -1544,58 +1537,12 @@ impl Client { } } - /// Spawn a background task to validate phash from server ack. - /// On mismatch, invalidates sender key device cache and group info cache. - fn spawn_phash_validation( - &self, - rx: futures::channel::oneshot::Receiver>, - our_phash: wacore_binary::CompactString, - jid: Jid, - invalidate_group_cache: bool, - message_id: String, - ) { - let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) else { - return; - }; - self.runtime - .spawn(Box::pin(async move { - let ack = match wacore::runtime::timeout( - &*client.runtime, - std::time::Duration::from_secs(10), - rx, - ) - .await - { - Ok(Ok(node)) => node, - _ => { - // Remove leaked waiter to prevent keepalive suppression - client.response_waiters_guard().remove(&message_id); - return; - } - }; - // Cold path: box the heavy mismatch handler so the common - // (phash matches) spawned future stays small instead of carrying - // all the invalidation/clear awaits inline. - if let Some(server) = ack.get().get_attr("phash").map(|v| v.as_str()) - && server != our_phash - { - Box::pin(client.handle_phash_mismatch( - &jid, - &our_phash, - &server, - invalidate_group_cache, - )) - .await; - } - })) - .detach(); - } - - /// Cold path of [`spawn_phash_validation`](Self::spawn_phash_validation): the - /// server's phash disagreed with ours, so invalidate the relevant - /// device/group caches and (for groups) force sender-key redistribution. + /// Cold path of the phash check: the server's phash disagreed with ours, so + /// invalidate the relevant device/group caches and (for groups) force + /// sender-key redistribution. Spawned only on a mismatch, which is why the + /// common path costs a string comparison on the read loop. #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.phash_mismatch", level = "debug", skip_all, fields(jid = %jid.observe())))] - async fn handle_phash_mismatch( + pub(crate) async fn handle_phash_mismatch( &self, jid: &Jid, our_phash: &str, @@ -1795,15 +1742,21 @@ impl Client { // keyed by outer stanza id, so it would overwrite the original send's // waiter (either ack could resolve the wrong send, and the older timeout // could remove the replacement). The edit's own ack is best-effort. - let ack = if !borrowed_message_id + // Registered before the stanza goes out: the ack can arrive while + // send_node is still returning, and a waiter installed afterwards would + // miss it. + let ack_message_id = if !borrowed_message_id && let Some(phash) = dm_phash && let Some(msg_id) = stanza_to_send .attrs() .optional_string("id") .map(|s| s.into_owned()) { - let rx = self.register_ack_waiter(&msg_id); - Some((rx, phash, msg_id)) + // Group sends also invalidate group cache on mismatch: the server's + // participant set diverged, so the next send needs a fresh query. + let invalidate_group = tc_issue_target.is_group(); + self.register_phash_waiter(&msg_id, phash, tc_issue_target.clone(), invalidate_group); + Some(msg_id) } else { None }; @@ -1819,7 +1772,7 @@ impl Client { } if let Err(e) = self.send_node(stanza_to_send).await { - if let Some((_, _, ref msg_id)) = ack { + if let Some(ref msg_id) = ack_message_id { self.response_waiters_guard().remove(msg_id); } return Err(e.into()); @@ -1847,19 +1800,6 @@ impl Client { } } - if let Some((rx, phash, msg_id)) = ack { - // Group sends also invalidate group cache on mismatch — server's - // participant set diverged, the next send needs a fresh query. - let invalidate_group = tc_issue_target.is_group(); - self.spawn_phash_validation( - rx, - phash, - tc_issue_target.clone(), - invalidate_group, - msg_id, - ); - } - if let Some(update) = skdm_update { self.update_sender_key_devices(&update.to_str, &update.devices) .await; diff --git a/src/test_utils.rs b/src/test_utils.rs index 909880019..66c48c22e 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -389,5 +389,8 @@ pub(crate) async fn answer_iq(client: &Arc, request_id: &str, response: .await .unwrap_or_else(|_| panic!("an IQ waiter should be registered for {request_id}")); - let _ = sender.send(node_to_owned_ref(response)); + // Test helper: the map only ever holds Iq waiters in these fixtures. + if let crate::client::ResponseWaiter::Iq(sender) = sender { + let _ = sender.send(node_to_owned_ref(response)); + } } diff --git a/wacore/src/request.rs b/wacore/src/request.rs index a791822dd..e1b465f78 100644 --- a/wacore/src/request.rs +++ b/wacore/src/request.rs @@ -196,23 +196,35 @@ impl RequestUtils { /// second, so a send that already sampled the clock for its own timestamps /// derives the id from that same instant instead of reading again. pub fn generate_message_id_at(&self, user_jid: Option<&Jid>, unix_secs: u64) -> String { - let mut data = Vec::with_capacity(8 + 20 + 16); + Self::message_id_at(user_jid, unix_secs) + } - data.extend_from_slice(&unix_secs.to_be_bytes()); + /// The id derivation itself, which reads nothing from `self`. Exposed + /// separately so a caller on the send path does not have to materialize a + /// `RequestUtils` (and clone the unique id inside it) just to name a + /// message. + pub fn message_id_at(user_jid: Option<&Jid>, unix_secs: u64) -> String { + // Fed straight into the digest instead of through a staging Vec: this + // runs once per outgoing message and the input is never needed as a + // contiguous buffer. + let mut hasher = Sha256::new(); + hasher.update(unix_secs.to_be_bytes()); if let Some(jid) = user_jid { - data.extend_from_slice(jid.user.as_bytes()); - data.extend_from_slice(b"@"); - data.extend_from_slice(LEGACY_USER_SERVER.as_bytes()); + hasher.update(jid.user.as_bytes()); + hasher.update(b"@"); + hasher.update(LEGACY_USER_SERVER.as_bytes()); } + // The thread-local generator directly: seeding a fresh StdRng per + // message ran a full ChaCha key schedule to produce 16 bytes. let mut random_bytes = [0u8; 16]; - rand::make_rng::().fill_bytes(&mut random_bytes); - data.extend_from_slice(&random_bytes); + rand::rng().fill_bytes(&mut random_bytes); + hasher.update(random_bytes); const HEX_UPPER: &[u8; 16] = b"0123456789ABCDEF"; - let hash = Sha256::digest(&data); + let hash = hasher.finalize(); let truncated = &hash[..9]; // WA Web message IDs are "3EB0" + 18 hex chars (9-byte truncated hash) @@ -390,3 +402,48 @@ mod iq_error_tests { } } } + +#[cfg(test)] +mod message_id_tests { + use super::RequestUtils; + use wacore_binary::Jid; + + fn jid() -> Jid { + "13135550100@s.whatsapp.net".parse().expect("valid jid") + } + + /// The wire format is what WA Web produces; a client that drifts from it is + /// identifiable as non-official, so it is pinned rather than left implicit. + #[test] + fn message_id_keeps_the_wa_web_shape() { + let id = RequestUtils::message_id_at(Some(&jid()), 1_700_000_000); + assert_eq!(id.len(), 22, "id must be 3EB0 plus 18 hex chars: {id}"); + assert!(id.starts_with("3EB0"), "id must carry the WA prefix: {id}"); + assert!( + id[4..] + .chars() + .all(|c| c.is_ascii_digit() || ('A'..='F').contains(&c)), + "id must be upper-case hex after the prefix: {id}" + ); + } + + /// Two ids from the same second and JID must still differ: the entropy comes + /// from the random block, not from the clock. + #[test] + fn message_id_is_unique_within_one_second() { + let jid = jid(); + let ids: std::collections::HashSet = (0..64) + .map(|_| RequestUtils::message_id_at(Some(&jid), 1_700_000_000)) + .collect(); + assert_eq!(ids.len(), 64, "ids repeated within the same second"); + } + + /// The JID is optional on this path (pre-pairing sends); dropping it must + /// not change the shape. + #[test] + fn message_id_without_jid_keeps_the_shape() { + let id = RequestUtils::message_id_at(None, 1_700_000_000); + assert_eq!(id.len(), 22); + assert!(id.starts_with("3EB0")); + } +}