perf(send): trim group-send warm-path CPU and cold fan-out allocations

src/client.rs

@@ -610,6 +610,17 @@ pub struct Client {
     pub(crate) group_devices_memo:
         Cache<Jid, Arc<crate::client::device_registry::GroupDevicesMemo>>,
 
+    /// Last `(devices, sender-key-device map)` Arc pair with an empty `needs_skdm`,
+    /// so a warm repeat send skips `filter_skdm_targets`. `Weak` keeps the pointer
+    /// comparison ABA-safe, matching `GroupDevicesMemo`.
+    pub(crate) skdm_warm_memo: Cache<
+        Jid,
+        (
+            std::sync::Weak<wacore::send::ResolvedGroupDevices>,
+            std::sync::Weak<crate::sender_key_device_cache::SenderKeyDeviceMap>,
+        ),
+    >,
+
     /// Router for dispatching stanzas to their appropriate handlers
     pub(crate) stanza_router: crate::handlers::router::StanzaRouter,
 

src/client/lifecycle.rs

@@ -243,6 +243,9 @@ impl Client {
             group_devices_memo: Cache::builder()
                 .max_capacity(GROUP_DEVICES_MEMO_CAPACITY)
                 .build(),
+            skdm_warm_memo: Cache::builder()
+                .max_capacity(GROUP_DEVICES_MEMO_CAPACITY)
+                .build(),
             stanza_router: Self::create_stanza_router(),
             synchronous_ack: false,
             http_client,

src/send/mod.rs

@@ -1048,12 +1048,34 @@ impl Client {
             .await
         {
             Ok(all_devices) => {
+                // Skip the O(devices) filter_skdm_targets scan when the same
+                // (devices, sender-key-map) Arc pair was already fully warm. Both
+                // Arcs swap on any warm-state or membership change, so a stale skip
+                // is impossible. Needs the device memo for a stable devices Arc.
+                if self.group_devices_memo_enabled
+                    && let Some((dw, cw)) = self.skdm_warm_memo.get(group).await
+                    && std::ptr::eq(dw.as_ptr(), std::sync::Arc::as_ptr(&all_devices))
+                    && std::ptr::eq(cw.as_ptr(), std::sync::Arc::as_ptr(&cached_map))
+                {
+                    return Some((all_devices, Vec::new()));
+                }
                 let needs_skdm = self.filter_skdm_targets(
                     group_jid,
                     all_devices.devices(),
                     &cached_map,
                     own_sending_jid,
                 );
+                if needs_skdm.is_empty() && self.group_devices_memo_enabled {
+                    self.skdm_warm_memo
+                        .insert(
+                            group.clone(),
+                            (
+                                std::sync::Arc::downgrade(&all_devices),
+                                std::sync::Arc::downgrade(&cached_map),
+                            ),
+                        )
+                        .await;
+                }
                 Some((all_devices, needs_skdm))
             }
             Err(e) => {

wacore/src/send/encrypt.rs

@@ -718,57 +718,71 @@ pub async fn encrypt_for_devices_with_sessions_raw(
         .await;
         push_raw_result(res, &mut encrypted, &mut includes_prekey_message);
     } else {
-        // Parallel encrypt fan-out across tokio tasks bounded by
-        // ENCRYPT_FANOUT_CONCURRENCY; collected in completion order so the
-        // fastest encrypts ship first.
+        // One task per chunk, not per device: the per-device fan-out allocated a
+        // task + oneshot + two store clones for every recipient. Same parallelism,
+        // spawns bounded by ENCRYPT_FANOUT_CONCURRENCY. Wire order is irrelevant
+        // (phash sorts before hashing on both ends).
         let plaintext_arc: std::sync::Arc<[u8]> = std::sync::Arc::from(plaintext_to_encrypt);
 
         let total = devices.len();
-        let mut next_spawn = 0usize;
+        let num_chunks = ENCRYPT_FANOUT_CONCURRENCY.min(total);
 
-        let make_encrypt_task = |idx: usize| {
-            let device_jid = devices[idx].clone();
-            // 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
-                .get(idx)
-                .and_then(|o| o.as_ref())
-                .unwrap_or(&devices[idx])
-                .to_protocol_address();
+        let mut in_flight: FuturesUnordered<_> = FuturesUnordered::new();
+        // Index partitioning gives exactly num_chunks slices (keeps the configured
+        // parallelism) and no-ops on an empty device set instead of dividing by zero.
+        for chunk_idx in 0..num_chunks {
+            let chunk_start = chunk_idx * total / num_chunks;
+            let chunk_end = (chunk_idx + 1) * total / num_chunks;
+            // The 'static task can't borrow devices/encryption_overrides.
+            let jobs: Vec<(ProtocolAddress, Jid)> = (chunk_start..chunk_end)
+                .map(|idx| {
+                    let addr = encryption_overrides
+                        .get(idx)
+                        .and_then(|o| o.as_ref())
+                        .unwrap_or(&devices[idx])
+                        .to_protocol_address();
+                    (addr, devices[idx].clone())
+                })
+                .collect();
             let plaintext = plaintext_arc.clone();
+            // clone_box shares the Arc-backed backend, so the sequential ratchet
+            // advances persist despite one clone serving the whole chunk.
             let mut session_store = stores.session_store.clone_box();
             let mut identity_store = stores.identity_store.clone_box();
 
-            spawn_oneshot(runtime, async move {
-                encrypt_one_device(
-                    &plaintext,
-                    &addr,
-                    &mut *session_store,
-                    &mut *identity_store,
-                    device_jid,
-                )
-                .await
-            })
-        };
-
-        let mut in_flight: FuturesUnordered<_> = FuturesUnordered::new();
-        while next_spawn < total && in_flight.len() < ENCRYPT_FANOUT_CONCURRENCY {
-            in_flight.push(make_encrypt_task(next_spawn));
-            next_spawn += 1;
+            in_flight.push(spawn_oneshot(runtime, async move {
+                let mut out = Vec::with_capacity(jobs.len());
+                for (addr, device_jid) in jobs {
+                    out.push(
+                        encrypt_one_device(
+                            &plaintext,
+                            &addr,
+                            &mut *session_store,
+                            &mut *identity_store,
+                            device_jid,
+                        )
+                        .await,
+                    );
+                }
+                out
+            }));
         }
         while let Some(spawn_result) = in_flight.next().await {
             match spawn_result {
-                Ok(res) => push_raw_result(res, &mut encrypted, &mut includes_prekey_message),
+                Ok(results) => {
+                    for res in results {
+                        push_raw_result(res, &mut encrypted, &mut includes_prekey_message);
+                    }
+                }
                 Err(SpawnCanceled) => {
-                    log::warn!("Encrypt task did not deliver a result; skipping device.");
+                    // A whole chunk drops (not one device); its members stay
+                    // un-warm and are re-targeted next send.
+                    log::warn!(
+                        "Encrypt chunk did not deliver a result; up to ~{} device(s) skipped this send.",
+                        total.div_ceil(num_chunks)
+                    );
                 }
             }
-
-            if next_spawn < total {
-                in_flight.push(make_encrypt_task(next_spawn));
-                next_spawn += 1;
-            }
         }
     }
 

wacore/src/send/group.rs

@@ -371,20 +371,31 @@ pub async fn prepare_group_stanza(
         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.
+    // Padding is chain-independent; compute it before the lock so the
+    // per-(group,sender) serialization point covers only the ratchet steps.
+    let plaintext = match &shared_content {
+        Some(content) => {
+            MessageUtils::pad_with_context_from_encoded(content, reporting_context.as_ref())
+        }
+        None => MessageUtils::encode_and_pad_with_context(message, reporting_context.as_ref()),
+    };
+
+    // The lock spans SKDM creation, the pairwise SKDM fan-out, and the skmsg
+    // encrypt. Creating the SKDM snapshots the sender key and the skmsg uses it,
+    // so those must be atomic. The fan-out also mutates the shared per-device
+    // Signal sessions, which the group path (unlike the DM path) does not lock,
+    // so it has to stay serialized here too or concurrent same-group sends race
+    // those sessions. Dropped after the encrypt so only the stanza build runs
+    // off the serialization point.
     let chain_lock = stores
         .sender_key_store
         .sender_key_lock(&sender_key_name)
         .await;
-    let _chain_guard = chain_lock.lock().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.
+        // record must exist so the skmsg below still encrypts.
         let axolotl_skdm_bytes = create_sender_key_distribution_message_for_group(
             stores.sender_key_store,
             &sender_key_name,
@@ -460,14 +471,6 @@ pub async fn prepare_group_stanza(
         }
     }
 
-    // Reuse the shared encode (also fed to the reporting token) for the skmsg plaintext
-    // instead of encoding the message a second time; the mci-hoist path re-encodes.
-    let plaintext = match &shared_content {
-        Some(content) => {
-            MessageUtils::pad_with_context_from_encoded(content, reporting_context.as_ref())
-        }
-        None => MessageUtils::encode_and_pad_with_context(message, reporting_context.as_ref()),
-    };
     let skmsg = encrypt_group_message(
         stores.sender_key_store,
         &sender_key_name,
@@ -476,6 +479,9 @@ pub async fn prepare_group_stanza(
     )
     .await?;
 
+    // Release before the chain-independent stanza build.
+    drop(chain_guard);
+
     let skmsg_ciphertext = skmsg.into_serialized();
 
     let mediatype = media_type_from_message(message);