From e0dc4a470048ef1a8aeecfe4356289fd4f4c037e Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 03:36:37 +0000 Subject: [PATCH 01/35] feat(recv)!: batch the inbound commit pipeline during the offline drain MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Production traces showed the offline drain paying every durability cost per message while globally serialized (1-permit semaphore): one hook commit, two pending-buffer transactions, one Signal-cache flush and one spawned event per message — ~2.8ms each on a Postgres-backed hook, so a 421-message backlog spent ~1.2s in per-message round-trips. WA Web batches all of this. Its MessageProcessorCache accumulates decrypted messages (size cap + timeout + end-of-drain triggers) and createSnapshot commits per batch, in order: bulk message-table write, bulk signal-store commit under lock, THEN the aggregated receipts for the batch. Our per-message hook was the parity divergence, not the batch. This introduces an InboundCommitBatcher with the same triggers (200 msgs / 4 MiB / 3s / end-of-drain) and the same commit ordering: pending buffer (one txn, single reused encode arena) -> Signal flush -> hook -> buffer clear (one txn) -> event -> acks. Nothing is acked or observable before it is durable; a failed batch stays unacked whole and the server redelivers it. Live traffic commits immediately as a batch of one (also WA Web behavior). Every flush acquires the global processing permit so the Signal flush can never interleave with a half-processed stanza, which would otherwise persist a ratchet advance for an uncommitted message and turn its redelivery into an unrecoverable duplicate. The end-of-drain flush runs before the semaphore widens, so the batcher is provably empty on the live path. Breaking (pre-1.0): - InboundDurabilityHook::on_message(info, message) is replaced by on_messages(&[InboundMessage]) — all-or-nothing, slice order, batch of one on live traffic. - Event::Message(Arc, Arc) is replaced by Event::Messages(MessageBatch { messages: Arc<[InboundMessage]>, origin: Live | OfflineDrain }) — Baileys' messages.upsert shape. The slice is the same allocation handed to the hook: what was committed is exactly what event consumers observe, in order. - EventKind::Message -> EventKind::Messages; Event::as_message() -> message_batch()/messages(); MessageContext::from_event -> from_inbound. Bot::on_message keeps its per-message signature by fanning out the batch in order. - ProtocolStore gains store/delete_pending_inbound_batch (defaults iterate the single-row methods, so third-party backends keep working); SqliteStore overrides both with one transaction per batch. Per-message costs for a 421-message drain: 421 hook commits -> ~3, 842 pending-buffer txns -> ~6, 421 Signal flushes -> ~3, 421 event spawns per handler -> ~3, 421 encode Vecs -> ~3 reused arenas. Delivery receipts and message secrets were already batched (offline receipt aggregation, MsgSecretWriteBuffer); bare stanzas stay 1:1, which matches WA Web (its pre-ack batcher batches persistence, not sends). --- agent_docs/e2e_testing.md | 4 +- examples/benchmark.rs | 14 +- examples/durability_hook.rs | 106 +++-- src/bot.rs | 32 +- src/client.rs | 3 + src/client/lifecycle.rs | 14 + src/client/sessions.rs | 15 +- src/handlers/ib.rs | 2 +- src/message.rs | 1 + src/message/commit_batch.rs | 460 ++++++++++++++++++++ src/message/dispatch.rs | 24 +- src/message/durability.rs | 177 ++++---- src/message/receive.rs | 17 +- src/message/tests.rs | 161 ++++--- src/pdo.rs | 13 +- src/reexports_test.rs | 5 +- src/send/mod.rs | 2 +- src/test_utils.rs | 7 + src/types/durability_hook.rs | 55 +-- storages/sqlite-storage/src/sqlite_store.rs | 78 ++++ tests/e2e/src/lib.rs | 10 +- tests/e2e/tests/app_state.rs | 10 +- tests/e2e/tests/concurrent_disconnect.rs | 7 +- tests/e2e/tests/groups.rs | 8 +- tests/e2e/tests/media.rs | 119 ++--- tests/e2e/tests/memory_soak.rs | 29 +- tests/e2e/tests/messaging.rs | 22 +- tests/e2e/tests/offline_groups.rs | 51 ++- tests/e2e/tests/offline_messages.rs | 22 +- tests/e2e/tests/prekey_sessions.rs | 6 +- tests/e2e/tests/privacy_tokens.rs | 5 +- tests/e2e/tests/receipts.rs | 24 +- tests/e2e/tests/session_reuse.rs | 13 +- wacore/src/store/traits.rs | 41 ++ wacore/src/types/events.rs | 71 ++- 35 files changed, 1177 insertions(+), 451 deletions(-) create mode 100644 src/message/commit_batch.rs diff --git a/agent_docs/e2e_testing.md b/agent_docs/e2e_testing.md index 56315812c..e0d1f216a 100644 --- a/agent_docs/e2e_testing.md +++ b/agent_docs/e2e_testing.md @@ -38,7 +38,7 @@ Use `wait_for_event()` with predicates instead of arbitrary sleeps. This is both ```rust // GOOD: event-driven — returns as soon as the event arrives let event = client_b - .wait_for_event(15, |e| matches!(e, Event::Message(msg, _) if msg.conversation.as_deref() == Some("hello"))) + .wait_for_event(15, |e| e.messages().any(|m| m.message.conversation.as_deref() == Some("hello"))) .await?; // BAD: arbitrary sleep — wastes time or causes flaky failures @@ -60,7 +60,7 @@ tokio::time::sleep(Duration::from_millis(100)).await; client_a.client.send_message(jid_b.clone(), message).await?; // Client reconnects automatically and receives from offline queue -let event = client_b.wait_for_event(30, |e| matches!(e, Event::Message(..))).await?; +let event = client_b.wait_for_event(30, |e| matches!(e, Event::Messages(_))).await?; ``` For full disconnects (no auto-reconnect): diff --git a/examples/benchmark.rs b/examples/benchmark.rs index de979bf1c..a25613729 100644 --- a/examples/benchmark.rs +++ b/examples/benchmark.rs @@ -79,7 +79,7 @@ fn main() { let bot = builder .on_event_for( &[ - EventKind::Message, + EventKind::Messages, EventKind::PairingQrCode, EventKind::Connected, EventKind::LoggedOut, @@ -88,12 +88,12 @@ fn main() { let admin_scan_url = admin_scan_url.clone(); async move { match &*event { - Event::Message(msg, info) => { - if let Some(text) = msg.text_content() - && text == "ping" - { - let ctx = - MessageContext::from_arc(Arc::clone(msg), info, client); + Event::Messages(batch) => { + for m in batch.messages.iter() { + if m.message.text_content() != Some("ping") { + continue; + } + let ctx = MessageContext::from_inbound(m, Arc::clone(&client)); info!("Received text ping, sending pong..."); let pong_text = format!("pong {}", ctx.info.id); diff --git a/examples/durability_hook.rs b/examples/durability_hook.rs index a77b63f35..6d05fb4cd 100644 --- a/examples/durability_hook.rs +++ b/examples/durability_hook.rs @@ -95,57 +95,73 @@ impl InboxArchiver { #[async_trait::async_trait] impl InboundDurabilityHook for InboxArchiver { - async fn on_message( + async fn on_messages( &self, _client: Arc, - info: &MessageInfo, - message: &wa::Message, + batch: &[whatsapp_rust::types::events::InboundMessage], ) -> anyhow::Result<()> { - let key: CommitKey = ( - info.source.chat.to_string(), - info.source.sender.to_string(), - info.id.clone(), - ); - - // Idempotency: a redelivery (or a replay after a crash between commit and - // ack) can hand us the same key more than once. Check, but only record it - // as committed AFTER the durable write below succeeds. - if self - .seen - .lock() - .map_err(|_| anyhow::anyhow!("seen lock poisoned"))? - .contains(&key) + // Live traffic arrives one message at a time; an offline drain hands + // over a whole batch. Either way the commit below is a single append + + // fsync, so the durability cost amortizes over the batch. + let mut lines = String::new(); + let mut keys: Vec = Vec::with_capacity(batch.len()); { - info!("[{}] already committed, skipping (dedup)", info.id); - return Ok(()); + // Idempotency: a redelivery (or a replay after a crash between + // commit and ack) can hand us the same keys more than once. Check, + // but only record them as committed AFTER the durable write below + // succeeds. + let seen = self + .seen + .lock() + .map_err(|_| anyhow::anyhow!("seen lock poisoned"))?; + for m in batch { + let key: CommitKey = ( + m.info.source.chat.to_string(), + m.info.source.sender.to_string(), + m.info.id.clone(), + ); + if seen.contains(&key) { + info!("[{}] already committed, skipping (dedup)", m.info.id); + continue; + } + // Sanitize so the tab-delimited archive stays parseable on restart. + let preview = m + .message + .conversation + .as_deref() + .unwrap_or("") + .replace(['\t', '\n'], " "); + lines.push_str(&format!("{}\t{}\t{}\t{preview}\n", key.0, key.1, key.2)); + keys.push(key); + } } - // Sanitize so the tab-delimited archive stays parseable on restart. - let preview = message - .conversation - .as_deref() - .unwrap_or("") - .replace(['\t', '\n'], " "); - let line = format!("{}\t{}\t{}\t{preview}\n", key.0, key.1, key.2); - - // Durable commit on a blocking thread: append then fsync. Returning Ok - // only after sync_all means "safe to ack"; any error returns Err, so the - // ack is suppressed and the server redelivers the message later. The hook - // is awaited on the receive path, so the disk I/O goes to spawn_blocking. - let file = Arc::clone(&self.file); - tokio::task::spawn_blocking(move || -> std::io::Result<()> { - let mut file = file.lock().expect("file lock poisoned"); - file.write_all(line.as_bytes())?; - file.sync_all() - }) - .await - .map_err(|e| anyhow::anyhow!("archive write task failed: {e}"))??; - - self.seen - .lock() - .map_err(|_| anyhow::anyhow!("seen lock poisoned"))? - .insert(key); - info!("[{}] committed durably: {preview}", info.id); + if !keys.is_empty() { + // Durable commit on a blocking thread: append then fsync — all-or- + // nothing for the batch. Returning Ok only after sync_all means + // "safe to ack every message"; any error returns Err, so the acks + // are suppressed and the server redelivers the batch later. The + // hook is awaited on the receive path, so disk I/O goes to + // spawn_blocking. + let file = Arc::clone(&self.file); + tokio::task::spawn_blocking(move || -> std::io::Result<()> { + let mut file = file.lock().expect("file lock poisoned"); + file.write_all(lines.as_bytes())?; + file.sync_all() + }) + .await + .map_err(|e| anyhow::anyhow!("archive write task failed: {e}"))??; + + let mut seen = self + .seen + .lock() + .map_err(|_| anyhow::anyhow!("seen lock poisoned"))?; + let count = keys.len(); + for key in keys { + seen.insert(key); + } + info!("committed {count} message(s) durably in one fsync"); + } Ok(()) } } diff --git a/src/bot.rs b/src/bot.rs index 3f7e8b1cc..c05e98e5a 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -143,7 +143,7 @@ pub struct MessageContext { impl MessageContext { /// Builds a context from borrowed parts, deep-cloning `message`. Prefer - /// [`MessageContext::from_arc`]/[`MessageContext::from_event`] when an + /// [`MessageContext::from_arc`]/[`MessageContext::from_inbound`] when an /// `Arc` is already at hand (the event bus always has one). pub fn from_parts(message: &wa::Message, info: &MessageInfo, client: Arc) -> Self { Self::from_arc(Arc::new(message.clone()), info, client) @@ -157,9 +157,11 @@ impl MessageContext { } } - pub fn from_event(event: &Event, client: Arc) -> Option { - let (msg, info) = event.as_message()?; - Some(Self::from_arc(Arc::clone(msg), info, client)) + pub fn from_inbound( + inbound: &wacore::types::events::InboundMessage, + client: Arc, + ) -> Self { + Self::from_arc(Arc::clone(&inbound.message), &inbound.info, client) } #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.bot.send_message", level = "debug", skip_all, fields(chat = %self.info.source.chat.observe()), err(Debug)))] @@ -671,16 +673,26 @@ impl BotBuilder { /// Run `handler` for every incoming message, with a ready /// [`MessageContext`] (reply/react/edit helpers included). + /// + /// [`Event::Messages`] batches (one per commit during an offline drain, + /// single-message on live traffic) are fanned out here in arrival order, + /// awaiting each handler before the next — per-message bots keep their + /// ergonomics and gain in-batch ordering. pub fn on_message(self, handler: F) -> Self where F: Fn(MessageContext) -> Fut + Send + Sync + 'static, Fut: Future + Send + 'static, { - self.on_event_for(&[EventKind::Message], move |event, client| { - let fut = MessageContext::from_event(&event, client).map(&handler); + let handler = Arc::new(handler); + self.on_event_for(&[EventKind::Messages], move |event, client| { + let contexts: Vec = event + .messages() + .map(|m| MessageContext::from_inbound(m, Arc::clone(&client))) + .collect(); + let handler = Arc::clone(&handler); async move { - if let Some(fut) = fut { - fut.await + for context in contexts { + handler(context).await; } } }) @@ -1495,7 +1507,7 @@ mod tests { let handlers = vec![ RegisteredHandler { callback: noop.clone(), - interest: EventInterest::of(&[EventKind::Message]), + interest: EventInterest::of(&[EventKind::Messages]), }, RegisteredHandler { callback: noop, @@ -1504,7 +1516,7 @@ mod tests { ]; let interest = combined_interest(&handlers); - assert!(interest.wants(EventKind::Message)); + assert!(interest.wants(EventKind::Messages)); assert!(interest.wants(EventKind::PairingQrCode)); assert!(!interest.wants(EventKind::Receipt)); } diff --git a/src/client.rs b/src/client.rs index 6a4d46f86..3248f8051 100644 --- a/src/client.rs +++ b/src/client.rs @@ -358,6 +358,9 @@ pub struct Client { /// Write-behind buffer for inbound messageSecret captures; readers check /// it before the backend so the durable write can leave the receive lane. pub(crate) msg_secret_buffer: Arc, + /// Accumulates decrypted messages during the offline drain for per-batch + /// commit (WA Web MessageProcessorCache parity). + pub(crate) inbound_commit_batch: crate::message::commit_batch::InboundCommitBatcher, pub(crate) media_conn: Arc>>, pub(crate) is_logged_in: Arc, diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 9cd7fb7d7..5dda72519 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -215,6 +215,7 @@ impl Client { offline_sync_notifier: Arc::new(event_listener::Event::new()), offline_sync_completed: Arc::new(AtomicBool::new(false)), offline_receipt_buffer: std::sync::Mutex::new(Vec::new()), + inbound_commit_batch: Default::default(), history_sync_tasks_in_flight: Arc::new(AtomicUsize::new(0)), history_sync_idle_notifier: Arc::new(event_listener::Event::new()), outbound_flush: Arc::new(crate::flush_scope::FlushScope::new()), @@ -455,6 +456,9 @@ impl Client { self.is_connected.store(false, Ordering::Relaxed); self.offline_sync_completed.store(false, Ordering::Relaxed); self.clear_offline_receipt_buffer(); + // Uncommitted batch entries were never acked; the server redelivers + // them on this fresh connection. + self.inbound_commit_batch.clear(); self.offline_batch.reset(); self.outbound_flush.reopen(); @@ -573,6 +577,11 @@ impl Client { // connection-state reset (clear_offline_receipt_buffer) and the server // redelivers their messages on the next connect, where they are // re-acked fresh. + // + // Commit any accumulated drain batch first so its acks land in this + // receipt drain; entries that cannot commit stay unacked and the + // server redelivers them. + self.flush_inbound_commits_acquiring_permit().await; self.flush_offline_receipts(); // Prevent late receipt producers from escaping the drain window. self.outbound_flush.close(); @@ -631,6 +640,7 @@ impl Client { self.auto_reconnect_errors .store(Self::RECONNECT_BACKOFF_STEP, Ordering::Relaxed); + self.flush_inbound_commits_acquiring_permit().await; self.flush_offline_receipts(); self.outbound_flush.close(); self.outbound_flush @@ -656,6 +666,7 @@ impl Client { info!("Reconnecting immediately (expected disconnect)."); self.expected_disconnect.store(true, Ordering::Relaxed); + self.flush_inbound_commits_acquiring_permit().await; self.flush_offline_receipts(); self.outbound_flush.close(); self.outbound_flush @@ -751,6 +762,9 @@ impl Client { // Reset offline sync state for next connection self.offline_sync_completed.store(false, Ordering::Relaxed); self.clear_offline_receipt_buffer(); + // Same rule as receipts: uncommitted entries drop here and the server + // redelivers them on the next connect. + self.inbound_commit_batch.clear(); self.offline_batch.reset(); self.offline_sync_metrics .active diff --git a/src/client/sessions.rs b/src/client/sessions.rs index 9d0189c10..7c8501ef2 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -14,7 +14,7 @@ impl Client { /// WA Web: `WAWebOfflineResumeConst.OFFLINE_STANZA_TIMEOUT_MS = 60000` pub(crate) const DEFAULT_OFFLINE_SYNC_TIMEOUT: Duration = Duration::from_secs(60); - pub(crate) fn complete_offline_sync(&self, count: i32) { + pub(crate) async fn complete_offline_sync(&self, count: i32) { self.offline_sync_metrics .active .store(false, Ordering::Release); @@ -36,6 +36,16 @@ impl Client { .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) .is_ok() { + // Commit the tail of the drain batch BEFORE widening the semaphore: + // acquiring the still-single permit serializes with the last + // in-flight stanza, and post-flip dispatches take the live path, so + // the batcher is provably empty from here on. Receipts flush after, + // so every receipt's message is durably committed first (WA Web's + // createSnapshot ordering). + if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { + client.flush_inbound_commits_acquiring_permit().await; + } + // Allow parallel message processing now that offline sync is done. // During offline sync, permits=1 serialized all message processing. // Replace with a new semaphore with 64 permits for concurrent processing. @@ -100,7 +110,8 @@ impl Client { processed, expected, ); - self.complete_offline_sync(i32::try_from(processed).unwrap_or(i32::MAX)); + self.complete_offline_sync(i32::try_from(processed).unwrap_or(i32::MAX)) + .await; } } diff --git a/src/handlers/ib.rs b/src/handlers/ib.rs index 9b43adac7..01598d8c4 100644 --- a/src/handlers/ib.rs +++ b/src/handlers/ib.rs @@ -184,7 +184,7 @@ async fn handle_ib_impl(client: Arc, node: &wacore_binary::NodeRef<'_>) let count = attrs.optional_u64("count").unwrap_or(0) as i32; debug!(target: "Client/OfflineSync", "Offline sync completed, received {} items", count); - client.complete_offline_sync(count); + client.complete_offline_sync(count).await; let client_clone = Arc::clone(&client); // Per-connection: the offline flush is tied to THIS connection. diff --git a/src/message.rs b/src/message.rs index 8018e50d7..d2a8f29ee 100644 --- a/src/message.rs +++ b/src/message.rs @@ -133,6 +133,7 @@ fn decrypt_fail_log_level(mode: crate::types::events::DecryptFailMode) -> log::L pub(crate) use wacore::protocol::retry::RetryReason; +pub(crate) mod commit_batch; mod dispatch; mod durability; mod msg_secret; diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs new file mode 100644 index 000000000..6d645d571 --- /dev/null +++ b/src/message/commit_batch.rs @@ -0,0 +1,460 @@ +//! Inbound commit batcher: accumulates decrypted messages during the offline +//! drain and commits them per batch — pending-inbound buffer, Signal-cache +//! flush, durability hook, event dispatch and acks all amortize over the +//! batch. Mirrors WA Web's `MessageProcessorCache` (`createSnapshot`: bulk +//! message-table write → bulk signal-store commit → aggregate receipts), with +//! the same flush triggers: batch size, timeout, and end-of-drain. +//! +//! Live messages bypass accumulation and commit as a batch of one, which is +//! also WA Web behavior (the same pipeline with an immediate flush). + +use super::*; +use portable_atomic::AtomicU64; +use std::sync::atomic::Ordering; +use wacore::store::traits::{PendingInboundKey, PendingInboundRow}; +use wacore::types::events::{BatchOrigin, InboundMessage, MessageBatch}; + +/// WA Web pulls the offline backlog in server batches of 200 +/// (`DEFAULT_MAX_BATCH_SIZE`); one commit per server batch is the natural +/// granularity. +const MAX_BATCH_MESSAGES: usize = 200; +/// Byte cap so a media-heavy backlog cannot hold multi-MB protos in memory; +/// WA Web caps by count only, we are stricter. +const MAX_BATCH_BYTES: usize = 4 * 1024 * 1024; +/// WA Web's offline pre-ack batcher uses `delayMs: 3000`; the message cache +/// timeout is an AB prop of the same magnitude. +const FLUSH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(3); + +#[derive(Default)] +struct BatchState { + entries: Vec, + /// Sum of `message_encoded_len` of the entries, for the byte cap. + bytes: usize, + timer_armed: bool, +} + +pub(crate) struct InboundCommitBatcher { + state: std::sync::Mutex, + /// Bumped on every take; a timer that observes a stale epoch stands down. + epoch: AtomicU64, + /// Serializes commit sequences so batches reach the hook in accumulation + /// order. The guard doubles as the reusable encode arena. + arena: async_lock::Mutex>, +} + +impl Default for InboundCommitBatcher { + fn default() -> Self { + Self { + state: std::sync::Mutex::new(BatchState::default()), + epoch: AtomicU64::new(0), + arena: async_lock::Mutex::new(Vec::new()), + } + } +} + +impl InboundCommitBatcher { + fn lock(&self) -> std::sync::MutexGuard<'_, BatchState> { + match self.state.lock() { + Ok(guard) => guard, + Err(poison) => poison.into_inner(), + } + } + + /// Take the accumulated batch, invalidating any armed timer. + fn take(&self) -> Vec { + let mut state = self.lock(); + self.epoch.fetch_add(1, Ordering::AcqRel); + state.bytes = 0; + state.timer_armed = false; + std::mem::take(&mut state.entries) + } + + /// Drop accumulated entries without committing (connection teardown). + /// Uncommitted messages were never acked, so the server redelivers them. + pub(crate) fn clear(&self) { + let dropped = self.take(); + if !dropped.is_empty() { + log::debug!( + "Dropping {} uncommitted inbound messages; the server will redeliver them", + dropped.len() + ); + } + } +} + +impl Client { + /// Queue a decrypted message for the next batch commit. Returns the armed + /// timer epoch when this push started a fresh batch (the caller spawns the + /// timeout flush), `None` otherwise. + fn enqueue_inbound_commit(&self, item: InboundMessage) -> Option { + let batcher = &self.inbound_commit_batch; + let mut state = batcher.lock(); + state.bytes += waproto::codec::message_encoded_len(&item.message); + state.entries.push(item); + if !state.timer_armed { + state.timer_armed = true; + Some(batcher.epoch.load(Ordering::Acquire)) + } else { + None + } + } + + /// Batch a message decrypted while the offline drain is active, or commit + /// immediately (batch of one) on the live path. + pub(crate) async fn commit_or_batch_inbound(self: &Arc, item: InboundMessage) { + if self.offline_sync_completed.load(Ordering::Relaxed) { + self.commit_inbound_batch(vec![item], BatchOrigin::Live, false) + .await; + return; + } + if let Some(epoch) = self.enqueue_inbound_commit(item) { + let client = self.clone(); + self.runtime + .spawn(Box::pin(async move { + client.runtime.sleep(FLUSH_TIMEOUT).await; + if client.inbound_commit_batch.epoch.load(Ordering::Acquire) == epoch { + client.flush_inbound_commits_acquiring_permit().await; + } + })) + .detach(); + } + } + + /// Size/byte-cap check, run at the end of stanza processing while the + /// global processing permit is still held (so the Signal flush inside the + /// commit cannot interleave with a half-processed stanza). + pub(crate) async fn maybe_flush_inbound_commits(self: &Arc) { + let over = { + let state = self.inbound_commit_batch.lock(); + state.entries.len() >= MAX_BATCH_MESSAGES || state.bytes >= MAX_BATCH_BYTES + }; + if over { + let batch = self.inbound_commit_batch.take(); + self.commit_inbound_batch(batch, BatchOrigin::OfflineDrain, true) + .await; + } + } + + /// Flush after acquiring a global processing permit, so no stanza is + /// mid-decrypt when the Signal cache is flushed: a crash could otherwise + /// persist a ratchet advance for a message no batch has committed, turning + /// its redelivery into an unrecoverable duplicate. During the drain the + /// semaphore holds a single permit, so this fully serializes with stanza + /// processing; after the drain the batcher is empty and this no-ops. + pub(crate) async fn flush_inbound_commits_acquiring_permit(self: &Arc) { + let _permit = loop { + let (generation, semaphore) = self.read_message_semaphore(); + let permit = semaphore.acquire_arc().await; + if generation + == self + .message_semaphore_generation + .load(std::sync::atomic::Ordering::SeqCst) + { + break permit; + } + drop(permit); + }; + let batch = self.inbound_commit_batch.take(); + if batch.is_empty() { + return; + } + self.commit_inbound_batch(batch, BatchOrigin::OfflineDrain, true) + .await; + } + + /// Commit one batch: durable buffer → Signal flush → hook → clear buffer → + /// event → acks. WA Web ordering (`createSnapshot`), so nothing is acked or + /// observable before it is durable. On any commit failure everything stays + /// unacked and the server redelivers the whole batch. + #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.recv.commit_batch", level = "debug", skip_all, fields(count = items.len())))] + pub(crate) async fn commit_inbound_batch( + self: &Arc, + items: Vec, + origin: BatchOrigin, + flush_signal: bool, + ) { + if items.is_empty() { + return; + } + + if let Some(hook) = self.inbound_durability_hook() { + // Key strings live for the whole commit; rows borrow them and the + // encode arena, so the batch write allocates nothing per row + // beyond these. + let keys: Vec<(String, String)> = items + .iter() + .map(|m| { + ( + m.info.source.chat.to_string(), + m.info.source.sender.to_string(), + ) + }) + .collect(); + + let mut arena = self.inbound_commit_batch.arena.lock().await; + arena.clear(); + let mut ranges = Vec::with_capacity(items.len()); + for item in &items { + let start = arena.len(); + waproto::codec::message_encode_into(&item.message, &mut arena); + ranges.push(start..arena.len()); + } + let rows: Vec> = items + .iter() + .zip(&keys) + .zip(&ranges) + .map(|((item, (chat, sender)), range)| PendingInboundRow { + chat, + sender, + id: &item.info.id, + message: &arena[range.clone()], + }) + .collect(); + + let backend = self.persistence_manager.backend(); + // Fail closed: without a durable buffered copy, do not run the hook + // and do not ack — the server redelivers once storage recovers. + if let Err(e) = backend.store_pending_inbound_batch(&rows).await { + log::error!( + "Failed to buffer inbound batch of {}; suppressing acks for redelivery: {e:?}", + items.len() + ); + return; + } + drop(rows); + + if flush_signal { + self.flush_signal_cache_logged("commit_batch", None).await; + } + + if let Err(e) = hook.on_messages(self.clone(), &items).await { + log::warn!( + "Inbound durability hook failed for batch of {}; suppressing acks for redelivery: {e:?}", + items.len() + ); + return; + } + + let delete_keys: Vec> = items + .iter() + .zip(&keys) + .map(|(item, (chat, sender))| PendingInboundKey { + chat, + sender, + id: &item.info.id, + }) + .collect(); + if let Err(e) = backend.delete_pending_inbound_batch(&delete_keys).await { + // Leftover rows replay as duplicates; the idempotent hook + // re-commits and the replay path clears them. + log::debug!( + "Failed to clear {} buffered inbound messages: {e:?}", + delete_keys.len() + ); + } + } else if flush_signal { + self.flush_signal_cache_logged("commit_batch", None).await; + } + + let batch = MessageBatch { + messages: items.into(), + origin, + }; + self.core.event_bus.dispatch(Event::Messages(batch.clone())); + for item in batch.messages.iter() { + self.ack_received_message(&item.info); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::create_test_client_with_failing_http; + use crate::types::durability_hook::InboundDurabilityHook; + use crate::types::message::{MessageInfo, MessageSource}; + use std::sync::Mutex; + use wacore::types::events::ChannelEventHandler; + + struct RecordingHook { + batches: Mutex>>, + } + + #[async_trait::async_trait] + impl InboundDurabilityHook for RecordingHook { + async fn on_messages( + &self, + _client: Arc, + batch: &[InboundMessage], + ) -> anyhow::Result<()> { + self.batches + .lock() + .expect("hook lock") + .push(batch.iter().map(|m| m.info.id.clone()).collect()); + Ok(()) + } + } + + fn item(id: &str) -> InboundMessage { + InboundMessage { + message: Arc::new(wa::Message { + conversation: Some(format!("text {id}")), + ..Default::default() + }), + info: Arc::new(MessageInfo { + id: id.to_string(), + source: MessageSource { + chat: "100@g.us".parse().unwrap(), + sender: "200@s.whatsapp.net".parse().unwrap(), + ..Default::default() + }, + ..Default::default() + }), + } + } + + // During the drain, messages accumulate and one flush commits them all in + // arrival order as a single hook call and a single OfflineDrain event. + #[tokio::test] + async fn drain_accumulates_then_commits_in_order() { + let client = create_test_client_with_failing_http("batch_drain").await; + let hook = Arc::new(RecordingHook { + batches: Mutex::new(Vec::new()), + }); + let _ = client.inbound_durability_hook.set(hook.clone()); + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + + client + .offline_sync_completed + .store(false, std::sync::atomic::Ordering::Relaxed); + for id in ["B1", "B2", "B3"] { + client.commit_or_batch_inbound(item(id)).await; + } + assert!( + hook.batches.lock().expect("hook lock").is_empty(), + "sub-threshold entries must accumulate, not commit" + ); + + client.flush_inbound_commits_acquiring_permit().await; + + let batches = hook.batches.lock().expect("hook lock").clone(); + assert_eq!(batches, vec![vec!["B1", "B2", "B3"]]); + + let event = rx.try_recv().expect("one batch event"); + let batch = event.message_batch().expect("Messages event"); + assert_eq!(batch.origin, BatchOrigin::OfflineDrain); + let ids: Vec<&str> = batch.messages.iter().map(|m| m.info.id.as_str()).collect(); + assert_eq!(ids, ["B1", "B2", "B3"]); + assert!(rx.try_recv().is_err(), "exactly one event for the batch"); + + // A committed batch leaves no buffered copies behind. + let backend = client.persistence_manager.backend(); + for id in ["B1", "B2", "B3"] { + assert!( + backend + .get_pending_inbound("100@g.us", "200@s.whatsapp.net", id) + .await + .unwrap() + .is_none() + ); + } + } + + // Live traffic commits immediately as a batch of one. + #[tokio::test] + async fn live_commits_as_batch_of_one() { + let client = create_test_client_with_failing_http("batch_live").await; + let hook = Arc::new(RecordingHook { + batches: Mutex::new(Vec::new()), + }); + let _ = client.inbound_durability_hook.set(hook.clone()); + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + client + .offline_sync_completed + .store(true, std::sync::atomic::Ordering::Relaxed); + + client.commit_or_batch_inbound(item("L1")).await; + + assert_eq!( + hook.batches.lock().expect("hook lock").clone(), + vec![vec!["L1"]] + ); + let event = rx.try_recv().expect("live event"); + let batch = event.message_batch().expect("Messages event"); + assert_eq!(batch.origin, BatchOrigin::Live); + assert_eq!(batch.messages.len(), 1); + } + + // The size trigger commits a full batch from the stanza-end check. + #[tokio::test] + async fn size_trigger_flushes_full_batch() { + let client = create_test_client_with_failing_http("batch_size").await; + client + .offline_sync_completed + .store(false, std::sync::atomic::Ordering::Relaxed); + let hook = Arc::new(RecordingHook { + batches: Mutex::new(Vec::new()), + }); + let _ = client.inbound_durability_hook.set(hook.clone()); + + for i in 0..MAX_BATCH_MESSAGES { + client.commit_or_batch_inbound(item(&format!("S{i}"))).await; + } + client.maybe_flush_inbound_commits().await; + + let batches = hook.batches.lock().expect("hook lock").clone(); + assert_eq!(batches.len(), 1, "one commit for the full batch"); + assert_eq!(batches[0].len(), MAX_BATCH_MESSAGES); + assert_eq!(batches[0][0], "S0"); + assert_eq!(batches[0][MAX_BATCH_MESSAGES - 1], "S199"); + } + + // Without a hook, the drain still batches the event dispatch. + #[tokio::test] + async fn drain_without_hook_batches_events() { + let client = create_test_client_with_failing_http("batch_no_hook").await; + client + .offline_sync_completed + .store(false, std::sync::atomic::Ordering::Relaxed); + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + + client.commit_or_batch_inbound(item("N1")).await; + client.commit_or_batch_inbound(item("N2")).await; + client.flush_inbound_commits_acquiring_permit().await; + + let event = rx.try_recv().expect("one batch event"); + assert_eq!( + event + .messages() + .map(|m| m.info.id.as_str()) + .collect::>(), + ["N1", "N2"] + ); + } + + // clear() drops uncommitted entries: no hook call, no event, and the + // pending buffer was never written (the server redelivers instead). + #[tokio::test] + async fn clear_drops_uncommitted_entries() { + let client = create_test_client_with_failing_http("batch_clear").await; + client + .offline_sync_completed + .store(false, std::sync::atomic::Ordering::Relaxed); + let hook = Arc::new(RecordingHook { + batches: Mutex::new(Vec::new()), + }); + let _ = client.inbound_durability_hook.set(hook.clone()); + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + + client.commit_or_batch_inbound(item("C1")).await; + client.inbound_commit_batch.clear(); + client.flush_inbound_commits_acquiring_permit().await; + + assert!(hook.batches.lock().expect("hook lock").is_empty()); + assert!(rx.try_recv().is_err()); + } +} diff --git a/src/message/dispatch.rs b/src/message/dispatch.rs index e28731140..d4e2aaa00 100644 --- a/src/message/dispatch.rs +++ b/src/message/dispatch.rs @@ -38,21 +38,15 @@ impl Client { } let dispatch_msg = Arc::new(decrypted.unwrap_or(msg)); - if let Some(hook) = self.inbound_durability_hook() { - // At-least-once: in-process handlers still fire, but the transport - // ack is deferred until the hook durably commits the message. - self.core - .event_bus - .dispatch(Event::Message(Arc::clone(&dispatch_msg), Arc::clone(&info))); - self.run_inbound_durability_hook(hook, &info, &dispatch_msg) - .await; - } else { - // Default at-most-once path (unchanged): ack, then dispatch. - self.ack_received_message(&info); - self.core - .event_bus - .dispatch(Event::Message(dispatch_msg, info)); - } + // Live traffic commits (and acks) as a batch of one; during the + // offline drain the message joins the accumulating commit batch and + // the event/ack fire only after its batch commits. Either way the + // hook (when registered) gates everything observable. + self.commit_or_batch_inbound(wacore::types::events::InboundMessage { + message: dispatch_msg, + info, + }) + .await; } /// Acknowledge a received message so the server drops it from the offline diff --git a/src/message/durability.rs b/src/message/durability.rs index e73c55038..434bf44bf 100644 --- a/src/message/durability.rs +++ b/src/message/durability.rs @@ -1,9 +1,15 @@ //! Inbound durability hook: opt-in at-least-once delivery by gating the //! transport ack on a consumer-provided durable commit. See //! [`crate::types::durability_hook::InboundDurabilityHook`] for the contract. +//! +//! The first-receipt path lives in [`super::commit_batch`]: messages commit +//! per batch (buffer → hook → ack). This module keeps the redelivery replay, +//! which is inherently per-message: each replayed stanza resolves against its +//! own buffered copy. use super::*; use crate::types::durability_hook::InboundDurabilityHook; +use wacore::types::events::InboundMessage; impl Client { /// The registered inbound durability hook, if any. `None` (default) keeps @@ -12,60 +18,6 @@ impl Client { self.inbound_durability_hook.get().cloned() } - /// First-receipt path: buffer the decrypted message durably, run the hook, - /// and ack only on success. On failure the buffered row is left for the - /// server's redelivery to retry. Replaces the plain ack for a freshly - /// dispatched user message when a hook is configured. - pub(crate) async fn run_inbound_durability_hook( - self: &Arc, - hook: Arc, - info: &Arc, - msg: &Arc, - ) { - let backend = self.persistence_manager.backend(); - let chat = info.source.chat.to_string(); - let sender = info.source.sender.to_string(); - // Persist before the Signal ratchet is flushed (which happens after this - // returns) so a crash mid-commit replays the message instead of losing it. - // `message_to_vec` is the shared non-generic encoder so this call does not - // monomorphize the whole `wa::Message` proto tree into this crate. - let bytes = waproto::codec::message_to_vec(msg); - // Fail closed: if we cannot durably buffer the message, do not run the - // hook and do not ack. The server keeps it queued and redelivers it once - // storage recovers, rather than us acking a message we cannot replay. - if let Err(e) = backend - .store_pending_inbound(&chat, &sender, &info.id, &bytes) - .await - { - log::error!( - "[msg:{}] failed to buffer inbound message; suppressing ack for redelivery: {e:?}", - info.id - ); - return; - } - - match hook.on_message(self.clone(), info, msg).await { - Ok(()) => { - if let Err(e) = backend - .delete_pending_inbound(&chat, &sender, &info.id) - .await - { - log::debug!( - "[msg:{}] failed to clear buffered inbound message: {e:?}", - info.id - ); - } - self.ack_received_message(info); - } - Err(e) => { - log::warn!( - "[msg:{}] inbound durability hook failed; suppressing ack for redelivery: {e:?}", - info.id - ); - } - } - } - /// Redelivery path: when the server replays an already-decrypted message /// (`DuplicatedMessage`), re-run the hook from the buffered copy instead of /// acking. A plain ack is sent only for a genuine duplicate (no buffered @@ -80,8 +32,14 @@ impl Client { Ok(Some(bytes)) => { match waproto::codec::message_decode(&bytes) { Ok(msg) => { - let msg = Arc::new(msg); - match hook.on_message(self.clone(), info, &msg).await { + let item = InboundMessage { + message: Arc::new(msg), + info: Arc::clone(info), + }; + match hook + .on_messages(self.clone(), std::slice::from_ref(&item)) + .await + { Ok(()) => { if let Err(e) = backend .delete_pending_inbound(&chat, &sender, &info.id) @@ -137,21 +95,23 @@ mod tests { use crate::test_utils::create_test_client_with_failing_http; use crate::types::message::MessageInfo; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use wacore::types::events::BatchOrigin; struct CountingHook { calls: AtomicUsize, + messages: AtomicUsize, succeed: AtomicBool, } #[async_trait::async_trait] impl InboundDurabilityHook for CountingHook { - async fn on_message( + async fn on_messages( &self, _client: Arc, - _info: &MessageInfo, - _message: &wa::Message, + batch: &[InboundMessage], ) -> anyhow::Result<()> { self.calls.fetch_add(1, Ordering::SeqCst); + self.messages.fetch_add(batch.len(), Ordering::SeqCst); if self.succeed.load(Ordering::SeqCst) { Ok(()) } else { @@ -160,6 +120,14 @@ mod tests { } } + fn counting_hook(succeed: bool) -> Arc { + Arc::new(CountingHook { + calls: AtomicUsize::new(0), + messages: AtomicUsize::new(0), + succeed: AtomicBool::new(succeed), + }) + } + fn test_info(id: &str) -> Arc { use crate::types::message::MessageSource; Arc::new(MessageInfo { @@ -173,66 +141,70 @@ mod tests { }) } - fn test_msg() -> Arc { - Arc::new(wa::Message { - conversation: Some("hello".to_string()), - ..Default::default() - }) + fn test_item(id: &str) -> InboundMessage { + InboundMessage { + message: Arc::new(wa::Message { + conversation: Some("hello".to_string()), + ..Default::default() + }), + info: test_info(id), + } } - // A successful hook acks the message and clears the buffered copy. + // A successful batch commit acks the messages and clears every buffered copy. #[tokio::test] - async fn hook_ok_clears_buffer() { + async fn commit_ok_clears_buffer() { let client = create_test_client_with_failing_http("durability_ok").await; - let hook = Arc::new(CountingHook { - calls: AtomicUsize::new(0), - succeed: AtomicBool::new(true), - }); + let hook = counting_hook(true); let _ = client.inbound_durability_hook.set(hook.clone()); - let info = test_info("MSG_OK"); + let items = vec![test_item("MSG_OK_1"), test_item("MSG_OK_2")]; + let infos: Vec<_> = items.iter().map(|i| Arc::clone(&i.info)).collect(); client - .run_inbound_durability_hook( - client.inbound_durability_hook().unwrap(), - &info, - &test_msg(), - ) + .commit_inbound_batch(items, BatchOrigin::OfflineDrain, false) .await; - assert_eq!(hook.calls.load(Ordering::SeqCst), 1); + assert_eq!(hook.calls.load(Ordering::SeqCst), 1, "one commit per batch"); + assert_eq!(hook.messages.load(Ordering::SeqCst), 2); let backend = client.persistence_manager.backend(); - assert!( - backend - .get_pending_inbound( - &info.source.chat.to_string(), - &info.source.sender.to_string(), - "MSG_OK", - ) - .await - .unwrap() - .is_none(), - "a committed message must not stay buffered" - ); + for info in &infos { + assert!( + backend + .get_pending_inbound( + &info.source.chat.to_string(), + &info.source.sender.to_string(), + &info.id, + ) + .await + .unwrap() + .is_none(), + "a committed message must not stay buffered" + ); + } } - // A failing hook suppresses the ack and keeps the buffered copy; a later - // redelivery re-runs the hook and, once it succeeds, clears the buffer. + // A failing batch commit suppresses the acks and keeps every buffered copy; + // later per-message redeliveries replay each one and, once the hook + // succeeds, clear them. #[tokio::test] - async fn hook_err_keeps_buffer_then_replays() { + async fn commit_err_keeps_buffer_then_replays() { let client = create_test_client_with_failing_http("durability_err").await; - let hook = Arc::new(CountingHook { - calls: AtomicUsize::new(0), - succeed: AtomicBool::new(false), - }); + let hook = counting_hook(false); let _ = client.inbound_durability_hook.set(hook.clone()); let backend = client.persistence_manager.backend(); let info = test_info("MSG_ERR"); client - .run_inbound_durability_hook( - client.inbound_durability_hook().unwrap(), - &info, - &test_msg(), + .commit_inbound_batch( + vec![InboundMessage { + message: Arc::new(wa::Message { + conversation: Some("hello".to_string()), + ..Default::default() + }), + info: Arc::clone(&info), + }], + BatchOrigin::OfflineDrain, + false, ) .await; @@ -292,10 +264,7 @@ mod tests { #[tokio::test] async fn replay_without_buffer_just_acks() { let client = create_test_client_with_failing_http("durability_dup").await; - let hook = Arc::new(CountingHook { - calls: AtomicUsize::new(0), - succeed: AtomicBool::new(true), - }); + let hook = counting_hook(true); let _ = client.inbound_durability_hook.set(hook.clone()); client.ack_or_replay_to_hook(&test_info("MSG_NONE")).await; diff --git a/src/message/receive.rs b/src/message/receive.rs index 37e3f3d39..e1bf3273c 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -528,9 +528,20 @@ impl Client { self.handle_msmsg_payload(&info, payload).await; } - // Flush cached Signal state to DB (matches WA Web's flushBufferToDiskIfNotMemOnlyMode) - self.flush_signal_cache_logged("message", Some(&info.id)) - .await; + // Live: flush cached Signal state per stanza (WA Web's + // flushBufferToDiskIfNotMemOnlyMode). During the offline drain the + // commit batcher owns the flush — one per batch, before any ack (WA + // Web's bulk signal-store snapshot) — so here only the batch size/byte + // triggers are checked, while the global permit is still held. + if self + .offline_sync_completed + .load(std::sync::atomic::Ordering::Relaxed) + { + self.flush_signal_cache_logged("message", Some(&info.id)) + .await; + } else { + self.maybe_flush_inbound_commits().await; + } } #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.recv.session_decrypt", level = "debug", skip_all, fields(chat = %info.source.chat.observe(), sender = %sender_encryption_jid.observe(), msg_id = %info.id)))] diff --git a/src/message/tests.rs b/src/message/tests.rs index e911b5d29..1d26d9d4c 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -5533,6 +5533,8 @@ async fn capturing_client( // other layers but not on this path. *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); seed_test_pn(&client).await; + // Live-path semantics by default; drain tests reset the flag themselves. + client.offline_sync_completed.store(true, Ordering::Relaxed); (client, transport) } @@ -6323,11 +6325,9 @@ fn message_events_for_id(rx: &async_channel::Receiver>, id: &str) -> let mut count = 0; let mut visible_content = 0; while let Ok(event) = rx.try_recv() { - if let Event::Message(msg, info) = event.as_ref() - && info.id == id - { + for m in event.messages().filter(|m| m.info.id == id) { count += 1; - if msg.conversation.is_some() { + if m.message.conversation.is_some() { visible_content += 1; } } @@ -6338,11 +6338,10 @@ fn message_events_for_id(rx: &async_channel::Receiver>, id: &str) -> fn message_texts_for_id(rx: &async_channel::Receiver>, id: &str) -> Vec { let mut texts = Vec::new(); while let Ok(event) = rx.try_recv() { - if let Event::Message(msg, info) = event.as_ref() - && info.id == id - && let Some(text) = &msg.conversation - { - texts.push(text.clone()); + for m in event.messages().filter(|m| m.info.id == id) { + if let Some(text) = &m.message.conversation { + texts.push(text.clone()); + } } } texts @@ -7687,10 +7686,12 @@ async fn secret_encrypted_message_edit_dispatches_legacy_edit() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == edit_id - && legacy_edit_text(msg.as_ref()) == Some("edited") - && msg.secret_encrypted_message.is_unset()) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == edit_id + && legacy_edit_text(msg.as_ref()) == Some("edited") + && msg.secret_encrypted_message.is_unset() + }) }, 500, ) @@ -7748,10 +7749,12 @@ async fn secret_encrypted_peer_edit_resolves_sender_from_envelope() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == edit_id - && legacy_edit_text(msg.as_ref()) == Some("edited") - && msg.secret_encrypted_message.is_unset()) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == edit_id + && legacy_edit_text(msg.as_ref()) == Some("edited") + && msg.secret_encrypted_message.is_unset() + }) }, 500, ) @@ -7812,10 +7815,12 @@ async fn run_secret_edit_with_window(test_id: &str, parent_ts: i64, edit_offset: &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == edit_id - && legacy_edit_text(msg.as_ref()) == Some("edited") - && msg.secret_encrypted_message.is_unset()) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == edit_id + && legacy_edit_text(msg.as_ref()) == Some("edited") + && msg.secret_encrypted_message.is_unset() + }) }, 500, ) @@ -7945,10 +7950,12 @@ async fn secret_encrypted_edit_decrypts_via_resolver_when_store_empty() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == edit_id - && legacy_edit_text(msg.as_ref()) == Some("edited via resolver") - && msg.secret_encrypted_message.is_unset()) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == edit_id + && legacy_edit_text(msg.as_ref()) == Some("edited via resolver") + && msg.secret_encrypted_message.is_unset() + }) }, 500, ) @@ -8037,9 +8044,10 @@ async fn decrypted_message_edit_recaptures_secret_for_next_edit() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == "EDIT_CHAIN_2" - && legacy_edit_text(msg.as_ref()) == Some("second")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == "EDIT_CHAIN_2" && legacy_edit_text(msg.as_ref()) == Some("second") + }) }, 500, ) @@ -8113,9 +8121,10 @@ async fn secret_encrypted_message_edit_uses_lid_pn_fallback_in_group() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == "GROUP_EDIT_1" - && legacy_edit_text(msg.as_ref()) == Some("group edited")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == "GROUP_EDIT_1" && legacy_edit_text(msg.as_ref()) == Some("group edited") + }) }, 500, ) @@ -8235,9 +8244,11 @@ async fn decrypted_message_edit_refreshes_alternate_secret_alias() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == "GROUP_EDIT_REFRESH_2" - && legacy_edit_text(msg.as_ref()) == Some("second")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == "GROUP_EDIT_REFRESH_2" + && legacy_edit_text(msg.as_ref()) == Some("second") + }) }, 500, ) @@ -8313,9 +8324,10 @@ async fn msmsg_decrypts_when_secret_is_stored() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("hi from bot")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == bot_reply_id && msg.conversation.as_deref() == Some("hi from bot") + }) }, 1500, ) @@ -8371,12 +8383,12 @@ async fn msmsg_without_stored_secret_nacks_495() { "missing messageSecret must nack with code 495" ); assert!( - collector - .events() - .iter() - .all(|e| !matches!(e.as_ref(), wacore::types::events::Event::Message(_, info) if info.id == bot_reply_id)), - "no Message event must be dispatched when decryption failed" - ); + collector + .events() + .iter() + .all(|e| !e.messages().any(|m| m.info.id == bot_reply_id)), + "no Message event must be dispatched when decryption failed" + ); } /// Tampered ciphertext → GCM tag fails → nack 495. @@ -8511,9 +8523,10 @@ async fn msmsg_bot_edit_uses_edit_target_id_for_hkdf() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == edit_reply_id - && msg.conversation.as_deref() == Some("edited content")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == edit_reply_id && msg.conversation.as_deref() == Some("edited content") + }) }, 1500, ) @@ -8651,7 +8664,7 @@ async fn msmsg_bot_edit_first_keeps_info_id() { let got = collect_event( &client, collector, - |e| matches!(e, wacore::types::events::Event::Message(_, info) if info.id == stanza_id), + |e| e.messages().any(|m| m.info.id == stanza_id), 1500, ) .await; @@ -8730,9 +8743,10 @@ async fn msmsg_falls_back_to_info_id_when_primary_uses_edit_target() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == stanza_id - && msg.conversation.as_deref() == Some("fallback ok")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == stanza_id && msg.conversation.as_deref() == Some("fallback ok") + }) }, 1500, ) @@ -8816,9 +8830,10 @@ async fn msmsg_falls_back_to_edit_target_when_primary_uses_info_id() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == stanza_id - && msg.conversation.as_deref() == Some("inverse fallback")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == stanza_id && msg.conversation.as_deref() == Some("inverse fallback") + }) }, 1500, ) @@ -9493,9 +9508,10 @@ async fn mixed_msmsg_and_unknown_enc_still_decrypts_msmsg() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("mixed ok")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == bot_reply_id && msg.conversation.as_deref() == Some("mixed ok") + }) }, 1500, ) @@ -9591,9 +9607,10 @@ async fn msmsg_alternate_lookup_resolves_lid_to_stored_pn() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("alt ok")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == bot_reply_id && msg.conversation.as_deref() == Some("alt ok") + }) }, 1500, ) @@ -9712,9 +9729,10 @@ async fn fanout_capture_lets_subsequent_msmsg_decrypt() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("bot reply")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == bot_reply_id && msg.conversation.as_deref() == Some("bot reply") + }) }, 1500, ) @@ -9808,9 +9826,10 @@ async fn msmsg_outbound_put_and_inbound_get_match_for_lid_bot() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("lid coherent")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == bot_reply_id && msg.conversation.as_deref() == Some("lid coherent") + }) }, 1500, ) @@ -9885,9 +9904,10 @@ async fn msmsg_with_bot_device_suffix_round_trips() { &client, collector, |e| { - matches!(e, wacore::types::events::Event::Message(msg, info) - if info.id == bot_reply_id - && msg.conversation.as_deref() == Some("with device")) + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.id == bot_reply_id && msg.conversation.as_deref() == Some("with device") + }) }, 1500, ) @@ -10108,9 +10128,8 @@ async fn enc_comment_inbound_dispatches_body_with_parent_link() { let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2); 'outer: while tokio::time::Instant::now() < deadline { while let Ok(event) = rx.try_recv() { - if let Event::Message(msg, info) = event.as_ref() - && info.id == COMMENT_ID - { + if let Some(m) = event.messages().find(|m| m.info.id == COMMENT_ID) { + let (msg, info) = (&m.message, &m.info); seen = true; assert_eq!( msg.extended_text_message diff --git a/src/pdo.rs b/src/pdo.rs index 40021e39e..724b9b346 100644 --- a/src/pdo.rs +++ b/src/pdo.rs @@ -417,11 +417,18 @@ impl Client { request_id ); + // PDO recovery is event-only (its ack runs on the PDO path, not the + // message pipeline), so this bypasses the commit batcher on purpose. self.core .event_bus - .dispatch(wacore::types::events::Event::Message( - Arc::from(message), - message_info, + .dispatch(wacore::types::events::Event::Messages( + wacore::types::events::MessageBatch { + messages: std::sync::Arc::from([wacore::types::events::InboundMessage { + message: Arc::from(message), + info: message_info, + }]), + origin: wacore::types::events::BatchOrigin::Live, + }, )); } diff --git a/src/reexports_test.rs b/src/reexports_test.rs index a1354497f..b54078785 100644 --- a/src/reexports_test.rs +++ b/src/reexports_test.rs @@ -41,11 +41,10 @@ struct NoopHook; #[whatsapp_rust::async_trait] impl whatsapp_rust::InboundDurabilityHook for NoopHook { - async fn on_message( + async fn on_messages( &self, _client: std::sync::Arc, - _info: &whatsapp_rust::types::message::MessageInfo, - _message: &wa::Message, + _batch: &[whatsapp_rust::types::events::InboundMessage], ) -> whatsapp_rust::anyhow::Result<()> { Ok(()) } diff --git a/src/send/mod.rs b/src/send/mod.rs index 8b679df0f..26157dc3d 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -4050,7 +4050,7 @@ mod tests { // The test client never connects, so the send's `ensure_e2e_sessions` // would otherwise block on `wait_for_offline_delivery_end` until timeout. - client.complete_offline_sync(0); + client.complete_offline_sync(0).await; // Seed a Signal session for the peer's LID device so the offline fanout // can encrypt without fetching prekeys over the (absent) socket. The diff --git a/src/test_utils.rs b/src/test_utils.rs index 6716d9a95..64e5a7cff 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -109,6 +109,13 @@ pub async fn create_test_client_with_config( ) .await; + // Tests exercise live-path semantics by default (a fresh client starts in + // drain mode, where inbound commits batch instead of dispatching + // immediately). Drain-specific tests reset this flag themselves. + client + .offline_sync_completed + .store(true, std::sync::atomic::Ordering::Relaxed); + client } diff --git a/src/types/durability_hook.rs b/src/types/durability_hook.rs index c96063880..6f0c2b72c 100644 --- a/src/types/durability_hook.rs +++ b/src/types/durability_hook.rs @@ -1,8 +1,7 @@ use crate::client::Client; -use crate::types::message::MessageInfo; use anyhow::Result; use std::sync::Arc; -use waproto::whatsapp as wa; +pub use wacore::types::events::InboundMessage; /// Hook invoked for every decrypted inbound user message before it is /// acknowledged to the server, turning the consumer from at-most-once into @@ -12,16 +11,18 @@ use waproto::whatsapp as wa; /// queue. By default the SDK acks as soon as a message is decrypted, so a crash /// (or a failed DB write) before the consumer persists the message loses it for /// good. When a hook is registered, the ack is deferred until the hook returns -/// `Ok`: the decrypted message is buffered durably first, the hook runs, and -/// only on success is the ack sent and the buffer cleared. On `Err` (or a -/// crash) the message stays unacked and the server redelivers it on the next -/// connect, where the hook runs again from the buffered copy. +/// `Ok`: the decrypted messages are buffered durably first, the hook runs, and +/// only on success are their acks sent and the buffer cleared. On `Err` (or a +/// crash) the messages stay unacked and the server redelivers them on the next +/// connect, where the hook runs again from the buffered copies. /// /// This is at-least-once, not exactly-once: a crash after the consumer commits /// but before the ack lands replays the message, so the hook MUST be idempotent. -/// Deduplicate by the message source AND id — `(info.source.chat, -/// info.source.sender, info.id)` — not `info.id` alone: stanza ids are only -/// unique within a `(chat, sender)`, so two chats can reuse the same id. +/// A failed batch is redelivered whole, so a partially-applied batch commit +/// must also be safe to re-run. Deduplicate by the message source AND id — +/// `(info.source.chat, info.source.sender, info.id)` — not `info.id` alone: +/// stanza ids are only unique within a `(chat, sender)`, so two chats can +/// reuse the same id. /// /// Durable replay across process crashes requires a backend that implements the /// `ProtocolStore` pending-inbound methods (the bundled `SqliteStore` does). @@ -30,32 +31,36 @@ use waproto::whatsapp as wa; /// /// The hook is awaited inside the receive pipeline, so a slow hook backpressures /// inbound processing (the same trade-off as whatsmeow's synchronous ack). Do -/// not perform blocking client operations for the same sender inside it (e.g. a -/// synchronous reply) — that can deadlock against the per-sender Signal lock held -/// while a 1:1 message is processed; persist and return, and spawn any reply. +/// not perform blocking client operations for a sender present in the batch +/// (e.g. a synchronous reply) — that can deadlock against the per-sender Signal +/// lock held while a 1:1 message is processed; persist and return, and spawn +/// any reply. /// /// Scope and known limitations: /// - Covers end-to-end encrypted messages (1:1 and group). Newsletter / broadcast /// channel messages are not encrypted and are acked on their own path, so the /// hook does not gate them. /// - If the durable buffer write itself fails (e.g. disk full, after retries), -/// the ack is suppressed, but if the process does not crash the Signal ratchet -/// still advances and that one message degrades to at-most-once on its next -/// redelivery (it can no longer be decrypted, and there is no buffered copy to -/// replay). The guarantee holds whenever the buffer write succeeds. +/// the acks are suppressed, but if the process does not crash the Signal +/// ratchet still advances and those messages degrade to at-most-once on their +/// next redelivery (they can no longer be decrypted, and there is no buffered +/// copy to replay). The guarantee holds whenever the buffer write succeeds. /// - On a redelivery replay the `info` is re-parsed from the stanza, so a few /// fields derived during the first dispatch (the ephemeral timer, encrypted /// comment threading) may be absent. The `message` body is always the original. #[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))] #[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)] pub trait InboundDurabilityHook: wacore::sync_marker::MaybeSendSync { - /// Durably commit `message` (e.g. INSERT into your DB, enqueue to a broker). - /// Return `Ok(())` only after the commit is durable; the SDK then acks the - /// message. Return `Err` to suppress the ack and have the server redeliver. - async fn on_message( - &self, - client: Arc, - info: &MessageInfo, - message: &wa::Message, - ) -> Result<()>; + /// Durably commit the whole batch, all-or-nothing, in slice order (e.g. one + /// multi-row INSERT transaction). Return `Ok(())` only after the commit is + /// durable; the SDK then acks every message in the batch. Return `Err` to + /// suppress all their acks and have the server redeliver them. + /// + /// Live messages arrive as batches of one. During the offline drain the + /// SDK accumulates and commits per batch (WA Web's MessageProcessorCache + /// granularity), so one round-trip covers the lot. + /// [`Event::Messages`](wacore::types::events::Event::Messages) then + /// carries the exact same items: what this method committed is what event + /// consumers observe. + async fn on_messages(&self, client: Arc, batch: &[InboundMessage]) -> Result<()>; } diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index 962636bd1..eceff2ec0 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -3093,6 +3093,84 @@ impl ProtocolStore for SqliteStore { .await .map_err(|e| StoreError::Database(Box::new(e)))? } + + async fn store_pending_inbound_batch( + &self, + rows: &[wacore::store::traits::PendingInboundRow<'_>], + ) -> Result<()> { + if rows.is_empty() { + return Ok(()); + } + // One owned copy shared across retry attempts; a single transaction + // amortizes the WAL commit over the whole batch. + let rows: Arc)>> = Arc::new( + rows.iter() + .map(|r| { + ( + r.chat.to_string(), + r.sender.to_string(), + r.id.to_string(), + r.message.to_vec(), + ) + }) + .collect(), + ); + let device_id = self.device_id; + self.with_retry("store_pending_inbound_batch", || { + let rows = Arc::clone(&rows); + Box::new(move |conn: &mut SqliteConnection| { + conn.transaction(|conn| { + for (chat, sender, id, message) in rows.iter() { + diesel::replace_into(pending_inbound_messages::table) + .values(( + pending_inbound_messages::chat.eq(chat), + pending_inbound_messages::sender.eq(sender), + pending_inbound_messages::id.eq(id), + pending_inbound_messages::message.eq(message.as_slice()), + pending_inbound_messages::device_id.eq(device_id), + )) + .execute(conn)?; + } + Ok(()) + }) + }) + }) + .await + } + + async fn delete_pending_inbound_batch( + &self, + keys: &[wacore::store::traits::PendingInboundKey<'_>], + ) -> Result<()> { + if keys.is_empty() { + return Ok(()); + } + let keys: Arc> = Arc::new( + keys.iter() + .map(|k| (k.chat.to_string(), k.sender.to_string(), k.id.to_string())) + .collect(), + ); + let device_id = self.device_id; + self.with_retry("delete_pending_inbound_batch", || { + let keys = Arc::clone(&keys); + Box::new(move |conn: &mut SqliteConnection| { + conn.transaction(|conn| { + for (chat, sender, id) in keys.iter() { + diesel::delete( + pending_inbound_messages::table + .filter(pending_inbound_messages::chat.eq(chat)) + .filter(pending_inbound_messages::sender.eq(sender)) + .filter(pending_inbound_messages::id.eq(id)) + .filter(pending_inbound_messages::device_id.eq(device_id)), + ) + .execute(conn)?; + } + Ok(()) + }) + }) + }) + .await + } } #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] diff --git a/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index dfef5fb74..2f716cc14 100644 --- a/tests/e2e/src/lib.rs +++ b/tests/e2e/src/lib.rs @@ -398,12 +398,10 @@ impl TestClient { let gid = group_jid.clone(); let text = text.to_string(); self.wait_for_event(timeout_secs, move |e| { - matches!( - e, - Event::Message(msg, info) - if info.source.chat == gid - && msg.conversation.as_deref() == Some(text.as_str()) - ) + e.messages().any(|m| { + m.info.source.chat == gid + && m.message.conversation.as_deref() == Some(text.as_str()) + }) }) .await } diff --git a/tests/e2e/tests/app_state.rs b/tests/e2e/tests/app_state.rs index 735cdc26c..df617c7e5 100644 --- a/tests/e2e/tests/app_state.rs +++ b/tests/e2e/tests/app_state.rs @@ -213,12 +213,16 @@ async fn test_star_received_message() -> anyhow::Result<()> { // B receives the message and extracts msg_id let event = client_b .wait_for_event(15, |e| { - matches!(e, Event::Message(msg, _) if msg.conversation.as_deref() == Some("Star me from the other side!")) + e.messages() + .any(|m| m.message.conversation.as_deref() == Some("Star me from the other side!")) }) .await?; - let msg_id = if let Event::Message(_, info) = &*event { - info.id.clone() + let msg_id = if let Some(m) = event + .messages() + .find(|m| m.message.conversation.as_deref() == Some("Star me from the other side!")) + { + m.info.id.clone() } else { panic!("Expected Message event"); }; diff --git a/tests/e2e/tests/concurrent_disconnect.rs b/tests/e2e/tests/concurrent_disconnect.rs index b5c30771e..cde56a036 100644 --- a/tests/e2e/tests/concurrent_disconnect.rs +++ b/tests/e2e/tests/concurrent_disconnect.rs @@ -18,7 +18,6 @@ use std::time::Duration; use wacore::time::Instant; use e2e_tests::{TestClient, text_msg}; -use wacore::types::events::Event; /// Baseline: 2 clients, multi-thread runtime. Should complete nearly /// instantly (the race window is microseconds in practice), so a generous @@ -125,13 +124,15 @@ async fn concurrent_disconnect_with_pending_receipts() -> anyhow::Result<()> { for i in 0..N { let expected_ab = format!("a->b #{i}"); bob.wait_for_event(10, |e| { - matches!(e, Event::Message(m, _) if m.conversation.as_deref() == Some(expected_ab.as_str())) + e.messages() + .any(|m| m.message.conversation.as_deref() == Some(expected_ab.as_str())) }) .await?; let expected_ba = format!("b->a #{i}"); alice .wait_for_event(10, |e| { - matches!(e, Event::Message(m, _) if m.conversation.as_deref() == Some(expected_ba.as_str())) + e.messages() + .any(|m| m.message.conversation.as_deref() == Some(expected_ba.as_str())) }) .await?; } diff --git a/tests/e2e/tests/groups.rs b/tests/e2e/tests/groups.rs index c6a22d621..744e564f0 100644 --- a/tests/e2e/tests/groups.rs +++ b/tests/e2e/tests/groups.rs @@ -171,7 +171,7 @@ async fn test_group_remove_member() -> anyhow::Result<()> { client_b .assert_no_event( 3, - |e| matches!(e, Event::Message(_, _)), + |e| matches!(e, Event::Messages(_)), "B should NOT receive messages after being removed", ) .await?; @@ -187,7 +187,7 @@ async fn test_group_remove_member() -> anyhow::Result<()> { client_b .assert_no_event( 3, - |e| matches!(e, Event::Message(_, _)), + |e| matches!(e, Event::Messages(_)), "B should NOT receive messages sent by C after being removed", ) .await?; @@ -584,7 +584,7 @@ async fn test_group_leave() -> anyhow::Result<()> { client_b .assert_no_event( 3, - |e| matches!(e, Event::Message(_, _)), + |e| matches!(e, Event::Messages(_)), "B should NOT receive messages after leaving", ) .await?; @@ -600,7 +600,7 @@ async fn test_group_leave() -> anyhow::Result<()> { client_b .assert_no_event( 3, - |e| matches!(e, Event::Message(_, _)), + |e| matches!(e, Event::Messages(_)), "B should NOT receive messages sent by C after leaving", ) .await?; diff --git a/tests/e2e/tests/media.rs b/tests/e2e/tests/media.rs index 131f033e8..44b44bd84 100644 --- a/tests/e2e/tests/media.rs +++ b/tests/e2e/tests/media.rs @@ -1,6 +1,5 @@ use e2e_tests::TestClient; use log::info; -use wacore::types::events::Event; use whatsapp_rust::download::{DownloadParams, Downloadable, MediaType}; use whatsapp_rust::upload::UploadResponse; use whatsapp_rust::waproto::whatsapp as wa; @@ -413,13 +412,13 @@ async fn test_send_image_message() -> anyhow::Result<()> { // B receives the image message let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(m, _) if m.image_message.is_set()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.image_message.is_set()) + }) .await?; - if let Event::Message(msg, info) = &*event { + if let Some(m) = event.messages().find(|m| m.message.image_message.is_set()) { + let (msg, info) = (&m.message, &m.info); let img = msg.image_message.as_option().unwrap(); assert_eq!(img.caption.as_deref(), Some(caption)); assert_eq!(img.mimetype.as_deref(), Some("image/jpeg")); @@ -467,13 +466,13 @@ async fn test_send_video_message() -> anyhow::Result<()> { client_a.client.send_message(jid_b.clone(), msg).await?; let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(m, _) if m.video_message.is_set()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.video_message.is_set()) + }) .await?; - if let Event::Message(msg, _) = &*event { + if let Some(m) = event.messages().find(|m| m.message.video_message.is_set()) { + let msg = &m.message; let vid = msg.video_message.as_option().unwrap(); assert_eq!(vid.caption.as_deref(), Some("Cool video")); assert_eq!(vid.seconds, Some(15)); @@ -513,13 +512,16 @@ async fn test_send_document_message() -> anyhow::Result<()> { client_a.client.send_message(jid_b.clone(), msg).await?; let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(m, _) if m.document_message.is_set()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.document_message.is_set()) + }) .await?; - if let Event::Message(msg, _) = &*event { + if let Some(m) = event + .messages() + .find(|m| m.message.document_message.is_set()) + { + let msg = &m.message; let doc = msg.document_message.as_option().unwrap(); assert_eq!(doc.file_name.as_deref(), Some("report.pdf")); assert_eq!(doc.mimetype.as_deref(), Some("application/pdf")); @@ -558,13 +560,13 @@ async fn test_send_audio_message() -> anyhow::Result<()> { client_a.client.send_message(jid_b.clone(), msg).await?; let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(m, _) if m.audio_message.is_set()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.audio_message.is_set()) + }) .await?; - if let Event::Message(msg, _) = &*event { + if let Some(m) = event.messages().find(|m| m.message.audio_message.is_set()) { + let msg = &m.message; let audio = msg.audio_message.as_option().unwrap(); assert_eq!(audio.seconds, Some(30)); assert_eq!(audio.ptt, Some(false)); @@ -603,13 +605,13 @@ async fn test_send_ptt_voice_message() -> anyhow::Result<()> { client_a.client.send_message(jid_b.clone(), msg).await?; let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(m, _) if m.audio_message.is_set()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.audio_message.is_set()) + }) .await?; - if let Event::Message(msg, _) = &*event { + if let Some(m) = event.messages().find(|m| m.message.audio_message.is_set()) { + let msg = &m.message; let audio = msg.audio_message.as_option().unwrap(); assert_eq!(audio.ptt, Some(true)); assert_eq!(audio.seconds, Some(5)); @@ -657,12 +659,12 @@ async fn test_send_image_bidirectional() -> anyhow::Result<()> { // B receives let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(m, _) if m.image_message.is_set()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.image_message.is_set()) + }) .await?; - if let Event::Message(msg, _) = &*event { + if let Some(m) = event.messages().find(|m| m.message.image_message.is_set()) { + let msg = &m.message; let img = msg.image_message.as_option().unwrap(); assert_eq!(img.caption.as_deref(), Some("From A")); let downloaded = client_b.client.download(img as &dyn Downloadable).await?; @@ -680,12 +682,12 @@ async fn test_send_image_bidirectional() -> anyhow::Result<()> { // A receives let event = client_a - .wait_for_event( - 30, - |e| matches!(e, Event::Message(m, _) if m.image_message.is_set()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.image_message.is_set()) + }) .await?; - if let Event::Message(msg, _) = &*event { + if let Some(m) = event.messages().find(|m| m.message.image_message.is_set()) { + let msg = &m.message; let img = msg.image_message.as_option().unwrap(); assert_eq!(img.caption.as_deref(), Some("From B")); let downloaded = client_a.client.download(img as &dyn Downloadable).await?; @@ -722,12 +724,12 @@ async fn test_send_multiple_media_types() -> anyhow::Result<()> { client_a.client.send_message(jid_b.clone(), img_msg).await?; let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(m, _) if m.image_message.is_set()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.image_message.is_set()) + }) .await?; - if let Event::Message(msg, _) = &*event { + if let Some(m) = event.messages().find(|m| m.message.image_message.is_set()) { + let msg = &m.message; let img = msg.image_message.as_option().unwrap(); let dl = client_b.client.download(img as &dyn Downloadable).await?; assert_eq!(dl, img_data); @@ -743,12 +745,15 @@ async fn test_send_multiple_media_types() -> anyhow::Result<()> { client_a.client.send_message(jid_b.clone(), doc_msg).await?; let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(m, _) if m.document_message.is_set()), - ) - .await?; - if let Event::Message(msg, _) = &*event { + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.document_message.is_set()) + }) + .await?; + if let Some(m) = event + .messages() + .find(|m| m.message.document_message.is_set()) + { + let msg = &m.message; let doc = msg.document_message.as_option().unwrap(); let dl = client_b.client.download(doc as &dyn Downloadable).await?; assert_eq!(dl, doc_data); @@ -764,12 +769,12 @@ async fn test_send_multiple_media_types() -> anyhow::Result<()> { client_a.client.send_message(jid_b.clone(), aud_msg).await?; let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(m, _) if m.audio_message.is_set()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.audio_message.is_set()) + }) .await?; - if let Event::Message(msg, _) = &*event { + if let Some(m) = event.messages().find(|m| m.message.audio_message.is_set()) { + let msg = &m.message; let audio = msg.audio_message.as_option().unwrap(); let dl = client_b.client.download(audio as &dyn Downloadable).await?; assert_eq!(dl, aud_data); @@ -876,13 +881,13 @@ async fn test_send_image_no_caption() -> anyhow::Result<()> { client_a.client.send_message(jid_b.clone(), msg).await?; let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(m, _) if m.image_message.is_set()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.image_message.is_set()) + }) .await?; - if let Event::Message(msg, _) = &*event { + if let Some(m) = event.messages().find(|m| m.message.image_message.is_set()) { + let msg = &m.message; let img = msg.image_message.as_option().unwrap(); assert!( img.caption.is_none() || img.caption.as_deref() == Some(""), diff --git a/tests/e2e/tests/memory_soak.rs b/tests/e2e/tests/memory_soak.rs index 3842f4a47..0b6a2a596 100644 --- a/tests/e2e/tests/memory_soak.rs +++ b/tests/e2e/tests/memory_soak.rs @@ -259,16 +259,16 @@ async fn send_and_recv( sender.client.send_message(to.clone(), msg).await?; let t = expected_text.to_string(); receiver - .wait_for_event(30, move |e| match e { - Event::Message(msg, _) => { + .wait_for_event(30, move |e| { + e.messages().any(|m| { + let msg = &m.message; msg.conversation.as_deref() == Some(t.as_str()) || msg .extended_text_message .as_option() .and_then(|ext| ext.text.as_deref()) .is_some_and(|txt| txt.starts_with(&t)) - } - _ => false, + }) }) .await?; Ok(()) @@ -282,16 +282,17 @@ async fn wait_for_group_msg( let gid = group_jid.clone(); let text = expected_text.to_string(); client - .wait_for_event(30, move |e| match e { - Event::Message(msg, info) if info.source.chat == gid => { - msg.conversation.as_deref() == Some(text.as_str()) - || msg - .extended_text_message - .as_option() - .and_then(|ext| ext.text.as_deref()) - .is_some_and(|txt| txt.starts_with(&text)) - } - _ => false, + .wait_for_event(30, move |e| { + e.messages().any(|m| { + let (msg, info) = (&m.message, &m.info); + info.source.chat == gid + && (msg.conversation.as_deref() == Some(text.as_str()) + || msg + .extended_text_message + .as_option() + .and_then(|ext| ext.text.as_deref()) + .is_some_and(|txt| txt.starts_with(&text))) + }) }) .await?; Ok(()) diff --git a/tests/e2e/tests/messaging.rs b/tests/e2e/tests/messaging.rs index 2167898df..567c1bfcc 100644 --- a/tests/e2e/tests/messaging.rs +++ b/tests/e2e/tests/messaging.rs @@ -1,6 +1,5 @@ use e2e_tests::{TestClient, text_msg}; use log::info; -use wacore::types::events::Event; use whatsapp_rust::waproto::whatsapp as wa; #[tokio::test] @@ -91,14 +90,16 @@ async fn test_message_revoke() -> anyhow::Result<()> { // B should receive the revoke as a protocol message let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(msg, _) if msg.protocol_message.is_set()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.protocol_message.is_set()) + }) .await?; - if let Event::Message(msg, _) = &*event { - let proto = msg.protocol_message.as_option().unwrap(); + if let Some(m) = event + .messages() + .find(|m| m.message.protocol_message.is_set()) + { + let proto = m.message.protocol_message.as_option().unwrap(); assert_eq!( proto.r#type, Some(wa::message::protocol_message::Type::REVOKE), @@ -139,9 +140,12 @@ async fn test_message_has_push_name() -> anyhow::Result<()> { // Assert the push_name field on the received event let event = client_b.wait_for_text(text, 15).await?; - if let Event::Message(_, info) = &*event { + if let Some(m) = event + .messages() + .find(|m| m.message.conversation.as_deref() == Some(text)) + { assert_eq!( - info.push_name, push_name, + m.info.push_name, push_name, "Received message push_name should match the sender's display name" ); } diff --git a/tests/e2e/tests/offline_groups.rs b/tests/e2e/tests/offline_groups.rs index 3a9d4e868..0d6cb2e92 100644 --- a/tests/e2e/tests/offline_groups.rs +++ b/tests/e2e/tests/offline_groups.rs @@ -145,16 +145,19 @@ async fn test_mixed_offline_event_ordering() -> anyhow::Result<()> { for _ in 0..5 { let result = client_c .wait_for_event(10, |e| { - matches!(e, Event::Message(msg, _) if msg.conversation.is_some()) + e.messages().any(|m| m.message.conversation.is_some()) || matches!(e, Event::Notification(node) if node.get_attr("type").is_some_and(|v| v.as_str() == "w:gp2")) }) .await; match result { - Ok(ref event) if let Some((msg, _)) = event.as_message() => { - let text = msg.conversation.clone().unwrap_or_default(); - info!("C received message: {text}"); - messages_received.push(text); + Ok(ref event) if event.message_batch().is_some() => { + for m in event.messages() { + if let Some(text) = &m.message.conversation { + info!("C received message: {text}"); + messages_received.push(text.clone()); + } + } } Ok(ref event) if matches!(**event, Event::Notification(_)) => { info!("C received group notification"); @@ -242,7 +245,11 @@ async fn test_offline_group_message_delivery() -> anyhow::Result<()> { // C should receive it after reconnecting (from offline queue) let event = client_c.wait_for_text(text, 30).await?; - if let Event::Message(msg, info) = &*event { + if let Some(m) = event + .messages() + .find(|m| m.message.conversation.as_deref() == Some(text)) + { + let (msg, info) = (&m.message, &m.info); assert_eq!(msg.conversation.as_deref(), Some(text)); assert!(info.source.is_group); assert_eq!(info.source.chat, group_jid); @@ -410,16 +417,20 @@ async fn test_offline_multi_sender_group_messages() -> anyhow::Result<()> { // Verify B receives all messages (sanity check that sends worked) let mut b_received = 0; - for _ in 0..expected_messages.len() { - if client_b - .wait_for_event( - 10, - |e| matches!(e, Event::Message(msg, _) if msg.conversation.is_some()), - ) + while b_received < expected_messages.len() { + match client_b + .wait_for_event(10, |e| { + e.messages().any(|m| m.message.conversation.is_some()) + }) .await - .is_ok() { - b_received += 1; + Ok(ref event) => { + b_received += event + .messages() + .filter(|m| m.message.conversation.is_some()) + .count(); + } + Err(_) => break, } } info!("B received {b_received} messages (online observer)"); @@ -437,16 +448,18 @@ async fn test_offline_multi_sender_group_messages() -> anyhow::Result<()> { let result = client_c .wait_for_event(timeout_secs, |e| { - matches!(e, Event::Message(msg, _) if msg.conversation.is_some()) + e.messages().any(|m| m.message.conversation.is_some()) || matches!(e, Event::Notification(node) if node.get_attr("type").is_some_and(|v| v.as_str() == "w:gp2")) }) .await; match result { - Ok(ref event) if let Some((msg, _)) = event.as_message() => { - if let Some(text) = &msg.conversation { - info!("C received: {text}"); - received_texts.insert(text.clone()); + Ok(ref event) if event.message_batch().is_some() => { + for m in event.messages() { + if let Some(text) = &m.message.conversation { + info!("C received: {text}"); + received_texts.insert(text.clone()); + } } } Ok(ref event) if matches!(**event, Event::Notification(_)) => {} diff --git a/tests/e2e/tests/offline_messages.rs b/tests/e2e/tests/offline_messages.rs index 275738118..7523ff635 100644 --- a/tests/e2e/tests/offline_messages.rs +++ b/tests/e2e/tests/offline_messages.rs @@ -1,6 +1,5 @@ use e2e_tests::{TestClient, text_msg}; use log::info; -use wacore::types::events::Event; #[tokio::test] async fn test_offline_message_delivery_on_reconnect() -> anyhow::Result<()> { @@ -56,20 +55,21 @@ async fn test_offline_message_ordering() -> anyhow::Result<()> { info!("Sent: {text}"); } - // Verify messages arrive in send order + // Verify messages arrive in send order. During offline drain a single + // event can carry several messages, so iterate each batch. let mut received = Vec::new(); - for _ in 0..messages.len() { + while received.len() < messages.len() { let event = client_b - .wait_for_event( - 30, - |e| matches!(e, Event::Message(msg, _) if msg.conversation.is_some()), - ) + .wait_for_event(30, |e| { + e.messages().any(|m| m.message.conversation.is_some()) + }) .await?; - if let Event::Message(msg, _) = &*event { - let text = msg.conversation.clone().unwrap(); - info!("Received: {text}"); - received.push(text); + for m in event.messages() { + if let Some(text) = &m.message.conversation { + info!("Received: {text}"); + received.push(text.clone()); + } } } diff --git a/tests/e2e/tests/prekey_sessions.rs b/tests/e2e/tests/prekey_sessions.rs index 719ca173d..9c033126e 100644 --- a/tests/e2e/tests/prekey_sessions.rs +++ b/tests/e2e/tests/prekey_sessions.rs @@ -63,10 +63,8 @@ async fn test_prekey_collision_regression() -> anyhow::Result<()> { .await?; recipient .wait_for_event(30, |e| { - matches!( - e, - Event::Message(msg, _) if msg.conversation.as_deref() == Some(text.as_str()) - ) + e.messages() + .any(|m| m.message.conversation.as_deref() == Some(text.as_str())) }) .await?; sender.disconnect().await; diff --git a/tests/e2e/tests/privacy_tokens.rs b/tests/e2e/tests/privacy_tokens.rs index 373d19a23..5872bf34e 100644 --- a/tests/e2e/tests/privacy_tokens.rs +++ b/tests/e2e/tests/privacy_tokens.rs @@ -78,7 +78,10 @@ async fn send_message_and_expect_463_with_id( recipient .assert_no_event( 5, - move |e| matches!(e, wacore::types::events::Event::Message(msg, _) if msg.conversation.as_deref() == Some(expected_text.as_str())), + move |e| { + e.messages() + .any(|m| m.message.conversation.as_deref() == Some(expected_text.as_str())) + }, "restricted recipient should not receive first-contact message without privacy token", ) .await?; diff --git a/tests/e2e/tests/receipts.rs b/tests/e2e/tests/receipts.rs index 2db102481..c8276597a 100644 --- a/tests/e2e/tests/receipts.rs +++ b/tests/e2e/tests/receipts.rs @@ -403,22 +403,26 @@ async fn test_delivery_receipts_flushed_on_disconnect() -> anyhow::Result<()> { while seen.len() < N { let event = client_b .wait_for_event(15, |e| { - matches!(e, Event::Message(m, _) if m - .conversation - .as_deref() - .and_then(|c| c.strip_prefix("flush burst ")) - .and_then(|s| s.parse::().ok()) - .is_some_and(|i| i < N)) + e.messages().any(|m| { + m.message + .conversation + .as_deref() + .and_then(|c| c.strip_prefix("flush burst ")) + .and_then(|s| s.parse::().ok()) + .is_some_and(|i| i < N) + }) }) .await?; - if let Event::Message(m, _) = &*event - && let Some(i) = m + for m in event.messages() { + if let Some(i) = m + .message .conversation .as_deref() .and_then(|c| c.strip_prefix("flush burst ")) .and_then(|s| s.parse::().ok()) - { - seen.insert(i); + { + seen.insert(i); + } } } info!("B saw all {N} message events"); diff --git a/tests/e2e/tests/session_reuse.rs b/tests/e2e/tests/session_reuse.rs index 022882a7e..752be9a07 100644 --- a/tests/e2e/tests/session_reuse.rs +++ b/tests/e2e/tests/session_reuse.rs @@ -3,7 +3,6 @@ use e2e_tests::{TestClient, send_and_expect_text}; use log::info; use wacore::libsignal::protocol::SessionRecord; -use wacore::types::events::Event; /// Scan backend for sessions matching a user across device IDs 0..=5. /// Returns Vec<(address, has_pending_pre_key)> for all found sessions. @@ -370,7 +369,11 @@ async fn test_message_info_fields() -> anyhow::Result<()> { let event = client_b.wait_for_text(text_ab, 30).await?; - if let Event::Message(msg, info) = &*event { + if let Some(m) = event + .messages() + .find(|m| m.message.conversation.as_deref() == Some(text_ab)) + { + let (msg, info) = (&m.message, &m.info); assert_eq!(msg.conversation.as_deref(), Some(text_ab)); assert!(!info.id.is_empty(), "Message ID must not be empty"); assert!( @@ -406,7 +409,11 @@ async fn test_message_info_fields() -> anyhow::Result<()> { let event = client_a.wait_for_text(text_ba, 30).await?; - if let Event::Message(_, info) = &*event { + if let Some(m) = event + .messages() + .find(|m| m.message.conversation.as_deref() == Some(text_ba)) + { + let info = &m.info; assert!(!info.source.is_from_me); assert!(!info.source.is_group); assert_eq!(info.source.sender.user, jid_b.user); diff --git a/wacore/src/store/traits.rs b/wacore/src/store/traits.rs index 76fe5d84d..917b0f8a4 100644 --- a/wacore/src/store/traits.rs +++ b/wacore/src/store/traits.rs @@ -509,6 +509,47 @@ pub trait ProtocolStore: Send + Sync { async fn delete_expired_pending_inbound(&self, _cutoff_timestamp: i64) -> Result { Ok(0) } + + /// Batched [`store_pending_inbound`](Self::store_pending_inbound): the + /// offline drain buffers one commit-batch of messages per call, so backends + /// should override this with a single transaction (the bundled SqliteStore + /// does). The default iterates the single-row method, preserving behavior + /// for third-party backends. + async fn store_pending_inbound_batch(&self, rows: &[PendingInboundRow<'_>]) -> Result<()> { + for row in rows { + self.store_pending_inbound(row.chat, row.sender, row.id, row.message) + .await?; + } + Ok(()) + } + + /// Batched [`delete_pending_inbound`](Self::delete_pending_inbound); same + /// override guidance as [`store_pending_inbound_batch`](Self::store_pending_inbound_batch). + async fn delete_pending_inbound_batch(&self, keys: &[PendingInboundKey<'_>]) -> Result<()> { + for key in keys { + self.delete_pending_inbound(key.chat, key.sender, key.id) + .await?; + } + Ok(()) + } +} + +/// One row of a pending-inbound batch write. Fields borrow from the in-flight +/// commit batch so building a batch allocates nothing per row. +#[derive(Debug, Clone, Copy)] +pub struct PendingInboundRow<'a> { + pub chat: &'a str, + pub sender: &'a str, + pub id: &'a str, + pub message: &'a [u8], +} + +/// Key of a buffered pending-inbound row, for batched deletes. +#[derive(Debug, Clone, Copy)] +pub struct PendingInboundKey<'a> { + pub chat: &'a str, + pub sender: &'a str, + pub id: &'a str, } /// Device data persistence operations. diff --git a/wacore/src/types/events.rs b/wacore/src/types/events.rs index 395369fac..638c2ca50 100755 --- a/wacore/src/types/events.rs +++ b/wacore/src/types/events.rs @@ -218,7 +218,7 @@ pub enum EventKind { PairingCode, QrScannedWithoutMultidevice, ClientOutdated, - Message, + Messages, Receipt, UndecryptableMessage, Notification, @@ -624,7 +624,13 @@ pub enum Event { QrScannedWithoutMultidevice(QrScannedWithoutMultidevice), ClientOutdated(ClientOutdated), - Message(Arc, Arc), + /// One or more decrypted inbound messages, in arrival order (Baileys' + /// `messages.upsert` shape). Live traffic arrives as single-message + /// batches; an offline drain delivers one batch per durable commit, so a + /// consumer never sees a message that a registered durability hook has + /// not committed. The `Arc` slice is shared with the hook call — same + /// items, same order, no copies. + Messages(MessageBatch), Receipt(Receipt), UndecryptableMessage(UndecryptableMessage), #[serde(skip)] @@ -771,7 +777,7 @@ impl Event { Event::PairingCode { .. } => EventKind::PairingCode, Event::QrScannedWithoutMultidevice(_) => EventKind::QrScannedWithoutMultidevice, Event::ClientOutdated(_) => EventKind::ClientOutdated, - Event::Message(_, _) => EventKind::Message, + Event::Messages(_) => EventKind::Messages, Event::Receipt(_) => EventKind::Receipt, Event::UndecryptableMessage(_) => EventKind::UndecryptableMessage, Event::Notification(_) => EventKind::Notification, @@ -820,20 +826,55 @@ impl Event { } } - pub fn as_message(&self) -> Option<(&Arc, &MessageInfo)> { - if let Event::Message(msg, info) = self { - Some((msg, &**info)) + pub fn message_batch(&self) -> Option<&MessageBatch> { + if let Event::Messages(batch) = self { + Some(batch) } else { None } } + /// The inbound messages carried by this event, in arrival order; empty for + /// every other event kind. + pub fn messages(&self) -> impl Iterator { + self.message_batch() + .map(|b| b.messages.iter()) + .into_iter() + .flatten() + } + pub fn message_text(&self) -> Option<&str> { - let (msg, _) = self.as_message()?; - msg.conversation.as_deref() + self.messages() + .find_map(|m| m.message.conversation.as_deref()) } } +/// One decrypted inbound message. The same items (and order) back both +/// consumer surfaces: the durability hook's batch and [`Event::Messages`]. +#[derive(Debug, Clone, Serialize)] +pub struct InboundMessage { + pub message: Arc, + pub info: Arc, +} + +/// Where a [`MessageBatch`] came from. Mirrors Baileys' `messages.upsert` +/// `type` field (`notify` / `append`). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +pub enum BatchOrigin { + /// A message received while connected; always a batch of one. + Live, + /// Part of the offline backlog drained on (re)connect; batched per + /// durable commit (WA Web's MessageProcessorCache snapshot granularity). + OfflineDrain, +} + +/// Payload of [`Event::Messages`]. +#[derive(Debug, Clone, Serialize)] +pub struct MessageBatch { + pub messages: Arc<[InboundMessage]>, + pub origin: BatchOrigin, +} + /// A newsletter live update notification, typically containing updated /// reaction counts for one or more messages. #[derive(Debug, Clone, Serialize)] @@ -1664,7 +1705,7 @@ mod tests { let bus = CoreEventBus::new(); let only_msg = Arc::new(Recorder { kinds: Mutex::new(Vec::new()), - interest: EventInterest::of(&[EventKind::Message]), + interest: EventInterest::of(&[EventKind::Messages]), }); let all = Arc::new(Recorder { kinds: Mutex::new(Vec::new()), @@ -1688,7 +1729,7 @@ mod tests { CALLS.fetch_add(1, Ordering::SeqCst); } fn interest(&self) -> EventInterest { - EventInterest::of(&[EventKind::Message]) + EventInterest::of(&[EventKind::Messages]) } } let bus2 = CoreEventBus::new(); @@ -1720,7 +1761,7 @@ mod tests { let bus = CoreEventBus::new(); let h = Arc::new(Dynamic { - interest: Mutex::new(EventInterest::of(&[EventKind::Message])), + interest: Mutex::new(EventInterest::of(&[EventKind::Messages])), hits: AtomicUsize::new(0), }); bus.add_handler(h.clone()); @@ -1754,17 +1795,17 @@ mod tests { let bus = CoreEventBus::new(); // Empty bus: nothing is wanted and there are no handlers. assert!(!bus.has_handlers()); - assert!(!bus.has_handler_for(EventKind::Message)); + assert!(!bus.has_handler_for(EventKind::Messages)); assert!(!bus.has_handler_for(EventKind::Receipt)); - bus.add_handler(Arc::new(Narrow(EventInterest::of(&[EventKind::Message])))); + bus.add_handler(Arc::new(Narrow(EventInterest::of(&[EventKind::Messages])))); assert!(bus.has_handlers()); - assert!(bus.has_handler_for(EventKind::Message)); + assert!(bus.has_handler_for(EventKind::Messages)); assert!(!bus.has_handler_for(EventKind::Receipt)); // has_handler_for is true once any registered handler wants the kind. bus.add_handler(Arc::new(Narrow(EventInterest::of(&[EventKind::Receipt])))); - assert!(bus.has_handler_for(EventKind::Message)); + assert!(bus.has_handler_for(EventKind::Messages)); assert!(bus.has_handler_for(EventKind::Receipt)); assert!(!bus.has_handler_for(EventKind::Connected)); } From d114f3a2b484f3042e3783677c838ed4389cfe83 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 03:46:43 +0000 Subject: [PATCH 02/35] fix(bot): keep MessageContext out of the batched on_message future MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MessageContext is not Send on wasm32 (the Client's trait objects carry no Send/Sync bounds there), so the batch fan-out must not hold contexts across an await. Build every handler future before the async block — only the futures, which the Fut: Send bound covers, cross the awaits. Same shape the pre-batch code used for its single future. --- src/bot.rs | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/bot.rs b/src/bot.rs index c05e98e5a..4b983cc92 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -683,16 +683,18 @@ impl BotBuilder { F: Fn(MessageContext) -> Fut + Send + Sync + 'static, Fut: Future + Send + 'static, { - let handler = Arc::new(handler); self.on_event_for(&[EventKind::Messages], move |event, client| { - let contexts: Vec = event + // Futures are built before the async block: `MessageContext` is + // not `Send` on wasm32 (the Client's trait objects aren't), so it + // must never be held across an await — only the handler futures + // (which the `Fut: Send` bound covers) may cross one. + let futures: Vec = event .messages() - .map(|m| MessageContext::from_inbound(m, Arc::clone(&client))) + .map(|m| handler(MessageContext::from_inbound(m, Arc::clone(&client)))) .collect(); - let handler = Arc::clone(&handler); async move { - for context in contexts { - handler(context).await; + for future in futures { + future.await; } } }) From 66ac89ae5416ede8b6aee1b5dc5054be4959f5f0 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 04:16:53 +0000 Subject: [PATCH 03/35] fix(recv): raceless drain-to-live transition; review fixes from cubic/codex/greptile/coderabbit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The drain/live gate moved off offline_sync_completed onto a dedicated batcher-owned `active` state that only the end-of-drain flush flips, while holding the single processing permit. Gating on the flag was racy in two confirmed ways: an in-flight stanza could enqueue pre-flip and then take the live per-stanza Signal flush at stanza end, persisting ratchet state for an uncommitted batch entry (crash => unrecoverable duplicate, cubic P1); and stanzas queued behind the permit could read the flipped flag at dispatch and commit as Live ahead of the still accumulated batch, inverting arrival order at the drain tail (codex P1). finish_inbound_commit_drain commits the tail and deactivates under the permit, so no stanza straddles the transition and queued stanzas provably commit after the tail. Regression test locks the sequence. Also from review: - Acks are sent before the event dispatch (all durability done by then): handle_event runs synchronously, so a panicking or blocking handler must not suppress acks for messages the consumer already owns — matches the old at-most-once ordering (codex). - The encode buffer is scoped to the durable write only and the live path uses a local buffer instead of the shared arena, so concurrent live commits never serialize behind a slow hook (greptile, coderabbit). - store_pending_inbound_batch uses multi-row VALUES inserts (chunked at 100 rows); deletes stay per-row inside the one transaction because Diesel's DSL cannot express a composite-key tuple IN (greptile). - The stale-permit re-acquire loop is one shared helper, acquire_message_processing_permit (cubic, coderabbit). - Teardown flushes are bounded (5s disconnect / 2s reconnect, like the outbound flush): a hung hook cannot wedge disconnect; entries stay unacked for redelivery (coderabbit). - complete_offline_sync logs loudly if the self_weak upgrade ever fails instead of silently skipping the tail commit (coderabbit). - Event::message_texts() iterator; message_text() documented as first-text convenience (coderabbit). - Test clients mirror full live state (flag, 64 permits, batcher live); batch example dedups within the batch; on_message documents the eager future construction wasm32 requires (cubic). --- examples/durability_hook.rs | 4 +- src/bot.rs | 6 + src/client/lifecycle.rs | 17 +- src/client/node_io.rs | 20 ++ src/client/sessions.rs | 17 +- src/message/commit_batch.rs | 228 ++++++++++++++------ src/message/receive.rs | 28 +-- src/message/tests.rs | 5 +- src/test_utils.rs | 8 +- storages/sqlite-storage/src/sqlite_store.rs | 27 ++- wacore/src/types/events.rs | 14 +- 11 files changed, 259 insertions(+), 115 deletions(-) diff --git a/examples/durability_hook.rs b/examples/durability_hook.rs index 6d05fb4cd..c16ee3e1f 100644 --- a/examples/durability_hook.rs +++ b/examples/durability_hook.rs @@ -120,7 +120,9 @@ impl InboundDurabilityHook for InboxArchiver { m.info.source.sender.to_string(), m.info.id.clone(), ); - if seen.contains(&key) { + // Dedup against the archive AND earlier entries of this same + // batch, so one fsync can never append a key twice. + if seen.contains(&key) || keys.contains(&key) { info!("[{}] already committed, skipping (dedup)", m.info.id); continue; } diff --git a/src/bot.rs b/src/bot.rs index 4b983cc92..b747e9816 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -678,6 +678,12 @@ impl BotBuilder { /// single-message on live traffic) are fanned out here in arrival order, /// awaiting each handler before the next — per-message bots keep their /// ergonomics and gain in-batch ordering. + /// + /// The handler is CALLED for every message in the batch up front and the + /// returned futures then run in order (an `async` closure runs no body + /// code at call time, so for the typical handler this is unobservable). + /// Interleaving call+await instead would hold a `MessageContext` across + /// an await, which is not `Send` on wasm32. pub fn on_message(self, handler: F) -> Self where F: Fn(MessageContext) -> Fut + Send + Sync + 'static, diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 5dda72519..05822858b 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -458,7 +458,7 @@ impl Client { self.clear_offline_receipt_buffer(); // Uncommitted batch entries were never acked; the server redelivers // them on this fresh connection. - self.inbound_commit_batch.clear(); + self.inbound_commit_batch.reset(); self.offline_batch.reset(); self.outbound_flush.reopen(); @@ -579,9 +579,10 @@ impl Client { // re-acked fresh. // // Commit any accumulated drain batch first so its acks land in this - // receipt drain; entries that cannot commit stay unacked and the - // server redelivers them. - self.flush_inbound_commits_acquiring_permit().await; + // receipt drain. Bounded like the outbound flush below: on timeout the + // entries simply stay unacked and the server redelivers them. + self.flush_inbound_commits_bounded(std::time::Duration::from_secs(5)) + .await; self.flush_offline_receipts(); // Prevent late receipt producers from escaping the drain window. self.outbound_flush.close(); @@ -640,7 +641,8 @@ impl Client { self.auto_reconnect_errors .store(Self::RECONNECT_BACKOFF_STEP, Ordering::Relaxed); - self.flush_inbound_commits_acquiring_permit().await; + self.flush_inbound_commits_bounded(std::time::Duration::from_secs(2)) + .await; self.flush_offline_receipts(); self.outbound_flush.close(); self.outbound_flush @@ -666,7 +668,8 @@ impl Client { info!("Reconnecting immediately (expected disconnect)."); self.expected_disconnect.store(true, Ordering::Relaxed); - self.flush_inbound_commits_acquiring_permit().await; + self.flush_inbound_commits_bounded(std::time::Duration::from_secs(2)) + .await; self.flush_offline_receipts(); self.outbound_flush.close(); self.outbound_flush @@ -764,7 +767,7 @@ impl Client { self.clear_offline_receipt_buffer(); // Same rule as receipts: uncommitted entries drop here and the server // redelivers them on the next connect. - self.inbound_commit_batch.clear(); + self.inbound_commit_batch.reset(); self.offline_batch.reset(); self.offline_sync_metrics .active diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 4fb3a78fa..d1879bd82 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -64,6 +64,26 @@ impl Client { .fetch_add(1, Ordering::SeqCst); } + /// Acquire one permit from the CURRENT message-processing semaphore. + /// + /// The semaphore can be swapped while a waiter sleeps (offline online + /// transition); a permit from the stale semaphore would be a no-op guard, + /// so re-acquire until generation and semaphore agree. Shared by stanza + /// processing and the commit batcher: both must serialize on the same + /// instance for the drain-flush safety argument to hold. + pub(crate) async fn acquire_message_processing_permit(&self) -> async_lock::SemaphoreGuardArc { + loop { + let (generation, semaphore) = self.read_message_semaphore(); + let permit = semaphore.acquire_arc().await; + if generation == self.message_semaphore_generation.load(Ordering::SeqCst) { + return permit; + } + // Generation changed while waiting: drop the stale permit and + // retry with the new semaphore. + drop(permit); + } + } + // err(...) stays at the default ERROR on purpose: with the routine server // recycle moved to Ok(ServerRecycle), an Err from this loop now always means // something genuinely wrong — so the automatic capture only ever reports diff --git a/src/client/sessions.rs b/src/client/sessions.rs index 7c8501ef2..590d9de3b 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -36,14 +36,21 @@ impl Client { .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) .is_ok() { - // Commit the tail of the drain batch BEFORE widening the semaphore: - // acquiring the still-single permit serializes with the last - // in-flight stanza, and post-flip dispatches take the live path, so - // the batcher is provably empty from here on. Receipts flush after, + // Commit the drain tail and flip the batcher to live mode, both + // under the still-single processing permit (see + // finish_inbound_commit_drain for the raceless-transition + // argument), BEFORE widening the semaphore. Receipts flush after, // so every receipt's message is durably committed first (WA Web's // createSnapshot ordering). if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { - client.flush_inbound_commits_acquiring_permit().await; + client.finish_inbound_commit_drain().await; + } else { + // Practically unreachable (the run loop owns a strong Arc), but + // a silent skip here would be exactly the acked-before-committed + // bug this ordering exists to prevent — make it loud. + log::error!( + "complete_offline_sync: self_weak upgrade failed; skipping drain-tail commit before widening the semaphore" + ); } // Allow parallel message processing now that offline sync is done. diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 6d645d571..afb1d0068 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -35,10 +35,20 @@ struct BatchState { pub(crate) struct InboundCommitBatcher { state: std::sync::Mutex, + /// Whether inbound commits accumulate here (offline drain) or commit + /// immediately (live). Flipped off ONLY by the end-of-drain flush while it + /// holds the single processing permit, so no stanza can straddle the + /// transition: gating on `offline_sync_completed` (which flips outside the + /// permit) would let an in-flight stanza persist Signal state for an + /// uncommitted batch entry, and let queued drain stanzas commit as Live + /// ahead of the accumulated batch. + active: std::sync::atomic::AtomicBool, /// Bumped on every take; a timer that observes a stale epoch stands down. epoch: AtomicU64, - /// Serializes commit sequences so batches reach the hook in accumulation - /// order. The guard doubles as the reusable encode arena. + /// Reusable encode arena for drain commits, which the processing permit + /// already serializes — so this lock is never contended there. Live + /// commits use a local buffer instead: sharing it would serialize + /// concurrent live-path hook calls that could previously overlap. arena: async_lock::Mutex>, } @@ -46,6 +56,7 @@ impl Default for InboundCommitBatcher { fn default() -> Self { Self { state: std::sync::Mutex::new(BatchState::default()), + active: std::sync::atomic::AtomicBool::new(true), epoch: AtomicU64::new(0), arena: async_lock::Mutex::new(Vec::new()), } @@ -69,9 +80,20 @@ impl InboundCommitBatcher { std::mem::take(&mut state.entries) } - /// Drop accumulated entries without committing (connection teardown). - /// Uncommitted messages were never acked, so the server redelivers them. - pub(crate) fn clear(&self) { + pub(crate) fn is_active(&self) -> bool { + self.active.load(Ordering::Acquire) + } + + /// Switch to immediate (live) commits. Only the end-of-drain flush calls + /// this, while holding the single processing permit. + fn deactivate(&self) { + self.active.store(false, Ordering::Release); + } + + /// Connection teardown/setup: drop uncommitted entries (they were never + /// acked, so the server redelivers them) and re-arm accumulation for the + /// next connection's drain. + pub(crate) fn reset(&self) { let dropped = self.take(); if !dropped.is_empty() { log::debug!( @@ -79,6 +101,13 @@ impl InboundCommitBatcher { dropped.len() ); } + self.active.store(true, Ordering::Release); + } + + /// Test-only: enter live mode without running a drain flush. + #[cfg(test)] + pub(crate) fn deactivate_for_tests(&self) { + self.deactivate(); } } @@ -102,7 +131,7 @@ impl Client { /// Batch a message decrypted while the offline drain is active, or commit /// immediately (batch of one) on the live path. pub(crate) async fn commit_or_batch_inbound(self: &Arc, item: InboundMessage) { - if self.offline_sync_completed.load(Ordering::Relaxed) { + if !self.inbound_commit_batch.is_active() { self.commit_inbound_batch(vec![item], BatchOrigin::Live, false) .await; return; @@ -142,18 +171,7 @@ impl Client { /// semaphore holds a single permit, so this fully serializes with stanza /// processing; after the drain the batcher is empty and this no-ops. pub(crate) async fn flush_inbound_commits_acquiring_permit(self: &Arc) { - let _permit = loop { - let (generation, semaphore) = self.read_message_semaphore(); - let permit = semaphore.acquire_arc().await; - if generation - == self - .message_semaphore_generation - .load(std::sync::atomic::Ordering::SeqCst) - { - break permit; - } - drop(permit); - }; + let _permit = self.acquire_message_processing_permit().await; let batch = self.inbound_commit_batch.take(); if batch.is_empty() { return; @@ -162,6 +180,44 @@ impl Client { .await; } + /// [`flush_inbound_commits_acquiring_permit`](Self::flush_inbound_commits_acquiring_permit) + /// with a deadline, for connection teardown: a stalled permit holder or a + /// hung hook must not wedge disconnect/reconnect. On timeout the entries + /// stay unacked and the server redelivers them on the next connect. + pub(crate) async fn flush_inbound_commits_bounded( + self: &Arc, + limit: std::time::Duration, + ) { + if wacore::runtime::timeout( + &*self.runtime, + limit, + self.flush_inbound_commits_acquiring_permit(), + ) + .await + .is_err() + { + log::warn!( + "Timed out committing the inbound drain batch during teardown; leaving entries for redelivery" + ); + } + } + + /// End-of-drain transition: commit the tail batch and switch the batcher + /// to live mode, all under the single processing permit. Holding the + /// permit across BOTH steps is what makes the transition raceless: no + /// stanza is mid-flight when the mode flips (so a stanza's enqueue and its + /// stanza-end flush always agree), and every stanza still queued behind + /// this permit commits as Live strictly AFTER the tail batch — arrival + /// order is preserved across the boundary. Runs before the semaphore + /// widens to the live permit count. + pub(crate) async fn finish_inbound_commit_drain(self: &Arc) { + let _permit = self.acquire_message_processing_permit().await; + let batch = self.inbound_commit_batch.take(); + self.inbound_commit_batch.deactivate(); + self.commit_inbound_batch(batch, BatchOrigin::OfflineDrain, true) + .await; + } + /// Commit one batch: durable buffer → Signal flush → hook → clear buffer → /// event → acks. WA Web ordering (`createSnapshot`), so nothing is acked or /// observable before it is durable. On any commit failure everything stays @@ -191,37 +247,52 @@ impl Client { }) .collect(); - let mut arena = self.inbound_commit_batch.arena.lock().await; - arena.clear(); - let mut ranges = Vec::with_capacity(items.len()); - for item in &items { - let start = arena.len(); - waproto::codec::message_encode_into(&item.message, &mut arena); - ranges.push(start..arena.len()); - } - let rows: Vec> = items - .iter() - .zip(&keys) - .zip(&ranges) - .map(|((item, (chat, sender)), range)| PendingInboundRow { - chat, - sender, - id: &item.info.id, - message: &arena[range.clone()], - }) - .collect(); - let backend = self.persistence_manager.backend(); - // Fail closed: without a durable buffered copy, do not run the hook - // and do not ack — the server redelivers once storage recovers. - if let Err(e) = backend.store_pending_inbound_batch(&rows).await { - log::error!( - "Failed to buffer inbound batch of {}; suppressing acks for redelivery: {e:?}", - items.len() - ); - return; + // Encode scope: the buffer lives only through the durable write, + // never across the slower flush/hook steps below. Drain reuses the + // shared arena (uncontended: the permit serializes drain flushes); + // live (batch of one) uses a local buffer so concurrent live + // commits never queue on a shared lock while a slow hook runs. + { + let mut local_arena; + let mut shared_arena; + let arena: &mut Vec = if matches!(origin, BatchOrigin::OfflineDrain) { + shared_arena = self.inbound_commit_batch.arena.lock().await; + &mut shared_arena + } else { + local_arena = Vec::new(); + &mut local_arena + }; + arena.clear(); + let mut ranges = Vec::with_capacity(items.len()); + for item in &items { + let start = arena.len(); + waproto::codec::message_encode_into(&item.message, arena); + ranges.push(start..arena.len()); + } + let rows: Vec> = items + .iter() + .zip(&keys) + .zip(&ranges) + .map(|((item, (chat, sender)), range)| PendingInboundRow { + chat, + sender, + id: &item.info.id, + message: &arena[range.clone()], + }) + .collect(); + + // Fail closed: without a durable buffered copy, do not run the + // hook and do not ack — the server redelivers once storage + // recovers. + if let Err(e) = backend.store_pending_inbound_batch(&rows).await { + log::error!( + "Failed to buffer inbound batch of {}; suppressing acks for redelivery: {e:?}", + items.len() + ); + return; + } } - drop(rows); if flush_signal { self.flush_signal_cache_logged("commit_batch", None).await; @@ -325,9 +396,7 @@ mod tests { let (handler, rx) = ChannelEventHandler::new(); client.core.event_bus.add_handler(handler); - client - .offline_sync_completed - .store(false, std::sync::atomic::Ordering::Relaxed); + client.inbound_commit_batch.reset(); for id in ["B1", "B2", "B3"] { client.commit_or_batch_inbound(item(id)).await; } @@ -371,10 +440,6 @@ mod tests { let _ = client.inbound_durability_hook.set(hook.clone()); let (handler, rx) = ChannelEventHandler::new(); client.core.event_bus.add_handler(handler); - client - .offline_sync_completed - .store(true, std::sync::atomic::Ordering::Relaxed); - client.commit_or_batch_inbound(item("L1")).await; assert_eq!( @@ -391,9 +456,7 @@ mod tests { #[tokio::test] async fn size_trigger_flushes_full_batch() { let client = create_test_client_with_failing_http("batch_size").await; - client - .offline_sync_completed - .store(false, std::sync::atomic::Ordering::Relaxed); + client.inbound_commit_batch.reset(); let hook = Arc::new(RecordingHook { batches: Mutex::new(Vec::new()), }); @@ -415,9 +478,7 @@ mod tests { #[tokio::test] async fn drain_without_hook_batches_events() { let client = create_test_client_with_failing_http("batch_no_hook").await; - client - .offline_sync_completed - .store(false, std::sync::atomic::Ordering::Relaxed); + client.inbound_commit_batch.reset(); let (handler, rx) = ChannelEventHandler::new(); client.core.event_bus.add_handler(handler); @@ -435,14 +496,51 @@ mod tests { ); } - // clear() drops uncommitted entries: no hook call, no event, and the + // End-of-drain transition: the tail batch commits first (as OfflineDrain), + // the batcher flips to live mode, and anything after commits as Live — + // never interleaved ahead of the tail (cubic/codex P1 regression). + #[tokio::test] + async fn finish_drain_commits_tail_then_switches_to_live() { + let client = create_test_client_with_failing_http("batch_transition").await; + client.inbound_commit_batch.reset(); + let hook = Arc::new(RecordingHook { + batches: Mutex::new(Vec::new()), + }); + let _ = client.inbound_durability_hook.set(hook.clone()); + let (handler, rx) = ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); + + client.commit_or_batch_inbound(item("T1")).await; + client.commit_or_batch_inbound(item("T2")).await; + client.finish_inbound_commit_drain().await; + assert!(!client.inbound_commit_batch.is_active()); + // A message arriving after the transition commits immediately as Live. + client.commit_or_batch_inbound(item("T3")).await; + + let batches = hook.batches.lock().expect("hook lock").clone(); + assert_eq!( + batches, + vec![vec!["T1", "T2"], vec!["T3"]], + "tail batch must commit before, and separately from, live traffic" + ); + let first = rx.try_recv().expect("tail event"); + assert_eq!( + first.message_batch().expect("Messages").origin, + BatchOrigin::OfflineDrain + ); + let second = rx.try_recv().expect("live event"); + assert_eq!( + second.message_batch().expect("Messages").origin, + BatchOrigin::Live + ); + } + + // reset() drops uncommitted entries: no hook call, no event, and the // pending buffer was never written (the server redelivers instead). #[tokio::test] async fn clear_drops_uncommitted_entries() { let client = create_test_client_with_failing_http("batch_clear").await; - client - .offline_sync_completed - .store(false, std::sync::atomic::Ordering::Relaxed); + client.inbound_commit_batch.reset(); let hook = Arc::new(RecordingHook { batches: Mutex::new(Vec::new()), }); @@ -451,7 +549,7 @@ mod tests { client.core.event_bus.add_handler(handler); client.commit_or_batch_inbound(item("C1")).await; - client.inbound_commit_batch.clear(); + client.inbound_commit_batch.reset(); client.flush_inbound_commits_acquiring_permit().await; assert!(hook.batches.lock().expect("hook lock").is_empty()); diff --git a/src/message/receive.rs b/src/message/receive.rs index e1bf3273c..957d356c1 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -344,24 +344,7 @@ impl Client { // can lose pkmsg messages carrying SKDM (sender key distribution). If the // SKDM is lost, ALL subsequent skmsg messages from that sender will fail // with "No sender key state". - let _global_permit = loop { - let (generation, semaphore) = self.read_message_semaphore(); - let permit = semaphore.acquire_arc().await; - if generation - == self - .message_semaphore_generation - .load(std::sync::atomic::Ordering::SeqCst) - { - break permit; - } - // Generation changed while waiting (e.g. offline→online transition). - // Drop the stale permit and retry with the new semaphore, which has - // more permits and will grant access quickly. - log::debug!( - "Semaphore generation changed during acquire, re-acquiring from new semaphore" - ); - drop(permit); - }; + let _global_permit = self.acquire_message_processing_permit().await; log::debug!( "Starting PASS 1: Processing {} session establishment messages (pkmsg/msg)", @@ -533,14 +516,11 @@ impl Client { // commit batcher owns the flush — one per batch, before any ack (WA // Web's bulk signal-store snapshot) — so here only the batch size/byte // triggers are checked, while the global permit is still held. - if self - .offline_sync_completed - .load(std::sync::atomic::Ordering::Relaxed) - { + if self.inbound_commit_batch.is_active() { + self.maybe_flush_inbound_commits().await; + } else { self.flush_signal_cache_logged("message", Some(&info.id)) .await; - } else { - self.maybe_flush_inbound_commits().await; } } diff --git a/src/message/tests.rs b/src/message/tests.rs index 1d26d9d4c..8ef229e81 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -5533,8 +5533,11 @@ async fn capturing_client( // other layers but not on this path. *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); seed_test_pn(&client).await; - // Live-path semantics by default; drain tests reset the flag themselves. + // Live-path semantics by default (flag, live permit count, batcher live); + // drain tests re-enter drain state themselves. client.offline_sync_completed.store(true, Ordering::Relaxed); + client.swap_message_semaphore(64); + client.inbound_commit_batch.deactivate_for_tests(); (client, transport) } diff --git a/src/test_utils.rs b/src/test_utils.rs index 64e5a7cff..d2a894331 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -110,11 +110,15 @@ pub async fn create_test_client_with_config( .await; // Tests exercise live-path semantics by default (a fresh client starts in - // drain mode, where inbound commits batch instead of dispatching - // immediately). Drain-specific tests reset this flag themselves. + // drain mode: 1-permit semaphore, inbound commits batch instead of + // dispatching immediately). Mirror a completed offline sync — flag, live + // permit count, batcher in live mode. Drain-specific tests re-enter drain + // state themselves. client .offline_sync_completed .store(true, std::sync::atomic::Ordering::Relaxed); + client.swap_message_semaphore(64); + client.inbound_commit_batch.deactivate_for_tests(); client } diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index eceff2ec0..4f5b7fd3f 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -3120,15 +3120,23 @@ impl ProtocolStore for SqliteStore { let rows = Arc::clone(&rows); Box::new(move |conn: &mut SqliteConnection| { conn.transaction(|conn| { - for (chat, sender, id, message) in rows.iter() { + // Multi-row VALUES per statement; chunked so binds stay far + // under SQLite's variable limit (5 columns × 100 rows). + for chunk in rows.chunks(100) { + let values: Vec<_> = chunk + .iter() + .map(|(chat, sender, id, message)| { + ( + pending_inbound_messages::chat.eq(chat), + pending_inbound_messages::sender.eq(sender), + pending_inbound_messages::id.eq(id), + pending_inbound_messages::message.eq(message.as_slice()), + pending_inbound_messages::device_id.eq(device_id), + ) + }) + .collect(); diesel::replace_into(pending_inbound_messages::table) - .values(( - pending_inbound_messages::chat.eq(chat), - pending_inbound_messages::sender.eq(sender), - pending_inbound_messages::id.eq(id), - pending_inbound_messages::message.eq(message.as_slice()), - pending_inbound_messages::device_id.eq(device_id), - )) + .values(&values) .execute(conn)?; } Ok(()) @@ -3154,6 +3162,9 @@ impl ProtocolStore for SqliteStore { self.with_retry("delete_pending_inbound_batch", || { let keys = Arc::clone(&keys); Box::new(move |conn: &mut SqliteConnection| { + // Per-row deletes stay: Diesel's DSL cannot express a composite + // `(chat, sender, id) IN (...)` tuple filter, and the single + // transaction already amortizes the WAL commit. conn.transaction(|conn| { for (chat, sender, id) in keys.iter() { diesel::delete( diff --git a/wacore/src/types/events.rs b/wacore/src/types/events.rs index 638c2ca50..b4f70f54b 100755 --- a/wacore/src/types/events.rs +++ b/wacore/src/types/events.rs @@ -843,9 +843,19 @@ impl Event { .flatten() } - pub fn message_text(&self) -> Option<&str> { + /// Every plain-text body in this event, in arrival order. Prefer this + /// over [`message_text`](Self::message_text) when a batch can carry more + /// than one message (offline drain). + pub fn message_texts(&self) -> impl Iterator { self.messages() - .find_map(|m| m.message.conversation.as_deref()) + .filter_map(|m| m.message.conversation.as_deref()) + } + + /// Convenience: the FIRST plain-text body in this event. A drain batch can + /// carry several texts — iterate [`message_texts`](Self::message_texts) + /// (or [`messages`](Self::messages)) to see them all. + pub fn message_text(&self) -> Option<&str> { + self.message_texts().next() } } From bbf931d94696ff8bd8f7599f026561e8c7f3b92c Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 04:31:17 +0000 Subject: [PATCH 04/35] fix(recv): commit the drain batch before cleanup's Signal flush; dispatch replayed messages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two Codex findings on the review round, both real: - cleanup_connection_state flushed the Signal cache and then reset the batcher, so an unexpected read-loop exit mid-drain persisted ratchet advances for entries it was about to drop — their redelivery would decrypt as a duplicate with no buffered copy and get acked without delivery. The cleanup now commits the batch (bounded, 5s) before that flush; acks/events are best-effort with the socket gone, but the durable hook commit is what prevents the loss. - A message whose hook only succeeds on redelivery replay never emitted Event::Messages (its original batch aborted before dispatch). The replay success path now dispatches the single-message batch after the ack, so consumers see it exactly once; test extended to lock this. Also: live batches build their Arc<[InboundMessage]> slice directly instead of round-tripping through a Vec — measured 2.0 -> 1.0 allocations and ~18ns (-24%) for that step per live message; commit_inbound_batch now takes the Arc slice, which the event reuses without conversion. --- src/client/lifecycle.rs | 13 +++++++++++++ src/message/commit_batch.rs | 18 +++++++++++------ src/message/durability.rs | 39 ++++++++++++++++++++++++++++++------- 3 files changed, 57 insertions(+), 13 deletions(-) diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 05822858b..c8bbd91a5 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -744,6 +744,19 @@ impl Client { .lock() .unwrap_or_else(|p| p.into_inner()) .clear(); + // Commit any accumulated drain batch BEFORE the Signal flush below: + // persisting ratchet advances while dropping their uncommitted batch + // entries would turn each redelivery into an ackable duplicate with no + // buffered copy — silent loss for hook consumers. Acks/events from + // this commit are best-effort (the socket is gone); the durable hook + // commit is what matters. Reached on every teardown path, including + // the run loop's unexpected read-loop exit, which never goes through + // disconnect(). + if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { + client + .flush_inbound_commits_bounded(std::time::Duration::from_secs(5)) + .await; + } // Flush before clear: clear() drops dirty entries, so a disconnect // racing an in-flight encrypt would lose the just-advanced sender-key // chain and force a full SKDM re-fanout. A disconnect is not a logout. diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index afb1d0068..5ab2f19da 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -132,7 +132,10 @@ impl Client { /// immediately (batch of one) on the live path. pub(crate) async fn commit_or_batch_inbound(self: &Arc, item: InboundMessage) { if !self.inbound_commit_batch.is_active() { - self.commit_inbound_batch(vec![item], BatchOrigin::Live, false) + // Arc::from([item]) builds the event/hook slice in one allocation; + // a Vec would add an alloc+dealloc per live message (measured + // ~18ns and 2x the allocations of this step). + self.commit_inbound_batch(std::sync::Arc::from([item]), BatchOrigin::Live, false) .await; return; } @@ -159,7 +162,7 @@ impl Client { }; if over { let batch = self.inbound_commit_batch.take(); - self.commit_inbound_batch(batch, BatchOrigin::OfflineDrain, true) + self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain, true) .await; } } @@ -176,7 +179,7 @@ impl Client { if batch.is_empty() { return; } - self.commit_inbound_batch(batch, BatchOrigin::OfflineDrain, true) + self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain, true) .await; } @@ -214,7 +217,10 @@ impl Client { let _permit = self.acquire_message_processing_permit().await; let batch = self.inbound_commit_batch.take(); self.inbound_commit_batch.deactivate(); - self.commit_inbound_batch(batch, BatchOrigin::OfflineDrain, true) + if batch.is_empty() { + return; + } + self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain, true) .await; } @@ -225,7 +231,7 @@ impl Client { #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.recv.commit_batch", level = "debug", skip_all, fields(count = items.len())))] pub(crate) async fn commit_inbound_batch( self: &Arc, - items: Vec, + items: std::sync::Arc<[InboundMessage]>, origin: BatchOrigin, flush_signal: bool, ) { @@ -265,7 +271,7 @@ impl Client { }; arena.clear(); let mut ranges = Vec::with_capacity(items.len()); - for item in &items { + for item in items.iter() { let start = arena.len(); waproto::codec::message_encode_into(&item.message, arena); ranges.push(start..arena.len()); diff --git a/src/message/durability.rs b/src/message/durability.rs index 434bf44bf..6825910c4 100644 --- a/src/message/durability.rs +++ b/src/message/durability.rs @@ -9,7 +9,7 @@ use super::*; use crate::types::durability_hook::InboundDurabilityHook; -use wacore::types::events::InboundMessage; +use wacore::types::events::{BatchOrigin, InboundMessage, MessageBatch}; impl Client { /// The registered inbound durability hook, if any. `None` (default) keeps @@ -51,6 +51,19 @@ impl Client { ); } self.ack_received_message(info); + // First successful commit of this message: + // its original batch never dispatched (the + // hook failed then), so consumers see it + // here or never. + let origin = if self.inbound_commit_batch.is_active() { + BatchOrigin::OfflineDrain + } else { + BatchOrigin::Live + }; + self.core.event_bus.dispatch(Event::Messages(MessageBatch { + messages: Arc::from([item]), + origin, + })); } Err(e) => { log::warn!( @@ -95,7 +108,6 @@ mod tests { use crate::test_utils::create_test_client_with_failing_http; use crate::types::message::MessageInfo; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; - use wacore::types::events::BatchOrigin; struct CountingHook { calls: AtomicUsize, @@ -158,10 +170,11 @@ mod tests { let hook = counting_hook(true); let _ = client.inbound_durability_hook.set(hook.clone()); - let items = vec![test_item("MSG_OK_1"), test_item("MSG_OK_2")]; + let items: Arc<[InboundMessage]> = + Arc::from([test_item("MSG_OK_1"), test_item("MSG_OK_2")]); let infos: Vec<_> = items.iter().map(|i| Arc::clone(&i.info)).collect(); client - .commit_inbound_batch(items, BatchOrigin::OfflineDrain, false) + .commit_inbound_batch(Arc::clone(&items), BatchOrigin::OfflineDrain, false) .await; assert_eq!(hook.calls.load(Ordering::SeqCst), 1, "one commit per batch"); @@ -196,13 +209,13 @@ mod tests { let info = test_info("MSG_ERR"); client .commit_inbound_batch( - vec![InboundMessage { + Arc::from([InboundMessage { message: Arc::new(wa::Message { conversation: Some("hello".to_string()), ..Default::default() }), info: Arc::clone(&info), - }], + }]), BatchOrigin::OfflineDrain, false, ) @@ -242,10 +255,22 @@ mod tests { "a still-failing hook keeps the buffered copy" ); - // Redelivery once the commit succeeds clears the buffer. + // Redelivery once the commit succeeds clears the buffer AND finally + // dispatches the event (the original batch never did). + let (handler, rx) = wacore::types::events::ChannelEventHandler::new(); + client.core.event_bus.add_handler(handler); hook.succeed.store(true, Ordering::SeqCst); client.ack_or_replay_to_hook(&info).await; assert_eq!(hook.calls.load(Ordering::SeqCst), 3); + let event = rx.try_recv().expect("successful replay must dispatch"); + assert_eq!( + event + .messages() + .map(|m| m.info.id.as_str()) + .collect::>(), + ["MSG_ERR"], + "consumers must observe a message whose hook only succeeded on replay" + ); assert!( backend .get_pending_inbound( From 9c9bee768c205ffad5c54a03076231e5557092a1 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 04:34:38 +0000 Subject: [PATCH 05/35] fix(recv): actually reorder acks before the batch event dispatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 66ac89a's message claimed this reorder but the patch hunk had silently failed to apply; the code still dispatched Event::Messages before sending acks (and with a redundant Arc conversion, caught by clippy's useless_conversion). Apply it for real: everything is durable by the time either happens, and acking first means a panicking or blocking synchronous handler cannot suppress acks for messages the consumer already owns — the ordering the pre-batch at-most-once path had. --- src/message/commit_batch.rs | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 5ab2f19da..989764f4b 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -333,14 +333,17 @@ impl Client { self.flush_signal_cache_logged("commit_batch", None).await; } - let batch = MessageBatch { - messages: items.into(), - origin, - }; - self.core.event_bus.dispatch(Event::Messages(batch.clone())); - for item in batch.messages.iter() { + // Acks first (everything durable by now): handle_event runs + // synchronously, so a handler that panics or blocks must not be able + // to suppress acks for messages the consumer already owns — the + // pre-batch at-most-once path acked before dispatching too. + for item in items.iter() { self.ack_received_message(&item.info); } + self.core.event_bus.dispatch(Event::Messages(MessageBatch { + messages: items, + origin, + })); } } From dcdc148c39a01e1363109dcc8537a1bfb2f16edd Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 04:48:36 +0000 Subject: [PATCH 06/35] fix(recv): flush Signal state on empty drain flushes; align docs; keep .text in budget MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - An SKDM-only drain stanza mutates Signal state without enqueueing a message, and its buffered receipt flushes at drain end / teardown — right after the batcher flush call sites. The empty-batch early return I added with the Arc refactor silently dropped the "flush Signal even when there is nothing to commit" guarantee, so a crash after the receipt flush could lose an acked sender key. Empty drain flushes now persist the Signal cache before returning (codex). - commit_inbound_batch's doc said event -> acks while the code acks first; the doc now states the ack -> event contract and why (a crash between them trades exactly like the old at-most-once path: the consumer's durable copy is the hook commit, not the event) (cubic). - A replayed message now always dispatches with BatchOrigin::Live, and BatchOrigin is documented as the delivery shape (immediate vs accumulated drain batch) — a stanza's provenance is info.is_offline. Inferring origin from the batcher's current state could mislabel replays (cubic). - complete_offline_sync documents why flag-gated waiters may unblock during the tail commit (same window in-flight stanzas always had; OfflineSyncCompleted still dispatches only after the commit). - Revert the multi-row VALUES insert to per-row statements inside the single batch transaction: the WAL commit is the amortized cost either way, and the multi-row form added ~4 KiB of monomorphized .text against a 32 KiB per-PR budget currently at 31.9 KiB. --- src/client/sessions.rs | 6 ++++- src/message/commit_batch.rs | 23 +++++++++++++--- storages/sqlite-storage/src/sqlite_store.rs | 29 +++++++++------------ wacore/src/types/events.rs | 13 +++++---- 4 files changed, 46 insertions(+), 25 deletions(-) diff --git a/src/client/sessions.rs b/src/client/sessions.rs index 590d9de3b..3c0ebf8d7 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -41,7 +41,11 @@ impl Client { // finish_inbound_commit_drain for the raceless-transition // argument), BEFORE widening the semaphore. Receipts flush after, // so every receipt's message is durably committed first (WA Web's - // createSnapshot ordering). + // createSnapshot ordering). Flag-gated waiters can unblock while + // this tail commit runs — the same window in-flight stanzas always + // had across the flip — but everything they unblock is safe to run + // concurrently, and OfflineSyncCompleted only dispatches below, + // after the commit. if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { client.finish_inbound_commit_drain().await; } else { diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 989764f4b..6bdcc37a5 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -175,8 +175,16 @@ impl Client { /// processing; after the drain the batcher is empty and this no-ops. pub(crate) async fn flush_inbound_commits_acquiring_permit(self: &Arc) { let _permit = self.acquire_message_processing_permit().await; + let was_draining = self.inbound_commit_batch.is_active(); let batch = self.inbound_commit_batch.take(); if batch.is_empty() { + // Even with nothing to commit, a drain-mode flush must persist the + // Signal cache: SKDM-only stanzas mutate Signal state without + // enqueueing a message, and their buffered receipts flush right + // after the teardown/drain-end call sites of this function. + if was_draining { + self.flush_signal_cache_logged("commit_batch", None).await; + } return; } self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain, true) @@ -218,6 +226,11 @@ impl Client { let batch = self.inbound_commit_batch.take(); self.inbound_commit_batch.deactivate(); if batch.is_empty() { + // Same rule as flush_inbound_commits_acquiring_permit: SKDM-only + // drain stanzas leave Signal state in the cache with their + // receipts buffered; those receipts flush right after this, so + // the state must be durable first. + self.flush_signal_cache_logged("commit_batch", None).await; return; } self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain, true) @@ -225,9 +238,13 @@ impl Client { } /// Commit one batch: durable buffer → Signal flush → hook → clear buffer → - /// event → acks. WA Web ordering (`createSnapshot`), so nothing is acked or - /// observable before it is durable. On any commit failure everything stays - /// unacked and the server redelivers the whole batch. + /// acks → event. Nothing is acked or observable before it is durable (WA + /// Web's `createSnapshot` ordering); acks precede the event dispatch so a + /// misbehaving synchronous handler cannot suppress them — the contract the + /// pre-batch at-most-once path had. A crash between ack and event trades + /// exactly like that old path: the consumer's durable copy is the hook + /// commit, not the event. On any commit failure everything stays unacked + /// and the server redelivers the whole batch. #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.recv.commit_batch", level = "debug", skip_all, fields(count = items.len())))] pub(crate) async fn commit_inbound_batch( self: &Arc, diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index 4f5b7fd3f..654b01abe 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -3119,24 +3119,21 @@ impl ProtocolStore for SqliteStore { self.with_retry("store_pending_inbound_batch", || { let rows = Arc::clone(&rows); Box::new(move |conn: &mut SqliteConnection| { + // Per-row statements inside ONE transaction: the WAL commit is + // the real per-message cost and it is already amortized. A + // multi-row VALUES insert was measurably faster per statement + // but cost ~4 KiB of extra monomorphized .text against a + // 32 KiB per-PR budget — not worth it for microseconds. conn.transaction(|conn| { - // Multi-row VALUES per statement; chunked so binds stay far - // under SQLite's variable limit (5 columns × 100 rows). - for chunk in rows.chunks(100) { - let values: Vec<_> = chunk - .iter() - .map(|(chat, sender, id, message)| { - ( - pending_inbound_messages::chat.eq(chat), - pending_inbound_messages::sender.eq(sender), - pending_inbound_messages::id.eq(id), - pending_inbound_messages::message.eq(message.as_slice()), - pending_inbound_messages::device_id.eq(device_id), - ) - }) - .collect(); + for (chat, sender, id, message) in rows.iter() { diesel::replace_into(pending_inbound_messages::table) - .values(&values) + .values(( + pending_inbound_messages::chat.eq(chat), + pending_inbound_messages::sender.eq(sender), + pending_inbound_messages::id.eq(id), + pending_inbound_messages::message.eq(message.as_slice()), + pending_inbound_messages::device_id.eq(device_id), + )) .execute(conn)?; } Ok(()) diff --git a/wacore/src/types/events.rs b/wacore/src/types/events.rs index b4f70f54b..d186bdafa 100755 --- a/wacore/src/types/events.rs +++ b/wacore/src/types/events.rs @@ -867,14 +867,17 @@ pub struct InboundMessage { pub info: Arc, } -/// Where a [`MessageBatch`] came from. Mirrors Baileys' `messages.upsert` -/// `type` field (`notify` / `append`). +/// How a [`MessageBatch`] was delivered. Mirrors Baileys' `messages.upsert` +/// `type` field (`notify` / `append`). This describes the delivery shape, +/// not a message's provenance: whether a stanza came from the offline queue +/// is `info.is_offline` on each [`InboundMessage`]. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] pub enum BatchOrigin { - /// A message received while connected; always a batch of one. + /// Delivered immediately as a batch of one: live traffic, and redelivery + /// replays that commit outside an accumulated batch. Live, - /// Part of the offline backlog drained on (re)connect; batched per - /// durable commit (WA Web's MessageProcessorCache snapshot granularity). + /// An accumulated batch from the offline drain, one per durable commit + /// (WA Web's MessageProcessorCache snapshot granularity). OfflineDrain, } From d1d00191906cb9189b20882c4a4ce9ff1129ecd8 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 05:00:59 +0000 Subject: [PATCH 07/35] fix(recv): route redelivery replays through the commit batcher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit During a drain, a buffered redelivery used to commit and dispatch immediately while earlier freshly-decrypted stanzas were still accumulating in the batcher, so consumers could observe the replayed message ahead of its predecessors (codex). The replay's Ok branch now just calls commit_or_batch_inbound: in drain mode it joins the accumulating batch (arrival order preserved, its pending row rewritten and cleared by the batch commit), live it commits as a batch of one — which is store -> hook -> delete -> ack -> event, exactly what the hand-rolled branch did, so 30 lines fold away. This also supersedes dcdc148's origin change for replays: that hunk had silently failed to apply (same failure mode as the one confessed in 9c9bee7 — patch scripts now assert), and the batcher routing makes it moot: a drain replay is genuinely part of an OfflineDrain batch and a live replay genuinely a Live one. --- src/message/durability.rs | 96 ++++++++++++++------------------------- 1 file changed, 33 insertions(+), 63 deletions(-) diff --git a/src/message/durability.rs b/src/message/durability.rs index 6825910c4..fe54e6cf3 100644 --- a/src/message/durability.rs +++ b/src/message/durability.rs @@ -9,7 +9,7 @@ use super::*; use crate::types::durability_hook::InboundDurabilityHook; -use wacore::types::events::{BatchOrigin, InboundMessage, MessageBatch}; +use wacore::types::events::InboundMessage; impl Client { /// The registered inbound durability hook, if any. `None` (default) keeps @@ -19,74 +19,43 @@ impl Client { } /// Redelivery path: when the server replays an already-decrypted message - /// (`DuplicatedMessage`), re-run the hook from the buffered copy instead of - /// acking. A plain ack is sent only for a genuine duplicate (no buffered - /// copy). A read failure fails closed (no ack) so a transient storage error - /// cannot drop a message that still needs its hook to commit. + /// (`DuplicatedMessage`), re-commit it from the buffered copy instead of + /// acking. The replay routes through the commit batcher: during a drain it + /// joins the accumulating batch, so its hook commit, ack and event keep + /// arrival order with the fresh stanzas around it; live it commits + /// immediately as a batch of one. Either way the batch commit rewrites and + /// then clears its pending row, and consumers observe the message there — + /// its original batch never dispatched (the hook failed then). A plain ack + /// is sent only for a genuine duplicate (no buffered copy). A read failure + /// fails closed (no ack) so a transient storage error cannot drop a + /// message that still needs its hook to commit. pub(crate) async fn ack_or_replay_to_hook(self: &Arc, info: &Arc) { - if let Some(hook) = self.inbound_durability_hook() { + if self.inbound_durability_hook().is_some() { let backend = self.persistence_manager.backend(); let chat = info.source.chat.to_string(); let sender = info.source.sender.to_string(); match backend.get_pending_inbound(&chat, &sender, &info.id).await { - Ok(Some(bytes)) => { - match waproto::codec::message_decode(&bytes) { - Ok(msg) => { - let item = InboundMessage { - message: Arc::new(msg), - info: Arc::clone(info), - }; - match hook - .on_messages(self.clone(), std::slice::from_ref(&item)) - .await - { - Ok(()) => { - if let Err(e) = backend - .delete_pending_inbound(&chat, &sender, &info.id) - .await - { - log::debug!( - "[msg:{}] failed to clear buffered inbound message: {e:?}", - info.id - ); - } - self.ack_received_message(info); - // First successful commit of this message: - // its original batch never dispatched (the - // hook failed then), so consumers see it - // here or never. - let origin = if self.inbound_commit_batch.is_active() { - BatchOrigin::OfflineDrain - } else { - BatchOrigin::Live - }; - self.core.event_bus.dispatch(Event::Messages(MessageBatch { - messages: Arc::from([item]), - origin, - })); - } - Err(e) => { - log::warn!( - "[msg:{}] inbound durability hook still failing on redelivery; keeping for retry: {e:?}", - info.id - ); - } - } - } - Err(e) => { - // Corrupt row (our own serialization): it can never be - // replayed, so drop it and ack to unstick the queue. - log::error!( - "[msg:{}] failed to decode buffered inbound message; acking to unstick queue: {e:?}", - info.id - ); - let _ = backend - .delete_pending_inbound(&chat, &sender, &info.id) - .await; - self.ack_received_message(info); - } + Ok(Some(bytes)) => match waproto::codec::message_decode(&bytes) { + Ok(msg) => { + self.commit_or_batch_inbound(InboundMessage { + message: Arc::new(msg), + info: Arc::clone(info), + }) + .await; } - } + Err(e) => { + // Corrupt row (our own serialization): it can never be + // replayed, so drop it and ack to unstick the queue. + log::error!( + "[msg:{}] failed to decode buffered inbound message; acking to unstick queue: {e:?}", + info.id + ); + let _ = backend + .delete_pending_inbound(&chat, &sender, &info.id) + .await; + self.ack_received_message(info); + } + }, // Genuine duplicate (never buffered, or already committed): ack it. Ok(None) => self.ack_received_message(info), // Fail closed: a transient read error must not ack a message whose @@ -108,6 +77,7 @@ mod tests { use crate::test_utils::create_test_client_with_failing_http; use crate::types::message::MessageInfo; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use wacore::types::events::BatchOrigin; struct CountingHook { calls: AtomicUsize, From a3accf312a49f79cfbed721ce27ad71822b4eac0 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 05:13:04 +0000 Subject: [PATCH 08/35] fix(recv): log the self_weak upgrade failure in cleanup's drain flush Same class of failure as the complete_offline_sync twin that already logs: silently skipping the drain-batch commit right before the Signal flush is the acked-before-committed loss scenario, so it must be loud. --- src/client/lifecycle.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index c8bbd91a5..8b0b3a83e 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -756,6 +756,12 @@ impl Client { client .flush_inbound_commits_bounded(std::time::Duration::from_secs(5)) .await; + } else { + // Same class of bug as the complete_offline_sync twin: a silent + // skip here is the acked-before-committed loss — make it loud. + log::error!( + "cleanup_connection_state: self_weak upgrade failed; skipping drain-batch commit before Signal flush — uncommitted entries will be dropped" + ); } // Flush before clear: clear() drops dirty entries, so a disconnect // racing an in-flight encrypt would lose the just-advanced sender-key From 8c44cffad73a0239c53443b0c5a943cc309f7ff1 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 06:23:25 +0000 Subject: [PATCH 09/35] fix(recv): close self-review findings on the inbound commit batcher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Self-review of the batching PR surfaced six confirmed issues; all fixed: - Newsletters no longer enter the commit pipeline: the plaintext stanza is already transport-acked and never redelivered, so gating it on the hook (or dropping it on a batcher reset) would lose it permanently. They dispatch event-only, as the hook docs always promised. - complete_offline_sync no longer parks the read loop: the ib end marker is processed inline, and awaiting the tail commit there starved IQ responses/pongs (a hook doing a server round-trip would deadlock; a slow hook tripped the keepalive at the end of every drain). The finisher now runs as a spawned task, ordered by the single processing permit and guarded by the connection generation; a dedicated once-flag replaces the completed-flag CAS so offline_sync_completed flips only AFTER the tail commit — restoring aggregate offline receipts for the tail batch, which previously fell back to 1:1 receipts on every reconnect. - Global Signal-cache flushes from paths unrelated to the batcher (retry receipts, sends, identity changes, session APIs) could persist ratchet advances for accumulated entries with no durable row; a crash then turned each redelivery into a plain-acked duplicate. Those sites now go through batch-safe variants that commit the pending batch under the permit first; the UntrustedIdentity path (already holding the permit) commits inline. - Teardown loss window: a bounded-flush timeout followed by the unconditional signal flush persisted ratchets for entries reset() then dropped. Taken batches are now restored by a guard if the commit is cancelled or fails before the durable write, and cleanup clears the cache without flushing when uncommitted entries remain. - Buffered offline receipts now flush only after a durable tail commit (commit_inbound_batch reports it); on failure they are dropped so the server redelivers everything unacked — an SKDM receipt must never outlive its sender-key state. - The 3s flush timer re-checks its epoch after acquiring the permit (a stale sleeper could commit a batch only milliseconds old) and holds a Weak client so stranded sleepers cannot delay drop by up to 3s. API/UX cleanups from the same review: Event::message_text/message_texts removed (first-of-batch footgun; wait_for_text in the e2e helper was already bitten), InboundMessage/MessageBatch/BatchOrigin added to the prelude, Event::Messages documents its at-least-once shape with a hook and the newsletter/PDO exceptions. The drain flush entry points collapsed into one flush_inbound_commits_under_permit, the pending-inbound row statements are shared between single-row and batch paths, tests enter live mode via a single enter_live_mode_for_tests helper, and the live-path encode buffer reserves its exact size up front. --- src/client.rs | 6 + src/client/adapters.rs | 55 ++++ src/client/device_registry.rs | 2 +- src/client/lifecycle.rs | 64 +++-- src/client/sender_keys.rs | 2 +- src/client/sessions.rs | 127 +++++++--- src/client/tests.rs | 11 + src/features/message_edit.rs | 2 +- src/features/signal.rs | 12 +- src/handlers/notification/device.rs | 2 +- src/lib.rs | 2 +- src/message/commit_batch.rs | 265 +++++++++++++++----- src/message/dispatch.rs | 19 ++ src/message/durability.rs | 16 +- src/message/receive.rs | 7 + src/message/tests.rs | 18 +- src/msg_secret_buffer.rs | 2 +- src/receipt.rs | 8 +- src/retry.rs | 22 +- src/send/mod.rs | 12 +- src/test_utils.rs | 11 +- src/types/durability_hook.rs | 4 +- src/voip/facade.rs | 2 +- storages/sqlite-storage/src/sqlite_store.rs | 80 +++--- tests/e2e/src/lib.rs | 5 +- tests/e2e/tests/receipts.rs | 2 +- wacore/src/types/events.rs | 23 +- 27 files changed, 559 insertions(+), 222 deletions(-) diff --git a/src/client.rs b/src/client.rs index 3248f8051..9df761528 100644 --- a/src/client.rs +++ b/src/client.rs @@ -525,7 +525,13 @@ pub struct Client { /// WhatsApp Web waits for this before sending passive tasks (prekey upload, active IQ, presence). pub(crate) offline_sync_notifier: Arc, /// Flag indicating offline sync has completed (received ib offline stanza). + /// Flips only AFTER the drain-tail commit, so the tail's acks still join + /// the aggregate offline-receipt drain. pub(crate) offline_sync_completed: Arc, + /// Once-guard for the drain finisher (the semaphore swap is not + /// idempotent). Separate from `offline_sync_completed` because the finish + /// runs off the read loop and the flag must flip only after its commit. + pub(crate) offline_sync_finish_started: Arc, /// Delivery receipts buffered during offline sync, flushed as aggregate /// `` stanzas at completion (WA Web `sendAggregateOfflineReceipts`). /// Empty (zero capacity) outside the offline window. diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 175ac9f41..08ee38dd3 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -85,6 +85,12 @@ impl Client { } /// [`flush_signal_cache`](Self::flush_signal_cache) with error logging instead of propagation. + /// + /// Both of these are safe only when the caller holds the message + /// processing permit or the batcher is known inactive: they persist the + /// WHOLE cache, including ratchet advances of drain entries that may not + /// have a durable buffered row yet. Everything else must go through the + /// `_batch_safe` variants below. pub(crate) async fn flush_signal_cache_logged(&self, context: &str, id: Option<&str>) { if let Err(e) = self.flush_signal_cache().await { if let Some(id) = id { @@ -94,4 +100,53 @@ impl Client { } } } + + /// Signal-cache flush that is safe while the offline drain is active. + /// + /// During the drain, decrypted messages accumulate in the commit batcher + /// with no durable buffered copy; flushing the cache from an unrelated + /// path (a retry receipt, a send, an identity change) would persist their + /// ratchet advances, and a crash/teardown that then drops the entries + /// turns each redelivery into an ackable duplicate — silent loss for hook + /// consumers. So in drain mode this routes through the batcher: commit + /// the pending entries (rows first) and flush under the processing + /// permit. Outside the drain it is exactly [`Self::flush_signal_cache`]. + /// + /// Must NOT be called while holding the processing permit (it acquires + /// it); permit-holding paths commit via the batcher directly. + pub(crate) async fn flush_signal_cache_batch_safe(&self) -> Result<(), anyhow::Error> { + if self.inbound_commit_batch.is_active() { + if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { + return if client.flush_inbound_commits_under_permit(false, None).await { + Ok(()) + } else { + Err(anyhow::anyhow!( + "inbound drain batch commit failed; Signal cache left unflushed so the server redelivers" + )) + }; + } + if self.inbound_commit_batch.has_entries() { + return Err(anyhow::anyhow!( + "client dropping with uncommitted drain entries; skipping Signal flush" + )); + } + } + self.flush_signal_cache().await + } + + /// [`flush_signal_cache_batch_safe`](Self::flush_signal_cache_batch_safe) + /// with error logging instead of propagation. + pub(crate) async fn flush_signal_cache_batch_safe_logged( + &self, + context: &str, + id: Option<&str>, + ) { + if let Err(e) = self.flush_signal_cache_batch_safe().await { + if let Some(id) = id { + log::error!("Failed to flush signal cache ({context} {id}): {e:?}"); + } else { + log::error!("Failed to flush signal cache ({context}): {e:?}"); + } + } + } } diff --git a/src/client/device_registry.rs b/src/client/device_registry.rs index 724432866..ba84957b7 100644 --- a/src/client/device_registry.rs +++ b/src/client/device_registry.rs @@ -673,7 +673,7 @@ impl Client { } } } - self.flush_signal_cache_logged("delete_sessions_for_devices", None) + self.flush_signal_cache_batch_safe_logged("delete_sessions_for_devices", None) .await; } diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 8b0b3a83e..96e57cfae 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -214,6 +214,7 @@ impl Client { prekey_upload_lock: Arc::new(async_lock::Mutex::new(())), offline_sync_notifier: Arc::new(event_listener::Event::new()), offline_sync_completed: Arc::new(AtomicBool::new(false)), + offline_sync_finish_started: Arc::new(AtomicBool::new(false)), offline_receipt_buffer: std::sync::Mutex::new(Vec::new()), inbound_commit_batch: Default::default(), history_sync_tasks_in_flight: Arc::new(AtomicUsize::new(0)), @@ -455,6 +456,8 @@ impl Client { self.is_ready.store(false, Ordering::Relaxed); self.is_connected.store(false, Ordering::Relaxed); self.offline_sync_completed.store(false, Ordering::Relaxed); + self.offline_sync_finish_started + .store(false, Ordering::Relaxed); self.clear_offline_receipt_buffer(); // Uncommitted batch entries were never acked; the server redelivers // them on this fresh connection. @@ -580,10 +583,16 @@ impl Client { // // Commit any accumulated drain batch first so its acks land in this // receipt drain. Bounded like the outbound flush below: on timeout the - // entries simply stay unacked and the server redelivers them. - self.flush_inbound_commits_bounded(std::time::Duration::from_secs(5)) - .await; - self.flush_offline_receipts(); + // entries simply stay unacked and the server redelivers them — and the + // buffered receipts stay unsent too, because their SKDM/session state + // may not be durable yet (receipting an SKDM whose sender key only + // lives in the cache would lose it to a crash with no redelivery). + if self + .flush_inbound_commits_bounded(std::time::Duration::from_secs(5)) + .await + { + self.flush_offline_receipts(); + } // Prevent late receipt producers from escaping the drain window. self.outbound_flush.close(); self.outbound_flush @@ -641,9 +650,13 @@ impl Client { self.auto_reconnect_errors .store(Self::RECONNECT_BACKOFF_STEP, Ordering::Relaxed); - self.flush_inbound_commits_bounded(std::time::Duration::from_secs(2)) - .await; - self.flush_offline_receipts(); + // Same durable-before-receipts gate as disconnect(). + if self + .flush_inbound_commits_bounded(std::time::Duration::from_secs(2)) + .await + { + self.flush_offline_receipts(); + } self.outbound_flush.close(); self.outbound_flush .flush(&*self.runtime, std::time::Duration::from_secs(2)) @@ -668,9 +681,13 @@ impl Client { info!("Reconnecting immediately (expected disconnect)."); self.expected_disconnect.store(true, Ordering::Relaxed); - self.flush_inbound_commits_bounded(std::time::Duration::from_secs(2)) - .await; - self.flush_offline_receipts(); + // Same durable-before-receipts gate as disconnect(). + if self + .flush_inbound_commits_bounded(std::time::Duration::from_secs(2)) + .await + { + self.flush_offline_receipts(); + } self.outbound_flush.close(); self.outbound_flush .flush(&*self.runtime, std::time::Duration::from_secs(2)) @@ -768,11 +785,26 @@ impl Client { // chain and force a full SKDM re-fanout. A disconnect is not a logout. // Only clear on a successful flush; on a backend error keep the cache so // the dirty state isn't dropped and the next operation can persist it. - match self.flush_signal_cache().await { - Ok(()) => self.signal_cache.clear().await, - Err(e) => log::error!( - "cleanup_connection_state: signal cache flush failed, keeping cache to avoid dropping Signal state: {e:?}" - ), + // + // Exception: if the bounded commit above left entries in the batcher + // (timed out or its durable write failed), the cache holds their + // ratchet advances. Persisting those would make each redelivery an + // ackable duplicate; keeping them cached would make it decrypt as one + // in-process. Clearing WITHOUT flushing un-advances them instead — + // everything they cover is unacked and redelivers fresh. Committed + // state is not at risk: every earlier commit flushed the cache itself. + if self.inbound_commit_batch.has_entries() { + log::warn!( + "cleanup_connection_state: dropping unflushed Signal state for uncommitted drain entries; the server redelivers them" + ); + self.signal_cache.clear().await; + } else { + match self.flush_signal_cache().await { + Ok(()) => self.signal_cache.clear().await, + Err(e) => log::error!( + "cleanup_connection_state: signal cache flush failed, keeping cache to avoid dropping Signal state: {e:?}" + ), + } } // Reset semaphore to 1 permit for next offline sync. self.swap_message_semaphore(1); @@ -783,6 +815,8 @@ impl Client { self.pending_device_sync.clear().await; // Reset offline sync state for next connection self.offline_sync_completed.store(false, Ordering::Relaxed); + self.offline_sync_finish_started + .store(false, Ordering::Relaxed); self.clear_offline_receipt_buffer(); // Same rule as receipts: uncommitted entries drop here and the server // redelivers them on the next connect. diff --git a/src/client/sender_keys.rs b/src/client/sender_keys.rs index 2058f81aa..e5e5aeadf 100644 --- a/src/client/sender_keys.rs +++ b/src/client/sender_keys.rs @@ -123,7 +123,7 @@ impl Client { .delete_sender_key(sk_name.cache_key()) .await; } - self.flush_signal_cache_logged("rotate_sender_key_on_participant_remove", None) + self.flush_signal_cache_batch_safe_logged("rotate_sender_key_on_participant_remove", None) .await; if let Err(e) = self diff --git a/src/client/sessions.rs b/src/client/sessions.rs index 3c0ebf8d7..8df93a47b 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -1,6 +1,7 @@ //! E2E Session management for Client. use anyhow::Result; +use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::Duration; use wacore::libsignal::store::SessionStore; @@ -23,57 +24,105 @@ impl Client { Err(poison) => *poison.into_inner() = None, } - // Signal that offline sync is complete - post-login tasks are waiting for this. - // This mimics WhatsApp Web's offlineDeliveryEnd event. - // Use compare_exchange to ensure we only run this once (add_permits is NOT idempotent). - // Readers that observe offline_sync_completed=true short-circuit without touching - // the semaphore (wait_for_offline_delivery_end returns early), so the ordering of - // flag flip vs. semaphore swap below is not observable: any in-flight worker keeps - // using its old 1-permit Arc and drains normally; newly-spawned workers pick up the - // 64-permit semaphore via read_message_semaphore(). + // Run the finisher once (the semaphore swap is not idempotent). The + // guard is a dedicated flag, NOT offline_sync_completed: that one only + // flips after the tail commit so the tail's acks still observe it as + // false and join the aggregate offline-receipt drain (WA Web + // `sendAggregateOfflineReceipts`) instead of going out 1:1. if self - .offline_sync_completed + .offline_sync_finish_started .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) - .is_ok() + .is_err() { - // Commit the drain tail and flip the batcher to live mode, both - // under the still-single processing permit (see - // finish_inbound_commit_drain for the raceless-transition - // argument), BEFORE widening the semaphore. Receipts flush after, - // so every receipt's message is durably committed first (WA Web's - // createSnapshot ordering). Flag-gated waiters can unblock while - // this tail commit runs — the same window in-flight stanzas always - // had across the flip — but everything they unblock is safe to run - // concurrently, and OfflineSyncCompleted only dispatches below, - // after the commit. - if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { - client.finish_inbound_commit_drain().await; - } else { - // Practically unreachable (the run loop owns a strong Arc), but - // a silent skip here would be exactly the acked-before-committed - // bug this ordering exists to prevent — make it loud. - log::error!( - "complete_offline_sync: self_weak upgrade failed; skipping drain-tail commit before widening the semaphore" - ); - } + return; + } - // Allow parallel message processing now that offline sync is done. - // During offline sync, permits=1 serialized all message processing. - // Replace with a new semaphore with 64 permits for concurrent processing. - // Old workers holding the previous semaphore Arc will finish normally. + let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) else { + // Practically unreachable (the run loop owns a strong Arc), but a + // silent skip would leave the client batching forever with a + // widened semaphore — leave consistent live state behind instead. + log::error!( + "complete_offline_sync: self_weak upgrade failed; dropping the drain tail and switching to live mode" + ); + self.inbound_commit_batch.force_live_dropping_entries(); + self.offline_sync_completed.store(true, Ordering::Release); self.swap_message_semaphore(64); + self.offline_sync_notifier.notify(usize::MAX); + self.core + .event_bus + .dispatch(Event::OfflineSyncCompleted(OfflineSyncCompleted { count })); + return; + }; + + // The `ib` offline end marker is processed INLINE on the read loop, and + // the tail commit awaits the processing permit plus the durable write, + // Signal flush and the consumer's durability hook. Parking the read + // loop on that would starve IQ responses and pongs — a hook awaiting + // any server round-trip would deadlock, and a merely slow hook would + // trip the keepalive at the end of every drain. Ordering does not need + // the inline await: the single permit already serializes the finisher + // against stanza processing, so run it off-loop. + let generation = self.connection_generation.load(Ordering::Acquire); + self.runtime + .spawn(Box::pin(async move { + client.finish_offline_sync(count, generation).await; + })) + .detach(); + } + /// Off-read-loop tail of [`complete_offline_sync`]: commit the drain tail, + /// then flip the completed flag, widen the semaphore and flush the + /// aggregate receipts. `generation` guards against a reconnect racing this + /// task — the new connection resets the drain state and must not have its + /// flag/semaphore/batcher touched by the old connection's finisher. + async fn finish_offline_sync(self: &Arc, count: i32, generation: u64) { + // Commit the drain tail and flip the batcher to live mode under the + // still-single processing permit (see flush_inbound_commits_under_permit + // for the raceless-transition argument), BEFORE widening the semaphore. + // Receipts flush after, so every receipt's message is durably committed + // first (WA Web's createSnapshot ordering). + let durable = self.finish_inbound_commit_drain().await; + + if self.connection_generation.load(Ordering::Acquire) != generation { + log::debug!( + "finish_offline_sync: connection generation changed during the tail commit; leaving the new connection's state alone" + ); + return; + } + + // Readers that observe offline_sync_completed=true short-circuit + // without touching the semaphore (wait_for_offline_delivery_end + // returns early), so the ordering of flag flip vs. semaphore swap is + // not observable: any in-flight worker keeps using its old 1-permit + // Arc and drains normally; newly-spawned workers pick up the 64-permit + // semaphore via read_message_semaphore(). + self.offline_sync_completed.store(true, Ordering::Release); + + // Allow parallel message processing now that offline sync is done. + self.swap_message_semaphore(64); + + if durable { // The flag flip above happens-before this drain takes the buffer // lock, so late offline receipts either land in this flush or // observe the flag and send 1:1 (see try_buffer_offline_receipt). self.flush_offline_receipts(); + } else { + // The tail's durable write failed: its entries are back in the + // batcher, unacked. Buffered receipts must not go out either — + // they may cover SKDM/session state that never became durable, and + // receipting those would trade a redeliverable failure for a + // crash-permanent one. Everything unacked redelivers next connect. + log::warn!( + "finish_offline_sync: tail commit not durable; dropping buffered offline receipts so the server redelivers" + ); + self.clear_offline_receipt_buffer(); + } - self.offline_sync_notifier.notify(usize::MAX); + self.offline_sync_notifier.notify(usize::MAX); - self.core - .event_bus - .dispatch(Event::OfflineSyncCompleted(OfflineSyncCompleted { count })); - } + self.core + .event_bus + .dispatch(Event::OfflineSyncCompleted(OfflineSyncCompleted { count })); } /// Wait for offline message delivery to complete (with timeout). diff --git a/src/client/tests.rs b/src/client/tests.rs index 648065d87..bbcb9209d 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -503,6 +503,17 @@ async fn test_wait_for_offline_delivery_end_times_out_when_flag_not_set() { .await; let elapsed = start.elapsed(); + // The drain finisher runs as a spawned task (off the read loop); wait for + // its completion signal before asserting on the state it flips. + let listener = client.offline_sync_notifier.listen(); + if !client + .offline_sync_completed + .load(std::sync::atomic::Ordering::Acquire) + { + tokio::time::timeout(std::time::Duration::from_secs(5), listener) + .await + .expect("drain finisher should complete"); + } // Count available permits by trying to acquire non-blockingly let semaphore = match client.message_processing_semaphore.lock() { Ok(guard) => guard.clone(), diff --git a/src/features/message_edit.rs b/src/features/message_edit.rs index db7e09c8b..77f13f94e 100644 --- a/src/features/message_edit.rs +++ b/src/features/message_edit.rs @@ -12,7 +12,7 @@ //! doing so requires a callback into the consumer's message store to //! fetch the parent's `messageContextInfo.messageSecret`. Consumers: //! -//! 1. Observe `Event::Message` for messages whose +//! 1. Observe `Event::Messages` for messages whose //! `message.secret_encrypted_message.secret_enc_type == MessageEdit`. //! 2. Detect the envelope with [`extract_envelope`]. //! 3. Look up the targeted message via `target_message_key`. diff --git a/src/features/signal.rs b/src/features/signal.rs index a0ed865ee..9be5d3d7b 100644 --- a/src/features/signal.rs +++ b/src/features/signal.rs @@ -70,7 +70,7 @@ impl<'a> Signal<'a> { .await?; drop(_guard); - self.client.flush_signal_cache().await?; + self.client.flush_signal_cache_batch_safe().await?; let (_, is_prekey, bytes) = wacore::send::extract_ciphertext(encrypted) .ok_or_else(|| SignalError::Unsupported("unexpected ciphertext variant".into()))?; @@ -144,7 +144,7 @@ impl<'a> Signal<'a> { } drop(_guard); - self.client.flush_signal_cache().await?; + self.client.flush_signal_cache_batch_safe().await?; Ok(decrypted.plaintext) } @@ -208,7 +208,7 @@ impl<'a> Signal<'a> { ) .await?; - self.client.flush_signal_cache().await?; + self.client.flush_signal_cache_batch_safe().await?; Ok((skdm_bytes, ciphertext.into_serialized().into_vec())) } @@ -238,7 +238,7 @@ impl<'a> Signal<'a> { ) .await?; - self.client.flush_signal_cache().await?; + self.client.flush_signal_cache_batch_safe().await?; Ok(plaintext.to_vec()) } @@ -279,7 +279,7 @@ impl<'a> Signal<'a> { self.client.signal_cache.delete_identity(&addr).await; } - self.client.flush_signal_cache().await?; + self.client.flush_signal_cache_batch_safe().await?; Ok(()) } @@ -323,7 +323,7 @@ impl<'a> Signal<'a> { .await?; drop(_session_guards); - self.client.flush_signal_cache().await?; + self.client.flush_signal_cache_batch_safe().await?; Ok((result.participant_nodes, result.includes_prekey_message)) } diff --git a/src/handlers/notification/device.rs b/src/handlers/notification/device.rs index 3bd3539c1..4ded278ee 100644 --- a/src/handlers/notification/device.rs +++ b/src/handlers/notification/device.rs @@ -278,7 +278,7 @@ pub(crate) async fn handle_identity_change(client: &Arc, node: &NodeRef< } client - .flush_signal_cache_logged("identity change", None) + .flush_signal_cache_batch_safe_logged("identity change", None) .await; } diff --git a/src/lib.rs b/src/lib.rs index 04f11ba69..908ea698c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -145,7 +145,7 @@ pub mod prelude { pub use crate::send::{SendError, SendOptions, SendResult}; #[cfg(feature = "sqlite-storage")] pub use crate::store::SqliteStore; - pub use crate::types::events::{Event, EventKind}; + pub use crate::types::events::{BatchOrigin, Event, EventKind, InboundMessage, MessageBatch}; pub use crate::types::message::MessageInfo; pub use crate::{Jid, Server}; pub use wacore::proto_helpers::{MessageBuilderExt, MessageExt}; diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 6bdcc37a5..d1e3f2551 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -84,12 +84,36 @@ impl InboundCommitBatcher { self.active.load(Ordering::Acquire) } + /// Whether uncommitted entries are still accumulated. Teardown uses this + /// to decide if flushing the Signal cache is safe: persisting ratchet + /// advances for entries about to be dropped would turn their redelivery + /// into an unrecoverable duplicate. + pub(crate) fn has_entries(&self) -> bool { + !self.lock().entries.is_empty() + } + /// Switch to immediate (live) commits. Only the end-of-drain flush calls /// this, while holding the single processing permit. fn deactivate(&self) { self.active.store(false, Ordering::Release); } + /// Last-resort drain exit for the unreachable-in-practice case where the + /// finisher cannot run (`self_weak` upgrade failure): drop any entries + /// (unacked, so the server redelivers them) and still switch to live mode + /// — staying active with a widened semaphore would batch live traffic + /// while flushes hold only 1 of its permits. + pub(crate) fn force_live_dropping_entries(&self) { + let dropped = self.take(); + if !dropped.is_empty() { + log::warn!( + "Dropping {} uncommitted inbound messages on forced drain exit; the server will redeliver them", + dropped.len() + ); + } + self.deactivate(); + } + /// Connection teardown/setup: drop uncommitted entries (they were never /// acked, so the server redelivers them) and re-arm accumulation for the /// next connection's drain. @@ -103,11 +127,63 @@ impl InboundCommitBatcher { } self.active.store(true, Ordering::Release); } +} - /// Test-only: enter live mode without running a drain flush. - #[cfg(test)] - pub(crate) fn deactivate_for_tests(&self) { - self.deactivate(); +#[cfg(test)] +impl Client { + /// Test-only mirror of a completed offline sync: the flag, live permit + /// count and batcher mode move together so tests can never run in a + /// hybrid drain/live state production cannot reach. Kept next to the + /// production transition (`finish_offline_sync`) so a new step gets added + /// to both. + pub(crate) fn enter_live_mode_for_tests(&self) { + self.offline_sync_completed.store(true, Ordering::Release); + self.offline_sync_finish_started + .store(true, Ordering::Release); + self.swap_message_semaphore(64); + self.inbound_commit_batch.deactivate(); + } +} + +/// Restores a taken drain batch if the commit is cancelled (bounded teardown +/// timeout) or fails before its durable point: taken entries must never +/// vanish while another path can still flush their ratchet advances — a +/// persisted ratchet without a buffered row makes the redelivery an ackable +/// duplicate (silent loss for hook consumers). +struct ReinsertGuard<'a> { + batcher: &'a InboundCommitBatcher, + items: Option>, +} + +impl ReinsertGuard<'_> { + fn disarm(&mut self) { + self.items = None; + } +} + +impl Drop for ReinsertGuard<'_> { + fn drop(&mut self) { + let Some(items) = self.items.take() else { + return; + }; + let mut state = self.batcher.lock(); + let mut restored: Vec = items.iter().cloned().collect(); + // Nothing newer can normally exist (the permit serializes drain + // flushes), but keep arrival order if it ever does. + restored.append(&mut state.entries); + state.bytes = restored + .iter() + .map(|i| waproto::codec::message_encoded_len(&i.message)) + .sum(); + // The taking flush bumped the epoch, so no sleeper owns this batch; + // the next enqueue arms a fresh one, and the drain-end/teardown + // flushes cover the gap regardless. + state.timer_armed = false; + state.entries = restored; + log::warn!( + "Restored {} uncommitted inbound messages to the batch after a failed or cancelled commit", + state.entries.len() + ); } } @@ -135,17 +211,28 @@ impl Client { // Arc::from([item]) builds the event/hook slice in one allocation; // a Vec would add an alloc+dealloc per live message (measured // ~18ns and 2x the allocations of this step). - self.commit_inbound_batch(std::sync::Arc::from([item]), BatchOrigin::Live, false) + let _ = self + .commit_inbound_batch(std::sync::Arc::from([item]), BatchOrigin::Live) .await; return; } if let Some(epoch) = self.enqueue_inbound_commit(item) { - let client = self.clone(); + // Weak: a sleeper must not keep the whole Client graph alive for + // up to 3s after the app drops its handle. + let client = Arc::downgrade(self); + let runtime = self.runtime.clone(); self.runtime .spawn(Box::pin(async move { - client.runtime.sleep(FLUSH_TIMEOUT).await; - if client.inbound_commit_batch.epoch.load(Ordering::Acquire) == epoch { - client.flush_inbound_commits_acquiring_permit().await; + runtime.sleep(FLUSH_TIMEOUT).await; + if let Some(client) = client.upgrade() + && client.inbound_commit_batch.epoch.load(Ordering::Acquire) == epoch + { + // The epoch re-checks after permit acquisition: a + // size-trigger flush racing this sleeper must not let + // it commit a batch that is only milliseconds old. + let _ = client + .flush_inbound_commits_under_permit(false, Some(epoch)) + .await; } })) .detach(); @@ -162,7 +249,8 @@ impl Client { }; if over { let batch = self.inbound_commit_batch.take(); - self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain, true) + let _ = self + .commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) .await; } } @@ -173,11 +261,34 @@ impl Client { /// its redelivery into an unrecoverable duplicate. During the drain the /// semaphore holds a single permit, so this fully serializes with stanza /// processing; after the drain the batcher is empty and this no-ops. - pub(crate) async fn flush_inbound_commits_acquiring_permit(self: &Arc) { + /// + /// With `deactivate`, this is the end-of-drain transition: commit the tail + /// batch and switch the batcher to live mode under the same permit hold. + /// That is what makes the transition raceless — no stanza is mid-flight + /// when the mode flips (a stanza's enqueue and its stanza-end flush always + /// agree), and every stanza still queued behind this permit commits as + /// Live strictly AFTER the tail batch, preserving arrival order across the + /// boundary. Runs before the semaphore widens to the live permit count. + /// + /// Returns whether the drain state is durable (see + /// [`commit_inbound_batch`](Self::commit_inbound_batch)); callers gate the + /// buffered offline-receipt flush on it. + pub(crate) async fn flush_inbound_commits_under_permit( + self: &Arc, + deactivate: bool, + expected_epoch: Option, + ) -> bool { let _permit = self.acquire_message_processing_permit().await; + if let Some(epoch) = expected_epoch + && self.inbound_commit_batch.epoch.load(Ordering::Acquire) != epoch + { + // Another flush took this sleeper's batch while it waited for the + // permit; whatever accumulates now belongs to a newer timer. + return true; + } let was_draining = self.inbound_commit_batch.is_active(); let batch = self.inbound_commit_batch.take(); - if batch.is_empty() { + let durable = if batch.is_empty() { // Even with nothing to commit, a drain-mode flush must persist the // Signal cache: SKDM-only stanzas mutate Signal state without // enqueueing a message, and their buffered receipts flush right @@ -185,56 +296,61 @@ impl Client { if was_draining { self.flush_signal_cache_logged("commit_batch", None).await; } - return; + true + } else { + self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) + .await + }; + if deactivate { + self.inbound_commit_batch.deactivate(); } - self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain, true) - .await; + durable } - /// [`flush_inbound_commits_acquiring_permit`](Self::flush_inbound_commits_acquiring_permit) + /// [`flush_inbound_commits_under_permit`](Self::flush_inbound_commits_under_permit) /// with a deadline, for connection teardown: a stalled permit holder or a - /// hung hook must not wedge disconnect/reconnect. On timeout the entries - /// stay unacked and the server redelivers them on the next connect. + /// hung hook must not wedge disconnect/reconnect. On timeout (or a failed + /// durable write) the entries stay in the batcher unacked — the caller + /// must not flush buffered receipts or persist the Signal cache for them, + /// and the server redelivers them on the next connect. pub(crate) async fn flush_inbound_commits_bounded( self: &Arc, limit: std::time::Duration, - ) { - if wacore::runtime::timeout( + ) -> bool { + match wacore::runtime::timeout( &*self.runtime, limit, - self.flush_inbound_commits_acquiring_permit(), + self.flush_inbound_commits_under_permit(false, None), ) .await - .is_err() { - log::warn!( - "Timed out committing the inbound drain batch during teardown; leaving entries for redelivery" - ); + Ok(durable) => durable, + Err(_) => { + log::warn!( + "Timed out committing the inbound drain batch during teardown; leaving entries for redelivery" + ); + false + } } } - /// End-of-drain transition: commit the tail batch and switch the batcher - /// to live mode, all under the single processing permit. Holding the - /// permit across BOTH steps is what makes the transition raceless: no - /// stanza is mid-flight when the mode flips (so a stanza's enqueue and its - /// stanza-end flush always agree), and every stanza still queued behind - /// this permit commits as Live strictly AFTER the tail batch — arrival - /// order is preserved across the boundary. Runs before the semaphore - /// widens to the live permit count. - pub(crate) async fn finish_inbound_commit_drain(self: &Arc) { - let _permit = self.acquire_message_processing_permit().await; + /// End-of-drain transition; see + /// [`flush_inbound_commits_under_permit`](Self::flush_inbound_commits_under_permit). + pub(crate) async fn finish_inbound_commit_drain(self: &Arc) -> bool { + self.flush_inbound_commits_under_permit(true, None).await + } + + /// Commit any accumulated entries while the caller ALREADY HOLDS the + /// processing permit (mid-stanza recovery paths like UntrustedIdentity), + /// so a full-cache Signal flush that follows cannot persist ratchet + /// advances for entries without a durable buffered row. + pub(crate) async fn commit_inbound_batch_holding_permit(self: &Arc) { let batch = self.inbound_commit_batch.take(); - self.inbound_commit_batch.deactivate(); - if batch.is_empty() { - // Same rule as flush_inbound_commits_acquiring_permit: SKDM-only - // drain stanzas leave Signal state in the cache with their - // receipts buffered; those receipts flush right after this, so - // the state must be durable first. - self.flush_signal_cache_logged("commit_batch", None).await; - return; + if !batch.is_empty() { + let _ = self + .commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) + .await; } - self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain, true) - .await; } /// Commit one batch: durable buffer → Signal flush → hook → clear buffer → @@ -245,16 +361,32 @@ impl Client { /// exactly like that old path: the consumer's durable copy is the hook /// commit, not the event. On any commit failure everything stays unacked /// and the server redelivers the whole batch. + /// + /// Drain commits also flush the Signal cache (bulk signal-store commit per + /// snapshot, WA Web ordering); live commits leave it to the per-stanza + /// flush at the end of processing. A drain batch whose durable write fails + /// (or whose future is cancelled by a bounded teardown flush before it) is + /// restored to the batcher, so "entries still batched" stays an accurate + /// signal for teardown's flush-vs-drop decision. + /// + /// Returns whether the durable state (buffer rows + Signal flush attempt) + /// committed — the gate for flushing buffered offline receipts. A hook + /// failure still returns `true`: its rows are durable and the replay path + /// retries it, so already-buffered receipts of other messages stay safe. #[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.recv.commit_batch", level = "debug", skip_all, fields(count = items.len())))] pub(crate) async fn commit_inbound_batch( self: &Arc, items: std::sync::Arc<[InboundMessage]>, origin: BatchOrigin, - flush_signal: bool, - ) { + ) -> bool { if items.is_empty() { - return; + return true; } + let is_drain = matches!(origin, BatchOrigin::OfflineDrain); + let mut reinsert = ReinsertGuard { + batcher: &self.inbound_commit_batch, + items: is_drain.then(|| std::sync::Arc::clone(&items)), + }; if let Some(hook) = self.inbound_durability_hook() { // Key strings live for the whole commit; rows borrow them and the @@ -279,11 +411,14 @@ impl Client { { let mut local_arena; let mut shared_arena; - let arena: &mut Vec = if matches!(origin, BatchOrigin::OfflineDrain) { + let arena: &mut Vec = if is_drain { shared_arena = self.inbound_commit_batch.arena.lock().await; &mut shared_arena } else { - local_arena = Vec::new(); + // Exact-size reservation: geometric growth would realloc + // several times per live message. + local_arena = + Vec::with_capacity(waproto::codec::message_encoded_len(&items[0].message)); &mut local_arena }; arena.clear(); @@ -306,18 +441,21 @@ impl Client { .collect(); // Fail closed: without a durable buffered copy, do not run the - // hook and do not ack — the server redelivers once storage - // recovers. + // hook and do not ack — the entries return to the batcher (via + // the guard) and the server redelivers once storage recovers. if let Err(e) = backend.store_pending_inbound_batch(&rows).await { log::error!( "Failed to buffer inbound batch of {}; suppressing acks for redelivery: {e:?}", items.len() ); - return; + return false; } } + // Rows are durable: from here on, redelivery replays from them, so + // a cancelled future must not restore the entries. + reinsert.disarm(); - if flush_signal { + if is_drain { self.flush_signal_cache_logged("commit_batch", None).await; } @@ -326,7 +464,7 @@ impl Client { "Inbound durability hook failed for batch of {}; suppressing acks for redelivery: {e:?}", items.len() ); - return; + return true; } let delete_keys: Vec> = items @@ -346,8 +484,12 @@ impl Client { delete_keys.len() ); } - } else if flush_signal { - self.flush_signal_cache_logged("commit_batch", None).await; + } else { + if is_drain { + self.flush_signal_cache_logged("commit_batch", None).await; + } + // No hook = at-most-once: the flush is the durable point. + reinsert.disarm(); } // Acks first (everything durable by now): handle_event runs @@ -361,6 +503,7 @@ impl Client { messages: items, origin, })); + true } } @@ -431,7 +574,7 @@ mod tests { "sub-threshold entries must accumulate, not commit" ); - client.flush_inbound_commits_acquiring_permit().await; + client.flush_inbound_commits_under_permit(false, None).await; let batches = hook.batches.lock().expect("hook lock").clone(); assert_eq!(batches, vec![vec!["B1", "B2", "B3"]]); @@ -510,7 +653,7 @@ mod tests { client.commit_or_batch_inbound(item("N1")).await; client.commit_or_batch_inbound(item("N2")).await; - client.flush_inbound_commits_acquiring_permit().await; + client.flush_inbound_commits_under_permit(false, None).await; let event = rx.try_recv().expect("one batch event"); assert_eq!( @@ -524,7 +667,7 @@ mod tests { // End-of-drain transition: the tail batch commits first (as OfflineDrain), // the batcher flips to live mode, and anything after commits as Live — - // never interleaved ahead of the tail (cubic/codex P1 regression). + // never interleaved ahead of the tail. #[tokio::test] async fn finish_drain_commits_tail_then_switches_to_live() { let client = create_test_client_with_failing_http("batch_transition").await; @@ -576,7 +719,7 @@ mod tests { client.commit_or_batch_inbound(item("C1")).await; client.inbound_commit_batch.reset(); - client.flush_inbound_commits_acquiring_permit().await; + client.flush_inbound_commits_under_permit(false, None).await; assert!(hook.batches.lock().expect("hook lock").is_empty()); assert!(rx.try_recv().is_err()); diff --git a/src/message/dispatch.rs b/src/message/dispatch.rs index d4e2aaa00..82911c8c4 100644 --- a/src/message/dispatch.rs +++ b/src/message/dispatch.rs @@ -38,6 +38,25 @@ impl Client { } let dispatch_msg = Arc::new(decrypted.unwrap_or(msg)); + // Newsletters never enter the commit pipeline: the plaintext stanza + // was already transport-acked at enqueue and the server never + // redelivers it, so a hook failure (or a batcher reset) would lose + // the message for good instead of trading on redelivery. They also + // reach here without the processing permit, so enqueueing could + // straddle the drain→live transition. + if info.source.chat.is_newsletter() { + self.core + .event_bus + .dispatch(Event::Messages(wacore::types::events::MessageBatch { + messages: Arc::from([wacore::types::events::InboundMessage { + message: dispatch_msg, + info, + }]), + origin: wacore::types::events::BatchOrigin::Live, + })); + return; + } + // Live traffic commits (and acks) as a batch of one; during the // offline drain the message joins the accumulating commit batch and // the event/ack fire only after its batch commits. Either way the diff --git a/src/message/durability.rs b/src/message/durability.rs index fe54e6cf3..723710ce3 100644 --- a/src/message/durability.rs +++ b/src/message/durability.rs @@ -24,11 +24,14 @@ impl Client { /// joins the accumulating batch, so its hook commit, ack and event keep /// arrival order with the fresh stanzas around it; live it commits /// immediately as a batch of one. Either way the batch commit rewrites and - /// then clears its pending row, and consumers observe the message there — - /// its original batch never dispatched (the hook failed then). A plain ack - /// is sent only for a genuine duplicate (no buffered copy). A read failure - /// fails closed (no ack) so a transient storage error cannot drop a - /// message that still needs its hook to commit. + /// then clears its pending row, and consumers observe the message there. + /// Usually its original batch never dispatched (the hook failed then); if + /// it did (post-commit row cleanup failed AND the ack was lost), event + /// consumers see it twice — the documented at-least-once shape of + /// `Event::Messages` with a hook registered. A plain ack is sent only for + /// a genuine duplicate (no buffered copy). A read failure fails closed + /// (no ack) so a transient storage error cannot drop a message that still + /// needs its hook to commit. pub(crate) async fn ack_or_replay_to_hook(self: &Arc, info: &Arc) { if self.inbound_durability_hook().is_some() { let backend = self.persistence_manager.backend(); @@ -144,7 +147,7 @@ mod tests { Arc::from([test_item("MSG_OK_1"), test_item("MSG_OK_2")]); let infos: Vec<_> = items.iter().map(|i| Arc::clone(&i.info)).collect(); client - .commit_inbound_batch(Arc::clone(&items), BatchOrigin::OfflineDrain, false) + .commit_inbound_batch(Arc::clone(&items), BatchOrigin::OfflineDrain) .await; assert_eq!(hook.calls.load(Ordering::SeqCst), 1, "one commit per batch"); @@ -187,7 +190,6 @@ mod tests { info: Arc::clone(&info), }]), BatchOrigin::OfflineDrain, - false, ) .await; diff --git a/src/message/receive.rs b/src/message/receive.rs index 957d356c1..e0e8a8c7b 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -741,6 +741,13 @@ impl Client { // archived (not deleted) when the new PreKeySignalMessage is processed, // allowing decryption of any in-flight messages encrypted with the old session. self.signal_cache.delete_identity(address).await; + // This stanza's processing permit is held here, and the + // full-cache flush below would otherwise persist ratchet + // advances for accumulated drain entries that have no + // durable row yet — commit them first. + if self.inbound_commit_batch.is_active() { + self.commit_inbound_batch_holding_permit().await; + } // Flush immediately so the backend is updated BEFORE the retry decrypt below. // Device::is_trusted_identity reads from backend, not cache. if let Err(e) = self.flush_signal_cache().await { diff --git a/src/message/tests.rs b/src/message/tests.rs index 8ef229e81..b87991d62 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -5533,11 +5533,9 @@ async fn capturing_client( // other layers but not on this path. *client.noise_socket.lock().await = Some(Arc::new(noise_socket)); seed_test_pn(&client).await; - // Live-path semantics by default (flag, live permit count, batcher live); - // drain tests re-enter drain state themselves. - client.offline_sync_completed.store(true, Ordering::Relaxed); - client.swap_message_semaphore(64); - client.inbound_commit_batch.deactivate_for_tests(); + // Live-path semantics by default; drain tests re-enter drain state + // themselves. + client.enter_live_mode_for_tests(); (client, transport) } @@ -6409,7 +6407,7 @@ async fn skdm_only_group_session_acknowledged_once_without_message_event() { assert_eq!( message_events_for_id(&rx, id), (0, 0), - "SKDM-only messages must not surface Event::Message" + "SKDM-only messages must not surface Event::Messages" ); } @@ -6459,7 +6457,7 @@ async fn session_plaintext_decode_error_is_not_acked_as_skdm_only() { assert_eq!( message_events_for_id(&rx, id), (0, 0), - "invalid plaintext must not surface Event::Message" + "invalid plaintext must not surface Event::Messages" ); let mut nack_code = None; for _ in 0..80 { @@ -6764,7 +6762,7 @@ async fn status_skdm_only_session_uses_one_status_receipt() { assert_eq!( message_events_for_id(&rx, id), (0, 0), - "status SKDM-only messages must not surface Event::Message" + "status SKDM-only messages must not surface Event::Messages" ); } @@ -8337,7 +8335,7 @@ async fn msmsg_decrypts_when_secret_is_stored() { .await; assert!( got.is_some(), - "msmsg decryption + dispatch must surface Event::Message" + "msmsg decryption + dispatch must surface Event::Messages" ); } @@ -9627,7 +9625,7 @@ async fn msmsg_alternate_lookup_resolves_lid_to_stored_pn() { /// End-to-end: phone fanout dispatches a wa::Message carrying the /// outbound `messageSecret`; later the Meta AI bot replies via msmsg /// referencing the same id. The captured secret must let the reply -/// decrypt and surface `Event::Message`. +/// decrypt and surface `Event::Messages`. #[tokio::test] async fn fanout_capture_lets_subsequent_msmsg_decrypt() { use crate::store::commands::DeviceCommand; diff --git a/src/msg_secret_buffer.rs b/src/msg_secret_buffer.rs index 3de4a9a5a..5e4a0b3e1 100644 --- a/src/msg_secret_buffer.rs +++ b/src/msg_secret_buffer.rs @@ -1,7 +1,7 @@ //! Write-behind buffer for inbound `messageSecret` persistence. //! //! Capturing a secret used to upsert SQLite synchronously inside the per-chat -//! receive lane, before the ack and the `Event::Message` dispatch. The buffer +//! receive lane, before the ack and the `Event::Messages` dispatch. The buffer //! splits visibility from durability: an insert is immediately readable //! through [`MsgSecretWriteBuffer::lookup`] (so an add-on referencing the //! secret of the stanza just processed always finds it), while the backend diff --git a/src/receipt.rs b/src/receipt.rs index 11f5a1141..f947585b0 100644 --- a/src/receipt.rs +++ b/src/receipt.rs @@ -554,10 +554,10 @@ impl Client { /// flush at offline-sync completion (WA Web `sendAggregateOfflineReceipts`). /// Returns `false` when the sync already completed, so the caller falls /// back to the live 1:1 receipt. The completed flag is re-checked under - /// the buffer lock: `complete_offline_sync` flips the flag before - /// draining, so a push that wins the lock either lands before the drain - /// (and is included) or observes the flag and goes 1:1 — a receipt can - /// never strand in the buffer. + /// the buffer lock: the drain finisher (`finish_offline_sync`) flips the + /// flag before draining, so a push that wins the lock either lands before + /// the drain (and is included) or observes the flag and goes 1:1 — a + /// receipt can never strand in the buffer. pub(crate) fn try_buffer_offline_receipt(&self, info: &Arc) -> bool { let mut buffer = self .offline_receipt_buffer diff --git a/src/retry.rs b/src/retry.rs index a0cab6610..587a26b78 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -451,8 +451,11 @@ impl Client { ); self.signal_cache.delete_session(&signal_address).await; drop(guard); - self.flush_signal_cache_logged("should_recreate_session", Some(&message_id)) - .await; + self.flush_signal_cache_batch_safe_logged( + "should_recreate_session", + Some(&message_id), + ) + .await; } } @@ -531,7 +534,7 @@ impl Client { .await?; self.send_node(stanza).await?; - self.flush_signal_cache().await?; + self.flush_signal_cache_batch_safe().await?; } else { // DM retry: pairwise resend to the requesting device only. // Use _resolved variant: resolved_jid is already in the correct @@ -566,7 +569,7 @@ impl Client { .await?; self.send_node(stanza).await?; - self.flush_signal_cache().await?; + self.flush_signal_cache_batch_safe().await?; } Ok(()) @@ -686,8 +689,11 @@ impl Client { let _guard = lock.lock().await; self.signal_cache.delete_session(&signal_address).await; drop(_guard); - self.flush_signal_cache_logged("reg ID mismatch session deletion", None) - .await; + self.flush_signal_cache_batch_safe_logged( + "reg ID mismatch session deletion", + None, + ) + .await; } } } @@ -759,7 +765,7 @@ impl Client { let _guard = lock.lock().await; self.signal_cache.delete_session(&signal_address).await; drop(_guard); - self.flush_signal_cache_logged( + self.flush_signal_cache_batch_safe_logged( "base key collision — forcing fresh session", None, ) @@ -1021,7 +1027,7 @@ impl Client { .await?; // Flush after session establishment - self.flush_signal_cache().await?; + self.flush_signal_cache_batch_safe().await?; if identity_change == wacore::libsignal::protocol::IdentityChange::ReplacedExisting { self.react_to_local_identity_change(requester_jid); diff --git a/src/send/mod.rs b/src/send/mod.rs index 26157dc3d..c87594c9c 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -927,7 +927,7 @@ impl Client { self.invalidate_device_cache(user).await; } - self.flush_signal_cache_logged("send_status_message", None) + self.flush_signal_cache_batch_safe_logged("send_status_message", None) .await; Ok(SendResult { @@ -1252,7 +1252,7 @@ impl Client { self.signal_cache.delete_sender_key(sk.cache_key()).await; } let _ = self - .flush_signal_cache_logged("phash-mismatch-fallback", None) + .flush_signal_cache_batch_safe_logged("phash-mismatch-fallback", None) .await; } } @@ -1935,7 +1935,7 @@ impl Client { drop(distribution_guard); // Flush cached Signal state to DB after encryption - self.flush_signal_cache_logged("send_message_impl", None) + self.flush_signal_cache_batch_safe_logged("send_message_impl", None) .await; // Issue new tc token after send if a bucket boundary was crossed. @@ -4049,8 +4049,10 @@ mod tests { } // The test client never connects, so the send's `ensure_e2e_sessions` - // would otherwise block on `wait_for_offline_delivery_end` until timeout. - client.complete_offline_sync(0).await; + // would otherwise block on `wait_for_offline_delivery_end` until + // timeout. Enter live state synchronously (the real finisher now runs + // as a spawned task). + client.enter_live_mode_for_tests(); // Seed a Signal session for the peer's LID device so the offline fanout // can encrypt without fetching prekeys over the (absent) socket. The diff --git a/src/test_utils.rs b/src/test_utils.rs index d2a894331..b1ecc2cd5 100644 --- a/src/test_utils.rs +++ b/src/test_utils.rs @@ -111,14 +111,9 @@ pub async fn create_test_client_with_config( // Tests exercise live-path semantics by default (a fresh client starts in // drain mode: 1-permit semaphore, inbound commits batch instead of - // dispatching immediately). Mirror a completed offline sync — flag, live - // permit count, batcher in live mode. Drain-specific tests re-enter drain - // state themselves. - client - .offline_sync_completed - .store(true, std::sync::atomic::Ordering::Relaxed); - client.swap_message_semaphore(64); - client.inbound_commit_batch.deactivate_for_tests(); + // dispatching immediately). Drain-specific tests re-enter drain state + // themselves. + client.enter_live_mode_for_tests(); client } diff --git a/src/types/durability_hook.rs b/src/types/durability_hook.rs index 6f0c2b72c..c576d4f1d 100644 --- a/src/types/durability_hook.rs +++ b/src/types/durability_hook.rs @@ -39,7 +39,9 @@ pub use wacore::types::events::InboundMessage; /// Scope and known limitations: /// - Covers end-to-end encrypted messages (1:1 and group). Newsletter / broadcast /// channel messages are not encrypted and are acked on their own path, so the -/// hook does not gate them. +/// hook does not gate them (they dispatch event-only). The same applies to PDO +/// placeholder recoveries (`info.unavailable_request_id` is set): their ack runs +/// on the PDO path, so the hook never sees them. /// - If the durable buffer write itself fails (e.g. disk full, after retries), /// the acks are suppressed, but if the process does not crash the Signal /// ratchet still advances and those messages degrade to at-most-once on their diff --git a/src/voip/facade.rs b/src/voip/facade.rs index 4339f5d2b..ef193c3ba 100644 --- a/src/voip/facade.rs +++ b/src/voip/facade.rs @@ -472,7 +472,7 @@ async fn place_call( .map_err(|e| CallError::Setup(e.to_string()))?; drop(_session_guards); client - .flush_signal_cache() + .flush_signal_cache_batch_safe() .await .map_err(|e| CallError::Setup(e.to_string()))?; raw diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index 654b01abe..467884e72 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -2156,6 +2156,46 @@ impl AppSyncStore for SqliteStore { } } +/// Single source of the pending-inbound row insert, shared by the single-row +/// and batch write paths so a schema or conflict-strategy change cannot +/// silently diverge between them. +fn insert_pending_inbound_row( + conn: &mut SqliteConnection, + device_id: i32, + chat: &str, + sender: &str, + id: &str, + message: &[u8], +) -> diesel::QueryResult { + diesel::replace_into(pending_inbound_messages::table) + .values(( + pending_inbound_messages::chat.eq(chat), + pending_inbound_messages::sender.eq(sender), + pending_inbound_messages::id.eq(id), + pending_inbound_messages::message.eq(message), + pending_inbound_messages::device_id.eq(device_id), + )) + .execute(conn) +} + +/// Batch/single-row shared delete; see [`insert_pending_inbound_row`]. +fn delete_pending_inbound_row( + conn: &mut SqliteConnection, + device_id: i32, + chat: &str, + sender: &str, + id: &str, +) -> diesel::QueryResult { + diesel::delete( + pending_inbound_messages::table + .filter(pending_inbound_messages::chat.eq(chat)) + .filter(pending_inbound_messages::sender.eq(sender)) + .filter(pending_inbound_messages::id.eq(id)) + .filter(pending_inbound_messages::device_id.eq(device_id)), + ) + .execute(conn) +} + #[cfg_attr(target_arch = "wasm32", async_trait(?Send))] #[cfg_attr(not(target_arch = "wasm32"), async_trait)] impl ProtocolStore for SqliteStore { @@ -2992,6 +3032,8 @@ impl ProtocolStore for SqliteStore { id: &str, message: &[u8], ) -> Result<()> { + // Row statement shared with store_pending_inbound_batch via + // insert_pending_inbound_row, so the two write paths cannot diverge. let chat = chat.to_string(); let sender = sender.to_string(); let id = id.to_string(); @@ -3004,15 +3046,7 @@ impl ProtocolStore for SqliteStore { let id = id.clone(); let message = Arc::clone(&message); Box::new(move |conn: &mut SqliteConnection| { - diesel::replace_into(pending_inbound_messages::table) - .values(( - pending_inbound_messages::chat.eq(&chat), - pending_inbound_messages::sender.eq(&sender), - pending_inbound_messages::id.eq(&id), - pending_inbound_messages::message.eq(message.as_slice()), - pending_inbound_messages::device_id.eq(device_id), - )) - .execute(conn)?; + insert_pending_inbound_row(conn, device_id, &chat, &sender, &id, &message)?; Ok(()) }) }) @@ -3060,14 +3094,7 @@ impl ProtocolStore for SqliteStore { let sender = sender.clone(); let id = id.clone(); Box::new(move |conn: &mut SqliteConnection| { - diesel::delete( - pending_inbound_messages::table - .filter(pending_inbound_messages::chat.eq(&chat)) - .filter(pending_inbound_messages::sender.eq(&sender)) - .filter(pending_inbound_messages::id.eq(&id)) - .filter(pending_inbound_messages::device_id.eq(device_id)), - ) - .execute(conn)?; + delete_pending_inbound_row(conn, device_id, &chat, &sender, &id)?; Ok(()) }) }) @@ -3126,15 +3153,7 @@ impl ProtocolStore for SqliteStore { // 32 KiB per-PR budget — not worth it for microseconds. conn.transaction(|conn| { for (chat, sender, id, message) in rows.iter() { - diesel::replace_into(pending_inbound_messages::table) - .values(( - pending_inbound_messages::chat.eq(chat), - pending_inbound_messages::sender.eq(sender), - pending_inbound_messages::id.eq(id), - pending_inbound_messages::message.eq(message.as_slice()), - pending_inbound_messages::device_id.eq(device_id), - )) - .execute(conn)?; + insert_pending_inbound_row(conn, device_id, chat, sender, id, message)?; } Ok(()) }) @@ -3164,14 +3183,7 @@ impl ProtocolStore for SqliteStore { // transaction already amortizes the WAL commit. conn.transaction(|conn| { for (chat, sender, id) in keys.iter() { - diesel::delete( - pending_inbound_messages::table - .filter(pending_inbound_messages::chat.eq(chat)) - .filter(pending_inbound_messages::sender.eq(sender)) - .filter(pending_inbound_messages::id.eq(id)) - .filter(pending_inbound_messages::device_id.eq(device_id)), - ) - .execute(conn)?; + delete_pending_inbound_row(conn, device_id, chat, sender, id)?; } Ok(()) }) diff --git a/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index 2f716cc14..ac68c3c88 100644 --- a/tests/e2e/src/lib.rs +++ b/tests/e2e/src/lib.rs @@ -382,8 +382,11 @@ impl TestClient { timeout_secs: u64, ) -> anyhow::Result> { let text = text.to_string(); + // Match anywhere in the batch: a drain event can carry the target + // text behind other messages. self.wait_for_event(timeout_secs, move |e| { - e.message_text() == Some(text.as_str()) + e.messages() + .any(|m| m.message.conversation.as_deref() == Some(text.as_str())) }) .await } diff --git a/tests/e2e/tests/receipts.rs b/tests/e2e/tests/receipts.rs index c8276597a..6406a7305 100644 --- a/tests/e2e/tests/receipts.rs +++ b/tests/e2e/tests/receipts.rs @@ -397,7 +397,7 @@ async fn test_delivery_receipts_flushed_on_disconnect() -> anyhow::Result<()> { info!("A sent {N} messages: {msg_ids:?}"); // Wait for every message event so later-arriving ones can't slip past the - // disconnect. Event::Message dispatches right after the receipt task is + // disconnect. Event::Messages dispatches right after the receipt task is // spawned, so by then the receipt may still be queued on the runtime. let mut seen = std::collections::HashSet::::new(); while seen.len() < N { diff --git a/wacore/src/types/events.rs b/wacore/src/types/events.rs index d186bdafa..ad088c7b9 100755 --- a/wacore/src/types/events.rs +++ b/wacore/src/types/events.rs @@ -630,6 +630,14 @@ pub enum Event { /// consumer never sees a message that a registered durability hook has /// not committed. The `Arc` slice is shared with the hook call — same /// items, same order, no copies. + /// + /// With a hook registered this event is at-least-once, like the hook: a + /// redelivery whose buffered copy survived (e.g. the post-commit cleanup + /// failed and the ack was lost) replays through the same commit and + /// dispatches again. Exceptions that bypass the hook: newsletter + /// messages (plaintext, acked on their own path, never redelivered) and + /// PDO placeholder recoveries (identified by + /// `info.unavailable_request_id`) dispatch event-only. Messages(MessageBatch), Receipt(Receipt), UndecryptableMessage(UndecryptableMessage), @@ -842,21 +850,6 @@ impl Event { .into_iter() .flatten() } - - /// Every plain-text body in this event, in arrival order. Prefer this - /// over [`message_text`](Self::message_text) when a batch can carry more - /// than one message (offline drain). - pub fn message_texts(&self) -> impl Iterator { - self.messages() - .filter_map(|m| m.message.conversation.as_deref()) - } - - /// Convenience: the FIRST plain-text body in this event. A drain batch can - /// carry several texts — iterate [`message_texts`](Self::message_texts) - /// (or [`messages`](Self::messages)) to see them all. - pub fn message_text(&self) -> Option<&str> { - self.message_texts().next() - } } /// One decrypted inbound message. The same items (and order) back both From a178372a33ae2fc49b0e28be345efb4b06f8c4d9 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 06:29:34 +0000 Subject: [PATCH 10/35] fix(recv): report Signal-flush failures as not-durable in drain commits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The drain flush swallowed a failed Signal-cache flush (log-and-continue) while still reporting the commit as durable, so buffered offline receipts — including receipts for SKDM-only stanzas whose sender-key state lives only in the cache — could go out with nothing persisted. All three drain flush points (empty-batch, hooked commit, hookless commit) now report the flush result; on failure the receipts are held for redelivery, the cache keeps its dirty entries for retry, and the hookless path restores its entries to the batcher. --- src/message/commit_batch.rs | 41 +++++++++++++++++++++++++++++-------- 1 file changed, 32 insertions(+), 9 deletions(-) diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index d1e3f2551..74d2f4358 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -292,11 +292,14 @@ impl Client { // Even with nothing to commit, a drain-mode flush must persist the // Signal cache: SKDM-only stanzas mutate Signal state without // enqueueing a message, and their buffered receipts flush right - // after the teardown/drain-end call sites of this function. + // after the teardown/drain-end call sites of this function — so a + // failed flush must report not-durable to hold those receipts + // back (the cache keeps its dirty entries for a later retry). if was_draining { - self.flush_signal_cache_logged("commit_batch", None).await; + self.drain_signal_flush_reporting().await + } else { + true } - true } else { self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) .await @@ -340,6 +343,23 @@ impl Client { self.flush_inbound_commits_under_permit(true, None).await } + /// Drain-mode Signal flush that reports success instead of swallowing it: + /// the drain call sites gate the buffered offline-receipt flush on this, + /// because an SKDM receipt must never be sent while its sender-key state + /// is only in the cache. On failure the cache keeps its dirty entries for + /// a later retry. + async fn drain_signal_flush_reporting(&self) -> bool { + match self.flush_signal_cache().await { + Ok(()) => true, + Err(e) => { + log::error!( + "Failed to flush signal cache (commit_batch): {e:?}; holding buffered receipts for redelivery" + ); + false + } + } + } + /// Commit any accumulated entries while the caller ALREADY HOLDS the /// processing permit (mid-stanza recovery paths like UntrustedIdentity), /// so a full-cache Signal flush that follows cannot persist ratchet @@ -369,7 +389,7 @@ impl Client { /// restored to the batcher, so "entries still batched" stays an accurate /// signal for teardown's flush-vs-drop decision. /// - /// Returns whether the durable state (buffer rows + Signal flush attempt) + /// Returns whether the durable state (buffer rows + Signal flush) /// committed — the gate for flushing buffered offline receipts. A hook /// failure still returns `true`: its rows are durable and the replay path /// retries it, so already-buffered receipts of other messages stay safe. @@ -455,8 +475,10 @@ impl Client { // a cancelled future must not restore the entries. reinsert.disarm(); - if is_drain { - self.flush_signal_cache_logged("commit_batch", None).await; + // A failed flush reports not-durable so buffered receipts are held + // back; the rows are in place, so redelivery replays the batch. + if is_drain && !self.drain_signal_flush_reporting().await { + return false; } if let Err(e) = hook.on_messages(self.clone(), &items).await { @@ -485,10 +507,11 @@ impl Client { ); } } else { - if is_drain { - self.flush_signal_cache_logged("commit_batch", None).await; + // No hook = at-most-once: the flush is the durable point, so a + // failure restores the entries (guard still armed) for a retry. + if is_drain && !self.drain_signal_flush_reporting().await { + return false; } - // No hook = at-most-once: the flush is the durable point. reinsert.disarm(); } From 1d95e3fc0c25fe42ae4f3e187c3cc991f84efb84 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 06:47:39 +0000 Subject: [PATCH 11/35] fix(recv): close review follow-ups on the drain finisher and lock order MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Guard the drain finisher by connection generation BEFORE it mutates the batcher: the check now runs inside flush_inbound_commits_under_permit, between permit acquisition and take(), so a stale finisher that only gets the permit after a reconnect stands down instead of taking or deactivating the new connection's drain. - The timeout arm of wait_for_offline_delivery_end now waits for the spawned finisher (ticked, with generation/shutdown guards), preserving the helper's contract that live state is in place when it returns. - commit_inbound_batch_holding_permit reports durability; the UntrustedIdentity recovery skips its identity flush when the drain commit failed, so restored entries' ratchets are never persisted without rows. - Release per-sender session locks (group/DM retry resends, retry key bundles) and the sender-key chain lock (group encrypt) before the batch-safe Signal flush: the flush acquires the processing permit, and a permit holder can need those same locks — holding them across the flush was an ABBA inversion. - The offline-sync test polls for the 64-permit semaphore instead of short-circuiting on the completed flag, which flips before the swap. - Shared log body between the two flush-logged variants. --- src/client/adapters.rs | 25 +++++++++++-------- src/client/sessions.rs | 18 ++++++++++++- src/client/tests.rs | 41 +++++++++++++++--------------- src/features/signal.rs | 4 +++ src/message/commit_batch.rs | 50 ++++++++++++++++++++++++++----------- src/message/receive.rs | 16 +++++++++--- src/retry.rs | 11 +++++++- 7 files changed, 114 insertions(+), 51 deletions(-) diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 08ee38dd3..da8e60308 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -93,11 +93,7 @@ impl Client { /// `_batch_safe` variants below. pub(crate) async fn flush_signal_cache_logged(&self, context: &str, id: Option<&str>) { if let Err(e) = self.flush_signal_cache().await { - if let Some(id) = id { - log::error!("Failed to flush signal cache ({context} {id}): {e:?}"); - } else { - log::error!("Failed to flush signal cache ({context}): {e:?}"); - } + log_signal_flush_error(context, id, &e); } } @@ -117,7 +113,10 @@ impl Client { pub(crate) async fn flush_signal_cache_batch_safe(&self) -> Result<(), anyhow::Error> { if self.inbound_commit_batch.is_active() { if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { - return if client.flush_inbound_commits_under_permit(false, None).await { + return if client + .flush_inbound_commits_under_permit(false, None, None) + .await + { Ok(()) } else { Err(anyhow::anyhow!( @@ -142,11 +141,15 @@ impl Client { id: Option<&str>, ) { if let Err(e) = self.flush_signal_cache_batch_safe().await { - if let Some(id) = id { - log::error!("Failed to flush signal cache ({context} {id}): {e:?}"); - } else { - log::error!("Failed to flush signal cache ({context}): {e:?}"); - } + log_signal_flush_error(context, id, &e); } } } + +fn log_signal_flush_error(context: &str, id: Option<&str>, e: &anyhow::Error) { + if let Some(id) = id { + log::error!("Failed to flush signal cache ({context} {id}): {e:?}"); + } else { + log::error!("Failed to flush signal cache ({context}): {e:?}"); + } +} diff --git a/src/client/sessions.rs b/src/client/sessions.rs index 8df93a47b..51ed48858 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -81,7 +81,7 @@ impl Client { // for the raceless-transition argument), BEFORE widening the semaphore. // Receipts flush after, so every receipt's message is durably committed // first (WA Web's createSnapshot ordering). - let durable = self.finish_inbound_commit_drain().await; + let durable = self.finish_inbound_commit_drain(generation).await; if self.connection_generation.load(Ordering::Acquire) != generation { log::debug!( @@ -172,6 +172,22 @@ impl Client { ); self.complete_offline_sync(i32::try_from(processed).unwrap_or(i32::MAX)) .await; + // The finisher runs as a spawned task; keep this helper's contract + // that live state is in place when it returns (callers start + // session/send work right after). Ticked so a reconnect or + // shutdown mid-commit cannot strand this waiter: the stale + // finisher stands down without notifying. + loop { + let listener = self.offline_sync_notifier.listen(); + if self.offline_sync_completed.load(Ordering::Acquire) + || self.connection_generation.load(Ordering::Acquire) != wait_generation + || self.expected_disconnect.load(Ordering::Relaxed) + { + return; + } + let _ = wacore::runtime::timeout(&*self.runtime, Duration::from_secs(1), listener) + .await; + } } } diff --git a/src/client/tests.rs b/src/client/tests.rs index bbcb9209d..a1985c3a8 100644 --- a/src/client/tests.rs +++ b/src/client/tests.rs @@ -503,28 +503,27 @@ async fn test_wait_for_offline_delivery_end_times_out_when_flag_not_set() { .await; let elapsed = start.elapsed(); - // The drain finisher runs as a spawned task (off the read loop); wait for - // its completion signal before asserting on the state it flips. - let listener = client.offline_sync_notifier.listen(); - if !client - .offline_sync_completed - .load(std::sync::atomic::Ordering::Acquire) - { - tokio::time::timeout(std::time::Duration::from_secs(5), listener) - .await - .expect("drain finisher should complete"); - } - // Count available permits by trying to acquire non-blockingly - let semaphore = match client.message_processing_semaphore.lock() { - Ok(guard) => guard.clone(), - Err(poisoned) => poisoned.into_inner().clone(), - }; - let mut guards = Vec::new(); - while let Some(guard) = semaphore.try_acquire() { - guards.push(guard); + // The drain finisher runs as a spawned task (off the read loop) and flips + // the flag BEFORE swapping the semaphore, so neither the flag nor the + // notifier alone proves the swap landed. Poll until the 64-permit + // semaphore is observable (counting by non-blocking acquire). + let mut permits = 0; + for _ in 0..100 { + let semaphore = match client.message_processing_semaphore.lock() { + Ok(guard) => guard.clone(), + Err(poisoned) => poisoned.into_inner().clone(), + }; + let mut guards = Vec::new(); + while let Some(guard) = semaphore.try_acquire() { + guards.push(guard); + } + permits = guards.len(); + drop(guards); + if permits == 64 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; } - let permits = guards.len(); - drop(guards); assert!( elapsed.as_millis() >= 45, // Allow small timing variance diff --git a/src/features/signal.rs b/src/features/signal.rs index 9be5d3d7b..9172b9d98 100644 --- a/src/features/signal.rs +++ b/src/features/signal.rs @@ -208,6 +208,10 @@ impl<'a> Signal<'a> { ) .await?; + // Chain mutation done; the batch-safe flush acquires the processing + // permit, and a permit holder can need this sender-key lock — release + // it first. + drop(_chain_guard); self.client.flush_signal_cache_batch_safe().await?; Ok((skdm_bytes, ciphertext.into_serialized().into_vec())) diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 74d2f4358..714900a0d 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -231,7 +231,7 @@ impl Client { // size-trigger flush racing this sleeper must not let // it commit a batch that is only milliseconds old. let _ = client - .flush_inbound_commits_under_permit(false, Some(epoch)) + .flush_inbound_commits_under_permit(false, Some(epoch), None) .await; } })) @@ -277,8 +277,17 @@ impl Client { self: &Arc, deactivate: bool, expected_epoch: Option, + expected_generation: Option, ) -> bool { let _permit = self.acquire_message_processing_permit().await; + if let Some(generation) = expected_generation + && self.connection_generation.load(Ordering::Acquire) != generation + { + // A reconnect reset the batcher while this (drain-finisher) call + // waited for the permit; the new connection owns the state now, so + // touching it here would take/deactivate the NEW drain. + return true; + } if let Some(epoch) = expected_epoch && self.inbound_commit_batch.epoch.load(Ordering::Acquire) != epoch { @@ -323,7 +332,7 @@ impl Client { match wacore::runtime::timeout( &*self.runtime, limit, - self.flush_inbound_commits_under_permit(false, None), + self.flush_inbound_commits_under_permit(false, None, None), ) .await { @@ -339,8 +348,12 @@ impl Client { /// End-of-drain transition; see /// [`flush_inbound_commits_under_permit`](Self::flush_inbound_commits_under_permit). - pub(crate) async fn finish_inbound_commit_drain(self: &Arc) -> bool { - self.flush_inbound_commits_under_permit(true, None).await + /// `generation` scopes it to the connection whose drain is ending, so a + /// stale finisher that only gets the permit after a reconnect stands down + /// instead of taking or deactivating the new connection's drain. + pub(crate) async fn finish_inbound_commit_drain(self: &Arc, generation: u64) -> bool { + self.flush_inbound_commits_under_permit(true, None, Some(generation)) + .await } /// Drain-mode Signal flush that reports success instead of swallowing it: @@ -363,14 +376,16 @@ impl Client { /// Commit any accumulated entries while the caller ALREADY HOLDS the /// processing permit (mid-stanza recovery paths like UntrustedIdentity), /// so a full-cache Signal flush that follows cannot persist ratchet - /// advances for entries without a durable buffered row. - pub(crate) async fn commit_inbound_batch_holding_permit(self: &Arc) { + /// advances for entries without a durable buffered row. Returns whether + /// that follow-up flush is safe: on a failed commit the entries are back + /// in the batcher (unbuffered), and the caller must skip its flush. + pub(crate) async fn commit_inbound_batch_holding_permit(self: &Arc) -> bool { let batch = self.inbound_commit_batch.take(); - if !batch.is_empty() { - let _ = self - .commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) - .await; + if batch.is_empty() { + return true; } + self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) + .await } /// Commit one batch: durable buffer → Signal flush → hook → clear buffer → @@ -597,7 +612,9 @@ mod tests { "sub-threshold entries must accumulate, not commit" ); - client.flush_inbound_commits_under_permit(false, None).await; + client + .flush_inbound_commits_under_permit(false, None, None) + .await; let batches = hook.batches.lock().expect("hook lock").clone(); assert_eq!(batches, vec![vec!["B1", "B2", "B3"]]); @@ -676,7 +693,9 @@ mod tests { client.commit_or_batch_inbound(item("N1")).await; client.commit_or_batch_inbound(item("N2")).await; - client.flush_inbound_commits_under_permit(false, None).await; + client + .flush_inbound_commits_under_permit(false, None, None) + .await; let event = rx.try_recv().expect("one batch event"); assert_eq!( @@ -704,7 +723,8 @@ mod tests { client.commit_or_batch_inbound(item("T1")).await; client.commit_or_batch_inbound(item("T2")).await; - client.finish_inbound_commit_drain().await; + let generation = client.connection_generation.load(Ordering::Acquire); + client.finish_inbound_commit_drain(generation).await; assert!(!client.inbound_commit_batch.is_active()); // A message arriving after the transition commits immediately as Live. client.commit_or_batch_inbound(item("T3")).await; @@ -742,7 +762,9 @@ mod tests { client.commit_or_batch_inbound(item("C1")).await; client.inbound_commit_batch.reset(); - client.flush_inbound_commits_under_permit(false, None).await; + client + .flush_inbound_commits_under_permit(false, None, None) + .await; assert!(hook.batches.lock().expect("hook lock").is_empty()); assert!(rx.try_recv().is_err()); diff --git a/src/message/receive.rs b/src/message/receive.rs index e0e8a8c7b..a721cd975 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -744,9 +744,19 @@ impl Client { // This stanza's processing permit is held here, and the // full-cache flush below would otherwise persist ratchet // advances for accumulated drain entries that have no - // durable row yet — commit them first. - if self.inbound_commit_batch.is_active() { - self.commit_inbound_batch_holding_permit().await; + // durable row yet — commit them first. A failed commit + // restores the entries, so the flush must be skipped + // too: the retry then fails and the message is + // redelivered, instead of stranding those ratchets. + if self.inbound_commit_batch.is_active() + && !self.commit_inbound_batch_holding_permit().await + { + log::warn!( + "Deferring identity-change flush for {}: the drain batch commit failed and its entries must stay unflushed", + wacore::types::jid::observe_protocol_address(address) + ); + outcome.had_failure = true; + continue; } // Flush immediately so the backend is updated BEFORE the retry decrypt below. // Device::is_trusted_identity reads from backend, not cache. diff --git a/src/retry.rs b/src/retry.rs index 587a26b78..08fcc89c9 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -533,6 +533,10 @@ impl Client { ) .await?; + // The Signal mutation is done; the batch-safe flush acquires the + // processing permit, and inbound processing holds that permit + // while taking this same session lock — release it first. + drop(_session_guard); self.send_node(stanza).await?; self.flush_signal_cache_batch_safe().await?; } else { @@ -568,6 +572,8 @@ impl Client { ) .await?; + // Same lock-ordering rule as the group branch above. + drop(_session_guard); self.send_node(stanza).await?; self.flush_signal_cache_batch_safe().await?; } @@ -1026,7 +1032,10 @@ impl Client { ) .await?; - // Flush after session establishment + // Flush after session establishment; release the session lock first + // (the batch-safe flush acquires the processing permit, whose holder + // may need this same lock). + drop(_session_guard); self.flush_signal_cache_batch_safe().await?; if identity_change == wacore::libsignal::protocol::IdentityChange::ReplacedExisting { From dacb7b4b5deaccc22aab6270e8782d082848ecf9 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 07:04:25 +0000 Subject: [PATCH 12/35] refactor(recv): consolidate the drain-to-live publication Share the flag/semaphore/receipts/notify/event tail between the drain finisher and its upgrade-failure fallback as one non-async helper (the duplicated codegen sat in two async state machines), and mark commit_inbound_batch_holding_permit #[must_use] so a future caller cannot silently ignore the flush-safety result. --- src/client/sessions.rs | 66 +++++++++++++++++++------------------ src/message/commit_batch.rs | 1 + 2 files changed, 35 insertions(+), 32 deletions(-) diff --git a/src/client/sessions.rs b/src/client/sessions.rs index 51ed48858..cc1b0c4ae 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -45,12 +45,7 @@ impl Client { "complete_offline_sync: self_weak upgrade failed; dropping the drain tail and switching to live mode" ); self.inbound_commit_batch.force_live_dropping_entries(); - self.offline_sync_completed.store(true, Ordering::Release); - self.swap_message_semaphore(64); - self.offline_sync_notifier.notify(usize::MAX); - self.core - .event_bus - .dispatch(Event::OfflineSyncCompleted(OfflineSyncCompleted { count })); + self.publish_offline_sync_live_state(count, None); return; }; @@ -90,36 +85,43 @@ impl Client { return; } - // Readers that observe offline_sync_completed=true short-circuit - // without touching the semaphore (wait_for_offline_delivery_end - // returns early), so the ordering of flag flip vs. semaphore swap is - // not observable: any in-flight worker keeps using its old 1-permit - // Arc and drains normally; newly-spawned workers pick up the 64-permit - // semaphore via read_message_semaphore(). - self.offline_sync_completed.store(true, Ordering::Release); + self.publish_offline_sync_live_state(count, Some(durable)); + } - // Allow parallel message processing now that offline sync is done. + /// The drain→live state publication, shared by the finisher and its + /// upgrade-failure fallback (non-async so the codegen stays out of their + /// state machines). + /// + /// Readers that observe offline_sync_completed=true short-circuit without + /// touching the semaphore (wait_for_offline_delivery_end returns early), + /// so the ordering of flag flip vs. semaphore swap is not observable: any + /// in-flight worker keeps using its old 1-permit Arc and drains normally; + /// newly-spawned workers pick up the 64-permit semaphore via + /// read_message_semaphore(). The flag flip happens-before the receipt + /// drain takes the buffer lock, so late offline receipts either land in + /// the flush or observe the flag and send 1:1 + /// (see try_buffer_offline_receipt). + /// + /// `durable`: `Some(true)` flushes the buffered offline receipts; + /// `Some(false)` drops them — the tail's durable write failed, its entries + /// are back in the batcher unacked, and receipting SKDM/session state that + /// never became durable would trade a redeliverable failure for a + /// crash-permanent one. `None` (upgrade-failure fallback) leaves the + /// buffer alone for the connection-state reset to clear. + fn publish_offline_sync_live_state(&self, count: i32, durable: Option) { + self.offline_sync_completed.store(true, Ordering::Release); self.swap_message_semaphore(64); - - if durable { - // The flag flip above happens-before this drain takes the buffer - // lock, so late offline receipts either land in this flush or - // observe the flag and send 1:1 (see try_buffer_offline_receipt). - self.flush_offline_receipts(); - } else { - // The tail's durable write failed: its entries are back in the - // batcher, unacked. Buffered receipts must not go out either — - // they may cover SKDM/session state that never became durable, and - // receipting those would trade a redeliverable failure for a - // crash-permanent one. Everything unacked redelivers next connect. - log::warn!( - "finish_offline_sync: tail commit not durable; dropping buffered offline receipts so the server redelivers" - ); - self.clear_offline_receipt_buffer(); + match durable { + Some(true) => self.flush_offline_receipts(), + Some(false) => { + log::warn!( + "finish_offline_sync: tail commit not durable; dropping buffered offline receipts so the server redelivers" + ); + self.clear_offline_receipt_buffer(); + } + None => {} } - self.offline_sync_notifier.notify(usize::MAX); - self.core .event_bus .dispatch(Event::OfflineSyncCompleted(OfflineSyncCompleted { count })); diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 714900a0d..cbf8c8084 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -379,6 +379,7 @@ impl Client { /// advances for entries without a durable buffered row. Returns whether /// that follow-up flush is safe: on a failed commit the entries are back /// in the batcher (unbuffered), and the caller must skip its flush. + #[must_use = "a false result means the entries are back in the batcher and the follow-up Signal flush must be skipped"] pub(crate) async fn commit_inbound_batch_holding_permit(self: &Arc) -> bool { let batch = self.inbound_commit_batch.take(); if batch.is_empty() { From d8ce860431da45b2e4d76114a085fb7759bc91f3 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 07:25:30 +0000 Subject: [PATCH 13/35] fix(recv): defer the drain-to-live transition when the tail commit fails MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A failed end-of-drain commit restored its entries to the batcher but still switched to live mode. In live mode the per-stanza full-cache Signal flush persists the restored entries' ratchet advances with no durable rows, so a crash or reconnect turned their redelivery into plain-acked duplicates — the tail messages were lost for hook consumers. The batcher now stays in drain mode on a failed tail: the restored entries keep batching (later traffic accumulates behind them in arrival order, and batch-safe flush paths keep routing through the batcher), the flush timer and threshold flushes retry the commit, and the first durable flush completes the deferred transition under the same permit-held, raceless flip as the normal end-of-drain path. Liveness is unchanged: the finisher still publishes completion (flag, semaphore, receipts dropped) — only the batcher mode lingers until storage recovers. Covered by a new test via a test-only commit-failure injection point (the real store offers no failure hook): failed tail → still active, entries restored, later message accumulates; durable flush → single ordered batch, transition completed. --- src/client/sessions.rs | 10 +-- src/message/commit_batch.rs | 135 ++++++++++++++++++++++++++++++++++-- 2 files changed, 134 insertions(+), 11 deletions(-) diff --git a/src/client/sessions.rs b/src/client/sessions.rs index cc1b0c4ae..da2a73a7d 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -104,10 +104,12 @@ impl Client { /// /// `durable`: `Some(true)` flushes the buffered offline receipts; /// `Some(false)` drops them — the tail's durable write failed, its entries - /// are back in the batcher unacked, and receipting SKDM/session state that - /// never became durable would trade a redeliverable failure for a - /// crash-permanent one. `None` (upgrade-failure fallback) leaves the - /// buffer alone for the connection-state reset to clear. + /// are back in the batcher unacked (which stays in drain mode until a + /// later durable flush completes the deferred live transition), and + /// receipting SKDM/session state that never became durable would trade a + /// redeliverable failure for a crash-permanent one. `None` + /// (upgrade-failure fallback) leaves the buffer alone for the + /// connection-state reset to clear. fn publish_offline_sync_live_state(&self, count: i32, durable: Option) { self.offline_sync_completed.store(true, Ordering::Release); self.swap_message_semaphore(64); diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index cbf8c8084..ee3d72b07 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -43,8 +43,21 @@ pub(crate) struct InboundCommitBatcher { /// uncommitted batch entry, and let queued drain stanzas commit as Live /// ahead of the accumulated batch. active: std::sync::atomic::AtomicBool, + /// Set when the end-of-drain flush could not commit its tail: the guard + /// restored the entries, and switching to live mode anyway would let the + /// per-stanza full-cache Signal flush persist their ratchet advances with + /// no durable rows — redelivery would then be acked as duplicates. The + /// batcher stays in drain mode instead (entries keep batching, the flush + /// timer retries the commit) and the next durable flush completes the + /// deferred transition. + pending_live: std::sync::atomic::AtomicBool, /// Bumped on every take; a timer that observes a stale epoch stands down. epoch: AtomicU64, + /// Test-only injection: fail the next drain commits before their durable + /// point, exercising the ReinsertGuard/deferred-transition paths (no + /// backend failure can be injected through the real store). + #[cfg(test)] + pub(crate) fail_commits: std::sync::atomic::AtomicBool, /// Reusable encode arena for drain commits, which the processing permit /// already serializes — so this lock is never contended there. Live /// commits use a local buffer instead: sharing it would serialize @@ -57,8 +70,11 @@ impl Default for InboundCommitBatcher { Self { state: std::sync::Mutex::new(BatchState::default()), active: std::sync::atomic::AtomicBool::new(true), + pending_live: std::sync::atomic::AtomicBool::new(false), epoch: AtomicU64::new(0), arena: async_lock::Mutex::new(Vec::new()), + #[cfg(test)] + fail_commits: std::sync::atomic::AtomicBool::new(false), } } } @@ -92,12 +108,24 @@ impl InboundCommitBatcher { !self.lock().entries.is_empty() } - /// Switch to immediate (live) commits. Only the end-of-drain flush calls - /// this, while holding the single processing permit. + /// Switch to immediate (live) commits. Only the end-of-drain flush (or a + /// later flush completing a deferred transition) calls this, while + /// holding the processing permit. fn deactivate(&self) { + self.pending_live.store(false, Ordering::Release); self.active.store(false, Ordering::Release); } + /// Record that the end-of-drain flush failed before its durable point; + /// see the `pending_live` field docs. + fn defer_live_transition(&self) { + self.pending_live.store(true, Ordering::Release); + } + + fn live_transition_pending(&self) -> bool { + self.pending_live.load(Ordering::Acquire) + } + /// Last-resort drain exit for the unreachable-in-practice case where the /// finisher cannot run (`self_weak` upgrade failure): drop any entries /// (unacked, so the server redelivers them) and still switch to live mode @@ -125,6 +153,7 @@ impl InboundCommitBatcher { dropped.len() ); } + self.pending_live.store(false, Ordering::Release); self.active.store(true, Ordering::Release); } } @@ -249,9 +278,15 @@ impl Client { }; if over { let batch = self.inbound_commit_batch.take(); - let _ = self + let durable = self .commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) .await; + // A durable commit may be the retry that clears a deferred + // drain→live transition (failed end-of-drain tail); the permit is + // held, so the flip is as raceless as the end-of-drain one. + if durable && self.inbound_commit_batch.live_transition_pending() { + self.inbound_commit_batch.deactivate(); + } } } @@ -313,8 +348,21 @@ impl Client { self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) .await }; - if deactivate { - self.inbound_commit_batch.deactivate(); + if deactivate || (was_draining && self.inbound_commit_batch.live_transition_pending()) { + if durable { + self.inbound_commit_batch.deactivate(); + } else if was_draining { + // The guard restored the entries; switching to live mode now + // would let per-stanza full-cache flushes persist their + // ratchet advances with no durable rows (acked-duplicate loss + // on redelivery). Keep batching — the flush timer and later + // flushes retry the commit, and the first durable one + // completes this transition. + self.inbound_commit_batch.defer_live_transition(); + log::warn!( + "End-of-drain commit failed; staying in drain mode until a durable flush completes the live transition" + ); + } } durable } @@ -385,8 +433,13 @@ impl Client { if batch.is_empty() { return true; } - self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) - .await + let durable = self + .commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) + .await; + if durable && self.inbound_commit_batch.live_transition_pending() { + self.inbound_commit_batch.deactivate(); + } + durable } /// Commit one batch: durable buffer → Signal flush → hook → clear buffer → @@ -423,6 +476,14 @@ impl Client { batcher: &self.inbound_commit_batch, items: is_drain.then(|| std::sync::Arc::clone(&items)), }; + #[cfg(test)] + if self + .inbound_commit_batch + .fail_commits + .load(Ordering::Acquire) + { + return false; + } if let Some(hook) = self.inbound_durability_hook() { // Key strings live for the whole commit; rows borrow them and the @@ -770,4 +831,64 @@ mod tests { assert!(hook.batches.lock().expect("hook lock").is_empty()); assert!(rx.try_recv().is_err()); } + + // A failed end-of-drain tail must NOT switch to live mode: the restored + // entries would sit in an inactive batcher while per-stanza full-cache + // flushes persist their ratchet advances rowless, turning redelivery into + // acked duplicates. The transition defers until a durable flush. + #[tokio::test] + async fn failed_tail_defers_live_transition_until_durable_flush() { + let client = create_test_client_with_failing_http("batch_defer").await; + client.inbound_commit_batch.reset(); + let hook = Arc::new(RecordingHook { + batches: Mutex::new(Vec::new()), + }); + let _ = client.inbound_durability_hook.set(hook.clone()); + + client.commit_or_batch_inbound(item("T1")).await; + + client + .inbound_commit_batch + .fail_commits + .store(true, Ordering::Release); + let generation = client.connection_generation.load(Ordering::Acquire); + let durable = client.finish_inbound_commit_drain(generation).await; + assert!(!durable, "injected failure must report not-durable"); + assert!( + client.inbound_commit_batch.is_active(), + "failed tail must keep the batcher in drain mode" + ); + assert!( + client.inbound_commit_batch.has_entries(), + "the guard must restore the taken entries" + ); + + // With the transition deferred, later traffic keeps batching instead + // of committing rowless around the restored tail. + client + .inbound_commit_batch + .fail_commits + .store(false, Ordering::Release); + client.commit_or_batch_inbound(item("T2")).await; + assert!( + hook.batches.lock().expect("hook lock").is_empty(), + "deferred mode must accumulate, not commit live" + ); + + let durable = client + .flush_inbound_commits_under_permit(false, None, None) + .await; + assert!(durable); + assert!( + !client.inbound_commit_batch.is_active(), + "the first durable flush completes the deferred transition" + ); + assert!(!client.inbound_commit_batch.has_entries()); + let batches = hook.batches.lock().expect("hook lock").clone(); + assert_eq!( + batches, + vec![vec!["T1", "T2"]], + "restored tail commits first, in arrival order" + ); + } } From 5ec1970711120aa9d90ee1f52d905d3b63c23015 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 07:42:38 +0000 Subject: [PATCH 14/35] fix(recv): hold the single permit while a deferred tail commit retries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the deferred drain-to-live transition: publishing completion still widened the semaphore to 64, so during the deferred window a timer/threshold flush (holding one of many permits) could persist the whole Signal cache while another live stanza was mid-decrypt with ratchet advances not yet enqueued — the same rowless-ratchet loss the deferral exists to prevent. The whole-cache flush is only safe while stanzas are serialized. The not-durable publication now keeps the one-permit semaphore, and completing the deferred transition (first durable flush, permit held) deactivates the batcher and widens the semaphore together. The test asserts both halves: one permit while deferred, 64 after completion. --- src/client/sessions.rs | 19 +++++++++----- src/message/commit_batch.rs | 52 ++++++++++++++++++++++++++++++++----- 2 files changed, 58 insertions(+), 13 deletions(-) diff --git a/src/client/sessions.rs b/src/client/sessions.rs index da2a73a7d..044ee7d7d 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -104,15 +104,20 @@ impl Client { /// /// `durable`: `Some(true)` flushes the buffered offline receipts; /// `Some(false)` drops them — the tail's durable write failed, its entries - /// are back in the batcher unacked (which stays in drain mode until a - /// later durable flush completes the deferred live transition), and - /// receipting SKDM/session state that never became durable would trade a - /// redeliverable failure for a crash-permanent one. `None` - /// (upgrade-failure fallback) leaves the buffer alone for the - /// connection-state reset to clear. + /// are back in the batcher unacked, and receipting SKDM/session state that + /// never became durable would trade a redeliverable failure for a + /// crash-permanent one. In that case the batcher stays in drain mode AND + /// the semaphore stays at one permit: the whole-cache flush inside a + /// retry commit is only safe while no other stanza can be mid-decrypt + /// with unenqueued ratchet advances. The first durable flush completes + /// the deferred transition (see `complete_deferred_live_transition`) and + /// widens the semaphore then. `None` (upgrade-failure fallback) leaves + /// the buffer alone for the connection-state reset to clear. fn publish_offline_sync_live_state(&self, count: i32, durable: Option) { self.offline_sync_completed.store(true, Ordering::Release); - self.swap_message_semaphore(64); + if durable != Some(false) { + self.swap_message_semaphore(64); + } match durable { Some(true) => self.flush_offline_receipts(), Some(false) => { diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index ee3d72b07..c64180281 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -282,10 +282,9 @@ impl Client { .commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) .await; // A durable commit may be the retry that clears a deferred - // drain→live transition (failed end-of-drain tail); the permit is - // held, so the flip is as raceless as the end-of-drain one. + // drain→live transition (failed end-of-drain tail). if durable && self.inbound_commit_batch.live_transition_pending() { - self.inbound_commit_batch.deactivate(); + self.complete_deferred_live_transition(); } } } @@ -348,7 +347,7 @@ impl Client { self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) .await }; - if deactivate || (was_draining && self.inbound_commit_batch.live_transition_pending()) { + if deactivate { if durable { self.inbound_commit_batch.deactivate(); } else if was_draining { @@ -360,13 +359,28 @@ impl Client { // completes this transition. self.inbound_commit_batch.defer_live_transition(); log::warn!( - "End-of-drain commit failed; staying in drain mode until a durable flush completes the live transition" + "End-of-drain commit failed; staying in drain mode (single permit) until a durable flush completes the live transition" ); } + } else if durable && was_draining && self.inbound_commit_batch.live_transition_pending() { + self.complete_deferred_live_transition(); } durable } + /// Finish a drain→live transition that a failed end-of-drain tail commit + /// deferred: the finisher already published completion but left the + /// batcher in drain mode AND the semaphore at one permit — the flush + /// invariant (no stanza mid-decrypt while the whole Signal cache is + /// persisted) only holds while stanzas are serialized. Called with the + /// permit held and a durable commit just done, so the flip is as raceless + /// as the normal end-of-drain one. + pub(crate) fn complete_deferred_live_transition(&self) { + self.inbound_commit_batch.deactivate(); + self.swap_message_semaphore(64); + log::info!("Deferred drain-to-live transition completed after a durable flush"); + } + /// [`flush_inbound_commits_under_permit`](Self::flush_inbound_commits_under_permit) /// with a deadline, for connection teardown: a stalled permit holder or a /// hung hook must not wedge disconnect/reconnect. On timeout (or a failed @@ -437,7 +451,7 @@ impl Client { .commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) .await; if durable && self.inbound_commit_batch.live_transition_pending() { - self.inbound_commit_batch.deactivate(); + self.complete_deferred_live_transition(); } durable } @@ -839,7 +853,10 @@ mod tests { #[tokio::test] async fn failed_tail_defers_live_transition_until_durable_flush() { let client = create_test_client_with_failing_http("batch_defer").await; + // Full drain state: batching AND the single-permit semaphore (the + // test client starts in live mode with 64 permits). client.inbound_commit_batch.reset(); + client.swap_message_semaphore(1); let hook = Arc::new(RecordingHook { batches: Mutex::new(Vec::new()), }); @@ -875,6 +892,12 @@ mod tests { "deferred mode must accumulate, not commit live" ); + assert_eq!( + available_permits(&client), + 1, + "the deferred transition must keep stanzas serialized" + ); + let durable = client .flush_inbound_commits_under_permit(false, None, None) .await; @@ -884,6 +907,11 @@ mod tests { "the first durable flush completes the deferred transition" ); assert!(!client.inbound_commit_batch.has_entries()); + assert_eq!( + available_permits(&client), + 64, + "completing the deferred transition widens the semaphore" + ); let batches = hook.batches.lock().expect("hook lock").clone(); assert_eq!( batches, @@ -891,4 +919,16 @@ mod tests { "restored tail commits first, in arrival order" ); } + + fn available_permits(client: &Client) -> usize { + let semaphore = match client.message_processing_semaphore.lock() { + Ok(guard) => guard.clone(), + Err(poisoned) => poisoned.into_inner().clone(), + }; + let mut guards = Vec::new(); + while let Some(guard) = semaphore.try_acquire() { + guards.push(guard); + } + guards.len() + } } From a93f243f95ea7d9844fae99f66057cc2f0b93aa4 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 07:53:57 +0000 Subject: [PATCH 15/35] fix(recv): arm a retry loop when the live transition is deferred MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The deferral left the retry to future traffic: the guard restores entries with their timer disarmed and the drain-end flush has already run, so on an idle connection nothing would ever recommit the tail (or flush dirty SKDM-only state) — it sat unacked until teardown while startup waiters were already released. A retry task armed at defer time now reflushes every 3s until the transition completes (any durable flush clears it) or a reconnect resets the batcher. Covered by a test that gets the deferred tail committed with no further inbound traffic. --- src/message/commit_batch.rs | 68 +++++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index c64180281..08a57c624 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -358,6 +358,7 @@ impl Client { // flushes retry the commit, and the first durable one // completes this transition. self.inbound_commit_batch.defer_live_transition(); + self.arm_deferred_transition_retry(); log::warn!( "End-of-drain commit failed; staying in drain mode (single permit) until a durable flush completes the live transition" ); @@ -381,6 +382,34 @@ impl Client { log::info!("Deferred drain-to-live transition completed after a durable flush"); } + /// Retry loop for a deferred transition, armed once at defer time: the + /// guard restored the entries with their timer disarmed and the drain-end + /// flush has already run, so on an idle connection NOTHING else would + /// retry the commit — the tail (or dirty SKDM-only state) would sit + /// uncommitted until teardown. Exits as soon as the transition completes + /// (here or via any other durable flush) or a reconnect resets the + /// batcher. + fn arm_deferred_transition_retry(self: &Arc) { + let client = Arc::downgrade(self); + let runtime = self.runtime.clone(); + self.runtime + .spawn(Box::pin(async move { + loop { + runtime.sleep(FLUSH_TIMEOUT).await; + let Some(client) = client.upgrade() else { + return; + }; + if !client.inbound_commit_batch.live_transition_pending() { + return; + } + let _ = client + .flush_inbound_commits_under_permit(false, None, None) + .await; + } + })) + .detach(); + } + /// [`flush_inbound_commits_under_permit`](Self::flush_inbound_commits_under_permit) /// with a deadline, for connection teardown: a stalled permit holder or a /// hung hook must not wedge disconnect/reconnect. On timeout (or a failed @@ -920,6 +949,45 @@ mod tests { ); } + // With no further inbound traffic, the retry loop armed at defer time + // must commit the restored tail and complete the transition on its own. + #[tokio::test] + async fn deferred_transition_retries_automatically() { + let client = create_test_client_with_failing_http("batch_defer_retry").await; + client.inbound_commit_batch.reset(); + client.swap_message_semaphore(1); + let hook = Arc::new(RecordingHook { + batches: Mutex::new(Vec::new()), + }); + let _ = client.inbound_durability_hook.set(hook.clone()); + + client.commit_or_batch_inbound(item("R1")).await; + client + .inbound_commit_batch + .fail_commits + .store(true, Ordering::Release); + let generation = client.connection_generation.load(Ordering::Acquire); + assert!(!client.finish_inbound_commit_drain(generation).await); + client + .inbound_commit_batch + .fail_commits + .store(false, Ordering::Release); + + for _ in 0..100 { + if !client.inbound_commit_batch.is_active() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + assert!( + !client.inbound_commit_batch.is_active(), + "the armed retry must complete the transition without new traffic" + ); + assert_eq!(available_permits(&client), 64); + let batches = hook.batches.lock().expect("hook lock").clone(); + assert_eq!(batches, vec![vec!["R1"]]); + } + fn available_permits(client: &Client) -> usize { let semaphore = match client.message_processing_semaphore.lock() { Ok(guard) => guard.clone(), From 9c02601f71082cfb3f67359d825f4f86a6c2adc2 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 08:03:37 +0000 Subject: [PATCH 16/35] fix(recv): route the session-establishment flush through the batch-safe wrapper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ensure_sessions_inner's post-establishment flush was the one site the batch-safe sweep missed: retry receipts reach it mid-drain, and post-timeout senders reach it during a deferred live transition — in both windows the raw whole-cache flush could persist rowless drain entries' ratchet advances, turning their redelivery into acked duplicates. The per-jid session guards are scoped to the loop iterations, so the wrapper's permit acquisition introduces no lock inversion here. --- src/client/sessions.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/client/sessions.rs b/src/client/sessions.rs index 044ee7d7d..4b7834c5e 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -378,9 +378,12 @@ impl Client { ); } - // Flush after all sessions established + // Flush after all sessions established. Batch-safe: retry receipts + // reach here mid-drain, and post-timeout senders reach here during a + // deferred live transition — in both windows a raw whole-cache flush + // would persist rowless drain entries' ratchet advances. if success_count > 0 { - self.flush_signal_cache().await?; + self.flush_signal_cache_batch_safe().await?; } Ok(success_count) From cf368ba4a8cc171e43a8cebc97e946e618a8cdfd Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 08:07:43 +0000 Subject: [PATCH 17/35] fix(recv): generation-scope the deferred-transition retry task Matches the finisher's scoping: after a reconnect the stale retry task stands down (checked in the loop and enforced by the flush's generation guard) instead of flushing the new connection's state when that connection also deferred its transition. The new connection arms its own retry. --- src/message/commit_batch.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 08a57c624..fa92d0ba2 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -392,6 +392,10 @@ impl Client { fn arm_deferred_transition_retry(self: &Arc) { let client = Arc::downgrade(self); let runtime = self.runtime.clone(); + // Generation-scoped like the finisher: after a reconnect this stale + // task stands down (a re-deferring new connection arms its own retry) + // instead of flushing the new connection's state. + let generation = self.connection_generation.load(Ordering::Acquire); self.runtime .spawn(Box::pin(async move { loop { @@ -399,11 +403,13 @@ impl Client { let Some(client) = client.upgrade() else { return; }; - if !client.inbound_commit_batch.live_transition_pending() { + if !client.inbound_commit_batch.live_transition_pending() + || client.connection_generation.load(Ordering::Acquire) != generation + { return; } let _ = client - .flush_inbound_commits_under_permit(false, None, None) + .flush_inbound_commits_under_permit(false, None, Some(generation)) .await; } })) From a5a0fd2d1d77048bf75680f7f1ee083c6f5a1ea3 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 08:22:14 +0000 Subject: [PATCH 18/35] fix(recv): bound the timeout-arm wait on an already-running finisher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit If the marker-triggered finisher was already started when the offline sync timeout fired, complete_offline_sync starts nothing and the wait loop had no deadline left — a stuck tail commit/hook turned the 'with timeout' helper into an unbounded wait for its callers. The loop is now bounded by a second timeout window; on expiry callers proceed and the finisher keeps running in the background. --- src/client/sessions.rs | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/src/client/sessions.rs b/src/client/sessions.rs index 4b7834c5e..0798f7fa3 100644 --- a/src/client/sessions.rs +++ b/src/client/sessions.rs @@ -184,8 +184,13 @@ impl Client { // The finisher runs as a spawned task; keep this helper's contract // that live state is in place when it returns (callers start // session/send work right after). Ticked so a reconnect or - // shutdown mid-commit cannot strand this waiter: the stale - // finisher stands down without notifying. + // shutdown mid-commit cannot strand this waiter, and bounded by a + // second `timeout` window: when the marker-triggered finisher was + // ALREADY running and its tail commit/hook is stuck, the + // complete_offline_sync above started nothing, and an unbounded + // wait here would defeat this helper's whole point — callers + // proceed and the finisher keeps running in the background. + let deadline = wacore::time::Instant::now() + timeout; loop { let listener = self.offline_sync_notifier.listen(); if self.offline_sync_completed.load(Ordering::Acquire) @@ -194,6 +199,14 @@ impl Client { { return; } + if wacore::time::Instant::now() >= deadline { + log::warn!( + target: "Client/OfflineSync", + "Drain finisher still running {:?} after the offline sync timeout; proceeding without it", + timeout, + ); + return; + } let _ = wacore::runtime::timeout(&*self.runtime, Duration::from_secs(1), listener) .await; } From 9a93e30458423c20a6859350c0ce884f4306a20b Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 08:37:37 +0000 Subject: [PATCH 19/35] fix(recv): keep hook batches retryable when the Signal flush fails after the rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In the hook path the reinsert guard was disarmed right after the durable row write, so a transient Signal-flush failure stranded the batch: entries gone from memory, rows buffered, but no hook, acks or event until a reconnect redelivery replay. The guard now stays armed across the flush — re-storing the rows is idempotent (replace-into), so the restored entries make the retry re-run the full commit this session. New test drives the fail-then-retry cycle via a flush-failure injection point and verifies the row lifecycle end to end. --- src/message/commit_batch.rs | 88 ++++++++++++++++++++++++++++++++++--- 1 file changed, 83 insertions(+), 5 deletions(-) diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index fa92d0ba2..a341bb0b6 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -58,6 +58,10 @@ pub(crate) struct InboundCommitBatcher { /// backend failure can be injected through the real store). #[cfg(test)] pub(crate) fail_commits: std::sync::atomic::AtomicBool, + /// Test-only injection: fail drain Signal flushes AFTER the durable row + /// write, exercising the rows-stored-but-unflushed retry path. + #[cfg(test)] + pub(crate) fail_flushes: std::sync::atomic::AtomicBool, /// Reusable encode arena for drain commits, which the processing permit /// already serializes — so this lock is never contended there. Live /// commits use a local buffer instead: sharing it would serialize @@ -75,6 +79,8 @@ impl Default for InboundCommitBatcher { arena: async_lock::Mutex::new(Vec::new()), #[cfg(test)] fail_commits: std::sync::atomic::AtomicBool::new(false), + #[cfg(test)] + fail_flushes: std::sync::atomic::AtomicBool::new(false), } } } @@ -459,6 +465,14 @@ impl Client { /// is only in the cache. On failure the cache keeps its dirty entries for /// a later retry. async fn drain_signal_flush_reporting(&self) -> bool { + #[cfg(test)] + if self + .inbound_commit_batch + .fail_flushes + .load(Ordering::Acquire) + { + return false; + } match self.flush_signal_cache().await { Ok(()) => true, Err(e) => { @@ -597,15 +611,19 @@ impl Client { return false; } } - // Rows are durable: from here on, redelivery replays from them, so - // a cancelled future must not restore the entries. - reinsert.disarm(); - // A failed flush reports not-durable so buffered receipts are held - // back; the rows are in place, so redelivery replays the batch. + // back. The guard is still armed: the rows are in place and + // re-storing them is idempotent (replace-into), so the restored + // entries make the batch retryable THIS session — the retry + // re-runs the full commit (rows → flush → hook → acks → event) + // instead of parking the messages until a reconnect replay. if is_drain && !self.drain_signal_flush_reporting().await { return false; } + // Rows durable and Signal flushed: from here on a cancelled + // future must not restore the entries — redelivery replays from + // the rows. + reinsert.disarm(); if let Err(e) = hook.on_messages(self.clone(), &items).await { log::warn!( @@ -994,6 +1012,66 @@ mod tests { assert_eq!(batches, vec![vec!["R1"]]); } + // A Signal-flush failure after the durable row write must keep the batch + // retryable this session: entries restored (re-storing rows is + // idempotent), hook not yet run, and the retry commits everything. + #[tokio::test] + async fn flush_failure_after_rows_keeps_batch_retryable() { + let client = create_test_client_with_failing_http("batch_flush_fail").await; + client.inbound_commit_batch.reset(); + let hook = Arc::new(RecordingHook { + batches: Mutex::new(Vec::new()), + }); + let _ = client.inbound_durability_hook.set(hook.clone()); + + client.commit_or_batch_inbound(item("F1")).await; + client + .inbound_commit_batch + .fail_flushes + .store(true, Ordering::Release); + let durable = client + .flush_inbound_commits_under_permit(false, None, None) + .await; + assert!(!durable); + assert!( + client.inbound_commit_batch.has_entries(), + "entries must be restored for a same-session retry" + ); + assert!( + hook.batches.lock().expect("hook lock").is_empty(), + "hook must not run before the Signal flush succeeds" + ); + let backend = client.persistence_manager.backend(); + assert!( + backend + .get_pending_inbound("100@g.us", "200@s.whatsapp.net", "F1") + .await + .unwrap() + .is_some(), + "the durable row from the failed attempt stays in place" + ); + + client + .inbound_commit_batch + .fail_flushes + .store(false, Ordering::Release); + let durable = client + .flush_inbound_commits_under_permit(false, None, None) + .await; + assert!(durable); + assert!(!client.inbound_commit_batch.has_entries()); + let batches = hook.batches.lock().expect("hook lock").clone(); + assert_eq!(batches, vec![vec!["F1"]]); + assert!( + backend + .get_pending_inbound("100@g.us", "200@s.whatsapp.net", "F1") + .await + .unwrap() + .is_none(), + "the retry commit clears the buffered row" + ); + } + fn available_permits(client: &Client) -> usize { let semaphore = match client.message_processing_semaphore.lock() { Ok(guard) => guard.clone(), From 3278ae66861e0c12bb70a08b6eab94f235876960 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 08:52:51 +0000 Subject: [PATCH 20/35] fix(recv): settle the teardown Signal cache under the processing permit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cleanup committed the drain batch (bounded, permit acquired and released) and only then ran the raw whole-cache flush. Old chat-lane workers keep draining until the connection generation changes, so one could take the freed permit and be mid-decrypt — ratchets advanced, entry not yet enqueued — while the flush persisted its rowless advances; the reset then dropped the entry and its redelivery was acked as a duplicate. teardown_inbound_commits_bounded now commits the batch and settles the cache (flush-or-drop, then clear) in one permit-held section; on timeout it drops the unflushed cache with the entries it covers. connect_internal additionally clears the cache next to its batcher reset: a worker that decrypts after cleanup settled leaves dirty rowless advances behind, and flushing those on the next connection would recreate the same acked-duplicate window. The pre-close receipt-gating flush (flush_inbound_commits_bounded) is unchanged. --- src/client/lifecycle.rs | 55 +++++++++++-------------------- src/message/commit_batch.rs | 65 ++++++++++++++++++++++++++++++++++--- 2 files changed, 79 insertions(+), 41 deletions(-) diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 96e57cfae..c95957c7b 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -460,7 +460,13 @@ impl Client { .store(false, Ordering::Relaxed); self.clear_offline_receipt_buffer(); // Uncommitted batch entries were never acked; the server redelivers - // them on this fresh connection. + // them on this fresh connection. The cache clear pairs with it: a + // late lane worker from the previous connection can decrypt AFTER + // cleanup settled the cache, leaving dirty ratchet advances whose + // entries are dropped right here — flushing those later would make + // their redeliveries ackable duplicates. Committed state was flushed + // by its commit, so this only drops rowless advances. + self.signal_cache.clear().await; self.inbound_commit_batch.reset(); self.offline_batch.reset(); self.outbound_flush.reopen(); @@ -761,50 +767,27 @@ impl Client { .lock() .unwrap_or_else(|p| p.into_inner()) .clear(); - // Commit any accumulated drain batch BEFORE the Signal flush below: + // Commit any accumulated drain batch and settle the Signal cache in + // ONE permit-held section (see teardown_inbound_commits_bounded): // persisting ratchet advances while dropping their uncommitted batch - // entries would turn each redelivery into an ackable duplicate with no - // buffered copy — silent loss for hook consumers. Acks/events from - // this commit are best-effort (the socket is gone); the durable hook - // commit is what matters. Reached on every teardown path, including - // the run loop's unexpected read-loop exit, which never goes through - // disconnect(). + // entries — or while an old lane worker is mid-decrypt — would turn + // redeliveries into ackable duplicates with no buffered copy. + // Acks/events from this commit are best-effort (the socket is gone); + // the durable hook commit is what matters. Reached on every teardown + // path, including the run loop's unexpected read-loop exit, which + // never goes through disconnect(). if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { client - .flush_inbound_commits_bounded(std::time::Duration::from_secs(5)) + .teardown_inbound_commits_bounded(std::time::Duration::from_secs(5)) .await; } else { // Same class of bug as the complete_offline_sync twin: a silent - // skip here is the acked-before-committed loss — make it loud. + // skip here is the acked-before-committed loss — make it loud, + // and drop the dirty cache with the entries it covers. log::error!( - "cleanup_connection_state: self_weak upgrade failed; skipping drain-batch commit before Signal flush — uncommitted entries will be dropped" - ); - } - // Flush before clear: clear() drops dirty entries, so a disconnect - // racing an in-flight encrypt would lose the just-advanced sender-key - // chain and force a full SKDM re-fanout. A disconnect is not a logout. - // Only clear on a successful flush; on a backend error keep the cache so - // the dirty state isn't dropped and the next operation can persist it. - // - // Exception: if the bounded commit above left entries in the batcher - // (timed out or its durable write failed), the cache holds their - // ratchet advances. Persisting those would make each redelivery an - // ackable duplicate; keeping them cached would make it decrypt as one - // in-process. Clearing WITHOUT flushing un-advances them instead — - // everything they cover is unacked and redelivers fresh. Committed - // state is not at risk: every earlier commit flushed the cache itself. - if self.inbound_commit_batch.has_entries() { - log::warn!( - "cleanup_connection_state: dropping unflushed Signal state for uncommitted drain entries; the server redelivers them" + "cleanup_connection_state: self_weak upgrade failed; dropping uncommitted drain entries and their unflushed Signal state" ); self.signal_cache.clear().await; - } else { - match self.flush_signal_cache().await { - Ok(()) => self.signal_cache.clear().await, - Err(e) => log::error!( - "cleanup_connection_state: signal cache flush failed, keeping cache to avoid dropping Signal state: {e:?}" - ), - } } // Reset semaphore to 1 permit for next offline sync. self.swap_message_semaphore(1); diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index a341bb0b6..e440005b6 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -422,12 +422,67 @@ impl Client { .detach(); } + /// Teardown twin of [`flush_inbound_commits_bounded`]: commit the batch + /// AND settle the Signal cache (flush-or-drop, then clear) in ONE + /// permit-held section. Doing these as separate steps released the permit + /// in between, and an old chat-lane worker (they drain until the + /// connection generation changes) could grab it and be mid-decrypt — + /// ratchets advanced, entry not yet enqueued — while the raw flush + /// persisted its rowless advances; the reset then dropped the entry and + /// its redelivery was acked as a duplicate. + /// + /// On timeout (stalled permit holder / hung hook) the cache is cleared + /// WITHOUT flushing: everything dirty then belongs to uncommitted + /// entries, and dropping both sides keeps redelivery consistent. + pub(crate) async fn teardown_inbound_commits_bounded( + self: &Arc, + limit: std::time::Duration, + ) { + let settle = async { + let _permit = self.acquire_message_processing_permit().await; + let batch = self.inbound_commit_batch.take(); + let durable = if batch.is_empty() { + true + } else { + self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) + .await + }; + // Permit still held: no stanza can be mid-decrypt while the cache + // is settled. Flush-before-clear preserves in-flight sender-key + // advances (a disconnect is not a logout); a failed or + // non-durable commit leaves entries restored, whose ratchets must + // drop with them instead of persisting rowless. + if durable && !self.inbound_commit_batch.has_entries() { + match self.flush_signal_cache().await { + Ok(()) => self.signal_cache.clear().await, + Err(e) => log::error!( + "cleanup_connection_state: signal cache flush failed, keeping cache to avoid dropping Signal state: {e:?}" + ), + } + } else { + log::warn!( + "cleanup_connection_state: dropping unflushed Signal state for uncommitted drain entries; the server redelivers them" + ); + self.signal_cache.clear().await; + } + }; + if wacore::runtime::timeout(&*self.runtime, limit, settle) + .await + .is_err() + { + log::warn!( + "Timed out committing the inbound drain batch during teardown; dropping unflushed Signal state so redelivery stays consistent" + ); + self.signal_cache.clear().await; + } + } + /// [`flush_inbound_commits_under_permit`](Self::flush_inbound_commits_under_permit) - /// with a deadline, for connection teardown: a stalled permit holder or a - /// hung hook must not wedge disconnect/reconnect. On timeout (or a failed - /// durable write) the entries stay in the batcher unacked — the caller - /// must not flush buffered receipts or persist the Signal cache for them, - /// and the server redelivers them on the next connect. + /// with a deadline, for pre-close receipt gating: a stalled permit holder + /// or a hung hook must not wedge disconnect/reconnect. On timeout (or a + /// failed durable write) the entries stay in the batcher unacked — the + /// caller must not flush buffered receipts or persist the Signal cache + /// for them, and the server redelivers them on the next connect. pub(crate) async fn flush_inbound_commits_bounded( self: &Arc, limit: std::time::Duration, From 9af8f0cac44bd71a9004dcdd9365da0629afafde Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 09:03:40 +0000 Subject: [PATCH 21/35] fix(recv): don't destroy teardown-retained Signal state at connect MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The connect-time cache clear (added to drop late lane-worker rowless advances) was unconditional, so it also destroyed the cache that a failed teardown flush deliberately RETAINED — committed/acked SKDM and session state the server never redelivers, making the loss permanent. A retention flag now distinguishes the two: the failed-flush branch sets it, the connect clear stands down when it is set (the next successful flush persists the state), and every dropping branch clears it. --- src/client.rs | 5 +++++ src/client/lifecycle.rs | 17 +++++++++++++++-- src/message/commit_batch.rs | 15 ++++++++++++--- 3 files changed, 32 insertions(+), 5 deletions(-) diff --git a/src/client.rs b/src/client.rs index 9df761528..e4d8497d6 100644 --- a/src/client.rs +++ b/src/client.rs @@ -532,6 +532,11 @@ pub struct Client { /// idempotent). Separate from `offline_sync_completed` because the finish /// runs off the read loop and the flag must flip only after its commit. pub(crate) offline_sync_finish_started: Arc, + /// Set when teardown could not flush the Signal cache and retained it + /// (committed/acked state that would be lost forever — the server never + /// redelivers acked stanzas). The connect-time cache clear skips when + /// this is set so the next successful flush can persist that state. + pub(crate) signal_cache_retained_dirty: AtomicBool, /// Delivery receipts buffered during offline sync, flushed as aggregate /// `` stanzas at completion (WA Web `sendAggregateOfflineReceipts`). /// Empty (zero capacity) outside the offline window. diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index c95957c7b..d8fa83729 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -215,6 +215,7 @@ impl Client { offline_sync_notifier: Arc::new(event_listener::Event::new()), offline_sync_completed: Arc::new(AtomicBool::new(false)), offline_sync_finish_started: Arc::new(AtomicBool::new(false)), + signal_cache_retained_dirty: AtomicBool::new(false), offline_receipt_buffer: std::sync::Mutex::new(Vec::new()), inbound_commit_batch: Default::default(), history_sync_tasks_in_flight: Arc::new(AtomicUsize::new(0)), @@ -465,8 +466,20 @@ impl Client { // cleanup settled the cache, leaving dirty ratchet advances whose // entries are dropped right here — flushing those later would make // their redeliveries ackable duplicates. Committed state was flushed - // by its commit, so this only drops rowless advances. - self.signal_cache.clear().await; + // by its commit, so this normally only drops rowless advances — the + // exception is a teardown whose flush FAILED and retained + // committed/acked state (never redelivered): keep that for the next + // successful flush instead of destroying it. + if self + .signal_cache_retained_dirty + .swap(false, Ordering::AcqRel) + { + log::warn!( + "connect: keeping Signal cache retained by a failed teardown flush; the next flush persists it" + ); + } else { + self.signal_cache.clear().await; + } self.inbound_commit_batch.reset(); self.offline_batch.reset(); self.outbound_flush.reopen(); diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index e440005b6..c063f1a4b 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -455,9 +455,16 @@ impl Client { if durable && !self.inbound_commit_batch.has_entries() { match self.flush_signal_cache().await { Ok(()) => self.signal_cache.clear().await, - Err(e) => log::error!( - "cleanup_connection_state: signal cache flush failed, keeping cache to avoid dropping Signal state: {e:?}" - ), + Err(e) => { + // Committed/acked state the server never redelivers: + // keep it resident and tell the connect-time clear to + // stand down so the next successful flush persists it. + self.signal_cache_retained_dirty + .store(true, Ordering::Release); + log::error!( + "cleanup_connection_state: signal cache flush failed, keeping cache to avoid dropping Signal state: {e:?}" + ); + } } } else { log::warn!( @@ -473,6 +480,8 @@ impl Client { log::warn!( "Timed out committing the inbound drain batch during teardown; dropping unflushed Signal state so redelivery stays consistent" ); + self.signal_cache_retained_dirty + .store(false, Ordering::Release); self.signal_cache.clear().await; } } From 9e15e3d99b1b049cb6faf0e5d7cd60e132c9aad4 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 09:11:31 +0000 Subject: [PATCH 22/35] fix(recv): clear the cache-retention flag only on a successful flush MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit swap(false) at connect consumed the flag on the first attempt, so a connect that failed before any flush left the second attempt free to clear the retained cache — destroying committed/acked Signal state the retention exists to protect. The connect check is now a plain load; the flag resets at the one choke point every flush goes through (flush_signal_cache success) and in the teardown branches that deliberately drop the cache. --- src/client/adapters.rs | 8 +++++++- src/client/lifecycle.rs | 8 ++++---- src/message/commit_batch.rs | 2 ++ 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/client/adapters.rs b/src/client/adapters.rs index da8e60308..04dbd6961 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -81,7 +81,13 @@ impl Client { self.signal_cache .flush(&*backend) .await - .map_err(|e| anyhow::anyhow!("Failed to flush signal cache: {e}")) + .map_err(|e| anyhow::anyhow!("Failed to flush signal cache: {e}"))?; + // Single choke point every flush goes through: state retained by a + // failed teardown flush is persisted now, so the connect-time clear + // may resume on future connections. + self.signal_cache_retained_dirty + .store(false, std::sync::atomic::Ordering::Release); + Ok(()) } /// [`flush_signal_cache`](Self::flush_signal_cache) with error logging instead of propagation. diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index d8fa83729..c5d54eeaf 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -470,10 +470,10 @@ impl Client { // exception is a teardown whose flush FAILED and retained // committed/acked state (never redelivered): keep that for the next // successful flush instead of destroying it. - if self - .signal_cache_retained_dirty - .swap(false, Ordering::AcqRel) - { + // load(), not swap(): a connect attempt that fails before any flush + // must not consume the flag, or the NEXT attempt would clear the + // retained cache. Only a successful flush_signal_cache resets it. + if self.signal_cache_retained_dirty.load(Ordering::Acquire) { log::warn!( "connect: keeping Signal cache retained by a failed teardown flush; the next flush persists it" ); diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index c063f1a4b..784d90724 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -470,6 +470,8 @@ impl Client { log::warn!( "cleanup_connection_state: dropping unflushed Signal state for uncommitted drain entries; the server redelivers them" ); + self.signal_cache_retained_dirty + .store(false, Ordering::Release); self.signal_cache.clear().await; } }; From c4cd5e18c8d145e4581c1afa7ddbcbdc694283ae Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 09:23:24 +0000 Subject: [PATCH 23/35] fix(recv): quiesce lane workers at teardown instead of patching around them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every teardown/cache hazard in the last review rounds traced back to one root cause: old chat-lane workers keep draining their queues until the connection generation changes — which only happened at the NEXT login — so they could decrypt after the teardown cache settle and leave rowless ratchet advances behind. Attack the root instead of the symptoms: - cleanup_connection_state bumps the connection generation FIRST. It is the this-connection-is-over signal every per-connection loop already polls: lane workers stop draining (their stanzas were never acked and redeliver), stale finishers/timers stand down earlier. - process_classified_message re-checks the generation after acquiring the processing permit, closing the one-message race where a worker passed its loop-top check before the bump: it bails unacked before any ratchet mutation, so the permit-held settle is provably the LAST Signal-cache activity of the connection. With no late writers possible, the connect-time cache clear and the retention flag (added to arbitrate between dropping late rowless advances and preserving flush-failure state) lose their reason to exist and are removed: teardown's settle owns the cache, and anything resident afterwards is deliberately retained committed state for the next successful flush. --- src/client.rs | 5 ----- src/client/adapters.rs | 8 +------- src/client/lifecycle.rs | 35 +++++++++++++++-------------------- src/message/commit_batch.rs | 22 ++++++++-------------- src/message/receive.rs | 18 ++++++++++++++++++ 5 files changed, 42 insertions(+), 46 deletions(-) diff --git a/src/client.rs b/src/client.rs index e4d8497d6..9df761528 100644 --- a/src/client.rs +++ b/src/client.rs @@ -532,11 +532,6 @@ pub struct Client { /// idempotent). Separate from `offline_sync_completed` because the finish /// runs off the read loop and the flag must flip only after its commit. pub(crate) offline_sync_finish_started: Arc, - /// Set when teardown could not flush the Signal cache and retained it - /// (committed/acked state that would be lost forever — the server never - /// redelivers acked stanzas). The connect-time cache clear skips when - /// this is set so the next successful flush can persist that state. - pub(crate) signal_cache_retained_dirty: AtomicBool, /// Delivery receipts buffered during offline sync, flushed as aggregate /// `` stanzas at completion (WA Web `sendAggregateOfflineReceipts`). /// Empty (zero capacity) outside the offline window. diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 04dbd6961..da8e60308 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -81,13 +81,7 @@ impl Client { self.signal_cache .flush(&*backend) .await - .map_err(|e| anyhow::anyhow!("Failed to flush signal cache: {e}"))?; - // Single choke point every flush goes through: state retained by a - // failed teardown flush is persisted now, so the connect-time clear - // may resume on future connections. - self.signal_cache_retained_dirty - .store(false, std::sync::atomic::Ordering::Release); - Ok(()) + .map_err(|e| anyhow::anyhow!("Failed to flush signal cache: {e}")) } /// [`flush_signal_cache`](Self::flush_signal_cache) with error logging instead of propagation. diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index c5d54eeaf..ffae685cc 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -215,7 +215,6 @@ impl Client { offline_sync_notifier: Arc::new(event_listener::Event::new()), offline_sync_completed: Arc::new(AtomicBool::new(false)), offline_sync_finish_started: Arc::new(AtomicBool::new(false)), - signal_cache_retained_dirty: AtomicBool::new(false), offline_receipt_buffer: std::sync::Mutex::new(Vec::new()), inbound_commit_batch: Default::default(), history_sync_tasks_in_flight: Arc::new(AtomicUsize::new(0)), @@ -461,25 +460,12 @@ impl Client { .store(false, Ordering::Relaxed); self.clear_offline_receipt_buffer(); // Uncommitted batch entries were never acked; the server redelivers - // them on this fresh connection. The cache clear pairs with it: a - // late lane worker from the previous connection can decrypt AFTER - // cleanup settled the cache, leaving dirty ratchet advances whose - // entries are dropped right here — flushing those later would make - // their redeliveries ackable duplicates. Committed state was flushed - // by its commit, so this normally only drops rowless advances — the - // exception is a teardown whose flush FAILED and retained - // committed/acked state (never redelivered): keep that for the next - // successful flush instead of destroying it. - // load(), not swap(): a connect attempt that fails before any flush - // must not consume the flag, or the NEXT attempt would clear the - // retained cache. Only a successful flush_signal_cache resets it. - if self.signal_cache_retained_dirty.load(Ordering::Acquire) { - log::warn!( - "connect: keeping Signal cache retained by a failed teardown flush; the next flush persists it" - ); - } else { - self.signal_cache.clear().await; - } + // them on this fresh connection. The Signal cache needs no clear + // here: teardown settled it under the permit, and its generation + // bump plus the post-permit re-check guarantee no late lane worker + // dirtied it afterwards. Anything still resident is state a failed + // teardown flush deliberately retained (committed/acked, never + // redelivered) for the next successful flush to persist. self.inbound_commit_batch.reset(); self.offline_batch.reset(); self.outbound_flush.reopen(); @@ -723,6 +709,15 @@ impl Client { tracing::instrument(name = "wa.conn.cleanup", level = "debug", skip_all) )] pub(crate) async fn cleanup_connection_state(&self) { + // Bump the generation FIRST: it is the "this connection is over" + // signal every per-connection loop already polls. Chat-lane workers + // stop draining their queues (their remaining stanzas were never + // acked and redeliver), stale finishers/timers stand down, and — + // combined with the post-permit generation re-check in + // process_classified_message — no decrypt can START after the + // permit-held cache settle below, so no rowless ratchet advances can + // dirty the cache behind teardown's back. + self.connection_generation.fetch_add(1, Ordering::SeqCst); // Note: node_waiters are intentionally NOT cleared here — they are // cross-connection (callers may register a waiter before an action that // completes on a subsequent connection, e.g. after 515 reconnect). diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 784d90724..e545bb5e3 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -455,23 +455,19 @@ impl Client { if durable && !self.inbound_commit_batch.has_entries() { match self.flush_signal_cache().await { Ok(()) => self.signal_cache.clear().await, - Err(e) => { - // Committed/acked state the server never redelivers: - // keep it resident and tell the connect-time clear to - // stand down so the next successful flush persists it. - self.signal_cache_retained_dirty - .store(true, Ordering::Release); - log::error!( - "cleanup_connection_state: signal cache flush failed, keeping cache to avoid dropping Signal state: {e:?}" - ); - } + // Committed/acked state the server never redelivers: keep + // it resident so the next successful flush persists it. + // Safe to carry across the reconnect — the teardown + // generation bump plus the post-permit re-check mean no + // late decrypt can mix rowless advances into it. + Err(e) => log::error!( + "cleanup_connection_state: signal cache flush failed, keeping cache to avoid dropping Signal state: {e:?}" + ), } } else { log::warn!( "cleanup_connection_state: dropping unflushed Signal state for uncommitted drain entries; the server redelivers them" ); - self.signal_cache_retained_dirty - .store(false, Ordering::Release); self.signal_cache.clear().await; } }; @@ -482,8 +478,6 @@ impl Client { log::warn!( "Timed out committing the inbound drain batch during teardown; dropping unflushed Signal state so redelivery stays consistent" ); - self.signal_cache_retained_dirty - .store(false, Ordering::Release); self.signal_cache.clear().await; } } diff --git a/src/message/receive.rs b/src/message/receive.rs index a721cd975..a05eeb475 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -344,7 +344,25 @@ impl Client { // can lose pkmsg messages carrying SKDM (sender key distribution). If the // SKDM is lost, ALL subsequent skmsg messages from that sender will fail // with "No sender key state". + let generation = self + .connection_generation + .load(std::sync::atomic::Ordering::Acquire); let _global_permit = self.acquire_message_processing_permit().await; + if self + .connection_generation + .load(std::sync::atomic::Ordering::Acquire) + != generation + { + // Teardown bumped the generation while this stanza waited for the + // permit; its cache settle must be the LAST Signal-cache activity + // of the connection. Decrypting now would advance ratchets with + // no committable entry — bail unacked, the server redelivers. + log::debug!( + "Connection torn down while awaiting the processing permit; leaving message {} for redelivery", + info.id + ); + return; + } log::debug!( "Starting PASS 1: Processing {} session establishment messages (pkmsg/msg)", From 4053f187e04190246d0444dd858e89649c364722 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 09:36:23 +0000 Subject: [PATCH 24/35] fix(recv): compare stale lane work against the lane's spawn generation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The post-permit re-check captured connection_generation inside process_classified_message, so a teardown bump landing while the worker was still classifying was absorbed into a fresher capture and the check passed — the stale stanza could still decrypt after the cache settle. The chat-lane worker now passes its spawn generation through handle_incoming_message_scoped, and the re-check compares against that stable value: any bump after the worker's loop-top validation trips it, before any ratchet mutation. --- src/handlers/message.rs | 4 +++- src/message/receive.rs | 33 +++++++++++++++++++++++------ src/message/tests.rs | 46 +++++++++++++++++++++++++---------------- 3 files changed, 58 insertions(+), 25 deletions(-) diff --git a/src/handlers/message.rs b/src/handlers/message.rs index 7c9b8cb85..1f46338bb 100644 --- a/src/handlers/message.rs +++ b/src/handlers/message.rs @@ -89,7 +89,9 @@ fn create_chat_lane(client: &Arc) -> ChatLane { // 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(msg_node).await; + 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!( diff --git a/src/message/receive.rs b/src/message/receive.rs index a05eeb475..849d64097 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -7,7 +7,26 @@ impl Client { feature = "tracing", tracing::instrument(name = "wa.recv.incoming", level = "debug", skip_all) )] + /// Test convenience: scope to the current generation. Production inbound + /// traffic always goes through the chat-lane worker, which passes its own + /// spawn generation. + #[cfg(test)] pub(crate) async fn handle_incoming_message(self: Arc, node: Arc) { + let generation = self + .connection_generation + .load(std::sync::atomic::Ordering::Acquire); + self.handle_incoming_message_scoped(node, generation).await + } + + /// `lane_generation` is the generation the CALLER validated (the chat-lane + /// worker's spawn generation) — not re-read here, so a teardown bump that + /// lands mid-classification still trips the post-permit re-check instead + /// of being absorbed into a fresher capture. + pub(crate) async fn handle_incoming_message_scoped( + self: Arc, + node: Arc, + lane_generation: u64, + ) { // Phase 1: classify borrows the node tree, extracts owned payloads, returns quickly. // Phase 2: process_classified_message holds no node borrows across heavy .await points, // keeping the async state machine small. @@ -17,7 +36,8 @@ impl Client { }; // node is no longer borrowed here -- drop it before the heavy phase drop(node); - self.process_classified_message(classified).await; + self.process_classified_message(classified, lane_generation) + .await; } #[cfg_attr( @@ -305,7 +325,11 @@ impl Client { feature = "tracing", tracing::instrument(name = "wa.recv.process", level = "debug", skip_all) )] - pub(crate) async fn process_classified_message(self: Arc, msg: ClassifiedMessage) { + pub(crate) async fn process_classified_message( + self: Arc, + msg: ClassifiedMessage, + lane_generation: u64, + ) { let ClassifiedMessage { info, sender_encryption_jid, @@ -344,14 +368,11 @@ impl Client { // can lose pkmsg messages carrying SKDM (sender key distribution). If the // SKDM is lost, ALL subsequent skmsg messages from that sender will fail // with "No sender key state". - let generation = self - .connection_generation - .load(std::sync::atomic::Ordering::Acquire); let _global_permit = self.acquire_message_processing_permit().await; if self .connection_generation .load(std::sync::atomic::Ordering::Acquire) - != generation + != lane_generation { // Teardown bumped the generation while this stanza waited for the // permit; its cache settle must be the LAST Signal-cache activity diff --git a/src/message/tests.rs b/src/message/tests.rs index b87991d62..2f928e06c 100644 --- a/src/message/tests.rs +++ b/src/message/tests.rs @@ -6210,15 +6210,20 @@ async fn process_session_ct( }); client .clone() - .process_classified_message(ClassifiedMessage { - info, - sender_encryption_jid: sender.clone(), - session_payloads: vec![payload], - group_payloads: vec![], - bot_payloads: vec![], - max_sender_retry_count: 0, - decrypt_fail_mode: crate::types::events::DecryptFailMode::Show, - }) + .process_classified_message( + ClassifiedMessage { + info, + sender_encryption_jid: sender.clone(), + session_payloads: vec![payload], + group_payloads: vec![], + bot_payloads: vec![], + max_sender_retry_count: 0, + decrypt_fail_mode: crate::types::events::DecryptFailMode::Show, + }, + client + .connection_generation + .load(std::sync::atomic::Ordering::Acquire), + ) .await; } @@ -6310,15 +6315,20 @@ async fn process_group_classified_with_payloads( ) { client .clone() - .process_classified_message(ClassifiedMessage { - info, - sender_encryption_jid: sender.clone(), - session_payloads, - group_payloads, - bot_payloads, - max_sender_retry_count: 0, - decrypt_fail_mode: crate::types::events::DecryptFailMode::Show, - }) + .process_classified_message( + ClassifiedMessage { + info, + sender_encryption_jid: sender.clone(), + session_payloads, + group_payloads, + bot_payloads, + max_sender_retry_count: 0, + decrypt_fail_mode: crate::types::events::DecryptFailMode::Show, + }, + client + .connection_generation + .load(std::sync::atomic::Ordering::Acquire), + ) .await; } From a5618bd528ebc94f77695da38a570de19a790825 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 09:51:06 +0000 Subject: [PATCH 25/35] fix(recv): gate classification on the lane generation; keep the tracing span in production Two follow-ups on the lane-generation guard: - classify_incoming_message is not side-effect-free (newsletter dispatch, unavailable-only acks, PDO scheduling), so a stale stanza is now dropped BEFORE classification when the generation moved past the lane's spawn value; the post-permit re-check still covers a bump landing between classification and the decrypt. - The wa.recv.incoming tracing span had moved onto the test-only wrapper, vanishing from production builds; it now instruments handle_incoming_message_scoped, the real entry point. --- src/message/receive.rs | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/src/message/receive.rs b/src/message/receive.rs index 849d64097..79417ecdd 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -3,10 +3,6 @@ use super::*; impl Client { - #[cfg_attr( - feature = "tracing", - tracing::instrument(name = "wa.recv.incoming", level = "debug", skip_all) - )] /// Test convenience: scope to the current generation. Production inbound /// traffic always goes through the chat-lane worker, which passes its own /// spawn generation. @@ -22,11 +18,29 @@ impl Client { /// worker's spawn generation) — not re-read here, so a teardown bump that /// lands mid-classification still trips the post-permit re-check instead /// of being absorbed into a fresher capture. + #[cfg_attr( + feature = "tracing", + tracing::instrument(name = "wa.recv.incoming", level = "debug", skip_all) + )] pub(crate) async fn handle_incoming_message_scoped( self: Arc, node: Arc, lane_generation: u64, ) { + // Classification is not side-effect-free (newsletter dispatch, + // unavailable-only acks, PDO scheduling), so a stale stanza must be + // dropped BEFORE it — this pairs with the post-permit re-check, which + // covers a bump landing between here and the decrypt. + if self + .connection_generation + .load(std::sync::atomic::Ordering::Acquire) + != lane_generation + { + log::debug!( + "Connection torn down before classification; leaving the stanza for redelivery" + ); + return; + } // Phase 1: classify borrows the node tree, extracts owned payloads, returns quickly. // Phase 2: process_classified_message holds no node borrows across heavy .await points, // keeping the async state machine small. From 371d4c954104fa6d08fbe94bfe94a24eec7ea03f Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 10:04:48 +0000 Subject: [PATCH 26/35] fix(recv): hold SKDM receipts through the deferred window; keep the cache on entryless teardown timeouts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two deferred-window edges: - try_buffer_offline_receipt short-circuited on offline_sync_completed, so an SKDM-only stanza processed during a deferred drain-to-live transition sent its delivery receipt 1:1 while its sender-key state was still cache-only; a crash before the retry flush made the loss permanent (the server does not redeliver receipted stanzas). Receipts now keep buffering while the batcher is active, and completing the deferred transition flushes the buffer — right after the durable flush that made it safe. - The teardown timeout cleared the Signal cache unconditionally, which also destroyed retained committed/acked state when the timeout fired after the durable point (e.g. a hung settle flush). Cancellation is synchronous, so has_entries() distinguishes the cases: with restored entries the cache holds their rowless advances and both drop together; without entries everything dirty is redeliverable SKDM residue or retained committed state, and the cache is kept for the next successful flush. --- src/message/commit_batch.rs | 28 ++++++++++++++++++++++++---- src/receipt.rs | 7 +++++++ 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index e545bb5e3..ac5d7f8c0 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -385,6 +385,11 @@ impl Client { pub(crate) fn complete_deferred_live_transition(&self) { self.inbound_commit_batch.deactivate(); self.swap_message_semaphore(64); + // Receipts buffered during the deferred window (SKDM-only stanzas + // keep buffering while the batcher is active) are safe to send now: + // the durable flush that triggered this completion persisted their + // Signal state. + self.flush_offline_receipts(); log::info!("Deferred drain-to-live transition completed after a durable flush"); } @@ -475,10 +480,25 @@ impl Client { .await .is_err() { - log::warn!( - "Timed out committing the inbound drain batch during teardown; dropping unflushed Signal state so redelivery stays consistent" - ); - self.signal_cache.clear().await; + // Cancellation is synchronous: a commit cut down mid-durable-write + // has already restored its entries via the guard by the time the + // timeout returns, so has_entries() accurately distinguishes the + // two kinds of dirty state. With entries, the cache holds their + // rowless ratchet advances — drop both sides so redelivery stays + // consistent. Without entries, everything dirty is either + // redeliverable SKDM residue or committed state a failed earlier + // flush retained (never redelivered) — keep it for the next + // successful flush instead of destroying it. + if self.inbound_commit_batch.has_entries() { + log::warn!( + "Timed out committing the inbound drain batch during teardown; dropping unflushed Signal state so redelivery stays consistent" + ); + self.signal_cache.clear().await; + } else { + log::warn!( + "Timed out settling the Signal cache during teardown; keeping it (no uncommitted entries) for the next successful flush" + ); + } } } diff --git a/src/receipt.rs b/src/receipt.rs index f947585b0..2f817915f 100644 --- a/src/receipt.rs +++ b/src/receipt.rs @@ -563,9 +563,16 @@ impl Client { .offline_receipt_buffer .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); + // Batcher still active covers the deferred drain→live window: the + // completion flag is already set there, but SKDM-only stanzas keep + // mutating cache-only sender-key state, and a 1:1 receipt for one of + // them before the deferred retry flushes would trade a redeliverable + // failure for a crash-permanent one. Completion of the deferred + // transition flushes this buffer (after its durable flush). if self .offline_sync_completed .load(std::sync::atomic::Ordering::Acquire) + && !self.inbound_commit_batch.is_active() { return false; } From 77cb458454260353b227a1f7b12712867ade8c05 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 10:08:59 +0000 Subject: [PATCH 27/35] test(recv): cover the deferred-window receipt-buffering contract The offline receipt protocol test flipped only the completion flag, which now (correctly) keeps buffering while the batcher is active; it asserts the deferred-window behavior explicitly and uses the full live transition for the 1:1 fallback assertion. --- src/receipt.rs | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/src/receipt.rs b/src/receipt.rs index 2f817915f..0a98c038a 100644 --- a/src/receipt.rs +++ b/src/receipt.rs @@ -2443,11 +2443,28 @@ mod tests { 2 ); - // Once the completion flag flips, late offline receipts fall back to - // 1:1 instead of stranding in the buffer (the exact race this guards). + // The completion flag alone is NOT enough to go 1:1: during a + // deferred drain-to-live transition the flag is set while the batcher + // stays active, and receipts must keep buffering (their SKDM state + // may still be cache-only until the deferred flush). client .offline_sync_completed .store(true, std::sync::atomic::Ordering::Release); + let deferred = offline_info( + "OFF2B", + "5511999990000@s.whatsapp.net", + "5511999990000@s.whatsapp.net", + false, + ); + assert!(client.try_buffer_offline_receipt(&deferred)); + assert_eq!( + client.offline_receipt_buffer.lock().expect("buffer").len(), + 3 + ); + + // Once the batcher goes live too, late offline receipts fall back to + // 1:1 instead of stranding in the buffer (the exact race this guards). + client.enter_live_mode_for_tests(); let late = offline_info( "OFF3", "5511999990000@s.whatsapp.net", @@ -2457,7 +2474,7 @@ mod tests { assert!(!client.try_buffer_offline_receipt(&late)); assert_eq!( client.offline_receipt_buffer.lock().expect("buffer").len(), - 2 + 3 ); // Flush drains everything and releases the backing capacity, so no From baf39523d66ac1036f4afdefc472fdb78b99135a Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 10:25:50 +0000 Subject: [PATCH 28/35] fix(recv): couple the Signal-cache drop to the batcher resets that drop entries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The teardown-timeout branch decided the cache's fate by sampling has_entries() at timeout time, which races the very permit holder the settle timed out on: a stanza can still be mid-decrypt with unenqueued advances (has_entries false), or enqueue after any clear taken there. The decision now lives where it cannot race: reset() reports whether it dropped entries, and both entry-dropping resets (cleanup's and connect's) drop the cache together with the entries — rowless advances never outlive their entries, no matter when the straggler lands. With nothing dropped, the cache survives for the next successful flush (redeliverable SKDM residue or retained committed state). The timeout branch itself just defers to that rule. --- src/client/lifecycle.rs | 29 ++++++++++++++++++--------- src/message/commit_batch.rs | 39 +++++++++++++++++-------------------- 2 files changed, 38 insertions(+), 30 deletions(-) diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index ffae685cc..741468db3 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -460,13 +460,20 @@ impl Client { .store(false, Ordering::Relaxed); self.clear_offline_receipt_buffer(); // Uncommitted batch entries were never acked; the server redelivers - // them on this fresh connection. The Signal cache needs no clear - // here: teardown settled it under the permit, and its generation - // bump plus the post-permit re-check guarantee no late lane worker - // dirtied it afterwards. Anything still resident is state a failed - // teardown flush deliberately retained (committed/acked, never - // redelivered) for the next successful flush to persist. - self.inbound_commit_batch.reset(); + // them on this fresh connection. The cache decision is coupled to the + // drop: entries present here mean their cache-only ratchet advances + // have no rows (e.g. a stanza that outlived the teardown settle and + // enqueued late), and flushing those later would make each redelivery + // an ackable duplicate — so the cache falls with them. With nothing + // dropped, anything resident is state a failed teardown flush + // deliberately retained (committed/acked, never redelivered) for the + // next successful flush to persist. + if self.inbound_commit_batch.reset() { + log::warn!( + "connect: dropping unflushed Signal state along with late uncommitted drain entries" + ); + self.signal_cache.clear().await; + } self.offline_batch.reset(); self.outbound_flush.reopen(); @@ -810,8 +817,12 @@ impl Client { .store(false, Ordering::Relaxed); self.clear_offline_receipt_buffer(); // Same rule as receipts: uncommitted entries drop here and the server - // redelivers them on the next connect. - self.inbound_commit_batch.reset(); + // redelivers them on the next connect. The cache falls with dropped + // entries (rowless advances — including a timed-out settle's restored + // batch); with nothing dropped it survives for the next flush. + if self.inbound_commit_batch.reset() { + self.signal_cache.clear().await; + } self.offline_batch.reset(); self.offline_sync_metrics .active diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index ac5d7f8c0..43ff6d476 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -150,8 +150,11 @@ impl InboundCommitBatcher { /// Connection teardown/setup: drop uncommitted entries (they were never /// acked, so the server redelivers them) and re-arm accumulation for the - /// next connection's drain. - pub(crate) fn reset(&self) { + /// next connection's drain. Returns whether entries were dropped — the + /// caller must drop the Signal cache with them (their cache-only ratchet + /// advances have no rows; flushing them later would make each redelivery + /// an ackable duplicate). + pub(crate) fn reset(&self) -> bool { let dropped = self.take(); if !dropped.is_empty() { log::debug!( @@ -161,6 +164,7 @@ impl InboundCommitBatcher { } self.pending_live.store(false, Ordering::Release); self.active.store(true, Ordering::Release); + !dropped.is_empty() } } @@ -480,25 +484,18 @@ impl Client { .await .is_err() { - // Cancellation is synchronous: a commit cut down mid-durable-write - // has already restored its entries via the guard by the time the - // timeout returns, so has_entries() accurately distinguishes the - // two kinds of dirty state. With entries, the cache holds their - // rowless ratchet advances — drop both sides so redelivery stays - // consistent. Without entries, everything dirty is either - // redeliverable SKDM residue or committed state a failed earlier - // flush retained (never redelivered) — keep it for the next - // successful flush instead of destroying it. - if self.inbound_commit_batch.has_entries() { - log::warn!( - "Timed out committing the inbound drain batch during teardown; dropping unflushed Signal state so redelivery stays consistent" - ); - self.signal_cache.clear().await; - } else { - log::warn!( - "Timed out settling the Signal cache during teardown; keeping it (no uncommitted entries) for the next successful flush" - ); - } + // Keep the cache: the decision is owned by the reset-coupled + // clear (cleanup's and connect's batcher resets drop the cache + // whenever they drop entries). Deciding here would race the + // permit holder we timed out on — it may still be mid-decrypt + // with unenqueued advances that has_entries() cannot see, and it + // may enqueue after any clear done here. Whatever is dirty when + // the resets find NO entries is redeliverable SKDM residue or + // retained committed state, which the next successful flush + // persists. + log::warn!( + "Timed out settling the inbound drain during teardown; deferring the Signal-cache decision to the batcher reset" + ); } } From 52c8a874f021ca28df05d8b584634f0be58ab6e2 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 10:44:05 +0000 Subject: [PATCH 29/35] fix(recv): re-scope the mode mutation after awaited commits; keep the identity path single-permit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two stale-state windows around awaited drain commits: - The finisher's generation check ran only before the take. A teardown that times out around the awaited commit (slow hook) resets the batcher, and the resuming finisher then deactivated the NEW connection's drain. The generation is now re-checked after the await, before any mode mutation; the commit itself stays sound (entries taken pre-reset, rows durable). - commit_inbound_batch_holding_permit no longer completes a pending deferred transition: the caller (UntrustedIdentity recovery) holds a permit from the old single-permit semaphore and follows up with a raw whole-cache flush, which is only safe while that permit excludes every other stanza — widening to 64 first would let new workers be mid-decrypt under the flush. The deferred-retry loop completes the transition moments later, outside any raw-flush window. --- src/message/commit_batch.rs | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 43ff6d476..d3fd78caf 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -357,6 +357,16 @@ impl Client { self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) .await }; + if let Some(generation) = expected_generation + && self.connection_generation.load(Ordering::Acquire) != generation + { + // The pre-take check only covers up to the take: a teardown can + // time out around the awaited commit above and reset the batcher, + // so mutating the mode now would deactivate (or defer) the NEW + // connection's drain. The commit itself was still sound — its + // entries were taken before the reset and its rows are durable. + return durable; + } if deactivate { if durable { self.inbound_commit_batch.deactivate(); @@ -573,13 +583,15 @@ impl Client { if batch.is_empty() { return true; } - let durable = self - .commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) - .await; - if durable && self.inbound_commit_batch.live_transition_pending() { - self.complete_deferred_live_transition(); - } - durable + // Deliberately NOT completing a pending deferred transition here even + // on a durable commit: the caller holds a permit from the old + // single-permit semaphore and follows up with a raw whole-cache + // flush, which is only safe while that permit excludes every other + // stanza — widening to 64 permits first would let new workers be + // mid-decrypt under it. The deferred-retry loop completes the + // transition moments later, outside any raw-flush window. + self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) + .await } /// Commit one batch: durable buffer → Signal flush → hook → clear buffer → From 68358c4fc642ccd997fb28344b15ef10d8fc140d Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 14:38:51 +0000 Subject: [PATCH 30/35] docs(recv): fix stale re-acquire comment at the permit call site MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The comment still described an inline generation-reacquire loop, but that logic moved into acquire_message_processing_permit; trim it to reference the helper while keeping the SKDM-loss rationale for why the 1→N re-acquire matters. --- src/message/receive.rs | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/src/message/receive.rs b/src/message/receive.rs index 79417ecdd..ff4ee6850 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -373,15 +373,12 @@ impl Client { ); } - // Acquire global processing permit (1 during offline sync, N after). - // Read generation + clone Arc under the same mutex so the pair is consistent. - // - // When the semaphore transitions from 1→N (offline→online), tasks waiting on - // the old 1-permit semaphore must re-acquire from the new N-permit semaphore. - // Without this re-acquire loop, those tasks would be silently dropped, which - // can lose pkmsg messages carrying SKDM (sender key distribution). If the - // SKDM is lost, ALL subsequent skmsg messages from that sender will fail - // with "No sender key state". + // Acquire the global processing permit (1 during offline sync, N after). + // The helper re-acquires across a 1→N semaphore swap (offline→online): + // without that, a task waiting on the old 1-permit semaphore would be + // silently dropped, losing pkmsg messages carrying SKDM (sender key + // distribution) — and a lost SKDM fails ALL subsequent skmsg from that + // sender with "No sender key state". let _global_permit = self.acquire_message_processing_permit().await; if self .connection_generation From 19eee50d48588e0eac68729aec27f73eef84c2d6 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 14:54:16 +0000 Subject: [PATCH 31/35] fix(recv): drop the dirty Signal cache on a teardown-settle timeout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The timeout branch deferred the cache decision to the batcher reset, but reset only clears when it drops entries — so a worker hung mid-decrypt (holding the permit the settle timed out waiting for) that advanced a ratchet WITHOUT enqueueing an entry (SKDM-only) left rowless state the reset could not see, and a later flush would persist it into an acked-duplicate loss. Clear unconditionally on timeout instead: it does not sample has_entries() (which races that permit holder — the concern from the earlier round), and it is the only rowless-safe action when a stuck worker may hold unenqueued advances. The common hook-timeout case is already safe (the ReinsertGuard restored the entries), so this only changes the rare hung-worker corner. The sole thing it can drop is committed state a prior failed flush retained across a total storage outage — a degraded corner where redeliver-on-reauth beats a silent rowless duplicate. --- src/message/commit_batch.rs | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index d3fd78caf..138cbe7c6 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -494,18 +494,25 @@ impl Client { .await .is_err() { - // Keep the cache: the decision is owned by the reset-coupled - // clear (cleanup's and connect's batcher resets drop the cache - // whenever they drop entries). Deciding here would race the - // permit holder we timed out on — it may still be mid-decrypt - // with unenqueued advances that has_entries() cannot see, and it - // may enqueue after any clear done here. Whatever is dirty when - // the resets find NO entries is redeliverable SKDM residue or - // retained committed state, which the next successful flush - // persists. + // Drop the whole dirty cache — do NOT sample has_entries() (it + // races the permit holder we timed out on, per the earlier + // review). A hook-timeout is the common case: the settle held + // the permit and its ReinsertGuard has restored the entries by + // now, and clearing drops their rowless advances safely (rows + // never persisted → server redelivers). The rare case is a + // worker hung mid-decrypt holding the permit; it may have + // advanced a ratchet WITHOUT enqueueing an entry (SKDM-only), + // which the reset-coupled clear cannot see — so the only + // rowless-safe action is to clear unconditionally here. The one + // thing this can drop is committed state a PRIOR connection's + // failed flush retained AND no flush since re-persisted: a + // total-storage-outage corner where the system is already + // degraded, and losing redeliverable-on-reauth state beats + // silently acking a rowless duplicate. log::warn!( - "Timed out settling the inbound drain during teardown; deferring the Signal-cache decision to the batcher reset" + "Timed out settling the inbound drain during teardown; dropping the dirty Signal cache so no rowless ratchet advance can persist" ); + self.signal_cache.clear().await; } } From 19bef0e35bf0ef6760b3e2278c40528ffee0b0eb Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 15:23:41 +0000 Subject: [PATCH 32/35] refactor(events): tighten the Messages event API surface (pre-1.0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit No production consumers yet, so shape the new batch API properly instead of carrying the accreted form: - Rename Event::message_batch() → Event::as_messages(): the idiomatic as_ prefix for a &self -> Option<&T> accessor, and it now tiers clearly against messages() (the flattening iterator) instead of two near-synonym names. - Make MessageBatch a first-class collection: iter/len/is_empty/first and IntoIterator for &MessageBatch, so `for m in &batch`, batch.iter(), batch.len() work without reaching through the .messages field. origin stays alongside as the delivery-shape metadata. - messages() is reimplemented on top of as_messages(); it stays because `event.messages().any(...)` / `for m in event.messages()` is the dominant consumer verb and reads better than the derived form. Callers migrated; the batcher tests and benchmark example dogfood the collection methods. --- examples/benchmark.rs | 2 +- src/message/commit_batch.rs | 12 +++---- tests/e2e/tests/offline_groups.rs | 4 +-- wacore/src/types/events.rs | 55 ++++++++++++++++++++++++------- 4 files changed, 52 insertions(+), 21 deletions(-) diff --git a/examples/benchmark.rs b/examples/benchmark.rs index a25613729..fa1aac172 100644 --- a/examples/benchmark.rs +++ b/examples/benchmark.rs @@ -89,7 +89,7 @@ fn main() { async move { match &*event { Event::Messages(batch) => { - for m in batch.messages.iter() { + for m in batch { if m.message.text_content() != Some("ping") { continue; } diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 138cbe7c6..c2e64746d 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -845,9 +845,9 @@ mod tests { assert_eq!(batches, vec![vec!["B1", "B2", "B3"]]); let event = rx.try_recv().expect("one batch event"); - let batch = event.message_batch().expect("Messages event"); + let batch = event.as_messages().expect("Messages event"); assert_eq!(batch.origin, BatchOrigin::OfflineDrain); - let ids: Vec<&str> = batch.messages.iter().map(|m| m.info.id.as_str()).collect(); + let ids: Vec<&str> = batch.iter().map(|m| m.info.id.as_str()).collect(); assert_eq!(ids, ["B1", "B2", "B3"]); assert!(rx.try_recv().is_err(), "exactly one event for the batch"); @@ -881,9 +881,9 @@ mod tests { vec![vec!["L1"]] ); let event = rx.try_recv().expect("live event"); - let batch = event.message_batch().expect("Messages event"); + let batch = event.as_messages().expect("Messages event"); assert_eq!(batch.origin, BatchOrigin::Live); - assert_eq!(batch.messages.len(), 1); + assert_eq!(batch.len(), 1); } // The size trigger commits a full batch from the stanza-end check. @@ -962,12 +962,12 @@ mod tests { ); let first = rx.try_recv().expect("tail event"); assert_eq!( - first.message_batch().expect("Messages").origin, + first.as_messages().expect("Messages").origin, BatchOrigin::OfflineDrain ); let second = rx.try_recv().expect("live event"); assert_eq!( - second.message_batch().expect("Messages").origin, + second.as_messages().expect("Messages").origin, BatchOrigin::Live ); } diff --git a/tests/e2e/tests/offline_groups.rs b/tests/e2e/tests/offline_groups.rs index 0d6cb2e92..6a89d275e 100644 --- a/tests/e2e/tests/offline_groups.rs +++ b/tests/e2e/tests/offline_groups.rs @@ -151,7 +151,7 @@ async fn test_mixed_offline_event_ordering() -> anyhow::Result<()> { .await; match result { - Ok(ref event) if event.message_batch().is_some() => { + Ok(ref event) if event.as_messages().is_some() => { for m in event.messages() { if let Some(text) = &m.message.conversation { info!("C received message: {text}"); @@ -454,7 +454,7 @@ async fn test_offline_multi_sender_group_messages() -> anyhow::Result<()> { .await; match result { - Ok(ref event) if event.message_batch().is_some() => { + Ok(ref event) if event.as_messages().is_some() => { for m in event.messages() { if let Some(text) = &m.message.conversation { info!("C received: {text}"); diff --git a/wacore/src/types/events.rs b/wacore/src/types/events.rs index ad088c7b9..5a82d89ed 100755 --- a/wacore/src/types/events.rs +++ b/wacore/src/types/events.rs @@ -834,21 +834,22 @@ impl Event { } } - pub fn message_batch(&self) -> Option<&MessageBatch> { - if let Event::Messages(batch) = self { - Some(batch) - } else { - None + /// This event as its [`MessageBatch`], or `None` for any other event kind. + /// Use this when you need the batch's [`origin`](MessageBatch::origin) or + /// want to treat the messages as a whole; to just iterate the messages, + /// prefer [`messages`](Self::messages). + pub fn as_messages(&self) -> Option<&MessageBatch> { + match self { + Event::Messages(batch) => Some(batch), + _ => None, } } - /// The inbound messages carried by this event, in arrival order; empty for - /// every other event kind. + /// The inbound messages carried by this event, in arrival order; an empty + /// iterator for every other event kind (so it drops cleanly into a + /// `for msg in event.messages()` scan over a mixed event stream). pub fn messages(&self) -> impl Iterator { - self.message_batch() - .map(|b| b.messages.iter()) - .into_iter() - .flatten() + self.as_messages().into_iter().flatten() } } @@ -874,13 +875,43 @@ pub enum BatchOrigin { OfflineDrain, } -/// Payload of [`Event::Messages`]. +/// Payload of [`Event::Messages`]: the decrypted messages of one durable +/// commit, in arrival order. Behaves as a collection of its messages — +/// `for msg in &batch`, `batch.iter()`, `batch.len()` — with `origin` +/// carrying the delivery shape alongside. #[derive(Debug, Clone, Serialize)] pub struct MessageBatch { pub messages: Arc<[InboundMessage]>, pub origin: BatchOrigin, } +impl MessageBatch { + pub fn iter(&self) -> std::slice::Iter<'_, InboundMessage> { + self.messages.iter() + } + + pub fn len(&self) -> usize { + self.messages.len() + } + + pub fn is_empty(&self) -> bool { + self.messages.is_empty() + } + + pub fn first(&self) -> Option<&InboundMessage> { + self.messages.first() + } +} + +impl<'a> IntoIterator for &'a MessageBatch { + type Item = &'a InboundMessage; + type IntoIter = std::slice::Iter<'a, InboundMessage>; + + fn into_iter(self) -> Self::IntoIter { + self.messages.iter() + } +} + /// A newsletter live update notification, typically containing updated /// reaction counts for one or more messages. #[derive(Debug, Clone, Serialize)] From 99fe8c79d5dfd86caf5d330fa75a404f56634f99 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 15:51:16 +0000 Subject: [PATCH 33/35] fix(recv): close batch-safe flush TOCTOU and log teardown drops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit flush_signal_cache_batch_safe entered on an is_active() check that the drain finisher can invalidate before the wrapper acquired the permit: the under-permit commit then took an empty batch in live mode and reported durable WITHOUT flushing the caller's out-of-band Signal advance, leaving it dirty until the next flush. Re-check is_active() after the commit and fall through to the raw flush when the batcher deactivated mid-wait — the same path a caller that found it already inactive takes. Covered by a deferred-transition test driven through the wrapper. Also log the teardown branch that drops uncommitted drain entries and clears the Signal cache, matching the connect-side twin and the upgrade-failure path so redelivery-triggering drops are visible. --- src/client/adapters.rs | 25 ++++++++++------ src/client/lifecycle.rs | 3 ++ src/message/commit_batch.rs | 57 +++++++++++++++++++++++++++++++++++++ 3 files changed, 77 insertions(+), 8 deletions(-) diff --git a/src/client/adapters.rs b/src/client/adapters.rs index da8e60308..15ce3c036 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -113,18 +113,27 @@ impl Client { pub(crate) async fn flush_signal_cache_batch_safe(&self) -> Result<(), anyhow::Error> { if self.inbound_commit_batch.is_active() { if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { - return if client + if !client .flush_inbound_commits_under_permit(false, None, None) .await { - Ok(()) - } else { - Err(anyhow::anyhow!( + return Err(anyhow::anyhow!( "inbound drain batch commit failed; Signal cache left unflushed so the server redelivers" - )) - }; - } - if self.inbound_commit_batch.has_entries() { + )); + } + // The is_active() check above races the drain finisher: if the + // batcher deactivated while we waited for the permit, the commit + // took an empty batch in live mode and reported durable WITHOUT + // flushing (it had no drain rows to tie an advance to). We are + // now live with no uncommitted drain entries, so fall through to + // the raw flush to persist our out-of-band advance — the same + // path a caller that found the batcher already inactive takes. + // While still draining, the commit already flushed under the + // permit, so return. + if self.inbound_commit_batch.is_active() { + return Ok(()); + } + } else if self.inbound_commit_batch.has_entries() { return Err(anyhow::anyhow!( "client dropping with uncommitted drain entries; skipping Signal flush" )); diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 741468db3..0db18d755 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -821,6 +821,9 @@ impl Client { // entries (rowless advances — including a timed-out settle's restored // batch); with nothing dropped it survives for the next flush. if self.inbound_commit_batch.reset() { + log::warn!( + "cleanup_connection_state: dropping unflushed Signal state along with late uncommitted drain entries" + ); self.signal_cache.clear().await; } self.offline_batch.reset(); diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index c2e64746d..8ec439e82 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -1168,6 +1168,63 @@ mod tests { ); } + // The batch-safe wrapper enters on an is_active() check that the drain + // finisher can invalidate before it acquires the permit. When the + // under-permit commit deactivates the batcher (here, by completing a + // deferred transition), the wrapper must fall through to the raw Signal + // flush instead of returning Ok on the stale check — otherwise the + // caller's out-of-band advance is silently skipped. This drives that + // completion through the wrapper and asserts it lands in live mode. + #[tokio::test] + async fn batch_safe_flush_completes_deferred_transition() { + let client = create_test_client_with_failing_http("batch_safe_defer").await; + client.inbound_commit_batch.reset(); + client.swap_message_semaphore(1); + let hook = Arc::new(RecordingHook { + batches: Mutex::new(Vec::new()), + }); + let _ = client.inbound_durability_hook.set(hook.clone()); + + client.commit_or_batch_inbound(item("B1")).await; + client + .inbound_commit_batch + .fail_commits + .store(true, Ordering::Release); + let generation = client.connection_generation.load(Ordering::Acquire); + assert!(!client.finish_inbound_commit_drain(generation).await); + assert!( + client.inbound_commit_batch.is_active() && client.inbound_commit_batch.has_entries(), + "failed tail defers the transition and restores the batch" + ); + client + .inbound_commit_batch + .fail_commits + .store(false, Ordering::Release); + + // Enters on is_active()==true, commits the restored tail under the + // permit, completes the deferred transition (deactivates), then falls + // through to the raw flush. + client + .flush_signal_cache_batch_safe() + .await + .expect("batch-safe flush must succeed once the tail commits"); + + assert!( + !client.inbound_commit_batch.is_active(), + "the batch-safe flush must complete the deferred transition" + ); + assert!(!client.inbound_commit_batch.has_entries()); + assert_eq!( + available_permits(&client), + 64, + "completing the transition widens the semaphore" + ); + assert_eq!( + hook.batches.lock().expect("hook lock").clone(), + vec![vec!["B1"]] + ); + } + fn available_permits(client: &Client) -> usize { let semaphore = match client.message_processing_semaphore.lock() { Ok(guard) => guard.clone(), From d82ccea0c0dd8a9001458d364e3e0ab3e689a4db Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 16:08:31 +0000 Subject: [PATCH 34/35] refactor(recv): make the under-permit commit always flush; drop stale recheck MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The batch-safe flush proved durability with a post-commit is_active() recheck, which is not a reliable proxy for whether the Signal cache was actually flushed. Instead make flush_inbound_commits_under_permit's empty path flush the cache unconditionally (it already did while draining; now also when the batcher went live between the caller's is_active() check and the permit). That makes "cache flushed under the permit" an unconditional post-condition, so the wrapper drops the racy recheck and collapses to a let-chain. Add a deterministic test (fail_flushes injection) proving the empty commit still flushes — the pre-fix empty+live path returned durable without flushing, so it now guards that regression observably. --- src/client/adapters.rs | 48 +++++++++++++----------------- src/message/commit_batch.rs | 59 ++++++++++++++++++++++++++++++------- 2 files changed, 69 insertions(+), 38 deletions(-) diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 15ce3c036..7de6a1099 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -111,33 +111,27 @@ impl Client { /// Must NOT be called while holding the processing permit (it acquires /// it); permit-holding paths commit via the batcher directly. pub(crate) async fn flush_signal_cache_batch_safe(&self) -> Result<(), anyhow::Error> { - if self.inbound_commit_batch.is_active() { - if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { - if !client - .flush_inbound_commits_under_permit(false, None, None) - .await - { - return Err(anyhow::anyhow!( - "inbound drain batch commit failed; Signal cache left unflushed so the server redelivers" - )); - } - // The is_active() check above races the drain finisher: if the - // batcher deactivated while we waited for the permit, the commit - // took an empty batch in live mode and reported durable WITHOUT - // flushing (it had no drain rows to tie an advance to). We are - // now live with no uncommitted drain entries, so fall through to - // the raw flush to persist our out-of-band advance — the same - // path a caller that found the batcher already inactive takes. - // While still draining, the commit already flushed under the - // permit, so return. - if self.inbound_commit_batch.is_active() { - return Ok(()); - } - } else if self.inbound_commit_batch.has_entries() { - return Err(anyhow::anyhow!( - "client dropping with uncommitted drain entries; skipping Signal flush" - )); - } + // Under the permit the commit ALWAYS flushes the Signal cache (an empty + // batch still flushes — see flush_inbound_commits_under_permit), so a + // successful call is proof the out-of-band advance is persisted; no + // stale is_active() re-check needed even if the drain finisher + // deactivated while we waited for the permit. + let drain_active = self.inbound_commit_batch.is_active(); + if drain_active && let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { + return if client + .flush_inbound_commits_under_permit(false, None, None) + .await + { + Ok(()) + } else { + Err(anyhow::anyhow!( + "inbound drain batch commit failed; Signal cache left unflushed so the server redelivers" + )) + }; + } else if drain_active && self.inbound_commit_batch.has_entries() { + return Err(anyhow::anyhow!( + "client dropping with uncommitted drain entries; skipping Signal flush" + )); } self.flush_signal_cache().await } diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index 8ec439e82..cd90d9d4b 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -342,17 +342,18 @@ impl Client { let was_draining = self.inbound_commit_batch.is_active(); let batch = self.inbound_commit_batch.take(); let durable = if batch.is_empty() { - // Even with nothing to commit, a drain-mode flush must persist the - // Signal cache: SKDM-only stanzas mutate Signal state without - // enqueueing a message, and their buffered receipts flush right - // after the teardown/drain-end call sites of this function — so a - // failed flush must report not-durable to hold those receipts - // back (the cache keeps its dirty entries for a later retry). - if was_draining { - self.drain_signal_flush_reporting().await - } else { - true - } + // Even with nothing to commit, flush the Signal cache under the + // permit. While draining this persists SKDM-only advances (stanzas + // that mutate Signal state without enqueueing a message), whose + // buffered receipts flush right after this function's + // drain-end/teardown call sites — a failed flush reports + // not-durable to hold them back (the cache keeps its dirty entries + // for a later retry). When the batcher deactivated between an + // out-of-band batch-safe caller's is_active() check and this + // permit, the SAME flush persists that caller's advance instead of + // silently no-oping, so is_active() never has to stand in for "was + // flushed". Idempotent and cheap when the cache is already clean. + self.drain_signal_flush_reporting().await } else { self.commit_inbound_batch(batch.into(), BatchOrigin::OfflineDrain) .await @@ -1225,6 +1226,42 @@ mod tests { ); } + // An empty commit must still flush the Signal cache under the permit — + // this is what lets flush_signal_cache_batch_safe persist an out-of-band + // advance when the drain finisher deactivated the batcher between its + // is_active() check and the permit. Observable via the injected flush + // failure: skipping the flush would wrongly report durable (the pre-fix + // empty+live path returned true without flushing). + #[tokio::test] + async fn empty_commit_still_flushes_under_permit() { + let client = create_test_client_with_failing_http("batch_empty_flush").await; + // Live mode, empty batcher: was_draining is false under the permit. + assert!(!client.inbound_commit_batch.is_active()); + assert!(!client.inbound_commit_batch.has_entries()); + + client + .inbound_commit_batch + .fail_flushes + .store(true, Ordering::Release); + assert!( + !client + .flush_inbound_commits_under_permit(false, None, None) + .await, + "an empty commit must still hit the Signal cache and surface its failure" + ); + + client + .inbound_commit_batch + .fail_flushes + .store(false, Ordering::Release); + assert!( + client + .flush_inbound_commits_under_permit(false, None, None) + .await, + "with the flush succeeding the empty commit reports durable" + ); + } + fn available_permits(client: &Client) -> usize { let semaphore = match client.message_processing_semaphore.lock() { Ok(guard) => guard.clone(), From 4b672d69045987f1ef182ba822fdada79e933b78 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 3 Jul 2026 16:16:51 +0000 Subject: [PATCH 35/35] fix(recv): fail closed on active drain without a live client; fix stale doc MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The batch-safe flush's upgrade-failure fallback gated on has_entries() before falling through to a raw flush. But an empty drain can still hold dirty SKDM-only Signal advances with no rows, so that raw flush could persist them rowless — the loss this path exists to prevent. Fail closed for any active drain when the weak self ref can't upgrade (practically unreachable; the run loop holds a strong Arc) so the cache stays unflushed and the server redelivers. Also refresh the flush_inbound_commits_under_permit doc: an empty batch is no longer a no-op — it still flushes the Signal cache under the permit (idempotent when clean). --- src/client/adapters.rs | 11 +++++++++-- src/message/commit_batch.rs | 5 ++++- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 7de6a1099..5bdbad020 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -128,9 +128,16 @@ impl Client { "inbound drain batch commit failed; Signal cache left unflushed so the server redelivers" )) }; - } else if drain_active && self.inbound_commit_batch.has_entries() { + } else if drain_active { + // Drain active but no live client to route the permit-held flush + // through (practically unreachable — the run loop holds a strong + // Arc). Fail closed regardless of has_entries(): an empty drain can + // still carry dirty SKDM-only advances with no rows, and a raw + // flush here would persist them rowless — the exact loss this + // batch-safe path exists to prevent. Leaving the cache unflushed + // makes the server redeliver. return Err(anyhow::anyhow!( - "client dropping with uncommitted drain entries; skipping Signal flush" + "client dropping while inbound drain is active; skipping Signal flush" )); } self.flush_signal_cache().await diff --git a/src/message/commit_batch.rs b/src/message/commit_batch.rs index cd90d9d4b..ad14732fb 100644 --- a/src/message/commit_batch.rs +++ b/src/message/commit_batch.rs @@ -304,7 +304,10 @@ impl Client { /// persist a ratchet advance for a message no batch has committed, turning /// its redelivery into an unrecoverable duplicate. During the drain the /// semaphore holds a single permit, so this fully serializes with stanza - /// processing; after the drain the batcher is empty and this no-ops. + /// processing. An empty batch is NOT a no-op: it still flushes the Signal + /// cache under the permit (idempotent when the cache is already clean), so + /// an out-of-band batch-safe caller that raced the drain finisher still + /// gets its advance persisted rather than silently skipped. /// /// With `deactivate`, this is the end-of-drain transition: commit the tail /// batch and switch the batcher to live mode under the same permit hold.