perf(signal): coalesce receive flushes and persist outbound state pre-wire

src/client.rs

@@ -836,6 +836,14 @@ pub struct Client {
     /// Initialized after `Arc::new(this)` in the constructor.
     pub(crate) self_weak: std::sync::OnceLock<std::sync::Weak<Client>>,
 
+    /// 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<Client>` ref count instead of the Bot wrapper's. Set once by
     /// `Bot::build`; on Client drop (last Arc), the handle drops and the saver

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

src/client/lifecycle.rs

@@ -240,6 +240,9 @@ 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),
+            #[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(),

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;

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;
         }
     }
 

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

src/signal_flush.rs

@@ -0,0 +1,229 @@
+//! 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_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).
+
+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.
+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
+    /// 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;
+        }
+        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 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 fire
+                    // 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) {
+                        return;
+                    }
+                }
+            }))
+            .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)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+    use std::time::Duration;
+
+    use wacore::libsignal::protocol::{ProtocolAddress, SessionRecord};
+
+    async fn backend_session(
+        client: &Arc<crate::client::Client>,
+        addr: &ProtocolAddress,
+    ) -> Option<bytes::Bytes> {
+        client
+            .persistence_manager
+            .backend()
+            .get_session(addr.as_str())
+            .await
+            .expect("backend read")
+    }
+
+    fn dirty_session(client: &Arc<crate::client::Client>, 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<crate::client::Client>, 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"
+        );
+    }
+
+    /// 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]
+    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;
+    }
+}

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).