From 62ff07f22f21a417d89e108cd9c74dfc5403e6f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Fri, 10 Jul 2026 22:58:38 -0300 Subject: [PATCH 01/24] perf(signal): coalesce hot-path Signal cache flushes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The live receive path and the send epilogue flushed the whole dirty Signal cache once per stanza: a session re-serialize plus a SQLite transaction per message, dominated by the record the ratchet re-dirties every time. Route both through a trailing-edge debounced scheduler (25 ms): one storage write covers a burst of messages. Ordering is unchanged — live acks already preceded the per-stanza flush — so this widens the existing crash-replay window from one stanza to at most one debounce, bounded and recoverable (receive chains re-derive forward; consumed prekeys stay buffered until their session is durable, the flush-internal atomicity is untouched). The offline drain, retry recovery, identity-change recovery and teardown keep their synchronous flushes: those gate acks, receipts or follow-up reads on durability. --- src/client.rs | 4 ++ src/client/adapters.rs | 6 -- src/client/lifecycle.rs | 1 + src/lib.rs | 1 + src/message/receive.rs | 16 +++-- src/send/mod.rs | 8 ++- src/signal_flush.rs | 156 ++++++++++++++++++++++++++++++++++++++++ 7 files changed, 176 insertions(+), 16 deletions(-) create mode 100644 src/signal_flush.rs diff --git a/src/client.rs b/src/client.rs index b9f0b6c9a..b298d20b2 100644 --- a/src/client.rs +++ b/src/client.rs @@ -836,6 +836,10 @@ pub struct Client { /// Initialized after `Arc::new(this)` in the constructor. pub(crate) self_weak: std::sync::OnceLock>, + /// Trailing-edge debounce flag for the coalesced Signal-cache flush + /// (see `signal_flush.rs`). True while a fire is armed. + pub(crate) signal_flush_pending: AtomicBool, + /// Holds the background saver's AbortHandle so the task lifetime follows /// `Arc` ref count instead of the Bot wrapper's. Set once by /// `Bot::build`; on Client drop (last Arc), the handle drops and the saver diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 5bdbad020..0e01f0c22 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -91,12 +91,6 @@ impl Client { /// 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 { - log_signal_flush_error(context, id, &e); - } - } - /// Signal-cache flush that is safe while the offline drain is active. /// /// During the drain, decrypted messages accumulate in the commit batcher diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 727185032..f0ece477a 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -240,6 +240,7 @@ impl Client { pair_code_state: Arc::new(Mutex::new(wacore::pair_code::PairCodeState::default())), passkey_state: Arc::new(Mutex::new(crate::passkey::flow::PasskeyFlowState::default())), passkey_opening: AtomicBool::new(false), + signal_flush_pending: AtomicBool::new(false), custom_enc_handlers: std::sync::OnceLock::new(), inbound_durability_hook: std::sync::OnceLock::new(), retry_admission: std::sync::OnceLock::new(), diff --git a/src/lib.rs b/src/lib.rs index 4d925b1ad..a2cc8a90a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -81,6 +81,7 @@ pub mod pair; pub mod pair_code; pub mod passkey; pub mod request; +pub(crate) mod signal_flush; pub use request::IqError; #[cfg(feature = "tokio-runtime")] pub mod runtime_impl; diff --git a/src/message/receive.rs b/src/message/receive.rs index 70421f904..47268d4f4 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -561,16 +561,18 @@ impl Client { self.handle_msmsg_payload(&info, payload).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. + // Live: schedule the coalesced Signal flush (WA Web flushes per + // stanza via flushBufferToDiskIfNotMemOnlyMode; we trade a bounded + // debounce window for one storage write per burst — the ack above + // already preceded the flush, so ordering is unchanged). 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.inbound_commit_batch.is_active() { self.maybe_flush_inbound_commits().await; } else { - self.flush_signal_cache_logged("message", Some(&info.id)) - .await; + self.schedule_signal_flush().await; } } diff --git a/src/send/mod.rs b/src/send/mod.rs index bd3e89601..912c1da06 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -1552,9 +1552,11 @@ impl Client { // Warm marking is visible; a waiting cold send may now re-resolve. drop(distribution_guard); - // Flush cached Signal state to DB after encryption - self.flush_signal_cache_batch_safe_logged("send_message_impl", None) - .await; + // Schedule the coalesced Signal flush for the encryption's ratchet + // advance (one storage write per debounce window instead of one per + // send; recovery paths that need read-after-write keep their own + // synchronous flushes). + self.schedule_signal_flush().await; // Issue new tc token after send if a bucket boundary was crossed. // Fire-and-forget so send_message returns without waiting for the IQ diff --git a/src/signal_flush.rs b/src/signal_flush.rs new file mode 100644 index 000000000..f450e76ec --- /dev/null +++ b/src/signal_flush.rs @@ -0,0 +1,156 @@ +//! Coalesced write-behind for the hot-path Signal cache flushes. +//! +//! The live receive path and the send epilogue used to flush the whole dirty +//! Signal cache to storage once per stanza — a serialize + SQLite transaction +//! per message, dominated by the session record the ratchet re-dirties every +//! time. Scheduling through here collapses those into one flush per debounce +//! window: under load, one storage write covers a burst of messages. +//! +//! Durability model (deliberate, bounded): +//! - Live acks already went out BEFORE the per-stanza flush, so coalescing +//! does not reorder acks vs durability — it widens the existing +//! crash-replay window from "one stanza" to at most [`SIGNAL_FLUSH_DEBOUNCE`] +//! plus one flush. Inbound receive chains re-derive forward after a lost +//! advance; consumed one-time prekeys stay buffered until their session is +//! durable (the flush-internal atomicity is untouched). +//! - The offline drain, retry recovery, identity-change recovery and +//! teardown keep their synchronous flushes: those paths gate acks, +//! receipts or follow-up reads on durability and are not routed here. +//! - Disconnect teardown settles the whole cache itself; a fire that lands +//! afterwards flushes an empty cache (no-op). + +use std::sync::atomic::Ordering; + +use crate::client::Client; + +/// Trailing-edge debounce for the coalesced flush. Small enough that the +/// widened crash window stays negligible next to network RTTs; large enough +/// to fold a full receive+reply cycle (and bursts) into one storage write. +const SIGNAL_FLUSH_DEBOUNCE: std::time::Duration = std::time::Duration::from_millis(25); + +impl Client { + /// Request a Signal-cache flush without paying one storage transaction + /// per stanza: the first request arms a trailing-edge timer and every + /// request inside the window rides the same fire. + /// + /// The pending flag clears BEFORE the fire's flush runs, so a request + /// that lands mid-flush arms a new fire instead of being absorbed by a + /// flush that may already have snapshotted the dirty set. + pub(crate) async fn schedule_signal_flush(&self) { + if self.signal_flush_pending.swap(true, Ordering::AcqRel) { + return; + } + let Some(weak) = self.self_weak.get() else { + // Constructor edge: no Arc identity to hold from the timer task. + // Flush inline so the request is never silently dropped. + self.signal_flush_pending.store(false, Ordering::Release); + self.flush_signal_cache_batch_safe_logged("coalesced-inline", None) + .await; + return; + }; + let client = weak.clone(); + let runtime = self.runtime.clone(); + self.runtime + .spawn(Box::pin(async move { + runtime.sleep(SIGNAL_FLUSH_DEBOUNCE).await; + let Some(client) = client.upgrade() else { + return; + }; + client.signal_flush_pending.store(false, Ordering::Release); + // Batch-safe: if an offline drain became active meanwhile, + // this routes under the processing permit like any + // out-of-band flush. + client + .flush_signal_cache_batch_safe_logged("coalesced", None) + .await; + })) + .detach(); + } + + #[cfg(test)] + pub(crate) fn signal_flush_is_pending(&self) -> bool { + self.signal_flush_pending.load(Ordering::Acquire) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::time::Duration; + + use wacore::libsignal::protocol::{ProtocolAddress, SessionRecord}; + + async fn backend_session( + client: &Arc, + addr: &ProtocolAddress, + ) -> Option { + client + .persistence_manager + .backend() + .get_session(addr.as_str()) + .await + .expect("backend read") + } + + fn dirty_session(client: &Arc, user: &str) -> ProtocolAddress { + let addr = ProtocolAddress::new(user.to_string(), 1.into()); + assert!( + client + .signal_cache + .try_put_session(&addr, SessionRecord::new_fresh()) + .is_ok() + ); + addr + } + + async fn wait_for_backend_session(client: &Arc, addr: &ProtocolAddress) { + let deadline = wacore::time::Instant::now() + Duration::from_secs(2); + while backend_session(client, addr).await.is_none() { + assert!( + wacore::time::Instant::now() < deadline, + "scheduled flush never persisted {addr}" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + } + + /// Requests inside one debounce window coalesce into a single armed fire, + /// and that fire persists every dirty entry written before it. + #[tokio::test] + async fn burst_of_requests_coalesces_and_persists() { + let client = crate::test_utils::create_test_client().await; + + let mut addrs = Vec::new(); + for i in 0..10 { + addrs.push(dirty_session(&client, &format!("155500011{i:02}"))); + client.schedule_signal_flush().await; + } + assert!( + client.signal_flush_is_pending(), + "burst must ride one armed fire" + ); + + for addr in &addrs { + wait_for_backend_session(&client, addr).await; + } + assert!( + !client.signal_flush_is_pending(), + "the fire must clear the pending flag" + ); + } + + /// A request after a completed fire arms a NEW fire — the flag round-trips + /// and later dirty state is not stranded behind an absorbed request. + #[tokio::test] + async fn reschedule_after_fire_flushes_again() { + let client = crate::test_utils::create_test_client().await; + + let first = dirty_session(&client, "15550002001"); + client.schedule_signal_flush().await; + wait_for_backend_session(&client, &first).await; + + let second = dirty_session(&client, "15550002002"); + client.schedule_signal_flush().await; + wait_for_backend_session(&client, &second).await; + } +} From f2d0212cd1153038ba69f08acbae5d0df8406d18 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Fri, 10 Jul 2026 23:56:00 -0300 Subject: [PATCH 02/24] test(e2e): settle the signal cache before the legacy-DB surgery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The on-the-fly migration test moved the LID session to PN in the backend right after a verification message, while that message's ratchet advance was still in the write-behind signal cache — the next reconnect's teardown flush then resurrected the LID session the test had just deleted, and the migration under test never ran. Settle with a reconnect before the surgery, like the durability test already does: the scenario this models (a legacy PN-only DB) arises with a quiescent client anyway. --- tests/e2e/tests/lid_sessions.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/e2e/tests/lid_sessions.rs b/tests/e2e/tests/lid_sessions.rs index 57023a742..f0e34b8b8 100644 --- a/tests/e2e/tests/lid_sessions.rs +++ b/tests/e2e/tests/lid_sessions.rs @@ -517,6 +517,13 @@ async fn test_pn_only_session_causes_undecryptable_on_lid_lookup() -> anyhow::Re .await?; info!("Verified messaging works after first reconnect"); + // Settle before the backend surgery below: the step-3 message left a + // ratchet advance in the (write-behind) signal cache, and a reconnect + // teardown would flush it AFTER the surgery — resurrecting the LID + // session this test deletes. Same pattern as the durability test. + client_a.reconnect_and_wait().await?; + info!("Settled signal cache before backend surgery"); + // Step 4: Simulate legacy DB — move session from LID to PN address. // Target B's connected device: under LID addressing a 1:1 peer sends from its // companion, never device 0, so the inbound session lives there (not at :0). From 430b2218b16670ceec6517f7b878047633cdd51a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Fri, 10 Jul 2026 23:57:01 -0300 Subject: [PATCH 03/24] fix(signal): re-arm the coalesced flush on error and fix window terminology MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A fire whose flush failed cleared the pending flag and exited, leaving dirty cache state unwritten until unrelated traffic scheduled again — the documented bounded window became unbounded under a transient storage failure. Re-arm on error so the window doubles as the retry backoff. Also rename the docs: this is a fixed coalescing window, not a trailing-edge debounce (deliberately — extending the deadline per request would defer the flush indefinitely under continuous traffic). --- src/signal_flush.rs | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/src/signal_flush.rs b/src/signal_flush.rs index f450e76ec..6749fb694 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -23,19 +23,25 @@ use std::sync::atomic::Ordering; use crate::client::Client; -/// Trailing-edge debounce for the coalesced flush. Small enough that the -/// widened crash window stays negligible next to network RTTs; large enough -/// to fold a full receive+reply cycle (and bursts) into one storage write. -const SIGNAL_FLUSH_DEBOUNCE: std::time::Duration = std::time::Duration::from_millis(25); +/// Fixed coalescing window for the flush. The fire runs this long after the +/// FIRST request; later requests inside the window ride the same fire (the +/// deadline is deliberately not extended — a true trailing-edge debounce +/// would defer the flush indefinitely under continuous traffic, while the +/// fixed window bounds the maximum deferral). Small enough that the widened +/// crash window stays negligible next to network RTTs; large enough to fold +/// a full receive+reply cycle (and bursts) into one storage write. +const SIGNAL_FLUSH_WINDOW: std::time::Duration = std::time::Duration::from_millis(25); impl Client { /// Request a Signal-cache flush without paying one storage transaction - /// per stanza: the first request arms a trailing-edge timer and every + /// per stanza: the first request arms a fixed-window timer and every /// request inside the window rides the same fire. /// /// The pending flag clears BEFORE the fire's flush runs, so a request /// that lands mid-flush arms a new fire instead of being absorbed by a - /// flush that may already have snapshotted the dirty set. + /// flush that may already have snapshotted the dirty set. A failed flush + /// re-arms the window, so dirty state is retried instead of sitting + /// unwritten until unrelated traffic schedules again. pub(crate) async fn schedule_signal_flush(&self) { if self.signal_flush_pending.swap(true, Ordering::AcqRel) { return; @@ -52,7 +58,7 @@ impl Client { let runtime = self.runtime.clone(); self.runtime .spawn(Box::pin(async move { - runtime.sleep(SIGNAL_FLUSH_DEBOUNCE).await; + runtime.sleep(SIGNAL_FLUSH_WINDOW).await; let Some(client) = client.upgrade() else { return; }; @@ -60,9 +66,12 @@ impl Client { // Batch-safe: if an offline drain became active meanwhile, // this routes under the processing permit like any // out-of-band flush. - client - .flush_signal_cache_batch_safe_logged("coalesced", None) - .await; + if let Err(e) = client.flush_signal_cache_batch_safe().await { + log::error!("Coalesced signal flush failed; re-arming for retry: {e:?}"); + // Re-arm with the same window as the retry backoff; the + // cache keeps its dirty entries until a flush succeeds. + client.schedule_signal_flush().await; + } })) .detach(); } From 1693785a0df4329a05697b9f55c3d3f4389d4e0e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Fri, 10 Jul 2026 23:58:34 -0300 Subject: [PATCH 04/24] fix(signal): retry the coalesced flush inline instead of recursing The re-arm-on-error path called schedule_signal_flush from inside its own fire, making the async fn recursive (not Send). Loop inside the fire task instead, holding only the Weak across each window; a concurrent request that re-armed meanwhile owns the retry. --- src/signal_flush.rs | 35 ++++++++++++++++++++++------------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/src/signal_flush.rs b/src/signal_flush.rs index 6749fb694..f8d36543a 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -54,23 +54,32 @@ impl Client { .await; return; }; - let client = weak.clone(); + let weak = weak.clone(); let runtime = self.runtime.clone(); self.runtime .spawn(Box::pin(async move { - runtime.sleep(SIGNAL_FLUSH_WINDOW).await; - let Some(client) = client.upgrade() else { - return; - }; - client.signal_flush_pending.store(false, Ordering::Release); - // Batch-safe: if an offline drain became active meanwhile, - // this routes under the processing permit like any - // out-of-band flush. - if let Err(e) = client.flush_signal_cache_batch_safe().await { + loop { + // Hold only the Weak across the sleep so an armed fire + // never extends the client's lifetime. + runtime.sleep(SIGNAL_FLUSH_WINDOW).await; + let Some(client) = weak.upgrade() else { + return; + }; + client.signal_flush_pending.store(false, Ordering::Release); + // Batch-safe: if an offline drain became active meanwhile, + // this routes under the processing permit like any + // out-of-band flush. + let Err(e) = client.flush_signal_cache_batch_safe().await else { + return; + }; log::error!("Coalesced signal flush failed; re-arming for retry: {e:?}"); - // Re-arm with the same window as the retry backoff; the - // cache keeps its dirty entries until a flush succeeds. - client.schedule_signal_flush().await; + // Re-arm inline with the same window as the retry backoff; + // the cache keeps its dirty entries until a flush + // succeeds. If a concurrent request re-armed already, its + // fire owns the retry. + if client.signal_flush_pending.swap(true, Ordering::AcqRel) { + return; + } } })) .detach(); From a75cfc4810bb73bcc2800f6acb7a678c267f16c3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sat, 11 Jul 2026 10:03:28 -0300 Subject: [PATCH 05/24] fix(signal): back off the failing-flush retry exponentially A persistent storage failure turned the re-arm loop into ~40 attempts and error logs per second. Double the backoff per consecutive failure up to a 5s ceiling, which rate-limits the log with it; a success (or a concurrent re-arm) still exits the loop. --- src/signal_flush.rs | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/src/signal_flush.rs b/src/signal_flush.rs index f8d36543a..43ae556f0 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -32,6 +32,12 @@ use crate::client::Client; /// a full receive+reply cycle (and bursts) into one storage write. const SIGNAL_FLUSH_WINDOW: std::time::Duration = std::time::Duration::from_millis(25); +/// Retry backoff ceiling for a failing flush. The backoff starts at the +/// window and doubles per consecutive failure, so a long-lived storage +/// outage settles at one attempt (and one error log) per ceiling instead of +/// ~40/s at the raw window. +const SIGNAL_FLUSH_RETRY_CEILING: std::time::Duration = std::time::Duration::from_secs(5); + impl Client { /// Request a Signal-cache flush without paying one storage transaction /// per stanza: the first request arms a fixed-window timer and every @@ -58,10 +64,11 @@ impl Client { let runtime = self.runtime.clone(); self.runtime .spawn(Box::pin(async move { + let mut backoff = SIGNAL_FLUSH_WINDOW; loop { // Hold only the Weak across the sleep so an armed fire // never extends the client's lifetime. - runtime.sleep(SIGNAL_FLUSH_WINDOW).await; + runtime.sleep(backoff).await; let Some(client) = weak.upgrade() else { return; }; @@ -72,11 +79,13 @@ impl Client { let Err(e) = client.flush_signal_cache_batch_safe().await else { return; }; - log::error!("Coalesced signal flush failed; re-arming for retry: {e:?}"); - // Re-arm inline with the same window as the retry backoff; - // the cache keeps its dirty entries until a flush - // succeeds. If a concurrent request re-armed already, its - // fire owns the retry. + // Exponential backoff doubles per consecutive failure and + // caps the error-log rate along with it; the cache keeps + // its dirty entries until a flush succeeds. + backoff = (backoff * 2).min(SIGNAL_FLUSH_RETRY_CEILING); + log::error!("Coalesced signal flush failed; retrying in {backoff:?}: {e:?}"); + // If a concurrent request re-armed already, its fire owns + // the retry. if client.signal_flush_pending.swap(true, Ordering::AcqRel) { return; } From 4e1a85115a808ae0a64523979edef2292f89fff5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sat, 11 Jul 2026 10:16:16 -0300 Subject: [PATCH 06/24] test(signal): cover the failing-flush retry path Inject flush failures via a cfg(test) counter (same pattern as the commit batcher's fail_flushes): with two injected failures the fire must consume both error attempts, re-arm through the backoff, and still persist the pre-existing dirty entry on the third attempt. --- src/client.rs | 6 ++++- src/client/lifecycle.rs | 2 ++ src/signal_flush.rs | 50 +++++++++++++++++++++++++++++++++++++++-- 3 files changed, 55 insertions(+), 3 deletions(-) diff --git a/src/client.rs b/src/client.rs index b298d20b2..c153db355 100644 --- a/src/client.rs +++ b/src/client.rs @@ -836,9 +836,13 @@ pub struct Client { /// Initialized after `Arc::new(this)` in the constructor. pub(crate) self_weak: std::sync::OnceLock>, - /// Trailing-edge debounce flag for the coalesced Signal-cache flush + /// Coalescing-window flag for the Signal-cache flush /// (see `signal_flush.rs`). True while a fire is armed. pub(crate) signal_flush_pending: AtomicBool, + /// Injected failures for the coalesced flush (consumed one per attempt), + /// so tests can exercise the retry/backoff path deterministically. + #[cfg(test)] + pub(crate) signal_flush_test_failures: AtomicU32, /// Holds the background saver's AbortHandle so the task lifetime follows /// `Arc` ref count instead of the Bot wrapper's. Set once by diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index f0ece477a..ea7031854 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -241,6 +241,8 @@ impl Client { passkey_state: Arc::new(Mutex::new(crate::passkey::flow::PasskeyFlowState::default())), passkey_opening: AtomicBool::new(false), signal_flush_pending: AtomicBool::new(false), + #[cfg(test)] + signal_flush_test_failures: AtomicU32::new(0), custom_enc_handlers: std::sync::OnceLock::new(), inbound_durability_hook: std::sync::OnceLock::new(), retry_admission: std::sync::OnceLock::new(), diff --git a/src/signal_flush.rs b/src/signal_flush.rs index 43ae556f0..89ef53df4 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -9,7 +9,7 @@ //! Durability model (deliberate, bounded): //! - Live acks already went out BEFORE the per-stanza flush, so coalescing //! does not reorder acks vs durability — it widens the existing -//! crash-replay window from "one stanza" to at most [`SIGNAL_FLUSH_DEBOUNCE`] +//! crash-replay window from "one stanza" to at most [`SIGNAL_FLUSH_WINDOW`] //! plus one flush. Inbound receive chains re-derive forward after a lost //! advance; consumed one-time prekeys stay buffered until their session is //! durable (the flush-internal atomicity is untouched). @@ -76,7 +76,7 @@ impl Client { // Batch-safe: if an offline drain became active meanwhile, // this routes under the processing permit like any // out-of-band flush. - let Err(e) = client.flush_signal_cache_batch_safe().await else { + let Err(e) = client.coalesced_flush_attempt().await else { return; }; // Exponential backoff doubles per consecutive failure and @@ -94,6 +94,22 @@ impl Client { .detach(); } + /// One flush attempt of the fire loop; tests can inject failures to + /// exercise the retry/backoff path (same pattern as the commit batcher's + /// `fail_flushes`). + async fn coalesced_flush_attempt(&self) -> Result<(), anyhow::Error> { + #[cfg(test)] + { + let remaining = self.signal_flush_test_failures.load(Ordering::Acquire); + if remaining > 0 { + self.signal_flush_test_failures + .store(remaining - 1, Ordering::Release); + anyhow::bail!("injected coalesced-flush failure"); + } + } + self.flush_signal_cache_batch_safe().await + } + #[cfg(test)] pub(crate) fn signal_flush_is_pending(&self) -> bool { self.signal_flush_pending.load(Ordering::Acquire) @@ -166,6 +182,36 @@ mod tests { ); } + /// Failed attempts re-arm and back off instead of dropping the dirty + /// state: with 2 injected failures, the fire must consume both error + /// attempts and still persist the pre-existing dirty entry on the third. + #[tokio::test] + async fn failed_fire_retries_until_the_dirty_entry_persists() { + use std::sync::atomic::Ordering; + + let client = crate::test_utils::create_test_client().await; + let addr = dirty_session(&client, "15550004001"); + client + .signal_flush_test_failures + .store(2, Ordering::Release); + + client.schedule_signal_flush().await; + + // Success only comes after both injected failures are consumed by the + // retry loop (25 + 50 ms of backoff), proving the error path re-armed + // rather than stranding the dirty entry. + wait_for_backend_session(&client, &addr).await; + assert_eq!( + client.signal_flush_test_failures.load(Ordering::Acquire), + 0, + "the retry loop must have consumed every injected failure" + ); + assert!( + !client.signal_flush_is_pending(), + "the successful retry must clear the pending flag" + ); + } + /// A request after a completed fire arms a NEW fire — the flag round-trips /// and later dirty state is not stranded behind an absorbed request. #[tokio::test] From 47b33d5c8db6aa0a46b195e373b777ef61400730 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sat, 11 Jul 2026 10:34:59 -0300 Subject: [PATCH 07/24] test(e2e): gate connect on the canonical is_ready signal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The shared connect helper waited for Event::Connected in a fixed 30s event-loop, then fell back to offline_sync_completed — an orthogonal signal. Under concurrent CI load the mock server was slow to serve the critical app-state IQs, so Connected raced past 30s and the fallback timed out on a signal that says nothing about readiness, failing 9 unrelated app-state tests at once. Gate on wait_for_connected instead: it resolves on is_ready (set at the same point Event::Connected is dispatched, after the critical sync) via a notifier with a TOCTOU re-check, so it does not depend on event arrival order, on who drains the channel, or on offline sync. The startup-sync drain is now best-effort — readiness is already guaranteed, and a message backlog must not fail the connect. The determinism comes from waiting on the right signal, not from a larger timeout. --- tests/e2e/src/lib.rs | 118 ++++++++++--------------------------------- 1 file changed, 28 insertions(+), 90 deletions(-) diff --git a/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index 53c79395f..af62b9785 100644 --- a/tests/e2e/src/lib.rs +++ b/tests/e2e/src/lib.rs @@ -159,102 +159,40 @@ impl TestClient { let run_handle = bot.spawn(); - // Wait for PairSuccess + Connected. + // Wait for the client to become fully ready: connected, logged in, and + // critical app-state sync complete. // - // PairSuccess arrives quickly (handshake only), but Connected is dispatched - // only after the critical app-state sync completes (sync_collections_batched). - // Under CI load with many concurrent clients, the mock server may be slow to - // serve app-state IQs, so Connected can take significantly longer than pairing. - // - // We use a two-phase timeout: 30s for pairing, then an additional 30s for - // Connected (which includes critical sync). This avoids a single shared timeout - // where a slow sync eats into the pairing budget. - let timeout = tokio::time::Duration::from_secs(30); - let mut got_pair = false; - let mut got_connected = false; - - let wait_result = tokio::time::timeout(timeout, async { - loop { - match event_rx.recv().await { - Ok(ref event) if matches!(**event, Event::PairSuccess(_)) => { - got_pair = true; - if got_connected { - break; - } - } - Ok(ref event) if matches!(**event, Event::Connected(_)) => { - got_connected = true; - if got_pair { - break; - } - } - Ok(_) => {} - Err(e) => { - return Err(anyhow::anyhow!("Event channel closed during connect: {e}")); - } - } - } - Ok(()) - }) - .await; - - match wait_result { - Err(_) => { - // If we got PairSuccess but not Connected, the critical sync is slow. - // Give it extra time via wait_for_startup_sync instead of failing immediately. - if got_pair && !got_connected { - eprintln!( - "WARN: Got PairSuccess but Connected timed out after {timeout:?}, \ - waiting for startup sync..." - ); - if let Err(e) = client - .wait_for_startup_sync(tokio::time::Duration::from_secs(30)) - .await - { - client.disconnect().await; - drop(run_handle); - return Err(anyhow::anyhow!( - "Timed out waiting for Connected after PairSuccess: {e}" - )); - } - // Drain the Connected event that should now be available - let connected_timeout = tokio::time::Duration::from_secs(5); - let _ = tokio::time::timeout(connected_timeout, async { - loop { - match event_rx.recv().await { - Ok(ref event) if matches!(**event, Event::Connected(_)) => break, - Ok(_) => continue, - Err(_) => break, - } - } - }) - .await; - } else { - client.disconnect().await; - drop(run_handle); - return Err(anyhow::anyhow!( - "Timed out waiting for PairSuccess + Connected \ - (got_pair={got_pair}, got_connected={got_connected})" - )); - } - } - Ok(Err(e)) => { - client.disconnect().await; - drop(run_handle); - return Err(e); - } - Ok(Ok(())) => {} - } - + // This gates on `wait_for_connected`, which resolves on the canonical + // `is_ready` signal (`dispatch_connected`, fired after the critical sync) + // via a notifier — not on the order events happen to arrive in the + // channel. That makes it immune to the earlier flake, where a fixed 30s + // event-loop wait for `Connected` raced how long the mock server took to + // serve app-state IQs under concurrent CI load, then fell back to + // `offline_sync_completed` — an orthogonal signal that could time out + // even though the client was otherwise ready. The single budget below is + // generous because it is a hard ceiling on a real hang, not a per-phase + // race window. PairSuccess/Connected still land in the unbounded + // `event_rx`; the predicate-filtered `wait_for_event` used by tests + // discards them. if let Err(e) = client - .wait_for_startup_sync(tokio::time::Duration::from_secs(15)) + .wait_for_connected(tokio::time::Duration::from_secs(60)) .await { client.disconnect().await; drop(run_handle); - return Err(anyhow::anyhow!( - "Timed out waiting for startup sync to become idle: {e}" - )); + return Err(anyhow::anyhow!("client never became ready after pairing: {e}")); + } + + // Drain the initial startup sync (offline messages + history) so tests + // start from a settled state. Best-effort: readiness is already + // guaranteed by `wait_for_connected` above, and a chat-action test does + // not depend on message backlog, so a slow mock server here must not + // fail the connect — it only means the (empty) backlog is still landing. + if let Err(e) = client + .wait_for_startup_sync(tokio::time::Duration::from_secs(30)) + .await + { + eprintln!("WARN: startup sync did not settle before the connect deadline: {e}"); } Ok(Self { From b46e161571fa0a776295c8f20a27e0b729dcc8fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sat, 11 Jul 2026 10:45:13 -0300 Subject: [PATCH 08/24] style(e2e): rustfmt the connect helper --- tests/e2e/src/lib.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index af62b9785..0cff24514 100644 --- a/tests/e2e/src/lib.rs +++ b/tests/e2e/src/lib.rs @@ -180,7 +180,9 @@ impl TestClient { { client.disconnect().await; drop(run_handle); - return Err(anyhow::anyhow!("client never became ready after pairing: {e}")); + return Err(anyhow::anyhow!( + "client never became ready after pairing: {e}" + )); } // Drain the initial startup sync (offline messages + history) so tests From 4e3e736585c2a9b0b95248c2efde1321d5b513a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Sat, 11 Jul 2026 10:57:12 -0300 Subject: [PATCH 09/24] test(e2e): settle the coalesced flush before inspecting durable sessions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A live send/receive now schedules a coalesced Signal flush rather than writing through synchronously, so tests that read session state from the backend right after messaging raced the debounce window — deterministic locally (the timer fires during network latency) but flaky under CI load where the timer slips. test_session_state_after_roundtrip hit this; several session_reuse/lid_sessions assertions had the same latent race. Add Client::flush_pending_signal_state (a public 'force durability now' that also fits pre-non-graceful-shutdown callers) and call it before every post-send backend read. Bootstrap-created sessions (primary device 0) and reads that already follow a reconnect are untouched — the reconnect teardown flushes the cache. --- src/client/adapters.rs | 13 +++++++++++++ tests/e2e/tests/lid_sessions.rs | 15 +++++++++++++-- tests/e2e/tests/session_reuse.rs | 18 +++++++++--------- 3 files changed, 35 insertions(+), 11 deletions(-) diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 0e01f0c22..ab0bac589 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -68,6 +68,19 @@ impl Client { .ok_or(ClientError::NotConnected) } + /// Force any pending write-behind Signal cache state to the backend and + /// wait for it to complete. + /// + /// Live sends and receives schedule a coalesced flush (see + /// `signal_flush.rs`) rather than writing synchronously, so durable + /// storage lags the in-memory cache by up to one debounce window. Callers + /// that need read-after-write durability on the backend — e.g. before + /// inspecting persisted session state, or ahead of a non-graceful + /// shutdown — use this to settle it deterministically. + pub async fn flush_pending_signal_state(&self) -> Result<(), anyhow::Error> { + self.flush_signal_cache_batch_safe().await + } + /// Flush the in-memory signal cache to the database backend. /// Called after each message is decrypted or after encryption operations. pub(crate) async fn flush_signal_cache(&self) -> Result<(), anyhow::Error> { diff --git a/tests/e2e/tests/lid_sessions.rs b/tests/e2e/tests/lid_sessions.rs index f0e34b8b8..59459dd68 100644 --- a/tests/e2e/tests/lid_sessions.rs +++ b/tests/e2e/tests/lid_sessions.rs @@ -236,6 +236,9 @@ async fn test_stale_pn_session_does_not_break_lid_messaging() -> anyhow::Result< ) .await?; + // Settle the coalesced flush before reading durable session state: the + // reply above created A's inbound session in the write-behind cache. + client_a.client.flush_pending_signal_state().await?; let backend_a = client_a.client.persistence_manager().backend(); // Read the existing LID session data at B's connected device (companion, not 0). @@ -279,6 +282,8 @@ async fn test_stale_pn_session_does_not_break_lid_messaging() -> anyhow::Result< .await?; info!("Messaging works despite stale PN session in DB"); + // Settle the coalesced flush from the post-inject sends before reading. + client_a.client.flush_pending_signal_state().await?; // LID session should still be authoritative assert!( backend_a.get_session(&lid_addr).await?.is_some(), @@ -339,7 +344,8 @@ async fn test_lid_session_survives_reconnect() -> anyhow::Result<()> { .await?; info!("Sessions established"); - // Verify LID-only before reconnect + // Verify LID-only before reconnect (settle the coalesced flush first). + client_a.client.flush_pending_signal_state().await?; let backend_a = client_a.client.persistence_manager().backend(); assert_lid_only_sessions(&*backend_a, &jid_b.user, &lid_b.user, "Before reconnect").await; @@ -370,7 +376,8 @@ async fn test_lid_session_survives_reconnect() -> anyhow::Result<()> { .await?; info!("Post-reconnect delivery confirmed (2 messages)"); - // Final check: still LID-only after post-reconnect sends + // Final check: still LID-only after post-reconnect sends (settle first). + client_a.client.flush_pending_signal_state().await?; assert_lid_only_sessions( &*backend_a, &jid_b.user, @@ -622,6 +629,8 @@ async fn test_inbound_1x1_session_keyed_at_companion_device() -> anyhow::Result< lid_b.device, 0, "a companion 1:1 peer must be addressed at a non-zero device" ); + // Settle the coalesced flush before reading durable session state. + client_a.client.flush_pending_signal_state().await?; let backend_a = client_a.client.persistence_manager().backend(); let addr = peer_session_addr(&lid_b.user, "lid", lid_b.device); assert!( @@ -675,6 +684,8 @@ async fn test_pn_migration_is_durable_across_followup_messages() -> anyhow::Resu .await?; } + // Settle the coalesced flush from the follow-up sends before reading. + client_a.client.flush_pending_signal_state().await?; assert!( backend_a.get_session(&lid_addr).await?.is_some(), "session stays under LID across follow-up messaging" diff --git a/tests/e2e/tests/session_reuse.rs b/tests/e2e/tests/session_reuse.rs index f1c2ac085..34a5c2693 100644 --- a/tests/e2e/tests/session_reuse.rs +++ b/tests/e2e/tests/session_reuse.rs @@ -140,15 +140,11 @@ async fn test_session_state_after_roundtrip() -> anyhow::Result<()> { send_and_expect_text(&client_b.client, &mut client_a, &jid_a, "Session reply", 30).await?; info!("Roundtrip complete"); - // Force cache flush by sending another message - send_and_expect_text( - &client_a.client, - &mut client_b, - &jid_b, - "Post-roundtrip flush", - 30, - ) - .await?; + // Settle A's write-behind Signal cache to the backend before inspecting + // it: a live send schedules a coalesced flush rather than writing + // synchronously, so a plain send-and-read races the debounce window (and + // its timer can slip arbitrarily under CI load). + client_a.client.flush_pending_signal_state().await?; // Inspect session state let backend = client_a.client.persistence_manager().backend(); @@ -236,6 +232,10 @@ async fn test_session_persistence() -> anyhow::Result<()> { .await?; info!("First message sent: {msg_id_1}"); + // Settle the coalesced send flush before reading durable state (see the + // roundtrip test): a live send does not write through synchronously. + client_a.client.flush_pending_signal_state().await?; + // Session may be under PN (c.us) or LID (lid) depending on whether // PN→LID mapping was resolved before encryption. let mut post_send = scan_sessions(&*backend, &jid_b.user, "c.us").await?; From 7be4107a0d13930a5669235bc7b6aa3129f882b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 15:36:37 -0300 Subject: [PATCH 10/24] docs(signal): document the settle API's permit and durability preconditions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit flush_pending_signal_state delegates to the batch-safe flush, which acquires the processing permit during an offline drain — so a caller inside an event handler or InboundDurabilityHook (which run under that permit) would self-deadlock. Document that, and qualify the lag bound: it holds on success; a backend outage extends it until the retry loop succeeds, so callers must check the Result. Also correct the receive comment to state WA Web flushes before the receipt (we ack before the flush, pre-dating this change) rather than implying parity. --- src/client/adapters.rs | 22 +++++++++++++++------- src/message/receive.rs | 17 +++++++++-------- 2 files changed, 24 insertions(+), 15 deletions(-) diff --git a/src/client/adapters.rs b/src/client/adapters.rs index ab0bac589..d0c95d5f4 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -68,15 +68,23 @@ impl Client { .ok_or(ClientError::NotConnected) } - /// Force any pending write-behind Signal cache state to the backend and - /// wait for it to complete. + /// Force any pending write-behind Signal cache state to the backend, + /// returning once the flush completes (or fails). /// /// Live sends and receives schedule a coalesced flush (see - /// `signal_flush.rs`) rather than writing synchronously, so durable - /// storage lags the in-memory cache by up to one debounce window. Callers - /// that need read-after-write durability on the backend — e.g. before - /// inspecting persisted session state, or ahead of a non-graceful - /// shutdown — use this to settle it deterministically. + /// `signal_flush.rs`) instead of writing through, so on success the backend + /// trails the in-memory cache by at most one coalescing window; a backend + /// outage extends that until the scheduler's retry loop succeeds. Use this + /// to settle durability deterministically before reading persisted state or + /// ahead of a non-graceful shutdown — and check the returned `Result`, as a + /// failure leaves state pending. + /// + /// Call from a control task, never from inside an event handler or an + /// [`InboundDurabilityHook`]: during an offline-sync drain those run while + /// the processing permit is held, and settling routes through that same + /// permit — re-entering it would deadlock. + /// + /// [`InboundDurabilityHook`]: crate::types::durability_hook::InboundDurabilityHook pub async fn flush_pending_signal_state(&self) -> Result<(), anyhow::Error> { self.flush_signal_cache_batch_safe().await } diff --git a/src/message/receive.rs b/src/message/receive.rs index 47268d4f4..1485b2f45 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -561,14 +561,15 @@ impl Client { self.handle_msmsg_payload(&info, payload).await; } - // Live: schedule the coalesced Signal flush (WA Web flushes per - // stanza via flushBufferToDiskIfNotMemOnlyMode; we trade a bounded - // debounce window for one storage write per burst — the ack above - // already preceded the flush, so ordering is unchanged). 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. + // Live: schedule the coalesced Signal flush. WA Web flushes per stanza + // (flushBufferToDiskIfNotMemOnlyMode, inside decrypt BEFORE the caller + // sends the receipt); our live path already acked BEFORE the per-stanza + // flush pre-dating this change, so we don't reorder ack-vs-flush — we + // only widen the crash-replay gap it already had to one coalescing + // window. 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.inbound_commit_batch.is_active() { self.maybe_flush_inbound_commits().await; } else { From 4d1b8fa91ef8b1e56c8b92ee00b68ca6434b0361 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:17:32 -0300 Subject: [PATCH 11/24] perf(signal): coalesce only the receive flush; keep sends synchronous MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reworks the coalescing to remove the outbound-durability hazard a reviewer flagged. The send path used to schedule its flush, so send_message could return success while the outbound ratchet advance was only in memory — a crash before the flush would reuse the counter, and with it the message key + IV, on the next send (a cryptographic regression, not just a duplicate). It now flushes synchronously before returning, matching WA Web (flush before the stanza hits the wire). Only the receive path coalesces, where a lost advance re-derives forward. The scheduler is now genuinely single-flight: a RUNNING/DIRTY atomic state means only the idle->running transition spawns a worker; requests mid-flush just mark DIRTY and the one worker runs another window. A failing flush is retried by that same worker with exponential backoff, so concurrent traffic can't reset it to the base delay, and it can't pile up detached tasks each holding an Arc. The worker is bound to connection_generation and the scheduler is reset at teardown, so a stale retry backoff never delays the next connection. --- src/client.rs | 6 +- src/client/lifecycle.rs | 7 +- src/message/receive.rs | 15 +- src/send/mod.rs | 13 +- src/signal_flush.rs | 314 +++++++++++++++++++++++++++------------- 5 files changed, 235 insertions(+), 120 deletions(-) diff --git a/src/client.rs b/src/client.rs index c153db355..a147af975 100644 --- a/src/client.rs +++ b/src/client.rs @@ -836,9 +836,9 @@ pub struct Client { /// Initialized after `Arc::new(this)` in the constructor. pub(crate) self_weak: std::sync::OnceLock>, - /// Coalescing-window flag for the Signal-cache flush - /// (see `signal_flush.rs`). True while a fire is armed. - pub(crate) signal_flush_pending: AtomicBool, + /// Single-flight state for the coalesced Signal-cache flush worker + /// (RUNNING/DIRTY bits; see `signal_flush.rs`). + pub(crate) signal_flush_state: AtomicU32, /// Injected failures for the coalesced flush (consumed one per attempt), /// so tests can exercise the retry/backoff path deterministically. #[cfg(test)] diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index ea7031854..48363b8f0 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -240,7 +240,7 @@ impl Client { pair_code_state: Arc::new(Mutex::new(wacore::pair_code::PairCodeState::default())), passkey_state: Arc::new(Mutex::new(crate::passkey::flow::PasskeyFlowState::default())), passkey_opening: AtomicBool::new(false), - signal_flush_pending: AtomicBool::new(false), + signal_flush_state: AtomicU32::new(0), #[cfg(test)] signal_flush_test_failures: AtomicU32::new(0), custom_enc_handlers: std::sync::OnceLock::new(), @@ -767,6 +767,11 @@ impl Client { // 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); + // Stand down the coalesced-flush worker: its generation guard exits it + // on the next wake, and clearing the arm lets the next connection's + // traffic spawn a fresh worker at the base window instead of waiting + // out a stale retry backoff. + self.reset_signal_flush_scheduler(); // 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/receive.rs b/src/message/receive.rs index 1485b2f45..94a690870 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -561,19 +561,14 @@ impl Client { self.handle_msmsg_payload(&info, payload).await; } - // Live: schedule the coalesced Signal flush. WA Web flushes per stanza - // (flushBufferToDiskIfNotMemOnlyMode, inside decrypt BEFORE the caller - // sends the receipt); our live path already acked BEFORE the per-stanza - // flush pre-dating this change, so we don't reorder ack-vs-flush — we - // only widen the crash-replay gap it already had to one coalescing - // window. 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. + // Live: coalesce the receive-side flush. A lost advance re-derives + // forward, so unlike the send path this tolerates the window (see + // `signal_flush.rs`). During the offline drain the commit batcher owns + // the flush instead, so only the batch size/byte triggers run here. if self.inbound_commit_batch.is_active() { self.maybe_flush_inbound_commits().await; } else { - self.schedule_signal_flush().await; + self.schedule_signal_flush(); } } diff --git a/src/send/mod.rs b/src/send/mod.rs index 912c1da06..b3af858d7 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -1552,11 +1552,14 @@ impl Client { // Warm marking is visible; a waiting cold send may now re-resolve. drop(distribution_guard); - // Schedule the coalesced Signal flush for the encryption's ratchet - // advance (one storage write per debounce window instead of one per - // send; recovery paths that need read-after-write keep their own - // synchronous flushes). - self.schedule_signal_flush().await; + // Flush the outbound ratchet advance synchronously before returning: + // a coalesced flush here would let send_message report success while + // the counter/chain-key advance is still only in memory, so a crash + // before the flush would reuse the counter (and its message key + IV) + // on the next send. Only the receive path — where a lost advance + // re-derives forward — coalesces. + self.flush_signal_cache_batch_safe_logged("send_message_impl", None) + .await; // Issue new tc token after send if a bucket boundary was crossed. // Fire-and-forget so send_message returns without waiting for the IQ diff --git a/src/signal_flush.rs b/src/signal_flush.rs index 89ef53df4..91feef399 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -1,102 +1,149 @@ -//! Coalesced write-behind for the hot-path Signal cache flushes. +//! Coalesced write-behind for the inbound Signal cache flush. //! -//! The live receive path and the send epilogue used to flush the whole dirty -//! Signal cache to storage once per stanza — a serialize + SQLite transaction -//! per message, dominated by the session record the ratchet re-dirties every -//! time. Scheduling through here collapses those into one flush per debounce -//! window: under load, one storage write covers a burst of messages. +//! The live receive path used to flush the whole dirty Signal cache to storage +//! once per stanza — a session re-serialize plus a SQLite transaction per +//! message, dominated by the record the ratchet re-dirties every time. Routing +//! it through here collapses a burst of receives into one flush per coalescing +//! window. //! -//! Durability model (deliberate, bounded): -//! - Live acks already went out BEFORE the per-stanza flush, so coalescing -//! does not reorder acks vs durability — it widens the existing -//! crash-replay window from "one stanza" to at most [`SIGNAL_FLUSH_WINDOW`] -//! plus one flush. Inbound receive chains re-derive forward after a lost -//! advance; consumed one-time prekeys stay buffered until their session is -//! durable (the flush-internal atomicity is untouched). -//! - The offline drain, retry recovery, identity-change recovery and -//! teardown keep their synchronous flushes: those paths gate acks, -//! receipts or follow-up reads on durability and are not routed here. -//! - Disconnect teardown settles the whole cache itself; a fire that lands -//! afterwards flushes an empty cache (no-op). +//! Scope and durability model (deliberate, bounded): +//! - Only the receive path coalesces. A lost receive-side advance re-derives +//! forward (the Double Ratchet receiving chain derives `CK_n → CK_n+1`), and +//! a consumed one-time prekey stays buffered until its session is durable, so +//! a crash inside the window is recoverable. The SEND path flushes +//! synchronously before returning: reusing an outbound counter would reuse +//! its message key + IV, so that advance must be durable before `send_message` +//! reports success. +//! - The offline drain, retry recovery, identity-change recovery and teardown +//! keep their own synchronous flushes: those gate acks, receipts or +//! follow-up reads on durability and are not routed here. +//! +//! Single-flight scheduler: at most one worker exists at a time. The first +//! request arms it; requests that arrive while it runs only mark it dirty +//! (never spawn a second worker), and it re-runs one more window if so. A +//! failing flush is retried by that same worker with exponential backoff, so a +//! backend outage cannot be reset to the base delay by concurrent traffic, and +//! a generation change (reconnect/teardown) stands the worker down. use std::sync::atomic::Ordering; use crate::client::Client; -/// Fixed coalescing window for the flush. The fire runs this long after the -/// FIRST request; later requests inside the window ride the same fire (the -/// deadline is deliberately not extended — a true trailing-edge debounce -/// would defer the flush indefinitely under continuous traffic, while the -/// fixed window bounds the maximum deferral). Small enough that the widened -/// crash window stays negligible next to network RTTs; large enough to fold -/// a full receive+reply cycle (and bursts) into one storage write. +/// Fixed coalescing window: the worker flushes this long after being armed. +/// Small enough that the widened crash-replay gap stays negligible next to +/// network RTTs; large enough to fold a burst of receives into one write. const SIGNAL_FLUSH_WINDOW: std::time::Duration = std::time::Duration::from_millis(25); -/// Retry backoff ceiling for a failing flush. The backoff starts at the -/// window and doubles per consecutive failure, so a long-lived storage -/// outage settles at one attempt (and one error log) per ceiling instead of -/// ~40/s at the raw window. +/// Backoff ceiling for a failing flush. The delay doubles per consecutive +/// failure up to this cap, which bounds the retry and error-log rate during a +/// long-lived backend outage. const SIGNAL_FLUSH_RETRY_CEILING: std::time::Duration = std::time::Duration::from_secs(5); +/// A worker is alive (armed or running/retrying). +const FLUSH_RUNNING: u32 = 0b01; +/// A request arrived while the worker was mid-flush; it must run one more window. +const FLUSH_DIRTY: u32 = 0b10; + impl Client { - /// Request a Signal-cache flush without paying one storage transaction - /// per stanza: the first request arms a fixed-window timer and every - /// request inside the window rides the same fire. - /// - /// The pending flag clears BEFORE the fire's flush runs, so a request - /// that lands mid-flush arms a new fire instead of being absorbed by a - /// flush that may already have snapshotted the dirty set. A failed flush - /// re-arms the window, so dirty state is retried instead of sitting - /// unwritten until unrelated traffic schedules again. - pub(crate) async fn schedule_signal_flush(&self) { - if self.signal_flush_pending.swap(true, Ordering::AcqRel) { - return; + /// Request a coalesced flush of the receive-path Signal cache. The first + /// request arms a single worker; concurrent requests only mark it dirty. + pub(crate) fn schedule_signal_flush(&self) { + loop { + let cur = self.signal_flush_state.load(Ordering::Acquire); + if cur & FLUSH_RUNNING == 0 { + // Idle → arm a worker. + if self + .signal_flush_state + .compare_exchange_weak(cur, FLUSH_RUNNING, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + self.spawn_signal_flush_worker(); + return; + } + } else { + // A worker is alive: mark dirty so it runs one more window. + if self + .signal_flush_state + .compare_exchange_weak( + cur, + cur | FLUSH_DIRTY, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + return; + } + } + // CAS lost a race; retry. } + } + + fn spawn_signal_flush_worker(&self) { let Some(weak) = self.self_weak.get() else { - // Constructor edge: no Arc identity to hold from the timer task. - // Flush inline so the request is never silently dropped. - self.signal_flush_pending.store(false, Ordering::Release); - self.flush_signal_cache_batch_safe_logged("coalesced-inline", None) - .await; + // Constructor edge: no Arc identity to hold from a timer task, so + // clear the arm and let the next post-construction request drive. + self.signal_flush_state.store(0, Ordering::Release); return; }; let weak = weak.clone(); let runtime = self.runtime.clone(); + let generation = self.connection_generation.load(Ordering::Acquire); self.runtime .spawn(Box::pin(async move { let mut backoff = SIGNAL_FLUSH_WINDOW; loop { - // Hold only the Weak across the sleep so an armed fire + // Hold only the Weak across the sleep so an armed worker // never extends the client's lifetime. runtime.sleep(backoff).await; let Some(client) = weak.upgrade() else { return; }; - client.signal_flush_pending.store(false, Ordering::Release); - // Batch-safe: if an offline drain became active meanwhile, - // this routes under the processing permit like any - // out-of-band flush. - let Err(e) = client.coalesced_flush_attempt().await else { - return; - }; - // Exponential backoff doubles per consecutive failure and - // caps the error-log rate along with it; the cache keeps - // its dirty entries until a flush succeeds. - backoff = (backoff * 2).min(SIGNAL_FLUSH_RETRY_CEILING); - log::error!("Coalesced signal flush failed; retrying in {backoff:?}: {e:?}"); - // If a concurrent request re-armed already, its fire owns - // the retry. - if client.signal_flush_pending.swap(true, Ordering::AcqRel) { + // A reconnect/teardown owns the state now (it reset the + // scheduler and the cache); stand down without touching the + // state so a fresh worker on the new connection is intact. + if client.connection_generation.load(Ordering::Acquire) != generation { return; } + match client.coalesced_flush_attempt().await { + Ok(()) => { + backoff = SIGNAL_FLUSH_WINDOW; + // Exit only if no request arrived mid-flush. If one + // did (DIRTY set), the RUNNING→IDLE CAS fails; clear + // DIRTY and run one more window. + if client + .signal_flush_state + .compare_exchange( + FLUSH_RUNNING, + 0, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + return; + } + client + .signal_flush_state + .fetch_and(!FLUSH_DIRTY, Ordering::AcqRel); + } + Err(e) => { + // Same worker retries with a growing backoff, so + // concurrent traffic cannot reset it to the base + // delay. The cache keeps its dirty entries. + backoff = (backoff * 2).min(SIGNAL_FLUSH_RETRY_CEILING); + log::error!( + "Coalesced signal flush failed; retrying in {backoff:?}: {e:?}" + ); + } + } } })) .detach(); } - /// One flush attempt of the fire loop; tests can inject failures to - /// exercise the retry/backoff path (same pattern as the commit batcher's - /// `fail_flushes`). + /// One flush attempt of the worker loop; tests inject failures via a + /// `cfg(test)` counter (same pattern as the commit batcher's `fail_flushes`). async fn coalesced_flush_attempt(&self) -> Result<(), anyhow::Error> { #[cfg(test)] { @@ -110,14 +157,23 @@ impl Client { self.flush_signal_cache_batch_safe().await } + /// Reset the scheduler at connection teardown so a worker stuck in a long + /// retry backoff on the old connection can't delay the next connection's + /// traffic: the worker's generation guard stands it down, and this clears + /// the arm so fresh traffic spawns a worker at the base window. + pub(crate) fn reset_signal_flush_scheduler(&self) { + self.signal_flush_state.store(0, Ordering::Release); + } + #[cfg(test)] - pub(crate) fn signal_flush_is_pending(&self) -> bool { - self.signal_flush_pending.load(Ordering::Acquire) + pub(crate) fn signal_flush_worker_alive(&self) -> bool { + self.signal_flush_state.load(Ordering::Acquire) & FLUSH_RUNNING != 0 } } #[cfg(test)] mod tests { + use super::*; use std::sync::Arc; use std::time::Duration; @@ -157,8 +213,8 @@ mod tests { } } - /// Requests inside one debounce window coalesce into a single armed fire, - /// and that fire persists every dirty entry written before it. + /// A burst of requests rides one armed worker and persists every dirty + /// entry written before it. #[tokio::test] async fn burst_of_requests_coalesces_and_persists() { let client = crate::test_utils::create_test_client().await; @@ -166,64 +222,120 @@ mod tests { let mut addrs = Vec::new(); for i in 0..10 { addrs.push(dirty_session(&client, &format!("155500011{i:02}"))); - client.schedule_signal_flush().await; + client.schedule_signal_flush(); } - assert!( - client.signal_flush_is_pending(), - "burst must ride one armed fire" - ); + assert!(client.signal_flush_worker_alive(), "burst arms one worker"); for addr in &addrs { wait_for_backend_session(&client, addr).await; } - assert!( - !client.signal_flush_is_pending(), - "the fire must clear the pending flag" - ); + // Worker exits back to idle once nothing is dirty. + let deadline = wacore::time::Instant::now() + Duration::from_secs(1); + while client.signal_flush_worker_alive() { + assert!( + wacore::time::Instant::now() < deadline, + "worker must return to idle" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + } + + /// A request after a completed worker arms a new one — later dirty state is + /// not stranded behind an exited worker. + #[tokio::test] + async fn reschedule_after_worker_exit_flushes_again() { + let client = crate::test_utils::create_test_client().await; + + let first = dirty_session(&client, "15550002001"); + client.schedule_signal_flush(); + wait_for_backend_session(&client, &first).await; + while client.signal_flush_worker_alive() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + + let second = dirty_session(&client, "15550002002"); + client.schedule_signal_flush(); + wait_for_backend_session(&client, &second).await; } - /// Failed attempts re-arm and back off instead of dropping the dirty - /// state: with 2 injected failures, the fire must consume both error - /// attempts and still persist the pre-existing dirty entry on the third. + /// While the worker is mid-flush, only one worker ever exists no matter how + /// many requests pile up; the dirty flag makes it run exactly one more + /// window rather than spawning a fleet. #[tokio::test] - async fn failed_fire_retries_until_the_dirty_entry_persists() { - use std::sync::atomic::Ordering; + async fn concurrent_requests_never_spawn_a_second_worker() { + let client = crate::test_utils::create_test_client().await; + // Block the first flush attempt so requests pile up against a running + // worker. + client + .signal_flush_test_failures + .store(3, Ordering::Release); + + for _ in 0..200 { + client.schedule_signal_flush(); + } + // State only ever carries the two defined bits; RUNNING stays set once, + // never "two workers" (which the single-flight design makes + // unrepresentable — this asserts the invariant holds under the pile-up). + for _ in 0..50 { + let s = client.signal_flush_state.load(Ordering::Acquire); + assert!(s & FLUSH_RUNNING != 0, "a worker stays armed"); + assert!( + s & !(FLUSH_RUNNING | FLUSH_DIRTY) == 0, + "no stray state bits" + ); + tokio::time::sleep(Duration::from_millis(2)).await; + } + let addr = dirty_session(&client, "15550005001"); + wait_for_backend_session(&client, &addr).await; + } + /// Failed attempts re-arm and back off in the SAME worker; the dirty entry + /// still persists once the injected failures drain. + #[tokio::test] + async fn failed_worker_retries_until_the_dirty_entry_persists() { let client = crate::test_utils::create_test_client().await; let addr = dirty_session(&client, "15550004001"); client .signal_flush_test_failures .store(2, Ordering::Release); - client.schedule_signal_flush().await; + client.schedule_signal_flush(); - // Success only comes after both injected failures are consumed by the - // retry loop (25 + 50 ms of backoff), proving the error path re-armed - // rather than stranding the dirty entry. + // Success comes only after both injected failures are consumed by the + // retry loop (25 + 50 + 100 ms of backoff), proving the same worker + // re-armed rather than stranding the entry. wait_for_backend_session(&client, &addr).await; assert_eq!( client.signal_flush_test_failures.load(Ordering::Acquire), 0, - "the retry loop must have consumed every injected failure" - ); - assert!( - !client.signal_flush_is_pending(), - "the successful retry must clear the pending flag" + "the retry loop consumed every injected failure" ); } - /// A request after a completed fire arms a NEW fire — the flag round-trips - /// and later dirty state is not stranded behind an absorbed request. + /// A generation bump (reconnect/teardown) stands the worker down instead of + /// imposing its retry backoff on the next connection. #[tokio::test] - async fn reschedule_after_fire_flushes_again() { + async fn generation_bump_stands_the_worker_down() { let client = crate::test_utils::create_test_client().await; + // Keep the worker failing so it stays in the retry loop across a bump. + client + .signal_flush_test_failures + .store(1_000, Ordering::Release); + dirty_session(&client, "15550006001"); + client.schedule_signal_flush(); + assert!(client.signal_flush_worker_alive()); - let first = dirty_session(&client, "15550002001"); - client.schedule_signal_flush().await; - wait_for_backend_session(&client, &first).await; + // Simulate teardown: bump the generation and reset the scheduler. + client.connection_generation.fetch_add(1, Ordering::SeqCst); + client.reset_signal_flush_scheduler(); - let second = dirty_session(&client, "15550002002"); - client.schedule_signal_flush().await; - wait_for_backend_session(&client, &second).await; + // The old worker exits on its next wake; the reset already cleared the + // arm, so a fresh request spawns a new worker immediately. + client + .signal_flush_test_failures + .store(0, Ordering::Release); + let addr = dirty_session(&client, "15550006002"); + client.schedule_signal_flush(); + wait_for_backend_session(&client, &addr).await; } } From 6549dcdf36a66226e740c1e3d4b2b2694d3a9598 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:17:32 -0300 Subject: [PATCH 12/24] docs(signal): fix orphaned rustdoc and document settle preconditions Reattach the batch-safe flush's doc (an earlier removal left a stray 'with error logging' line on it), and on the public settle API document the permit precondition (self-deadlock if called from an event handler / durability hook under the drain permit) and qualify the lag bound (holds on success; a backend outage extends it until the retry loop succeeds). --- src/client/adapters.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/client/adapters.rs b/src/client/adapters.rs index d0c95d5f4..e087cc670 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -105,15 +105,14 @@ impl Client { .map_err(|e| anyhow::anyhow!("Failed to flush signal cache: {e}")) } - /// [`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. /// Signal-cache flush that is safe while the offline drain is active. /// + /// [`flush_signal_cache`](Self::flush_signal_cache) is safe only when the + /// caller holds the message processing permit or the batcher is known + /// inactive: it persists the WHOLE cache, including ratchet advances of + /// drain entries that may not have a durable buffered row yet. Everything + /// else must go through this `_batch_safe` variant. + /// /// 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 From fea31e0957e24a3124f86418171afa8c6188918d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:17:32 -0300 Subject: [PATCH 13/24] test(e2e): require startup-sync quiescence; prove outbound flush is durable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Restore wait_for_startup_sync as a hard requirement in the shared connect helper (keeping the canonical wait_for_connected gate) so a real startup hang or a mid-sync client fails the connect instead of racing assertions. Add a test that each send leaves a strictly higher sender-chain counter durable in the backend with no explicit settle — the synchronous-flush invariant a crash-after-send depends on. --- tests/e2e/src/lib.rs | 34 ++++++------- tests/e2e/tests/session_reuse.rs | 87 ++++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 19 deletions(-) diff --git a/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index 0cff24514..521f43829 100644 --- a/tests/e2e/src/lib.rs +++ b/tests/e2e/src/lib.rs @@ -159,20 +159,12 @@ impl TestClient { let run_handle = bot.spawn(); - // Wait for the client to become fully ready: connected, logged in, and - // critical app-state sync complete. - // - // This gates on `wait_for_connected`, which resolves on the canonical - // `is_ready` signal (`dispatch_connected`, fired after the critical sync) - // via a notifier — not on the order events happen to arrive in the - // channel. That makes it immune to the earlier flake, where a fixed 30s - // event-loop wait for `Connected` raced how long the mock server took to - // serve app-state IQs under concurrent CI load, then fell back to - // `offline_sync_completed` — an orthogonal signal that could time out - // even though the client was otherwise ready. The single budget below is - // generous because it is a hard ceiling on a real hang, not a per-phase - // race window. PairSuccess/Connected still land in the unbounded - // `event_rx`; the predicate-filtered `wait_for_event` used by tests + // Readiness gate: `wait_for_connected` resolves on the canonical + // `is_ready` signal (`dispatch_connected`, after the critical sync) via + // a notifier, so it does not race event arrival order or fall back to an + // orthogonal signal — the earlier flake, where a fixed 30s wait for + // `Connected` timed out under CI load. PairSuccess/Connected still land + // in the unbounded `event_rx`; the predicate-filtered `wait_for_event` // discards them. if let Err(e) = client .wait_for_connected(tokio::time::Duration::from_secs(60)) @@ -186,15 +178,19 @@ impl TestClient { } // Drain the initial startup sync (offline messages + history) so tests - // start from a settled state. Best-effort: readiness is already - // guaranteed by `wait_for_connected` above, and a chat-action test does - // not depend on message backlog, so a slow mock server here must not - // fail the connect — it only means the (empty) backlog is still landing. + // start from a quiescent state. This is a hard requirement, not + // best-effort: a timeout here means a real startup hang or a mid-sync + // client that would make assertions race changing state, so it fails + // the connect. if let Err(e) = client .wait_for_startup_sync(tokio::time::Duration::from_secs(30)) .await { - eprintln!("WARN: startup sync did not settle before the connect deadline: {e}"); + client.disconnect().await; + drop(run_handle); + return Err(anyhow::anyhow!( + "startup sync did not settle before the connect deadline: {e}" + )); } Ok(Self { diff --git a/tests/e2e/tests/session_reuse.rs b/tests/e2e/tests/session_reuse.rs index 34a5c2693..61c6b6b6a 100644 --- a/tests/e2e/tests/session_reuse.rs +++ b/tests/e2e/tests/session_reuse.rs @@ -37,6 +37,93 @@ async fn scan_sessions( Ok(results) } +/// Read the sender-chain counter of the first established session for `user` +/// straight from the backend (no cache, no settle) — the durable outbound +/// ratchet position. +async fn durable_sender_chain_index( + backend: &dyn wacore::store::traits::SignalStore, + user: &str, + server: &str, +) -> anyhow::Result> { + for device_id in 0..=99u16 { + let addr = if device_id == 0 { + format!("{user}@{server}.0") + } else { + format!("{user}:{device_id}@{server}.0") + }; + if let Some(data) = backend.get_session(&addr).await? + && let Some(state) = SessionRecord::deserialize(&data)?.session_state() + && let Ok(chain) = state.get_sender_chain_key() + { + return Ok(Some(chain.index())); + } + } + Ok(None) +} + +/// The outbound ratchet advance must be durable BEFORE send_message returns: +/// a coalesced send could report success while the counter/chain-key advance +/// is still only in the in-memory cache, so a crash before the flush would +/// reuse the counter (and its message key + IV) on the next send. Read the +/// backend directly after each send — with no explicit settle — and require +/// the sender-chain counter to have already advanced, which is exactly what a +/// crash at that instant would preserve. +#[tokio::test] +async fn test_outbound_ratchet_is_durable_before_send_returns() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let mut client_a = TestClient::connect("e2e_sig_durable_a").await?; + let mut client_b = TestClient::connect("e2e_sig_durable_b").await?; + let jid_a = client_a.jid().await; + let jid_b = client_b.jid().await; + let lid_b = client_b.client.get_lid(); + + // Establish the outbound session A→B. + send_and_expect_text(&client_a.client, &mut client_b, &jid_b, "establish", 30).await?; + send_and_expect_text(&client_b.client, &mut client_a, &jid_a, "reply", 30).await?; + + let backend_a = client_a.client.persistence_manager().backend(); + let read_index = async |user: &str, server: &str| { + durable_sender_chain_index(&*backend_a, user, server).await + }; + + // The session may be keyed under LID (modern) or PN. + let (user, server) = match lid_b { + Some(ref lid) if read_index(&lid.user, "lid").await?.is_some() => (lid.user.clone(), "lid"), + _ => (jid_b.user.clone(), "c.us"), + }; + + let mut last = read_index(&user, server) + .await? + .expect("an outbound session must exist after the roundtrip"); + + // Each send must leave a strictly higher counter durable in the backend + // WITHOUT any explicit flush — proving the synchronous outbound flush. + for i in 0..3 { + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + &format!("m{i}"), + 30, + ) + .await?; + let now = read_index(&user, server) + .await? + .expect("session persists across sends"); + assert!( + now > last, + "send #{i} must persist the advanced sender-chain counter before returning \ + (durable {last} -> {now}); a coalesced send would leave it stale" + ); + last = now; + } + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + /// Multiple sequential sends without a reply should all be delivered. #[tokio::test] async fn test_one_way_multiple_sends() -> anyhow::Result<()> { From 9cf00c572ecdcbbed5decbfa27d9de5cf3767fbf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:54:20 -0300 Subject: [PATCH 14/24] fix(send): flush the outbound ratchet before the stanza hits the wire MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The synchronous send flush ran after send_node and swallowed its error, so a crash between wire and flush — or a backend failure — could still reuse the outbound counter (with its message key + IV). Move it ahead of send_node with error propagation (matching WA Web's flush-before-send): persist first, and abort the send if persistence fails rather than transmitting an advance we couldn't save. Same fix for send_status_message. --- src/send/mod.rs | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/src/send/mod.rs b/src/send/mod.rs index b3af858d7..845d34328 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -962,6 +962,10 @@ impl Client { .ensure_status_participants(prepared.node, &group_info) .await?; + // Persist the sender-key ratchet advance before the stanza hits the + // wire (same rule as the DM/group send path); a failure aborts the send. + self.flush_signal_cache_batch_safe().await?; + let ack = if let Some(phash) = stanza .attrs() .optional_string("phash") @@ -991,9 +995,6 @@ impl Client { self.invalidate_device_cache(user).await; } - self.flush_signal_cache_batch_safe_logged("send_status_message", None) - .await; - Ok(SendResult { message_id: request_id, to, @@ -1481,6 +1482,15 @@ impl Client { .await? }; + // Persist the outbound ratchet advance BEFORE the stanza hits the wire + // (WA Web flushes the Signal store ahead of send). Reusing an outbound + // counter reuses its message key + IV, so the advance must be durable + // before anyone can act on the ciphertext — and a persistence failure + // must abort the send rather than transmit an advance we couldn't save. + // Only the receive path, where a lost advance re-derives forward, + // coalesces. + self.flush_signal_cache_batch_safe().await?; + let ack = if let Some(phash) = dm_phash && let Some(msg_id) = stanza_to_send .attrs() @@ -1552,15 +1562,6 @@ impl Client { // Warm marking is visible; a waiting cold send may now re-resolve. drop(distribution_guard); - // Flush the outbound ratchet advance synchronously before returning: - // a coalesced flush here would let send_message report success while - // the counter/chain-key advance is still only in memory, so a crash - // before the flush would reuse the counter (and its message key + IV) - // on the next send. Only the receive path — where a lost advance - // re-derives forward — coalesces. - self.flush_signal_cache_batch_safe_logged("send_message_impl", None) - .await; - // Issue new tc token after send if a bucket boundary was crossed. // Fire-and-forget so send_message returns without waiting for the IQ if should_issue_tc_token_after_send { From 10a06dbe2a17cba565617a267f90f2776770a4de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:54:20 -0300 Subject: [PATCH 15/24] fix(signal): make the flush scheduler generation-scoped MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The worker checked connection_generation only before its flush await, then mutated the shared state after it. A reconnect during an in-flight flush let a stale worker clear a new-generation worker's RUNNING bit, breaking single-flight and untracking the new worker. Embed the generation in the atomic state (generation << 2 | flags) so every CAS is generation-scoped: a stale worker cannot mutate state a new generation owns, and the teardown reset is no longer needed — the generation bump hands ownership to the next request's worker. Adds a test that holds a worker inside the flush across a bump and proves it cannot clobber. --- src/client.rs | 14 ++- src/client/lifecycle.rs | 14 +-- src/signal_flush.rs | 217 +++++++++++++++++++++++++++++----------- 3 files changed, 176 insertions(+), 69 deletions(-) diff --git a/src/client.rs b/src/client.rs index a147af975..43a5c5025 100644 --- a/src/client.rs +++ b/src/client.rs @@ -836,13 +836,21 @@ pub struct Client { /// Initialized after `Arc::new(this)` in the constructor. pub(crate) self_weak: std::sync::OnceLock>, - /// Single-flight state for the coalesced Signal-cache flush worker - /// (RUNNING/DIRTY bits; see `signal_flush.rs`). - pub(crate) signal_flush_state: AtomicU32, + /// Single-flight state for the coalesced Signal-cache flush worker: + /// `(connection_generation << 2) | RUNNING/DIRTY bits` (see `signal_flush.rs`). + pub(crate) signal_flush_state: AtomicU64, /// Injected failures for the coalesced flush (consumed one per attempt), /// so tests can exercise the retry/backoff path deterministically. #[cfg(test)] pub(crate) signal_flush_test_failures: AtomicU32, + /// Blocks each coalesced flush attempt while set, so a test can hold a + /// worker inside the flush and drive a concurrent generation change. + #[cfg(test)] + pub(crate) signal_flush_test_block: AtomicBool, + /// Counts entries into the coalesced flush attempt, so a test can wait + /// until a worker is actually inside the (blocked) flush. + #[cfg(test)] + pub(crate) signal_flush_test_in_attempt: AtomicU32, /// Holds the background saver's AbortHandle so the task lifetime follows /// `Arc` ref count instead of the Bot wrapper's. Set once by diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 48363b8f0..3bc93b14e 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -240,9 +240,13 @@ impl Client { pair_code_state: Arc::new(Mutex::new(wacore::pair_code::PairCodeState::default())), passkey_state: Arc::new(Mutex::new(crate::passkey::flow::PasskeyFlowState::default())), passkey_opening: AtomicBool::new(false), - signal_flush_state: AtomicU32::new(0), + signal_flush_state: AtomicU64::new(0), #[cfg(test)] signal_flush_test_failures: AtomicU32::new(0), + #[cfg(test)] + signal_flush_test_block: AtomicBool::new(false), + #[cfg(test)] + signal_flush_test_in_attempt: AtomicU32::new(0), custom_enc_handlers: std::sync::OnceLock::new(), inbound_durability_hook: std::sync::OnceLock::new(), retry_admission: std::sync::OnceLock::new(), @@ -767,11 +771,9 @@ impl Client { // 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); - // Stand down the coalesced-flush worker: its generation guard exits it - // on the next wake, and clearing the arm lets the next connection's - // traffic spawn a fresh worker at the base window instead of waiting - // out a stale retry backoff. - self.reset_signal_flush_scheduler(); + // The coalesced-flush scheduler needs no explicit reset: its state is + // generation-scoped, so the bump above already hands ownership to the + // next connection's first request and retires any stale worker. // 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/signal_flush.rs b/src/signal_flush.rs index 91feef399..ecc3b0c8f 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -39,29 +39,51 @@ const SIGNAL_FLUSH_WINDOW: std::time::Duration = std::time::Duration::from_milli /// long-lived backend outage. const SIGNAL_FLUSH_RETRY_CEILING: std::time::Duration = std::time::Duration::from_secs(5); -/// A worker is alive (armed or running/retrying). -const FLUSH_RUNNING: u32 = 0b01; +// Scheduler state is `(connection_generation << 2) | flags`. Embedding the +// generation makes worker ownership generation-scoped: a worker from an old +// connection cannot mutate state a new-connection worker owns, because its CAS +// targets its own generation's exact value. So a reconnect during an in-flight +// flush needs no teardown reset — the next request on the new generation takes +// over via CAS, and the stale worker retires when it sees a foreign generation. +const FLUSH_RUNNING: u64 = 0b01; /// A request arrived while the worker was mid-flush; it must run one more window. -const FLUSH_DIRTY: u32 = 0b10; +const FLUSH_DIRTY: u64 = 0b10; +#[cfg(test)] +const FLUSH_FLAGS: u64 = FLUSH_RUNNING | FLUSH_DIRTY; + +#[inline] +fn pack_flush_state(generation: u64, flags: u64) -> u64 { + (generation << 2) | flags +} impl Client { /// Request a coalesced flush of the receive-path Signal cache. The first - /// request arms a single worker; concurrent requests only mark it dirty. + /// request for the current connection generation arms a single worker; + /// concurrent requests only mark it dirty. pub(crate) fn schedule_signal_flush(&self) { + let generation = self.connection_generation.load(Ordering::Acquire); loop { let cur = self.signal_flush_state.load(Ordering::Acquire); - if cur & FLUSH_RUNNING == 0 { - // Idle → arm a worker. + let running_this_generation = (cur >> 2) == generation && cur & FLUSH_RUNNING != 0; + if !running_this_generation { + // Idle, or only a stale-generation worker is present → take + // over for this generation. if self .signal_flush_state - .compare_exchange_weak(cur, FLUSH_RUNNING, Ordering::AcqRel, Ordering::Acquire) + .compare_exchange_weak( + cur, + pack_flush_state(generation, FLUSH_RUNNING), + Ordering::AcqRel, + Ordering::Acquire, + ) .is_ok() { - self.spawn_signal_flush_worker(); + self.spawn_signal_flush_worker(generation); return; } } else { - // A worker is alive: mark dirty so it runs one more window. + // A worker for this generation is alive: mark dirty so it runs + // one more window. if self .signal_flush_state .compare_exchange_weak( @@ -79,16 +101,16 @@ impl Client { } } - fn spawn_signal_flush_worker(&self) { + fn spawn_signal_flush_worker(&self, generation: u64) { let Some(weak) = self.self_weak.get() else { // Constructor edge: no Arc identity to hold from a timer task, so - // clear the arm and let the next post-construction request drive. - self.signal_flush_state.store(0, Ordering::Release); + // release the arm and let the next post-construction request drive. + self.signal_flush_state + .store(pack_flush_state(generation, 0), Ordering::Release); return; }; let weak = weak.clone(); let runtime = self.runtime.clone(); - let generation = self.connection_generation.load(Ordering::Acquire); self.runtime .spawn(Box::pin(async move { let mut backoff = SIGNAL_FLUSH_WINDOW; @@ -99,35 +121,48 @@ impl Client { let Some(client) = weak.upgrade() else { return; }; - // A reconnect/teardown owns the state now (it reset the - // scheduler and the cache); stand down without touching the - // state so a fresh worker on the new connection is intact. - if client.connection_generation.load(Ordering::Acquire) != generation { - return; - } match client.coalesced_flush_attempt().await { Ok(()) => { backoff = SIGNAL_FLUSH_WINDOW; - // Exit only if no request arrived mid-flush. If one - // did (DIRTY set), the RUNNING→IDLE CAS fails; clear - // DIRTY and run one more window. - if client - .signal_flush_state - .compare_exchange( - FLUSH_RUNNING, - 0, - Ordering::AcqRel, - Ordering::Acquire, - ) - .is_ok() - { - return; + // Settle ownership under a CAS scoped to our + // generation: exit to idle if nothing is dirty, run + // one more window if it is, or stand down if a + // reconnect handed the state to a new generation. + loop { + let cur = client.signal_flush_state.load(Ordering::Acquire); + if (cur >> 2) != generation { + return; + } + let next = if cur & FLUSH_DIRTY != 0 { + pack_flush_state(generation, FLUSH_RUNNING) + } else { + pack_flush_state(generation, 0) + }; + if client + .signal_flush_state + .compare_exchange_weak( + cur, + next, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + if next & FLUSH_RUNNING == 0 { + return; + } + break; + } } - client - .signal_flush_state - .fetch_and(!FLUSH_DIRTY, Ordering::AcqRel); } Err(e) => { + // A reconnect handed the state to a new generation: + // stand down instead of imposing a stale backoff. + if (client.signal_flush_state.load(Ordering::Acquire) >> 2) + != generation + { + return; + } // Same worker retries with a growing backoff, so // concurrent traffic cannot reset it to the base // delay. The cache keeps its dirty entries. @@ -147,6 +182,11 @@ impl Client { async fn coalesced_flush_attempt(&self) -> Result<(), anyhow::Error> { #[cfg(test)] { + self.signal_flush_test_in_attempt + .fetch_add(1, Ordering::AcqRel); + while self.signal_flush_test_block.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } let remaining = self.signal_flush_test_failures.load(Ordering::Acquire); if remaining > 0 { self.signal_flush_test_failures @@ -157,14 +197,6 @@ impl Client { self.flush_signal_cache_batch_safe().await } - /// Reset the scheduler at connection teardown so a worker stuck in a long - /// retry backoff on the old connection can't delay the next connection's - /// traffic: the worker's generation guard stands it down, and this clears - /// the arm so fresh traffic spawns a worker at the base window. - pub(crate) fn reset_signal_flush_scheduler(&self) { - self.signal_flush_state.store(0, Ordering::Release); - } - #[cfg(test)] pub(crate) fn signal_flush_worker_alive(&self) -> bool { self.signal_flush_state.load(Ordering::Acquire) & FLUSH_RUNNING != 0 @@ -273,16 +305,13 @@ mod tests { for _ in 0..200 { client.schedule_signal_flush(); } - // State only ever carries the two defined bits; RUNNING stays set once, - // never "two workers" (which the single-flight design makes - // unrepresentable — this asserts the invariant holds under the pile-up). + // RUNNING stays set for generation 0 the whole time — never a second + // worker; only the two flag bits are ever set (generation 0 leaves the + // high bits clear), so the pile-up never mints stray state. for _ in 0..50 { let s = client.signal_flush_state.load(Ordering::Acquire); assert!(s & FLUSH_RUNNING != 0, "a worker stays armed"); - assert!( - s & !(FLUSH_RUNNING | FLUSH_DIRTY) == 0, - "no stray state bits" - ); + assert_eq!(s & !FLUSH_FLAGS, 0, "generation 0: no stray high bits"); tokio::time::sleep(Duration::from_millis(2)).await; } let addr = dirty_session(&client, "15550005001"); @@ -312,12 +341,15 @@ mod tests { ); } - /// A generation bump (reconnect/teardown) stands the worker down instead of - /// imposing its retry backoff on the next connection. + /// A generation bump (reconnect/teardown) hands scheduler ownership to the + /// next connection without an explicit reset: the new request takes over + /// via a generation-scoped CAS while the stale worker is still looping, and + /// the stale worker cannot clobber the new worker's state. #[tokio::test] - async fn generation_bump_stands_the_worker_down() { + async fn generation_bump_hands_off_without_clobbering() { let client = crate::test_utils::create_test_client().await; - // Keep the worker failing so it stays in the retry loop across a bump. + // Keep the old worker failing so it stays in the retry loop across the + // bump (a stale worker that could still reach its exit CAS). client .signal_flush_test_failures .store(1_000, Ordering::Release); @@ -325,17 +357,82 @@ mod tests { client.schedule_signal_flush(); assert!(client.signal_flush_worker_alive()); - // Simulate teardown: bump the generation and reset the scheduler. + // Bump the generation as teardown does — no explicit scheduler reset. client.connection_generation.fetch_add(1, Ordering::SeqCst); - client.reset_signal_flush_scheduler(); - - // The old worker exits on its next wake; the reset already cleared the - // arm, so a fresh request spawns a new worker immediately. + // The next connection's first request takes over for the new + // generation; stop injecting failures so its worker succeeds. client .signal_flush_test_failures .store(0, Ordering::Release); let addr = dirty_session(&client, "15550006002"); client.schedule_signal_flush(); wait_for_backend_session(&client, &addr).await; + + // The state must be tagged with the new generation, proving the + // hand-off (and that the stale worker did not reclaim it). + let new_gen = client.connection_generation.load(Ordering::Acquire); + let deadline = wacore::time::Instant::now() + Duration::from_secs(1); + loop { + let s = client.signal_flush_state.load(Ordering::Acquire); + if s >> 2 == new_gen { + break; + } + assert!( + wacore::time::Instant::now() < deadline, + "scheduler state must carry the new generation, got {s:#x}" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + } + + /// The race the generation-scoped CAS closes: an old worker held INSIDE the + /// flush while a reconnect hands the scheduler to a new-generation worker + /// must not clobber the new worker's state when it finally completes. + #[tokio::test] + async fn stale_worker_inside_flush_cannot_clobber_new_generation() { + let client = crate::test_utils::create_test_client().await; + + // Hold every flush attempt inside the flush until we release it. + client + .signal_flush_test_block + .store(true, Ordering::Release); + + // Arm worker A on generation 0 and wait until it is inside the flush. + dirty_session(&client, "15550007001"); + client.schedule_signal_flush(); + let deadline = wacore::time::Instant::now() + Duration::from_secs(1); + while client.signal_flush_test_in_attempt.load(Ordering::Acquire) == 0 { + assert!( + wacore::time::Instant::now() < deadline, + "worker A must enter the flush" + ); + tokio::time::sleep(Duration::from_millis(2)).await; + } + + // Reconnect: bump the generation, then a new request takes over and + // arms worker B on the new generation (also blocked in the flush). + let new_gen = client.connection_generation.fetch_add(1, Ordering::SeqCst) + 1; + dirty_session(&client, "15550007002"); + client.schedule_signal_flush(); + let s = client.signal_flush_state.load(Ordering::Acquire); + assert_eq!(s >> 2, new_gen, "worker B took over for the new generation"); + + // Release both. The stale worker A completes its flush and tries to + // settle ownership — its generation-scoped CAS must fail, leaving B's + // state intact rather than reverting to generation 0. + client + .signal_flush_test_block + .store(false, Ordering::Release); + for _ in 0..100 { + let s = client.signal_flush_state.load(Ordering::Acquire); + assert!( + s >> 2 >= new_gen, + "stale worker A clobbered the new generation's state: {s:#x}" + ); + tokio::time::sleep(Duration::from_millis(2)).await; + } + // And B still makes progress: its dirty entry lands. + let addr = ProtocolAddress::new("15550007002".to_string(), 1.into()); + wait_for_backend_session(&client, &addr).await; } } From 14b9cac35e3d63c2fb8ee57872174e9f262586b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:54:20 -0300 Subject: [PATCH 16/24] test(e2e),docs: prove the send flush ordering; fix stale coalescing docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rework the outbound-durability test to read the backend immediately after send_message (no delivery wait, no settle) so a coalesced/deferred flush would fail it — the previous version read after delivery, where even a 25 ms window had elapsed. Drop the now-redundant settle before a post-send backend read (send is synchronous; a missing session now catches a regression). Correct the docs that still said sends coalesce (only receives do) and that promised a hard one-window bound. --- src/client/adapters.rs | 12 +++++---- tests/e2e/tests/lid_sessions.rs | 5 ++-- tests/e2e/tests/session_reuse.rs | 45 ++++++++++++++------------------ 3 files changed, 30 insertions(+), 32 deletions(-) diff --git a/src/client/adapters.rs b/src/client/adapters.rs index e087cc670..947fac6bf 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -71,11 +71,13 @@ impl Client { /// Force any pending write-behind Signal cache state to the backend, /// returning once the flush completes (or fails). /// - /// Live sends and receives schedule a coalesced flush (see - /// `signal_flush.rs`) instead of writing through, so on success the backend - /// trails the in-memory cache by at most one coalescing window; a backend - /// outage extends that until the scheduler's retry loop succeeds. Use this - /// to settle durability deterministically before reading persisted state or + /// The live receive path schedules a coalesced flush (see `signal_flush.rs`) + /// instead of writing through (sends flush synchronously). On success the + /// backend normally trails the cache by about the coalescing window, but + /// that is not a hard wall-clock bound — the timer can slip under runtime + /// starvation and the flush can wait on locks or slow/failing storage (a + /// backend outage extends it until the retry loop succeeds). Use this to + /// settle durability deterministically before reading persisted state or /// ahead of a non-graceful shutdown — and check the returned `Result`, as a /// failure leaves state pending. /// diff --git a/tests/e2e/tests/lid_sessions.rs b/tests/e2e/tests/lid_sessions.rs index 59459dd68..faa7f889f 100644 --- a/tests/e2e/tests/lid_sessions.rs +++ b/tests/e2e/tests/lid_sessions.rs @@ -282,7 +282,8 @@ async fn test_stale_pn_session_does_not_break_lid_messaging() -> anyhow::Result< .await?; info!("Messaging works despite stale PN session in DB"); - // Settle the coalesced flush from the post-inject sends before reading. + // Settle the coalesced receive-path flush (A received the post-inject + // reply) before reading. client_a.client.flush_pending_signal_state().await?; // LID session should still be authoritative assert!( @@ -684,7 +685,7 @@ async fn test_pn_migration_is_durable_across_followup_messages() -> anyhow::Resu .await?; } - // Settle the coalesced flush from the follow-up sends before reading. + // Settle the coalesced receive-path flush (A received the follow-up messages) before reading. client_a.client.flush_pending_signal_state().await?; assert!( backend_a.get_session(&lid_addr).await?.is_some(), diff --git a/tests/e2e/tests/session_reuse.rs b/tests/e2e/tests/session_reuse.rs index 61c6b6b6a..6c430510c 100644 --- a/tests/e2e/tests/session_reuse.rs +++ b/tests/e2e/tests/session_reuse.rs @@ -61,15 +61,14 @@ async fn durable_sender_chain_index( Ok(None) } -/// The outbound ratchet advance must be durable BEFORE send_message returns: -/// a coalesced send could report success while the counter/chain-key advance -/// is still only in the in-memory cache, so a crash before the flush would -/// reuse the counter (and its message key + IV) on the next send. Read the -/// backend directly after each send — with no explicit settle — and require -/// the sender-chain counter to have already advanced, which is exactly what a -/// crash at that instant would preserve. +/// The outbound ratchet advance must be durable by the time `send_message` +/// returns: reusing an outbound counter reuses its message key + IV, so a crash +/// after a successful send must never leave the advance only in memory. This +/// reads the backend IMMEDIATELY after `send_message` (no wait for delivery, no +/// explicit settle): a coalesced send would still be inside its window and the +/// counter would be stale, so only the synchronous outbound flush passes. #[tokio::test] -async fn test_outbound_ratchet_is_durable_before_send_returns() -> anyhow::Result<()> { +async fn test_outbound_ratchet_is_durable_when_send_returns() -> anyhow::Result<()> { let _ = env_logger::builder().is_test(true).try_init(); let mut client_a = TestClient::connect("e2e_sig_durable_a").await?; @@ -97,17 +96,14 @@ async fn test_outbound_ratchet_is_durable_before_send_returns() -> anyhow::Resul .await? .expect("an outbound session must exist after the roundtrip"); - // Each send must leave a strictly higher counter durable in the backend - // WITHOUT any explicit flush — proving the synchronous outbound flush. + // send_message returns only after the synchronous pre-wire flush, so the + // advanced counter is already durable — read it with no delivery wait and + // no settle. A coalesced (window-deferred) flush would leave it unchanged. for i in 0..3 { - send_and_expect_text( - &client_a.client, - &mut client_b, - &jid_b, - &format!("m{i}"), - 30, - ) - .await?; + client_a + .client + .send_message(jid_b.clone(), e2e_tests::text_msg(&format!("m{i}"))) + .await?; let now = read_index(&user, server) .await? .expect("session persists across sends"); @@ -227,10 +223,9 @@ async fn test_session_state_after_roundtrip() -> anyhow::Result<()> { send_and_expect_text(&client_b.client, &mut client_a, &jid_a, "Session reply", 30).await?; info!("Roundtrip complete"); - // Settle A's write-behind Signal cache to the backend before inspecting - // it: a live send schedules a coalesced flush rather than writing - // synchronously, so a plain send-and-read races the debounce window (and - // its timer can slip arbitrarily under CI load). + // Settle A's write-behind Signal cache before inspecting it: A's last op + // here is receiving B's reply, and the receive-path flush is coalesced, so + // a plain read races the coalescing window. client_a.client.flush_pending_signal_state().await?; // Inspect session state @@ -319,9 +314,9 @@ async fn test_session_persistence() -> anyhow::Result<()> { .await?; info!("First message sent: {msg_id_1}"); - // Settle the coalesced send flush before reading durable state (see the - // roundtrip test): a live send does not write through synchronously. - client_a.client.flush_pending_signal_state().await?; + // No settle: the send path flushes synchronously, so the session is durable + // in the backend by the time send_message returned. (A missing session here + // would catch a regression back to a coalesced/deferred send flush.) // Session may be under PN (c.us) or LID (lid) depending on whether // PN→LID mapping was resolved before encryption. From 8a38d465944cc1a3695344e154d3d26f6a50d410 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 16:58:23 -0300 Subject: [PATCH 17/24] fix(signal): skip the stale worker's flush after a generation change Re-add the pre-flush generation check the generation-scoped rewrite dropped: the exit CAS already rejects a stale worker's mutation, but checking before the flush also avoids a needless stale write after teardown settled the cache. --- src/signal_flush.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/signal_flush.rs b/src/signal_flush.rs index ecc3b0c8f..269b78e8a 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -121,6 +121,13 @@ impl Client { let Some(client) = weak.upgrade() else { return; }; + // A reconnect handed ownership to a new generation: stand + // down before flushing. The generation-scoped exit CAS would + // reject our mutation anyway, but skipping the flush avoids a + // stale write after teardown settled the cache. + if client.connection_generation.load(Ordering::Acquire) != generation { + return; + } match client.coalesced_flush_attempt().await { Ok(()) => { backoff = SIGNAL_FLUSH_WINDOW; From 4b6522c509d4ba6617eb93c2436a8580fd4048fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 17:53:49 -0300 Subject: [PATCH 18/24] fix(signal): reject stale-generation schedule calls (no scheduler regression) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The generation-scoped worker CAS protected against a stale WORKER, but schedule_signal_flush itself read the generation once and its take-over CAS accepted any state — so a call carrying a torn-down connection's generation could CAS a newer generation's RUNNING|DIRTY state back down to its own, dropping the pending DIRTY and orphaning the state (both the new worker and the demoted one then stand down). Pass the caller's already-validated lane_generation, no-op if it is not the live generation, and never move the state to an older generation. Two tests: a stale schedule leaves gen-1 state bit-for-bit unchanged, and a stale schedule during a new worker's flush cannot clear the pending DIRTY. --- src/message/receive.rs | 2 +- src/signal_flush.rs | 133 +++++++++++++++++++++++++++++++++++------ 2 files changed, 115 insertions(+), 20 deletions(-) diff --git a/src/message/receive.rs b/src/message/receive.rs index 94a690870..31c0c82f5 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -568,7 +568,7 @@ impl Client { if self.inbound_commit_batch.is_active() { self.maybe_flush_inbound_commits().await; } else { - self.schedule_signal_flush(); + self.schedule_signal_flush(lane_generation); } } diff --git a/src/signal_flush.rs b/src/signal_flush.rs index 269b78e8a..1b7f14e96 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -57,28 +57,40 @@ fn pack_flush_state(generation: u64, flags: u64) -> u64 { } impl Client { - /// Request a coalesced flush of the receive-path Signal cache. The first - /// request for the current connection generation arms a single worker; - /// concurrent requests only mark it dirty. - pub(crate) fn schedule_signal_flush(&self) { - let generation = self.connection_generation.load(Ordering::Acquire); + /// Request a coalesced flush of the receive-path Signal cache for the + /// caller's already-validated `lane_generation`. The first request for the + /// live generation arms a single worker; concurrent requests only mark it + /// dirty. A request from a torn-down connection is a no-op — it can neither + /// arm a worker for a dead generation nor move the scheduler backwards. + pub(crate) fn schedule_signal_flush(&self, lane_generation: u64) { loop { + // A stale caller (its connection was torn down) must not touch the + // scheduler; the live generation is the source of truth. + if self.connection_generation.load(Ordering::Acquire) != lane_generation { + return; + } let cur = self.signal_flush_state.load(Ordering::Acquire); - let running_this_generation = (cur >> 2) == generation && cur & FLUSH_RUNNING != 0; + let cur_gen = cur >> 2; + // The state never moves to an older generation: a newer worker + // already owns it, so this request is redundant. + if cur_gen > lane_generation { + return; + } + let running_this_generation = cur_gen == lane_generation && cur & FLUSH_RUNNING != 0; if !running_this_generation { - // Idle, or only a stale-generation worker is present → take - // over for this generation. + // Idle, or only an older generation's leftover state → take over + // for the live generation. if self .signal_flush_state .compare_exchange_weak( cur, - pack_flush_state(generation, FLUSH_RUNNING), + pack_flush_state(lane_generation, FLUSH_RUNNING), Ordering::AcqRel, Ordering::Acquire, ) .is_ok() { - self.spawn_signal_flush_worker(generation); + self.spawn_signal_flush_worker(lane_generation); return; } } else { @@ -208,6 +220,13 @@ impl Client { pub(crate) fn signal_flush_worker_alive(&self) -> bool { self.signal_flush_state.load(Ordering::Acquire) & FLUSH_RUNNING != 0 } + + /// Schedule for the live generation, as the receive path does with its + /// validated `lane_generation`. + #[cfg(test)] + pub(crate) fn schedule_signal_flush_live(&self) { + self.schedule_signal_flush(self.connection_generation.load(Ordering::Acquire)); + } } #[cfg(test)] @@ -261,7 +280,7 @@ mod tests { let mut addrs = Vec::new(); for i in 0..10 { addrs.push(dirty_session(&client, &format!("155500011{i:02}"))); - client.schedule_signal_flush(); + client.schedule_signal_flush_live(); } assert!(client.signal_flush_worker_alive(), "burst arms one worker"); @@ -286,14 +305,14 @@ mod tests { let client = crate::test_utils::create_test_client().await; let first = dirty_session(&client, "15550002001"); - client.schedule_signal_flush(); + client.schedule_signal_flush_live(); wait_for_backend_session(&client, &first).await; while client.signal_flush_worker_alive() { tokio::time::sleep(Duration::from_millis(5)).await; } let second = dirty_session(&client, "15550002002"); - client.schedule_signal_flush(); + client.schedule_signal_flush_live(); wait_for_backend_session(&client, &second).await; } @@ -310,7 +329,7 @@ mod tests { .store(3, Ordering::Release); for _ in 0..200 { - client.schedule_signal_flush(); + client.schedule_signal_flush_live(); } // RUNNING stays set for generation 0 the whole time — never a second // worker; only the two flag bits are ever set (generation 0 leaves the @@ -335,7 +354,7 @@ mod tests { .signal_flush_test_failures .store(2, Ordering::Release); - client.schedule_signal_flush(); + client.schedule_signal_flush_live(); // Success comes only after both injected failures are consumed by the // retry loop (25 + 50 + 100 ms of backoff), proving the same worker @@ -361,7 +380,7 @@ mod tests { .signal_flush_test_failures .store(1_000, Ordering::Release); dirty_session(&client, "15550006001"); - client.schedule_signal_flush(); + client.schedule_signal_flush_live(); assert!(client.signal_flush_worker_alive()); // Bump the generation as teardown does — no explicit scheduler reset. @@ -372,7 +391,7 @@ mod tests { .signal_flush_test_failures .store(0, Ordering::Release); let addr = dirty_session(&client, "15550006002"); - client.schedule_signal_flush(); + client.schedule_signal_flush_live(); wait_for_backend_session(&client, &addr).await; // The state must be tagged with the new generation, proving the @@ -406,7 +425,7 @@ mod tests { // Arm worker A on generation 0 and wait until it is inside the flush. dirty_session(&client, "15550007001"); - client.schedule_signal_flush(); + client.schedule_signal_flush_live(); let deadline = wacore::time::Instant::now() + Duration::from_secs(1); while client.signal_flush_test_in_attempt.load(Ordering::Acquire) == 0 { assert!( @@ -420,7 +439,7 @@ mod tests { // arms worker B on the new generation (also blocked in the flush). let new_gen = client.connection_generation.fetch_add(1, Ordering::SeqCst) + 1; dirty_session(&client, "15550007002"); - client.schedule_signal_flush(); + client.schedule_signal_flush_live(); let s = client.signal_flush_state.load(Ordering::Acquire); assert_eq!(s >> 2, new_gen, "worker B took over for the new generation"); @@ -442,4 +461,80 @@ mod tests { let addr = ProtocolAddress::new("15550007002".to_string(), 1.into()); wait_for_backend_session(&client, &addr).await; } + + /// A schedule call carrying a torn-down connection's generation must not + /// move the scheduler backwards — the exact bit-state must be untouched. + #[tokio::test] + async fn stale_schedule_cannot_move_the_generation_backwards() { + let client = crate::test_utils::create_test_client().await; + + // Live generation 1, with generation 1's running+dirty state installed. + client.connection_generation.store(1, Ordering::SeqCst); + let installed = pack_flush_state(1, FLUSH_RUNNING | FLUSH_DIRTY); + client + .signal_flush_state + .store(installed, Ordering::Release); + + // A stale gen-0 schedule is a no-op (live-generation guard). + client.schedule_signal_flush(0); + assert_eq!( + client.signal_flush_state.load(Ordering::Acquire), + installed, + "a stale (gen-0) schedule must not alter gen-1 state" + ); + + // And the no-regress guard holds even if the live generation matched + // the lane while the state already carried a newer generation (the + // window between the live-generation check and the CAS). + client.connection_generation.store(0, Ordering::SeqCst); + client.schedule_signal_flush(0); + assert_eq!( + client.signal_flush_state.load(Ordering::Acquire), + installed, + "a schedule must never CAS gen-1 state down to gen-0" + ); + } + + /// A stale schedule arriving while a new-generation worker is mid-flush must + /// not steal the DIRTY bit — the pending update still gets its own window. + #[tokio::test] + async fn stale_schedule_does_not_starve_the_new_generation() { + let client = crate::test_utils::create_test_client().await; + client + .signal_flush_test_block + .store(true, Ordering::Release); + + // Arm worker B on generation 1 and wait until it is inside the flush. + client.connection_generation.fetch_add(1, Ordering::SeqCst); + dirty_session(&client, "15550008001"); + client.schedule_signal_flush_live(); + let deadline = wacore::time::Instant::now() + Duration::from_secs(1); + while client.signal_flush_test_in_attempt.load(Ordering::Acquire) == 0 { + assert!( + wacore::time::Instant::now() < deadline, + "worker B must enter" + ); + tokio::time::sleep(Duration::from_millis(2)).await; + } + + // A live request during the flush marks DIRTY; a stale gen-0 schedule + // must not clear it. + let second = dirty_session(&client, "15550008002"); + client.schedule_signal_flush_live(); + assert!( + client.signal_flush_state.load(Ordering::Acquire) & FLUSH_DIRTY != 0, + "the live request set DIRTY" + ); + client.schedule_signal_flush(0); + assert!( + client.signal_flush_state.load(Ordering::Acquire) & FLUSH_DIRTY != 0, + "a stale schedule must not clear the pending DIRTY" + ); + + // Release: B sees DIRTY and runs another window that persists it. + client + .signal_flush_test_block + .store(false, Ordering::Release); + wait_for_backend_session(&client, &second).await; + } } From 935c83c75deed7a4d1dd53ee34bd205bd0c982e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 17:53:49 -0300 Subject: [PATCH 19/24] docs,test(e2e): fix stale per-message-flush rustdoc; settle without reconnect MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit signal_cache and flush_signal_cache no longer flush unconditionally per message (send is synchronous, receive coalesces) — correct their docs. In the migration test, settle via flush_pending_signal_state instead of a full reconnect: faster and isolates the scenario. --- src/client.rs | 5 +++-- src/client/adapters.rs | 5 +++-- tests/e2e/tests/lid_sessions.rs | 8 ++++---- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/client.rs b/src/client.rs index 43a5c5025..731ec5e29 100644 --- a/src/client.rs +++ b/src/client.rs @@ -560,8 +560,9 @@ pub struct Client { pub(crate) unified_session: crate::unified_session::UnifiedSessionManager, /// In-memory cache for Signal protocol state (sessions, identities, sender keys). - /// Matches WhatsApp Web's SignalStoreCache pattern: crypto ops read/write this cache, - /// and DB writes are deferred to flush() after each message is processed. + /// Matches WhatsApp Web's SignalStoreCache pattern: crypto ops read/write this + /// cache, and DB writes are flushed out of it — synchronously on the send path + /// and coalesced on the receive path (see `signal_flush.rs`). pub(crate) signal_cache: Arc, /// Limits message processing concurrency (1 permit during offline sync, N after). diff --git a/src/client/adapters.rs b/src/client/adapters.rs index 947fac6bf..dce5a9ddc 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -91,8 +91,9 @@ impl Client { self.flush_signal_cache_batch_safe().await } - /// Flush the in-memory signal cache to the database backend. - /// Called after each message is decrypted or after encryption operations. + /// Flush the in-memory signal cache to the database backend. Invoked by the + /// send path (synchronously, pre-wire), the receive-path coalescer, and the + /// drain/retry/teardown recovery paths — not unconditionally per message. pub(crate) async fn flush_signal_cache(&self) -> Result<(), anyhow::Error> { // Hold no device guard across the flush: this per-message batched SQLite // write would otherwise block every concurrent Device write for its duration. diff --git a/tests/e2e/tests/lid_sessions.rs b/tests/e2e/tests/lid_sessions.rs index faa7f889f..6c2bbdd96 100644 --- a/tests/e2e/tests/lid_sessions.rs +++ b/tests/e2e/tests/lid_sessions.rs @@ -526,10 +526,10 @@ async fn test_pn_only_session_causes_undecryptable_on_lid_lookup() -> anyhow::Re info!("Verified messaging works after first reconnect"); // Settle before the backend surgery below: the step-3 message left a - // ratchet advance in the (write-behind) signal cache, and a reconnect - // teardown would flush it AFTER the surgery — resurrecting the LID - // session this test deletes. Same pattern as the durability test. - client_a.reconnect_and_wait().await?; + // receive-side ratchet advance in the coalesced signal cache, which a later + // flush would write AFTER the surgery — resurrecting the LID session this + // test deletes. Settle it directly (no full reconnect needed). + client_a.client.flush_pending_signal_state().await?; info!("Settled signal cache before backend surgery"); // Step 4: Simulate legacy DB — move session from LID to PN address. From ef1cf76d53cc568dd1af0834b62e81607ec6ab29 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 18:11:56 -0300 Subject: [PATCH 20/24] test(e2e): prove send aborts before the wire when persistence fails Adds two InMemoryBackend test hooks: a put_sessions_batch call counter and a fail switch. The new e2e test fails session persistence, then asserts the send returns Err, the flush was attempted (counter moved, so the flush runs pre-wire), the peer receives nothing, and delivery resumes once persistence recovers. This is the deterministic proof of the send-path durability ordering the coalescer relies on. TestClient retains the concrete backend to reach the hooks. --- tests/e2e/src/lib.rs | 7 +++- tests/e2e/tests/session_reuse.rs | 69 ++++++++++++++++++++++++++++++++ wacore/src/store/in_memory.rs | 27 ++++++++++++- 3 files changed, 101 insertions(+), 2 deletions(-) diff --git a/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index 521f43829..f56102609 100644 --- a/tests/e2e/src/lib.rs +++ b/tests/e2e/src/lib.rs @@ -97,6 +97,9 @@ pub struct TestClient { pub client: Arc, pub event_rx: async_channel::Receiver>, pub run_handle: whatsapp_rust::bot::BotHandle, + /// The concrete backend, retained for its test hooks + /// (`session_batch_write_count`, `set_fail_session_writes`). + pub backend: Arc, } impl TestClient { @@ -126,8 +129,9 @@ impl TestClient { let transport_factory = TokioWebSocketTransportFactory::new().with_url(mock_server_url()); let (event_handler, event_rx) = ChannelEventHandler::new(); + let backend = Arc::new(InMemoryBackend::new()); let mut builder = Bot::builder() - .with_backend(InMemoryBackend::new()) + .with_backend_arc(backend.clone()) .with_transport_factory(transport_factory) .with_http_client(UreqHttpClient::new()) .with_runtime(whatsapp_rust::TokioRuntime) @@ -197,6 +201,7 @@ impl TestClient { client, event_rx, run_handle, + backend, }) } diff --git a/tests/e2e/tests/session_reuse.rs b/tests/e2e/tests/session_reuse.rs index 6c430510c..67038f52b 100644 --- a/tests/e2e/tests/session_reuse.rs +++ b/tests/e2e/tests/session_reuse.rs @@ -120,6 +120,75 @@ async fn test_outbound_ratchet_is_durable_when_send_returns() -> anyhow::Result< Ok(()) } +/// A send whose outbound-ratchet persistence fails must abort BEFORE the stanza +/// reaches the wire: the flush precedes the send on the send path, so if the +/// advance cannot be stored, `send_message` returns `Err` and the peer receives +/// nothing. Otherwise a crash after a wire-committed send would leave the +/// advance only in memory and the next send would reuse that counter's key + IV. +#[tokio::test] +async fn test_send_aborts_before_wire_when_persist_fails() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let mut client_a = TestClient::connect("e2e_sig_abort_a").await?; + let mut client_b = TestClient::connect("e2e_sig_abort_b").await?; + let jid_a = client_a.jid().await; + let jid_b = client_b.jid().await; + + // Establish the session both ways so the next A→B send is a steady-state + // encrypt (its only new durable write is the ratchet advance we fail). + send_and_expect_text(&client_a.client, &mut client_b, &jid_b, "establish", 30).await?; + send_and_expect_text(&client_b.client, &mut client_a, &jid_a, "reply", 30).await?; + + // Persisting the outbound advance now fails. + client_a.backend.set_fail_session_writes(true); + let writes_before = client_a.backend.session_batch_write_count(); + let result = client_a + .client + .send_message( + jid_b.clone(), + e2e_tests::text_msg("must not reach the wire"), + ) + .await; + assert!( + result.is_err(), + "send must fail when the ratchet advance cannot be persisted, got {result:?}" + ); + // The send reached the (failing) persistence step, proving the flush runs on + // the send path before the wire rather than being skipped or deferred. + assert!( + client_a.backend.session_batch_write_count() > writes_before, + "the send path must attempt to persist the ratchet advance before the wire" + ); + + // The stanza never went out: B must not see it. + client_b + .assert_no_event( + 3, + |e| { + e.messages() + .any(|m| m.message.conversation.as_deref() == Some("must not reach the wire")) + }, + "a send that failed to persist must not deliver", + ) + .await?; + + // Recovery: once persistence works again, sends deliver normally, proving + // the failure aborted cleanly rather than wedging the session. + client_a.backend.set_fail_session_writes(false); + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + "after recovery", + 30, + ) + .await?; + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + /// Multiple sequential sends without a reply should all be delivered. #[tokio::test] async fn test_one_way_multiple_sends() -> anyhow::Result<()> { diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index 9f7f2e9e4..745e50a40 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -6,7 +6,7 @@ use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::{AtomicI32, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicI32, AtomicU32, Ordering}; use crate::appstate::hash::HashState; use crate::store::Device; @@ -91,6 +91,13 @@ const MAX_SENT_MESSAGES: usize = 4096; pub struct InMemoryBackend { state: Mutex, next_device_id: AtomicI32, + /// Count of `put_sessions_batch` calls. Test hook: lets a harness prove the + /// receive-path flush coalesces (N receives collapse to fewer batch writes). + session_batch_writes: AtomicU32, + /// When set, `put_sessions_batch` fails. Test hook: lets a harness prove the + /// send path aborts (and never hits the wire) when the ratchet advance + /// cannot be persisted. + fail_session_writes: AtomicBool, } impl InMemoryBackend { @@ -99,8 +106,20 @@ impl InMemoryBackend { Self { state: Mutex::new(InMemoryState::default()), next_device_id: AtomicI32::new(1), + session_batch_writes: AtomicU32::new(0), + fail_session_writes: AtomicBool::new(false), } } + + /// Number of `put_sessions_batch` calls since construction. + pub fn session_batch_write_count(&self) -> u32 { + self.session_batch_writes.load(Ordering::Relaxed) + } + + /// Make every subsequent `put_sessions_batch` fail (or stop failing). + pub fn set_fail_session_writes(&self, fail: bool) { + self.fail_session_writes.store(fail, Ordering::Relaxed); + } } impl Default for InMemoryBackend { @@ -148,6 +167,12 @@ impl SignalStore for InMemoryBackend { } async fn put_sessions_batch(&self, sessions: &[(Arc, Bytes)]) -> Result<()> { + self.session_batch_writes.fetch_add(1, Ordering::Relaxed); + if self.fail_session_writes.load(Ordering::Relaxed) { + return Err(crate::store::error::StoreError::Io(std::io::Error::other( + "put_sessions_batch failing (test hook)", + ))); + } let mut state = self.state.lock().await; state.sessions.reserve(sessions.len()); for (address, session) in sessions { From 63fe36b1c05bf6c5a840b0938c65777114f2c341 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:16:11 -0300 Subject: [PATCH 21/24] refactor(wacore): feature-gate InMemoryBackend test hooks behind test-util The put_sessions_batch call counter and fault switch were unconditional fields plus a fetch_add/load on every batch write, so every production build carried them and any in-memory benchmark saw instrumentation the SQLite path does not. Move both behind a new off-by-default test-util feature (and cfg(test) for wacore own tests); the e2e crate enables it. Normal builds regain a zero-cost, two-field struct with no fault- injection surface. --- tests/e2e/Cargo.toml | 2 +- wacore/Cargo.toml | 4 ++++ wacore/src/store/in_memory.rs | 35 ++++++++++++++++++++++++----------- 3 files changed, 29 insertions(+), 12 deletions(-) diff --git a/tests/e2e/Cargo.toml b/tests/e2e/Cargo.toml index 45b6fc303..c9521b1b3 100644 --- a/tests/e2e/Cargo.toml +++ b/tests/e2e/Cargo.toml @@ -17,7 +17,7 @@ dhat = { version = "0.3", optional = true } futures = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread", "sync", "time"] } uuid = { workspace = true, features = ["v4"] } -wacore = { path = "../../wacore" } +wacore = { path = "../../wacore", features = ["test-util"] } wacore-binary = { path = "../../wacore/binary" } whatsapp-rust = { path = "../..", default-features = false, features = [ "danger-skip-cert-chain-verify", diff --git a/wacore/Cargo.toml b/wacore/Cargo.toml index f9fd522e5..17126c5ba 100644 --- a/wacore/Cargo.toml +++ b/wacore/Cargo.toml @@ -19,6 +19,10 @@ ignored = ["getrandom"] default = ["simd"] simd = ["wacore-appstate/simd"] debug-snapshots = [] +# Expose the InMemoryBackend fault-injection / call-count hooks used by the e2e +# suite. Off by default so normal builds carry no extra fields or per-call +# bookkeeping. Enabled only from test crates. +test-util = [] # Optional observability: emit tracing spans/events. Off by default (no dep). tracing = ["dep:tracing"] # Optional metrics via the `metrics` facade. Off by default (no dep). diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index 745e50a40..dac9139da 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -6,7 +6,9 @@ use std::collections::HashMap; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicI32, AtomicU32, Ordering}; +#[cfg(any(test, feature = "test-util"))] +use std::sync::atomic::{AtomicBool, AtomicU32}; +use std::sync::atomic::{AtomicI32, Ordering}; use crate::appstate::hash::HashState; use crate::store::Device; @@ -91,12 +93,16 @@ const MAX_SENT_MESSAGES: usize = 4096; pub struct InMemoryBackend { state: Mutex, next_device_id: AtomicI32, - /// Count of `put_sessions_batch` calls. Test hook: lets a harness prove the - /// receive-path flush coalesces (N receives collapse to fewer batch writes). + /// Count of `put_sessions_batch` calls. Test hook (see `test-util`): lets a + /// harness prove receive-path flush coalescing (N receives collapse to fewer + /// batch writes). Gated so normal builds carry neither the field nor the + /// per-call bookkeeping. + #[cfg(any(test, feature = "test-util"))] session_batch_writes: AtomicU32, - /// When set, `put_sessions_batch` fails. Test hook: lets a harness prove the - /// send path aborts (and never hits the wire) when the ratchet advance - /// cannot be persisted. + /// When set, `put_sessions_batch` fails. Test hook (see `test-util`): lets a + /// harness prove the send path aborts (and never hits the wire) when the + /// ratchet advance cannot be persisted. + #[cfg(any(test, feature = "test-util"))] fail_session_writes: AtomicBool, } @@ -106,17 +112,21 @@ impl InMemoryBackend { Self { state: Mutex::new(InMemoryState::default()), next_device_id: AtomicI32::new(1), + #[cfg(any(test, feature = "test-util"))] session_batch_writes: AtomicU32::new(0), + #[cfg(any(test, feature = "test-util"))] fail_session_writes: AtomicBool::new(false), } } /// Number of `put_sessions_batch` calls since construction. + #[cfg(any(test, feature = "test-util"))] pub fn session_batch_write_count(&self) -> u32 { self.session_batch_writes.load(Ordering::Relaxed) } /// Make every subsequent `put_sessions_batch` fail (or stop failing). + #[cfg(any(test, feature = "test-util"))] pub fn set_fail_session_writes(&self, fail: bool) { self.fail_session_writes.store(fail, Ordering::Relaxed); } @@ -167,11 +177,14 @@ impl SignalStore for InMemoryBackend { } async fn put_sessions_batch(&self, sessions: &[(Arc, Bytes)]) -> Result<()> { - self.session_batch_writes.fetch_add(1, Ordering::Relaxed); - if self.fail_session_writes.load(Ordering::Relaxed) { - return Err(crate::store::error::StoreError::Io(std::io::Error::other( - "put_sessions_batch failing (test hook)", - ))); + #[cfg(any(test, feature = "test-util"))] + { + self.session_batch_writes.fetch_add(1, Ordering::Relaxed); + if self.fail_session_writes.load(Ordering::Relaxed) { + return Err(crate::store::error::StoreError::Io(std::io::Error::other( + "put_sessions_batch failing (test hook)", + ))); + } } let mut state = self.state.lock().await; state.sessions.reserve(sessions.len()); From 1b476b7df6921581e430cbc9ad1e65fbb0aa08e8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:16:11 -0300 Subject: [PATCH 22/24] test,docs(signal): prove the second flush window and the pre-wire abort The stale-schedule starvation test only checked that the pending DIRTY survived; the first (unblocked) flush would persist the second session on its own, so it never proved a second window ran. Require the attempt counter to reach 2 after release. The pre-wire abort e2e now registers a sent-node waiter and asserts it stays pending after the failed send: no message node was marshaled, so send_node (and the wire) was never reached, replacing the sole reliance on a negative timed assertion. Also correct the flush_signal_cache comment (no longer per-message, and the backend is generic, not SQLite). --- src/client/adapters.rs | 4 ++-- src/signal_flush.rs | 14 +++++++++++++- tests/e2e/tests/session_reuse.rs | 12 +++++++++++- 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/src/client/adapters.rs b/src/client/adapters.rs index dce5a9ddc..380e9ae54 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -95,8 +95,8 @@ impl Client { /// send path (synchronously, pre-wire), the receive-path coalescer, and the /// drain/retry/teardown recovery paths — not unconditionally per message. pub(crate) async fn flush_signal_cache(&self) -> Result<(), anyhow::Error> { - // Hold no device guard across the flush: this per-message batched SQLite - // write would otherwise block every concurrent Device write for its duration. + // Clone the backend before awaiting so a slow write cannot retain the + // device guard and stall every concurrent Device write. let backend = self .persistence_manager .get_device_snapshot() diff --git a/src/signal_flush.rs b/src/signal_flush.rs index 1b7f14e96..9e8975d4c 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -531,10 +531,22 @@ mod tests { "a stale schedule must not clear the pending DIRTY" ); - // Release: B sees DIRTY and runs another window that persists it. + // Release: B's first (blocked) attempt completes, then — because DIRTY is + // still set — it runs a SECOND attempt. Proving the second window ran means + // requiring the attempt counter to reach 2, not just that `second` landed + // (the first attempt alone would persist it, since it was dirtied before + // the blocked flush actually executed). client .signal_flush_test_block .store(false, Ordering::Release); + let deadline = wacore::time::Instant::now() + Duration::from_secs(2); + while client.signal_flush_test_in_attempt.load(Ordering::Acquire) < 2 { + assert!( + wacore::time::Instant::now() < deadline, + "the pending DIRTY must trigger a second flush attempt" + ); + tokio::time::sleep(Duration::from_millis(2)).await; + } wait_for_backend_session(&client, &second).await; } } diff --git a/tests/e2e/tests/session_reuse.rs b/tests/e2e/tests/session_reuse.rs index 67038f52b..ae9547436 100644 --- a/tests/e2e/tests/session_reuse.rs +++ b/tests/e2e/tests/session_reuse.rs @@ -142,6 +142,10 @@ async fn test_send_aborts_before_wire_when_persist_fails() -> anyhow::Result<()> // Persisting the outbound advance now fails. client_a.backend.set_fail_session_writes(true); let writes_before = client_a.backend.session_batch_write_count(); + // Resolves the instant any `message` node is marshaled for the wire, which + // `send_node` does BEFORE `send_raw_bytes`; if the send aborts before that, + // it never fires. + let mut sent_waiter = client_a.next_sent_message_waiter(); let result = client_a .client .send_message( @@ -159,8 +163,14 @@ async fn test_send_aborts_before_wire_when_persist_fails() -> anyhow::Result<()> client_a.backend.session_batch_write_count() > writes_before, "the send path must attempt to persist the ratchet advance before the wire" ); + // Deterministic: no `message` node was ever marshaled, so `send_node` (and + // thus the wire) was never reached. `Ok(None)` == pending, sender still alive. + assert!( + matches!(sent_waiter.try_recv(), Ok(None)), + "the send must abort before send_node marshals the stanza for the wire" + ); - // The stanza never went out: B must not see it. + // End-to-end corroboration: the stanza never went out, so B never sees it. client_b .assert_no_event( 3, From 43561b8643004b0f40094ed30ad77e22ec2d4357 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:57:32 -0300 Subject: [PATCH 23/24] fix(signal): gate coalesced flush writes against teardown cache settle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The generation-scoped atomic ordered only signal_flush_state, not the backend writes. A worker that passed its pre-flush generation check could be preempted (e.g. blocked on the sessions lock teardown holds), and resume its flush after teardown settled the cache and the next connection's drain dirtied it — persisting that drain's rowless ratchet advances out of band, the silent-loss class the commit batcher prevents. Add a signal_flush_lifecycle mutex: the worker holds it only across the flush (never across sleep/backoff) and re-checks the generation under it; teardown holds it across the whole cache settle. So a worker either wins the gate and flushes its own generation before any settle, or acquires it after teardown and stands down on the generation check. Lock order is always gate -> permit/sessions-lock, so no inversion. Two deterministic tests: the gate blocks a worker flush while held, and a worker reaching the gate after the bump stands down without writing. --- src/client.rs | 10 +++++ src/client/lifecycle.rs | 10 +++++ src/signal_flush.rs | 97 +++++++++++++++++++++++++++++++++++------ 3 files changed, 104 insertions(+), 13 deletions(-) diff --git a/src/client.rs b/src/client.rs index 731ec5e29..f2c0044c3 100644 --- a/src/client.rs +++ b/src/client.rs @@ -840,6 +840,16 @@ pub struct Client { /// Single-flight state for the coalesced Signal-cache flush worker: /// `(connection_generation << 2) | RUNNING/DIRTY bits` (see `signal_flush.rs`). pub(crate) signal_flush_state: AtomicU64, + /// Barrier between a coalesced-flush worker's backend write and teardown's + /// Signal-cache settle. The generation-scoped atomic only orders + /// `signal_flush_state`, not the writes themselves: a worker that passed its + /// pre-flush generation check could still be mid-flush when teardown settles + /// the cache and the next connection's drain dirties it, persisting rowless + /// advances out of band. The worker holds this only across the flush (never + /// across sleep/backoff) and re-checks the generation under it; teardown + /// holds it around the settle. Lock order is always this-gate → processing + /// permit / sessions lock, so no inversion. + pub(crate) signal_flush_lifecycle: async_lock::Mutex<()>, /// Injected failures for the coalesced flush (consumed one per attempt), /// so tests can exercise the retry/backoff path deterministically. #[cfg(test)] diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 3bc93b14e..f958d03f2 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -241,6 +241,7 @@ impl Client { passkey_state: Arc::new(Mutex::new(crate::passkey::flow::PasskeyFlowState::default())), passkey_opening: AtomicBool::new(false), signal_flush_state: AtomicU64::new(0), + signal_flush_lifecycle: async_lock::Mutex::new(()), #[cfg(test)] signal_flush_test_failures: AtomicU32::new(0), #[cfg(test)] @@ -840,6 +841,13 @@ impl Client { // 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(). + // + // Hold the coalesced-flush barrier across the whole settle: a stale flush + // worker that already passed its generation check must not interleave a + // backend write between our commit and the next connection's drain, or it + // could persist that drain's rowless advances. The worker re-checks the + // generation (bumped above) once it gets the gate, so it stands down. + let flush_gate = self.signal_flush_lifecycle.lock().await; if let Some(client) = self.self_weak.get().and_then(|w| w.upgrade()) { client .teardown_inbound_commits_bounded(std::time::Duration::from_secs(5)) @@ -874,6 +882,8 @@ impl Client { ); self.signal_cache.clear().await; } + // Cache is settled and any dropped entries cleared; a worker may run again. + drop(flush_gate); self.offline_batch.reset(); self.offline_sync_metrics .active diff --git a/src/signal_flush.rs b/src/signal_flush.rs index 9e8975d4c..f50c79465 100644 --- a/src/signal_flush.rs +++ b/src/signal_flush.rs @@ -18,12 +18,16 @@ //! keep their own synchronous flushes: those gate acks, receipts or //! follow-up reads on durability and are not routed here. //! -//! Single-flight scheduler: at most one worker exists at a time. The first +//! Single-flight scheduler: at most one worker per generation. The first //! request arms it; requests that arrive while it runs only mark it dirty -//! (never spawn a second worker), and it re-runs one more window if so. A -//! failing flush is retried by that same worker with exponential backoff, so a -//! backend outage cannot be reset to the base delay by concurrent traffic, and -//! a generation change (reconnect/teardown) stands the worker down. +//! (never spawn a second worker for that generation), and it re-runs one more +//! window if so. A failing flush is retried by that same worker with +//! exponential backoff, so a backend outage cannot be reset to the base delay +//! by concurrent traffic. Across a reconnect a stale worker and the new +//! generation's worker can briefly coexist; the stale one stands down at its +//! generation check. Actual backend writes are additionally serialized against +//! teardown's cache settle by the `signal_flush_lifecycle` gate, so at most one +//! flush is ever mid-write — see `Client::signal_flush_lifecycle`. use std::sync::atomic::Ordering; @@ -133,14 +137,23 @@ impl Client { let Some(client) = weak.upgrade() else { return; }; - // A reconnect handed ownership to a new generation: stand - // down before flushing. The generation-scoped exit CAS would - // reject our mutation anyway, but skipping the flush avoids a - // stale write after teardown settled the cache. - if client.connection_generation.load(Ordering::Acquire) != generation { - return; - } - match client.coalesced_flush_attempt().await { + // Take the lifecycle barrier for the flush, then re-check the + // generation under it. Teardown holds this same gate while it + // settles the cache, so either we got here first (and flush + // this generation's state before any settle) or teardown ran + // and the generation moved (stand down). Without the gate the + // bare generation check races: we could pass it, get preempted, + // and resume the flush after teardown settled and the next + // connection's drain dirtied the cache — persisting rowless + // advances out of band. + let flush_result = { + let _gate = client.signal_flush_lifecycle.lock().await; + if client.connection_generation.load(Ordering::Acquire) != generation { + return; + } + client.coalesced_flush_attempt().await + }; + match flush_result { Ok(()) => { backoff = SIGNAL_FLUSH_WINDOW; // Settle ownership under a CAS scoped to our @@ -549,4 +562,62 @@ mod tests { } wait_for_backend_session(&client, &second).await; } + + /// The lifecycle gate is a hard barrier: while it is held (as teardown holds + /// it across the cache settle), a worker cannot write to the backend; once + /// released, it flushes. This is what stops a stale flush from interleaving + /// a backend write into teardown's settle. + #[tokio::test] + async fn lifecycle_gate_blocks_the_worker_flush_until_released() { + let client = crate::test_utils::create_test_client().await; + let s1 = dirty_session(&client, "15550009001"); + + // Hold the gate as teardown's settle would. + let gate = client.signal_flush_lifecycle.lock().await; + + client.schedule_signal_flush_live(); + // The worker arms, wakes after the window, and blocks acquiring the gate. + // It genuinely cannot proceed, so the dirty session never reaches the + // backend no matter how long we wait here. + tokio::time::sleep(SIGNAL_FLUSH_WINDOW * 3).await; + assert!( + backend_session(&client, &s1).await.is_none(), + "the worker must not flush while the lifecycle gate is held" + ); + + // Release as teardown would after settling; the worker proceeds. + drop(gate); + wait_for_backend_session(&client, &s1).await; + } + + /// A worker that only reaches the gate after teardown bumped the generation + /// re-checks under the gate and stands down: it never enters the flush, so + /// the attempt counter stays 0 and the session is not persisted out of band. + #[tokio::test] + async fn worker_stands_down_if_it_reaches_the_gate_after_the_bump() { + let client = crate::test_utils::create_test_client().await; + let s1 = dirty_session(&client, "15550009101"); + + // Teardown-first: hold the gate, arm a worker for the current generation, + // then bump under the gate (mirrors cleanup_connection_state bumping the + // generation and holding the gate across the settle). + let gate = client.signal_flush_lifecycle.lock().await; + client.schedule_signal_flush_live(); + client.connection_generation.fetch_add(1, Ordering::SeqCst); + + // Let the worker wake and block on the gate, then release it. + tokio::time::sleep(SIGNAL_FLUSH_WINDOW * 3).await; + drop(gate); + tokio::time::sleep(SIGNAL_FLUSH_WINDOW * 3).await; + + assert_eq!( + client.signal_flush_test_in_attempt.load(Ordering::Acquire), + 0, + "a worker reaching the gate after the bump must stand down before flushing" + ); + assert!( + backend_session(&client, &s1).await.is_none(), + "the stale worker must not persist S1 after teardown" + ); + } } From d6b29977e54927609547de3b191d5b36c7159496 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= <55464917+jlucaso1@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:57:32 -0300 Subject: [PATCH 24/24] docs(e2e): correct the sent-node waiter ordering comment send_node resolves the waiter before marshaling the node, not "when marshaled"; a pending waiter proves a pre-wire abort. --- tests/e2e/tests/session_reuse.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/e2e/tests/session_reuse.rs b/tests/e2e/tests/session_reuse.rs index ae9547436..9dc29c635 100644 --- a/tests/e2e/tests/session_reuse.rs +++ b/tests/e2e/tests/session_reuse.rs @@ -142,9 +142,8 @@ async fn test_send_aborts_before_wire_when_persist_fails() -> anyhow::Result<()> // Persisting the outbound advance now fails. client_a.backend.set_fail_session_writes(true); let writes_before = client_a.backend.session_batch_write_count(); - // Resolves the instant any `message` node is marshaled for the wire, which - // `send_node` does BEFORE `send_raw_bytes`; if the send aborts before that, - // it never fires. + // `send_node` resolves this before marshaling the node, so a still-pending + // waiter proves the send aborted before reaching the wire. let mut sent_waiter = client_a.next_sent_message_waiter(); let result = client_a .client