perf(send): establish sessions before taking the sender-key chain lock

src/store/signal_adapter.rs

@@ -309,6 +309,13 @@ impl wacore::libsignal::protocol::SenderKeyStore for SenderKeyAdapter {
     ) -> std::sync::Arc<async_lock::Mutex<()>> {
         self.0.cache.sender_key_lock(sender_key_name).await
     }
+
+    async fn session_setup_lock(
+        &self,
+        sender_key_name: &SenderKeyName,
+    ) -> std::sync::Arc<async_lock::Mutex<()>> {
+        self.0.cache.session_setup_lock(sender_key_name).await
+    }
 }
 
 #[cfg(test)]

wacore/libsignal/src/protocol/storage/traits.rs

@@ -172,6 +172,19 @@ pub trait SenderKeyStore: ThreadSafe {
     ) -> std::sync::Arc<async_lock::Mutex<()>> {
         std::sync::Arc::new(async_lock::Mutex::new(()))
     }
+
+    /// Serializes per-group session setup (prekey fetch + X3DH) so concurrent
+    /// cold sends to the same group can't race writes to the same per-device
+    /// sessions. Distinct from [`sender_key_lock`](Self::sender_key_lock):
+    /// held only across setup — never the chain-advancing critical section —
+    /// so it may span network I/O without blocking warm sends. Default is
+    /// uncontended; stores over shared state override it.
+    async fn session_setup_lock(
+        &self,
+        _sender_key_name: &SenderKeyName,
+    ) -> std::sync::Arc<async_lock::Mutex<()>> {
+        std::sync::Arc::new(async_lock::Mutex::new(()))
+    }
 }
 
 /// Mixes in all the store interfaces defined in this module.

wacore/src/send/encrypt.rs

@@ -286,9 +286,14 @@ fn push_encrypt_result(
 /// by [`ENCRYPT_FANOUT_CONCURRENCY`]. Each task clones the store handles
 /// (Arc bumps under the hood); the shared cache provides interior mutability.
 ///
+/// Composition of [`ensure_sessions_for_devices`] (network: prekey fetch +
+/// X3DH for missing sessions) and [`encrypt_for_devices_with_sessions`]
+/// (CPU: the pairwise encrypt fan-out). Callers that must not hold a lock
+/// across network I/O (the group sender-key chain lock) call the two phases
+/// directly with the lock taken only around the second.
+///
 /// Callers must hold per-device session locks before calling this function —
 /// concurrent ratchet mutations will corrupt Signal session state.
-#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.encrypt_fanout", level = "debug", skip_all, fields(count = devices.len()), err(Debug)))]
 pub async fn encrypt_for_devices<'a, S, I, P, SP>(
     runtime: &dyn Runtime,
     stores: &mut SignalStores<'a, S, I, P, SP>,
@@ -298,6 +303,46 @@ pub async fn encrypt_for_devices<'a, S, I, P, SP>(
     hide_decrypt_fail: bool,
     mediatype: Option<&str>,
 ) -> Result<EncryptResult>
+where
+    S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static,
+    I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static,
+    P: crate::libsignal::protocol::PreKeyStore + Send + Sync,
+    SP: crate::libsignal::protocol::SignedPreKeyStore + Send + Sync,
+{
+    let plan = ensure_sessions_for_devices(runtime, stores, resolver, devices).await?;
+    encrypt_for_devices_with_sessions(
+        runtime,
+        stores,
+        devices,
+        plaintext_to_encrypt,
+        hide_decrypt_fail,
+        mediatype,
+        plan,
+    )
+    .await
+}
+
+/// Session material prepared for one encrypt fan-out: per-index LID
+/// encryption overrides (mirroring the `devices` slice it was built from)
+/// plus whether any device 406'd during prekey fetch. Produced only by
+/// [`ensure_sessions_for_devices`]; consumed by
+/// [`encrypt_for_devices_with_sessions`] over the same `devices` slice.
+pub struct SessionPlan {
+    encryption_overrides: Vec<Option<Jid>>,
+    pub had_unregistered_device: bool,
+}
+
+/// Resolve LID overrides and establish missing Signal sessions (prekey
+/// fetch + X3DH) for `devices`. This is the network half of the encrypt
+/// fan-out and touches only session/identity state — never a sender-key
+/// chain — so group sends run it before taking the chain lock.
+#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.ensure_sessions", level = "debug", skip_all, fields(count = devices.len()), err(Debug)))]
+pub async fn ensure_sessions_for_devices<'a, S, I, P, SP>(
+    runtime: &dyn Runtime,
+    stores: &mut SignalStores<'a, S, I, P, SP>,
+    resolver: &dyn SendContextResolver,
+    devices: &[Jid],
+) -> Result<SessionPlan>
 where
     S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static,
     I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static,
