perf: parallelize remaining serial/blocking hot paths (startup, media, send/recv)

http_clients/ureq-client/src/lib.rs

@@ -23,8 +23,9 @@ const MAX_IDLE_CONNECTIONS: u64 = 3;
 #[derive(Debug, Clone)]
 pub struct UreqHttpClient {
     agent: ureq::Agent,
-    /// Cap for [`UreqHttpClient::execute`]. Streaming is unbounded — the
-    /// caller owns the sink.
+    /// Total-bytes cap for both [`UreqHttpClient::execute`] and the reader from
+    /// [`UreqHttpClient::execute_streaming`]. Bounds an in-memory sink so a
+    /// hostile CDN can't drive it to OOM; defaults to WA's 2 GiB max file size.
     max_body_bytes: u64,
     /// Best-effort pool footprint for `resource_report`. `None` when a custom
     /// agent is supplied (its buffer/pool config is opaque to us).
@@ -64,8 +65,9 @@ impl UreqHttpClient {
         }
     }
 
-    /// Override the per-response cap for [`UreqHttpClient::execute`]. Set to
-    /// `u64::MAX` to disable; a hostile server can then exhaust memory.
+    /// Override the per-response cap for [`UreqHttpClient::execute`] and
+    /// [`UreqHttpClient::execute_streaming`]. Set to `u64::MAX` to disable; a
+    /// hostile server can then exhaust memory.
     pub fn with_max_body_bytes(mut self, max_body_bytes: u64) -> Self {
         self.max_body_bytes = max_body_bytes;
         self
@@ -172,7 +174,13 @@ impl HttpClient for UreqHttpClient {
         };
 
         let status_code = response.status().as_u16();
-        let reader = response.into_body().into_reader();
+        // Bound the streaming reader to the same cap `execute` enforces: an
+        // in-memory sink (`Client::download` buffers into a `Vec`) must not be
+        // driveable to OOM by a CDN that streams past the declared length. Over
+        // the cap the reader hits EOF and the downstream MAC/SHA check fails,
+        // rather than growing the sink unbounded. `DOWNLOAD_PREALLOC_CAP` only
+        // sizes the initial allocation, not the total read.
+        let reader = std::io::Read::take(response.into_body().into_reader(), self.max_body_bytes);
 
         Ok(StreamingHttpResponse {
             status_code,
@@ -302,6 +310,32 @@ mod tests {
             .expect_err("1 KiB cap must reject a 4 MiB body");
     }
 
+    // The streaming reader must honor the same cap: an over-cap body is
+    // truncated at EOF (the caller's decrypt/MAC check then rejects it) instead
+    // of growing an in-memory sink to OOM.
+    #[tokio::test(flavor = "current_thread")]
+    async fn execute_streaming_bounds_body_at_cap() {
+        const SIZE: usize = 4 * 1024 * 1024;
+        const CAP: u64 = 1024;
+        let url = spawn_fixed_size_server(SIZE);
+        let read = tokio::task::spawn_blocking(move || {
+            let mut resp = UreqHttpClient::new()
+                .with_max_body_bytes(CAP)
+                .execute_streaming(HttpRequest {
+                    method: "GET".into(),
+                    url,
+                    headers: std::collections::HashMap::new(),
+                    body: None,
+                })
+                .expect("streaming GET should start");
+            let mut sink = std::io::sink();
+            std::io::copy(&mut resp.body, &mut sink).expect("draining the reader should not error")
+        })
+        .await
+        .unwrap();
+        assert_eq!(read, CAP, "streaming body must stop at the cap");
+    }
+
     /// Captures the raw request headers and body of a single POST, then replies 200.
     fn spawn_capture_server() -> (String, std::sync::mpsc::Receiver<(String, Vec<u8>)>) {
         let listener = TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port");

src/bot.rs

@@ -467,7 +467,17 @@ impl Bot {
         } = self;
 
         if let Some(receiver) = sync_task_receiver {
+            // This channel carries only HistorySync tasks: app-state sync runs via
+            // its own direct path (fetch_app_state_with_retry), nothing enqueues
+            // AppStateSync here. Chunks are independent (order-free upserts; the
+            // event carries chunk_order), so ingest concurrently, bounded low — each
+            // in-flight chunk decompresses a blob and the connect path is peak-
+            // memory-conscious (WA Web caps at histSyncChunk=3). Taking the permit in
+            // the recv loop backpressures history intake on a burst; since no
+            // app-state task flows here, that can't head-of-line block one.
+            const HISTORY_SYNC_CONCURRENCY: usize = 2;
             let worker_client = Arc::downgrade(&client);
+            let history_permits = Arc::new(async_lock::Semaphore::new(HISTORY_SYNC_CONCURRENCY));
             client
                 .runtime
                 .spawn(Box::pin(async move {
@@ -476,9 +486,25 @@ impl Bot {
                             break;
                         };
 
-                        worker_client.process_sync_task(task).await;
+                        if matches!(task, crate::sync_task::MajorSyncTask::HistorySync { .. }) {
+                            let permit = history_permits.acquire_arc().await;
+                            let task_client = worker_client.clone();
+                            worker_client
+                                .runtime
+                                .spawn(Box::pin(async move {
+                                    let _permit = permit;
+                                    task_client.process_sync_task(task).await;
+                                }))
+                                .detach();
+                        } else {
+                            // Defensive: nothing enqueues AppStateSync today, but if
+                            // that changes it must run serially (ordered patches).
+                            worker_client.process_sync_task(task).await;
+                        }
                     }
-                    info!("Sync worker shutting down.");
+                    info!(
+                        "Sync worker intake loop finished (detached history-sync tasks may still be running)."
+                    );
                 }))
                 .detach();
         }

src/client.rs

@@ -646,6 +646,11 @@ pub struct Client {
     /// Single-flights signed pre-key rotation so overlapping post-login tasks
     /// (from reconnect churn) can't run the rotate/upload/prune flow concurrently.
     pub(crate) signed_pre_key_rotation_lock: Arc<async_lock::Mutex<()>>,
+    /// Serializes the history-sync tc-token get-then-store: chunks ingest
+    /// concurrently, but the newer-wins guard is a non-atomic read-then-write,
+    /// so without this two same-contact candidates could interleave and let an
+    /// older privacy token overwrite a fresher one.
+    pub(crate) tc_token_lock: Arc<async_lock::Mutex<()>>,
     /// Notifier for when offline sync (ib offline stanza) is received.
     /// WhatsApp Web waits for this before sending passive tasks (prekey upload, active IQ, presence).
     pub(crate) offline_sync_notifier: Arc<event_listener::Event>,

src/client/app_state.rs

@@ -2,6 +2,12 @@
 
 use super::*;
 
+/// Concurrency cap for pre-downloading app-state external blobs (independent CDN
+/// GETs, keyed by directPath — LTHash ordering is in patch application, not blob
+/// fetching). WA Web fans these out under `Promise.all` (`Syncd/CollectionHandler`);
+/// bounded here because a snapshot can be multi-MB and a batch carries several.
+const APPSTATE_BLOB_DOWNLOAD_CONCURRENCY: usize = 4;
+
 impl Client {
     pub(crate) async fn get_app_state_processor(&self) -> Arc<AppStateProcessor> {
         let mut guard = self.app_state_processor.lock().await;
@@ -17,6 +23,94 @@ impl Client {
         proc
     }
 
+    /// Pre-download every external blob (snapshots + patch external mutations)
+    /// referenced by `patch_lists`, keyed by directPath, fetching concurrently
+    /// (bounded by [`APPSTATE_BLOB_DOWNLOAD_CONCURRENCY`]). A failed download is
+    /// logged and omitted; the later inline step surfaces the missing blob as
+    /// before. Mirrors WA Web's parallel syncd blob fetch.
+    async fn pre_download_external_blobs(
+        &self,
+        patch_lists: &[wacore::appstate::patch_decode::PatchList],
+    ) -> std::collections::HashMap<String, Vec<u8>> {
+        use futures::StreamExt;
+
+        // Kept only so a failed download logs the right message (snapshot vs patch).
+        enum BlobKind {
+            Snapshot(WAPatchName),
+            Mutation(u64),
+        }
+
+        // Clone the (small) blob ref into each job so the task owns its input and
+        // captures only `&self` (keeps the future Send); the directPath is
+        // recovered from the moved `ext` after the fetch. Dedup by directPath so
+        // patches sharing a blob don't fetch it twice into the same map key.
+        let mut jobs: Vec<(wa::ExternalBlobReference, BlobKind)> = Vec::new();
+        let mut seen_paths: std::collections::HashSet<&str> = std::collections::HashSet::new();
+        for pl in patch_lists {
+            if let Some(ext) = &pl.snapshot_ref
+                && let Some(path) = ext.direct_path.as_deref()
+                && seen_paths.insert(path)
+            {
+                jobs.push((ext.clone(), BlobKind::Snapshot(pl.name)));
+            }
+            for patch in &pl.patches {
+                if let Some(ext) = patch.external_mutations.as_option()
+                    && let Some(path) = ext.direct_path.as_deref()
+                    && seen_paths.insert(path)
+                {
+                    let v = patch
+                        .version
+                        .as_option()
+                        .and_then(|v| v.version)
+                        .unwrap_or(0);
+                    jobs.push((ext.clone(), BlobKind::Mutation(v)));
+                }
+            }
+        }
+
+        if jobs.is_empty() {
+            return std::collections::HashMap::new();
+        }
+
+        let mut pre_downloaded = std::collections::HashMap::with_capacity(jobs.len());
+        let results = futures::stream::iter(jobs.into_iter().map(|(ext, kind)| async move {
+            let bytes = self.download(&ext).await;
+            // directPath presence was checked when the job was built.
+            (ext.direct_path, kind, bytes)
+        }))
+        .buffer_unordered(APPSTATE_BLOB_DOWNLOAD_CONCURRENCY)
+        .collect::<Vec<_>>()
+        .await;
+
+        for (path, kind, res) in results {
+            match res {
+                Ok(bytes) => {
+                    if let BlobKind::Mutation(v) = kind {
+                        debug!(target: "Client/AppState", "Downloaded external mutations for patch v{} ({} bytes)", v, bytes.len());
+                    } else {
+                        debug!(target: "Client/AppState", "Downloaded external snapshot ({} bytes)", bytes.len());
+                    }
+                    if let Some(path) = path {
+                        pre_downloaded.insert(path, bytes);
+                    }
+                }
+                Err(e) => match kind {
+                    BlobKind::Snapshot(name) => {
+                        warn!("Failed to download external snapshot for {:?}: {e}", name)
+                    }
+                    BlobKind::Mutation(v) => {
+                        warn!(
+                            "Failed to download external mutations for patch v{}: {e}",
+                            v
+                        )
+                    }
+                },
+            }
+        }
+
+        pre_downloaded
+    }
+
     /// Public entry point for processing [`MajorSyncTask`] from the sync channel.
     #[cfg_attr(
         feature = "tracing",
@@ -218,60 +312,15 @@ impl Client {
 
             let resp = self.send_iq(iq).await?;
 
-            // Pre-download all external blobs for all collections in the response
-            let mut pre_downloaded: std::collections::HashMap<String, Vec<u8>> =
-                std::collections::HashMap::new();
-
             // Parse the response once here for pre-download; the same parsed
             // lists are handed to the processor below (no second parse).
             let mut patch_lists =
                 wacore::appstate::patch_decode::parse_patch_lists_ref(resp.get())?;
 
             let proc = self.get_app_state_processor().await;
-            {
-                for pl in &patch_lists {
-                    // Download external snapshot
-                    if let Some(ext) = &pl.snapshot_ref
-                        && let Some(path) = &ext.direct_path
-                    {
-                        match self.download(ext).await {
-                            Ok(bytes) => {
-                                pre_downloaded.insert(path.clone(), bytes);
-                            }
-                            Err(e) => {
-                                warn!(
-                                    "Failed to download external snapshot for {:?}: {e}",
-                                    pl.name
-                                );
-                            }
-                        }
-                    }
-
-                    // Download external mutations
-                    for patch in &pl.patches {
-                        if let Some(ext) = patch.external_mutations.as_option()
-                            && let Some(path) = &ext.direct_path
-                        {
-                            match self.download(ext).await {
-                                Ok(bytes) => {
-                                    pre_downloaded.insert(path.clone(), bytes);
-                                }
-                                Err(e) => {
-                                    let v = patch
-                                        .version
-                                        .as_option()
-                                        .and_then(|v| v.version)
-                                        .unwrap_or(0);
-                                    warn!(
-                                        "Failed to download external mutations for patch v{}: {e}",
-                                        v
-                                    );
-                                }
-                            }
-                        }
-                    }
-                }
-            }
+            // Pre-download all external blobs for all collections in the response,
+            // concurrently (independent CDN GETs, keyed by directPath).
+            let pre_downloaded = self.pre_download_external_blobs(&patch_lists).await;
 
             let download = |ext: &wa::ExternalBlobReference| -> anyhow::Result<Vec<u8>> {
                 if let Some(path) = &ext.direct_path {
@@ -477,51 +526,11 @@ impl Client {
 
             let proc = self.get_app_state_processor().await;
 
-            // Pre-download all external blobs (snapshot and patch mutations); keyed by
-            // directPath.
-            let mut pre_downloaded: std::collections::HashMap<String, Vec<u8>> =
-                std::collections::HashMap::new();
-            {
-                // Download external snapshot if present
-                if let Some(ext) = &pl.snapshot_ref
-                    && let Some(path) = &ext.direct_path
-                {
-                    match self.download(ext).await {
-                        Ok(bytes) => {
-                            debug!(target: "Client/AppState", "Downloaded external snapshot ({} bytes)", bytes.len());
-                            pre_downloaded.insert(path.clone(), bytes);
-                        }
-                        Err(e) => {
-                            warn!("Failed to download external snapshot: {e}");
-                        }
-                    }
-                }
-
-                // Download external mutations for each patch that has them
-                for patch in &pl.patches {
-                    if let Some(ext) = patch.external_mutations.as_option()
-                        && let Some(path) = &ext.direct_path
-                    {
-                        let patch_version = patch
-                            .version
-                            .as_option()
-                            .and_then(|v| v.version)
-                            .unwrap_or(0);
-                        match self.download(ext).await {
-                            Ok(bytes) => {
-                                debug!(target: "Client/AppState", "Downloaded external mutations for patch v{} ({} bytes)", patch_version, bytes.len());
-                                pre_downloaded.insert(path.clone(), bytes);
-                            }
-                            Err(e) => {
-                                warn!(
-                                    "Failed to download external mutations for patch v{}: {e}",
-                                    patch_version
-                                );
-                            }
-                        }
-                    }
-                }
-            }
+            // Pre-download all external blobs (snapshot and patch mutations),
+            // concurrently, keyed by directPath.
+            let pre_downloaded = self
+                .pre_download_external_blobs(std::slice::from_ref(&pl))
+                .await;
 
             let download = |ext: &wa::ExternalBlobReference| -> anyhow::Result<Vec<u8>> {
                 if let Some(path) = &ext.direct_path {

src/client/lifecycle.rs

@@ -218,6 +218,7 @@ impl Client {
             initial_app_state_keys_received: Arc::new(AtomicBool::new(false)),
             prekey_upload_lock: Arc::new(async_lock::Mutex::new(())),
             signed_pre_key_rotation_lock: Arc::new(async_lock::Mutex::new(())),
+            tc_token_lock: Arc::new(async_lock::Mutex::new(())),
             offline_sync_notifier: Arc::new(event_listener::Event::new()),
             offline_sync_completed: Arc::new(AtomicBool::new(false)),
             offline_sync_finish_started: Arc::new(AtomicBool::new(false)),