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>>,
 
+    /// 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)]
+    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

@@ -68,6 +68,27 @@ impl Client {
             .ok_or(ClientError::NotConnected)
     }
 
+    /// 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
+    /// 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.
     /// Called after each message is decrypted or after encryption operations.
     pub(crate) async fn flush_signal_cache(&self) -> Result<(), anyhow::Error> {
@@ -84,21 +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.
-    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,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_state: AtomicU32::new(0),
+            #[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(),
@@ -764,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).

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

src/send/mod.rs

@@ -1552,7 +1552,12 @@ 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
+        // 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;