@@ -489,6 +534,44 @@ where
         }
     }
 
+    Ok(SessionPlan {
+        encryption_overrides,
+        had_unregistered_device: had_406,
+    })
+}
+
+/// CPU half of the encrypt fan-out: pairwise-encrypt `plaintext_to_encrypt`
+/// for each device using sessions prepared by [`ensure_sessions_for_devices`]
+/// over the same `devices` slice. No resolver, no network — safe to run
+/// under locks that must not span I/O. A device whose session is still
+/// missing (e.g. its bundle was absent) fails its encrypt and is skipped,
+/// matching the combined path's behavior.
+#[cfg_attr(feature = "tracing", tracing::instrument(name = "wa.send.encrypt_fanout", level = "debug", skip_all, fields(count = devices.len()), err(Debug)))]
+pub async fn encrypt_for_devices_with_sessions<'a, S, I, P, SP>(
+    runtime: &dyn Runtime,
+    stores: &mut SignalStores<'a, S, I, P, SP>,
+    devices: &[Jid],
+    plaintext_to_encrypt: &[u8],
+    hide_decrypt_fail: bool,
+    mediatype: Option<&str>,
+    plan: SessionPlan,
+) -> Result<EncryptResult>
+where
+    S: crate::libsignal::protocol::SessionStore + Clone + Send + Sync + 'static,
+    I: crate::libsignal::protocol::IdentityKeyStore + Clone + Send + Sync + 'static,
+    P: crate::libsignal::protocol::PreKeyStore + Send + Sync,
+    SP: crate::libsignal::protocol::SignedPreKeyStore + Send + Sync,
+{
+    debug_assert_eq!(
+        plan.encryption_overrides.len(),
+        devices.len(),
+        "SessionPlan built for a different device list"
+    );
+    let SessionPlan {
+        encryption_overrides,
+        had_unregistered_device,
+    } = plan;
+
     let mut participant_nodes = Vec::with_capacity(devices.len());
     let mut includes_prekey_message = false;
     let mut encrypted_devices = Vec::with_capacity(devices.len());
@@ -502,8 +585,9 @@ where
         // a FuturesUnordered, and two store clones), with no parallelism to gain.
         // Encrypt inline.
         let device_jid = devices[0].clone();
-        let addr = encryption_overrides[0]
-            .as_ref()
+        let addr = encryption_overrides
+            .first()
+            .and_then(|o| o.as_ref())
             .unwrap_or(&devices[0])
             .to_protocol_address();
         let res = encrypt_one_device(
@@ -536,8 +620,9 @@ where
             // The encryption JID is only needed to build the Signal address, so
             // derive it here from a borrow rather than cloning the whole Jid into
             // the task (device_jid is still cloned because it's returned).
-            let addr = encryption_overrides[idx]
-                .as_ref()
+            let addr = encryption_overrides
+                .get(idx)
+                .and_then(|o| o.as_ref())
                 .unwrap_or(&devices[idx])
                 .to_protocol_address();
             let plaintext = plaintext_arc.clone();
@@ -587,6 +672,6 @@ where
         participant_nodes,
         includes_prekey_message,
         encrypted_devices,
-        had_unregistered_device: had_406,
+        had_unregistered_device,
     })
 }

