perf(signal): lease outbound counters in batches instead of flushing every send

src/client/adapters.rs

@@ -72,7 +72,8 @@ impl Client {
     /// 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
+    /// instead of writing through, and lease-covered sends do the same (only a
+    /// send that raises its counter lease flushes 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
@@ -160,6 +161,26 @@ impl Client {
         self.flush_signal_cache().await
     }
 
+    /// Pre-wire durability gate for the send path. Flushes synchronously only
+    /// when an outbound crypto advance actually demands it — a raised session
+    /// counter lease or a sender-key chain advance not yet persisted (see
+    /// `SignalStoreCache::needs_pre_wire_flush`). Otherwise the dirty state is
+    /// covered by an existing durable lease, so it only needs to land
+    /// eventually: it rides the same coalesced write-behind as the receive
+    /// path instead of paying a serialize + storage transaction per message.
+    /// A failure must abort the send — transmitting a ciphertext whose lease
+    /// could not be persisted reintroduces the counter-reuse window.
+    pub(crate) async fn persist_signal_state_pre_wire(&self) -> Result<(), anyhow::Error> {
+        if self.signal_cache.needs_pre_wire_flush().await {
+            return self.flush_signal_cache_batch_safe().await;
+        }
+        self.schedule_signal_flush(
+            self.connection_generation
+                .load(std::sync::atomic::Ordering::Acquire),
+        );
+        Ok(())
+    }
+
     /// [`flush_signal_cache_batch_safe`](Self::flush_signal_cache_batch_safe)
     /// with error logging instead of propagation.
     pub(crate) async fn flush_signal_cache_batch_safe_logged(

src/features/signal.rs

@@ -70,7 +70,9 @@ impl<'a> Signal<'a> {
         .await?;
 
         drop(_guard);
-        self.client.flush_signal_cache_batch_safe().await?;
+        // Same pre-wire gate as the send path: the caller transmits these
+        // bytes, so a raised lease must be durable before they leave here.
+        self.client.persist_signal_state_pre_wire().await?;
 
         let (_, is_prekey, bytes) = wacore::send::extract_ciphertext(encrypted)
             .ok_or_else(|| SignalError::Unsupported("unexpected ciphertext variant".into()))?;

src/send/mod.rs

@@ -962,9 +962,9 @@ 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?;
+        // Gate the stanza on the sender-key ratchet advance being durable
+        // (same rule as the DM/group send path); a failure aborts the send.
+        self.persist_signal_state_pre_wire().await?;
 
         let ack = if let Some(phash) = stanza
             .attrs()
@@ -1482,14 +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?;
+        // The outbound advance must be durable BEFORE the stanza hits the wire:
+        // reusing an outbound counter reuses its message key + IV. Counters are
+        // leased in batches (see `SessionRecord::reserve_sender_chain_counters`),
+        // so most sends are already covered by a durable lease and only
+        // schedule the coalesced write-behind; a send that raised the lease (or
+        // advanced a sender-key chain) flushes synchronously, and a persistence
+        // failure must abort the send rather than transmit an advance we
+        // couldn't save.
+        self.persist_signal_state_pre_wire().await?;
 
         let ack = if let Some(phash) = dm_phash
             && let Some(msg_id) = stanza_to_send

src/signal_flush.rs

@@ -7,13 +7,16 @@
 //! window.
 //!
 //! Scope and durability model (deliberate, bounded):
-//! - Only the receive path coalesces. A lost receive-side advance re-derives
+//! - 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.
+//!   a crash inside the window is recoverable. The SEND path coalesces too
+//!   whenever its advance is covered by a durable counter lease (see
+//!   `SessionRecord::reserve_sender_chain_counters` and
+//!   `Client::persist_signal_state_pre_wire`); only a send that raises the
+//!   lease — or advances a group sender-key chain — still flushes
+//!   synchronously before the wire, because reusing an outbound counter would
+//!   reuse its message key + IV.
 //! - 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.
@@ -284,6 +287,36 @@ mod tests {
         }
     }
 
+    /// The send-path gate: a raised counter lease flushes synchronously
+    /// before returning; lease-covered dirty state only schedules the
+    /// coalesced worker and still lands within its window.
+    #[tokio::test]
+    async fn pre_wire_gate_flushes_leases_synchronously_and_coalesces_the_rest() {
+        let client = crate::test_utils::create_test_client().await;
+
+        // Lease-covered advance (no raised reservation): no synchronous write —
+        // the session is still absent from the backend when the gate returns.
+        let covered = dirty_session(&client, "15550003001");
+        client.persist_signal_state_pre_wire().await.unwrap();
+        assert!(
+            backend_session(&client, &covered).await.is_none(),
+            "a covered advance must not flush synchronously"
+        );
+        // ...but it rides the coalescer and still becomes durable.
+        wait_for_backend_session(&client, &covered).await;
+
+        // A raised lease must be durable when the call returns.
+        let addr = ProtocolAddress::new("15550003002".to_string(), 1.into());
+        let mut record = SessionRecord::new_fresh();
+        record.reserve_sender_chain_counters(0);
+        assert!(client.signal_cache.try_put_session(&addr, record).is_ok());
+        client.persist_signal_state_pre_wire().await.unwrap();
+        assert!(
+            backend_session(&client, &addr).await.is_some(),
+            "the leased session must be durable when the gate returns"
+        );
+    }
+
     /// A burst of requests rides one armed worker and persists every dirty
     /// entry written before it.
     #[tokio::test]

tests/e2e/tests/session_reuse.rs

@@ -61,25 +61,31 @@ async fn durable_sender_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.
+/// The durable snapshot must always be able to resume PAST every counter that
+/// may have hit the wire: reusing an outbound counter reuses its message key +
+/// IV. Counters are leased in batches (`SENDER_CHAIN_RESERVATION_BATCH`): the
+/// send that raises the lease flushes synchronously before the wire, and every
+/// lease-covered send may defer its advance to the coalesced write-behind
+/// because `SessionRecord::deserialize` fast-forwards the reloaded chain past
+/// the whole lease. This reads the backend IMMEDIATELY after `send_message`
+/// (no delivery wait, no settle): the resume position — what a crash restore
+/// would actually use — must already cover every counter spent so far.
 #[tokio::test]
-async fn test_outbound_ratchet_is_durable_when_send_returns() -> anyhow::Result<()> {
+async fn test_durable_resume_position_always_covers_spent_counters() -> 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 client_a = TestClient::connect("e2e_sig_durable_a").await?;
+    let client_b = TestClient::connect("e2e_sig_durable_b").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?;
+    // The first send raises the lease, so its flush is synchronous: the
+    // durable resume position must already be past counter 0 the moment
+    // send_message returns, with no settle.
+    client_a
+        .client
+        .send_message(jid_b.clone(), e2e_tests::text_msg("establish"))
+        .await?;
 
     let backend_a = client_a.client.persistence_manager().backend();
     let read_index = async |user: &str, server: &str| {
@@ -92,26 +98,35 @@ async fn test_outbound_ratchet_is_durable_when_send_returns() -> anyhow::Result<
         _ => (jid_b.user.clone(), "c.us"),
     };
 
+    // `durable_sender_chain_index` deserializes the stored record, which
+    // applies the crash-restore fast-forward: this IS the resume position.
     let mut last = read_index(&user, server)
         .await?
-        .expect("an outbound session must exist after the roundtrip");
+        .expect("the lease raise must persist the session before the wire");
+    let mut spent = 1u32; // counter 0 went out with "establish"
+    assert!(
+        last >= spent,
+        "resume position {last} must cover the {spent} spent counter(s)"
+    );
 
-    // 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.
+    // Lease-covered sends may leave the durable snapshot trailing (that is
+    // the optimization), but the resume position must never fall behind the
+    // wire and never regress.
     for i in 0..3 {
         client_a
             .client
             .send_message(jid_b.clone(), e2e_tests::text_msg(&format!("m{i}")))
             .await?;
+        spent += 1;
         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"
+            now >= spent,
+            "send #{i}: resume position {now} fell behind the {spent} spent counter(s); \
+             a crash here would re-derive a (key, IV) pair"
         );
+        assert!(now >= last, "resume position must never regress");
         last = now;
     }
 
@@ -120,26 +135,21 @@ 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
+/// A send that RAISES the counter lease gates the wire on persisting it: if
+/// the flush fails, `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.
+/// lease only in memory and a reload would re-derive that counter's key + IV.
+/// The first send on a fresh session always raises the lease, so the failure
+/// is injected before it.
 #[tokio::test]
-async fn test_send_aborts_before_wire_when_persist_fails() -> anyhow::Result<()> {
+async fn test_send_aborts_before_wire_when_lease_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 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.
+    // Persisting the outbound lease 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
@@ -154,13 +164,13 @@ async fn test_send_aborts_before_wire_when_persist_fails() -> anyhow::Result<()>
         .await;
     assert!(
         result.is_err(),
-        "send must fail when the ratchet advance cannot be persisted, got {result:?}"
+        "send must fail when the raised lease 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.
+    // The send reached the (failing) persistence step, proving the lease 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 send path must attempt to persist the raised lease 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.
@@ -198,6 +208,43 @@ async fn test_send_aborts_before_wire_when_persist_fails() -> anyhow::Result<()>
     Ok(())
 }
 
+/// The counterpart of the abort test: a send COVERED by an already-durable
+/// lease does not depend on this flush for safety — a crash would reload the
+/// durable snapshot and fast-forward past the whole lease, so its counter can
+/// never be re-derived. Such a send must therefore succeed even while the
+/// backend is refusing session writes (the advance lands later via the
+/// coalescer's retry), instead of turning a storage hiccup into message loss.
+#[tokio::test]
+async fn test_lease_covered_send_survives_persist_failure() -> anyhow::Result<()> {
+    let _ = env_logger::builder().is_test(true).try_init();
+
+    let mut client_a = TestClient::connect("e2e_sig_covered_a").await?;
+    let mut client_b = TestClient::connect("e2e_sig_covered_b").await?;
+    let jid_a = client_a.jid().await;
+    let jid_b = client_b.jid().await;
+
+    // Establish both ways: the lease raise happens (and is persisted) here.
+    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?;
+
+    // Storage starts refusing session writes; the next send is lease-covered,
+    // so it must still deliver.
+    client_a.backend.set_fail_session_writes(true);
+    send_and_expect_text(
+        &client_a.client,
+        &mut client_b,
+        &jid_b,
+        "covered by the lease",
+        30,
+    )
+    .await?;
+
+    client_a.backend.set_fail_session_writes(false);
+    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<()> {

wacore/libsignal/src/protocol/consts.rs

@@ -24,3 +24,19 @@ pub const MAX_SENDER_KEY_STATES: usize = 5;
 /// Eviction only triggers when buffer exceeds MAX_MESSAGE_KEYS + PRUNE_THRESHOLD,
 /// reducing O(n) drain() calls from every insert to once every PRUNE_THRESHOLD inserts.
 pub const MESSAGE_KEY_PRUNE_THRESHOLD: usize = 50;
+
+/// Sender-chain counters leased per durable reservation (see
+/// `SessionRecord::reserve_sender_chain_counters`). Message keys and IVs are
+/// derived deterministically from the counter, so an outbound counter must
+/// never repeat across a crash; instead of persisting every advance before it
+/// hits the wire, the record durably reserves this many counters ahead and a
+/// reloaded snapshot fast-forwards past them. Bounds both the sync-flush
+/// amortization (one per this many sends) and the worst-case counter gap a
+/// receiver sees after a crash — keep it well under MAX_FORWARD_JUMPS.
+pub const SENDER_CHAIN_RESERVATION_BATCH: u32 = 64;
+
+/// Upper bound for the reservation fast-forward on load. A legitimate lease
+/// gap is < SENDER_CHAIN_RESERVATION_BATCH; anything past this ceiling means
+/// a corrupt record, and refusing it keeps a bogus reserved index from
+/// turning the load into an unbounded KDF loop.
+pub const MAX_RESERVATION_FAST_FORWARD: u32 = MAX_FORWARD_JUMPS as u32;

wacore/libsignal/src/protocol/group_cipher.rs

@@ -111,6 +111,11 @@ pub async fn group_encrypt<R: Rng + CryptoRng>(
 
     sender_key_state.set_sender_chain_key(next_sender_chain_key);
 
+    // Outbound advance: this iteration's (key, IV) must never be re-derivable,
+    // so the store must gate the ciphertext on durability. Decrypt-side
+    // advances stay ungated (they re-derive forward).
+    record.mark_wire_gated();
+
     sender_key_store
         .store_sender_key(sender_key_name, record)
         .await?;