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

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<crate::store::signal_cache::SignalStoreCache>,
 
     /// 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<std::sync::Weak<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)]
+    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<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

@@ -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

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

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

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 {