wacore/src/send/group.rs

@@ -172,17 +172,9 @@ pub async fn prepare_group_stanza<
     let mut phash_for_stanza: Option<String> = None;
     let mut skdm_encrypted_devices: Vec<Jid> = Vec::new();
 
-    // Build the chain name once and hold its lock across SKDM creation + the
-    // skmsg encrypt, so concurrent same-(group, sender) sends can't split the
-    // key between the SKDM and the skmsg (nor reuse a chain iteration).
-    let sender_key_name = make_sender_key_name(&to_jid, &own_sending_jid.to_protocol_address());
-    let chain_lock = stores
-        .sender_key_store
-        .sender_key_lock(&sender_key_name)
-        .await;
-    let _chain_guard = chain_lock.lock().await;
-
-    // Determine if we need to distribute SKDM and to which devices
+    // Determine if we need to distribute SKDM and to which devices.
+    // Resolved before the chain lock below: device resolution is network I/O
+    // on state independent of the sender-key chain.
     let distribution_list: Option<Vec<Jid>> = if let Some(target_devices) = skdm_target_devices {
         // Use the specific list of devices that need SKDM
         if target_devices.is_empty() {
@@ -338,69 +330,119 @@ pub async fn prepare_group_stanza<
 
     let mut had_unregistered_devices = false;
 
+    let sender_key_name = make_sender_key_name(&to_jid, &own_sending_jid.to_protocol_address());
+
+    // Establish missing pairwise sessions (prekey fetch + X3DH) for the SKDM
+    // targets before taking the chain lock, so the chain critical section
+    // below never spans a network RTT — concurrent sends to the same group
+    // would otherwise serialize behind it. The setup lock serializes this
+    // phase per group instead: two cold sends can't race fetch + X3DH writes
+    // to the same per-device sessions, while warm sends (no SKDM) never take
+    // it. WA Web's GroupSkmsgJob wraps ensureE2ESessions in try/catch — logs
+    // but does NOT rethrow: a session setup failure must not prevent the
+    // group message from being sent.
+    let session_plan = match distribution_list.as_deref() {
+        Some(list) => {
+            let setup_lock = stores
+                .sender_key_store
+                .session_setup_lock(&sender_key_name)
+                .await;
+            let _setup_guard = setup_lock.lock().await;
+            match ensure_sessions_for_devices(runtime, stores, resolver, list).await {
+                Ok(plan) => Some(plan),
+                Err(e) => {
+                    log::warn!(
+                        "SKDM session setup failed for group {}, continuing without distribution: {e}",
+                        to_jid.observe()
+                    );
+                    if is_device_unregistered_error(&e) {
+                        had_unregistered_devices = true;
+                    }
+                    None
+                }
+            }
+        }
+        None => None,
+    };
+
+    // Hold the chain lock across SKDM creation + the skmsg encrypt, so
+    // concurrent same-(group, sender) sends can't split the key between the
+    // SKDM and the skmsg (nor reuse a chain iteration). The lock covers only
+    // chain-touching steps; session setup and device resolution stay outside.
+    let chain_lock = stores
+        .sender_key_store
+        .sender_key_lock(&sender_key_name)
+        .await;
+    let _chain_guard = chain_lock.lock().await;
+
     if let Some(ref distribution_list) = distribution_list {
+        // Created even when session setup failed (plan None): the sender-key
+        // record must exist so the skmsg below still encrypts, preserving the
+        // continue-without-distribution semantics.
         let axolotl_skdm_bytes = create_sender_key_distribution_message_for_group(
             stores.sender_key_store,
             &sender_key_name,
         )
         .await?;
 
-        let skdm_wrapper_msg = wa::Message {
-            sender_key_distribution_message: Some(wa::message::SenderKeyDistributionMessage {
-                group_id: Some(to_jid.to_string()),
-                axolotl_sender_key_distribution_message: Some(axolotl_skdm_bytes),
-            }),
-            ..Default::default()
-        };
-        let skdm_plaintext_to_encrypt = MessageUtils::encode_and_pad(&skdm_wrapper_msg);
-
-        // WA Web's GroupSkmsgJob wraps ensureE2ESessions in try/catch — logs error
-        // but does NOT rethrow. SKDM distribution failure must not prevent the group
-        // message from being sent. Only successfully encrypted devices are tracked.
-        // Must match the rule applied to the main skmsg payload below: if SKDM carries
-        // `decrypt-fail="hide"` but the payload does not (e.g. AdminRevoke), recipients
-        // without a sender key never decrypt the skmsg and the revoke is silently dropped.
-        let skdm_hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit.as_ref(), message);
-        match encrypt_for_devices(
-            runtime,
-            stores,
-            resolver,
-            distribution_list,
-            &skdm_plaintext_to_encrypt,
-            skdm_hide_decrypt_fail,
-            None,
-        )
-        .await
-        {
-            Ok(result) => {
-                includes_prekey_message = includes_prekey_message || result.includes_prekey_message;
-                if result.had_unregistered_device {
-                    had_unregistered_devices = true;
-                }
-                skdm_encrypted_devices = result.encrypted_devices;
+        if let Some(plan) = session_plan {
+            let skdm_wrapper_msg = wa::Message {
+                sender_key_distribution_message: Some(wa::message::SenderKeyDistributionMessage {
+                    group_id: Some(to_jid.to_string()),
+                    axolotl_sender_key_distribution_message: Some(axolotl_skdm_bytes),
+                }),
+                ..Default::default()
+            };
+            let skdm_plaintext_to_encrypt = MessageUtils::encode_and_pad(&skdm_wrapper_msg);
+
+            // SKDM distribution failure must not prevent the group message from
+            // being sent. Only successfully encrypted devices are tracked.
+            // Must match the rule applied to the main skmsg payload below: if SKDM carries
+            // `decrypt-fail="hide"` but the payload does not (e.g. AdminRevoke), recipients
+            // without a sender key never decrypt the skmsg and the revoke is silently dropped.
+            let skdm_hide_decrypt_fail = should_hide_decrypt_fail_for_send(edit.as_ref(), message);
+            match encrypt_for_devices_with_sessions(
+                runtime,
+                stores,
+                distribution_list,
+                &skdm_plaintext_to_encrypt,
+                skdm_hide_decrypt_fail,
+                None,
+                plan,
+            )
+            .await
+            {
+                Ok(result) => {
+                    includes_prekey_message =
+                        includes_prekey_message || result.includes_prekey_message;
+                    if result.had_unregistered_device {
+                        had_unregistered_devices = true;
+                    }
+                    skdm_encrypted_devices = result.encrypted_devices;
 
-                if !result.participant_nodes.is_empty() {
-                    message_children.push(
-                        NodeBuilder::new("participants")
-                            .children(result.participant_nodes)
-                            .build(),
-                    );
-                    if includes_prekey_message && let Some(acc) = account {
+                    if !result.participant_nodes.is_empty() {
                         message_children.push(
-                            NodeBuilder::new("device-identity")
-                                .bytes(acc.encode_to_vec())
+                            NodeBuilder::new("participants")
+                                .children(result.participant_nodes)
                                 .build(),
                         );
+                        if includes_prekey_message && let Some(acc) = account {
+                            message_children.push(
+                                NodeBuilder::new("device-identity")
+                                    .bytes(acc.encode_to_vec())
+                                    .build(),
+                            );
+                        }
                     }
                 }
-            }
-            Err(e) => {
-                log::warn!(
-                    "SKDM distribution failed for group {}, continuing without it: {e}",
-                    to_jid.observe()
-                );
-                if is_device_unregistered_error(&e) {
-                    had_unregistered_devices = true;
+                Err(e) => {
+                    log::warn!(
+                        "SKDM distribution failed for group {}, continuing without it: {e}",
+                        to_jid.observe()
+                    );
+                    if is_device_unregistered_error(&e) {
+                        had_unregistered_devices = true;
+                    }
                 }
             }
         }