diff --git a/src/client.rs b/src/client.rs index b9f0b6c9a..f2c0044c3 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). @@ -836,6 +837,32 @@ 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: + /// `(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)] + 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 /// `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..380e9ae54 100644 --- a/src/client/adapters.rs +++ b/src/client/adapters.rs @@ -68,11 +68,35 @@ impl Client { .ok_or(ClientError::NotConnected) } - /// Flush the in-memory signal cache to the database backend. - /// Called after each message is decrypted or after encryption operations. + /// Force any pending write-behind Signal cache state to the backend, + /// returning once the flush completes (or fails). + /// + /// 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. + /// + /// 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 + } + + /// 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. + // 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() @@ -84,21 +108,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. - 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. /// + /// [`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 diff --git a/src/client/lifecycle.rs b/src/client/lifecycle.rs index 727185032..f958d03f2 100644 --- a/src/client/lifecycle.rs +++ b/src/client/lifecycle.rs @@ -240,6 +240,14 @@ 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: AtomicU64::new(0), + signal_flush_lifecycle: async_lock::Mutex::new(()), + #[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(), @@ -764,6 +772,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); + // 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). @@ -830,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)) @@ -864,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/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..31c0c82f5 100644 --- a/src/message/receive.rs +++ b/src/message/receive.rs @@ -561,16 +561,14 @@ 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: 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.flush_signal_cache_logged("message", Some(&info.id)) - .await; + self.schedule_signal_flush(lane_generation); } } diff --git a/src/send/mod.rs b/src/send/mod.rs index bd3e89601..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,10 +1562,6 @@ 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; - // 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 { diff --git a/src/signal_flush.rs b/src/signal_flush.rs new file mode 100644 index 000000000..f50c79465 --- /dev/null +++ b/src/signal_flush.rs @@ -0,0 +1,623 @@ +//! Coalesced write-behind for the inbound Signal cache flush. +//! +//! 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. +//! +//! 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 per generation. The first +//! request arms it; requests that arrive while it runs only mark it dirty +//! (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; + +use crate::client::Client; + +/// 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); + +/// 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); + +// 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: 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 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 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 an older generation's leftover state → take over + // for the live generation. + if self + .signal_flush_state + .compare_exchange_weak( + cur, + pack_flush_state(lane_generation, FLUSH_RUNNING), + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + self.spawn_signal_flush_worker(lane_generation); + return; + } + } else { + // A worker for this generation 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, generation: u64) { + let Some(weak) = self.self_weak.get() else { + // Constructor edge: no Arc identity to hold from a timer task, so + // 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(); + self.runtime + .spawn(Box::pin(async move { + let mut backoff = SIGNAL_FLUSH_WINDOW; + loop { + // 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; + }; + // 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 + // 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; + } + } + } + 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. + backoff = (backoff * 2).min(SIGNAL_FLUSH_RETRY_CEILING); + log::error!( + "Coalesced signal flush failed; retrying in {backoff:?}: {e:?}" + ); + } + } + } + })) + .detach(); + } + + /// 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)] + { + 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 + .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_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)] +mod tests { + use super::*; + 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; + } + } + + /// 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; + + let mut addrs = Vec::new(); + for i in 0..10 { + addrs.push(dirty_session(&client, &format!("155500011{i:02}"))); + client.schedule_signal_flush_live(); + } + assert!(client.signal_flush_worker_alive(), "burst arms one worker"); + + for addr in &addrs { + wait_for_backend_session(&client, addr).await; + } + // 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_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_live(); + wait_for_backend_session(&client, &second).await; + } + + /// 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 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_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 + // 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_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"); + 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_live(); + + // 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 consumed every injected failure" + ); + } + + /// 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_hands_off_without_clobbering() { + let client = crate::test_utils::create_test_client().await; + // 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); + dirty_session(&client, "15550006001"); + client.schedule_signal_flush_live(); + assert!(client.signal_flush_worker_alive()); + + // Bump the generation as teardown does — no explicit scheduler reset. + client.connection_generation.fetch_add(1, Ordering::SeqCst); + // 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_live(); + 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_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 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_live(); + 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; + } + + /// 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'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; + } + + /// 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" + ); + } +} 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/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index 53c79395f..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) @@ -159,101 +163,37 @@ impl TestClient { let run_handle = bot.spawn(); - // Wait for PairSuccess + Connected. - // - // 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(())) => {} + // 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)) + .await + { + client.disconnect().await; + drop(run_handle); + return Err(anyhow::anyhow!( + "client never became ready after pairing: {e}" + )); } + // Drain the initial startup sync (offline messages + history) so tests + // 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(15)) + .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 startup sync to become idle: {e}" + "startup sync did not settle before the connect deadline: {e}" )); } @@ -261,6 +201,7 @@ impl TestClient { client, event_rx, run_handle, + backend, }) } diff --git a/tests/e2e/tests/lid_sessions.rs b/tests/e2e/tests/lid_sessions.rs index 57023a742..6c2bbdd96 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,9 @@ 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 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!( backend_a.get_session(&lid_addr).await?.is_some(), @@ -339,7 +345,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 +377,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, @@ -517,6 +525,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 + // 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. // 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). @@ -615,6 +630,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!( @@ -668,6 +685,8 @@ async fn test_pn_migration_is_durable_across_followup_messages() -> anyhow::Resu .await?; } + // 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(), "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..9dc29c635 100644 --- a/tests/e2e/tests/session_reuse.rs +++ b/tests/e2e/tests/session_reuse.rs @@ -37,6 +37,167 @@ 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 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_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?; + 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"); + + // 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 { + 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"); + 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(()) +} + +/// 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(); + // `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 + .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" + ); + // 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" + ); + + // End-to-end corroboration: the stanza never went out, so B never sees 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<()> { @@ -140,15 +301,10 @@ 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 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 let backend = client_a.client.persistence_manager().backend(); @@ -236,6 +392,10 @@ async fn test_session_persistence() -> anyhow::Result<()> { .await?; info!("First message sent: {msg_id_1}"); + // 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. let mut post_send = scan_sessions(&*backend, &jid_b.user, "c.us").await?; 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 9f7f2e9e4..dac9139da 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -6,6 +6,8 @@ use std::collections::HashMap; use std::sync::Arc; +#[cfg(any(test, feature = "test-util"))] +use std::sync::atomic::{AtomicBool, AtomicU32}; use std::sync::atomic::{AtomicI32, Ordering}; use crate::appstate::hash::HashState; @@ -91,6 +93,17 @@ const MAX_SENT_MESSAGES: usize = 4096; pub struct InMemoryBackend { state: Mutex, next_device_id: AtomicI32, + /// 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 (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, } impl InMemoryBackend { @@ -99,8 +112,24 @@ 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); + } } impl Default for InMemoryBackend { @@ -148,6 +177,15 @@ impl SignalStore for InMemoryBackend { } async fn put_sessions_batch(&self, sessions: &[(Arc, Bytes)]) -> Result<()> { + #[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()); for (address, session) in sessions {