Fix proxy batch fallback ignoring API concurrency (#614)

crates/socket-patch-core/src/api/client.rs

@@ -290,9 +290,12 @@ pub struct ApiClient {
     /// In-flight slots for the public proxy's batch path — the
     /// `/patch/batch` POSTs and the legacy per-package GETs they degrade
     /// to. Shared by clones, so concurrent [`Self::search_patches_batch`]
-    /// calls on one client (scan's batch windows) stay within
-    /// [`PROXY_BATCH_PATH_CONCURRENCY`] requests in total: the peak the
-    /// serial batch loop reached, never that peak times the window.
+    /// calls on one client (scan's batch windows) stay within the public
+    /// proxy's in-flight cap in total, never that cap times the window.
+    /// Sized when a proxy client is built from
+    /// [`crate::utils::concurrent::api_concurrency`], so
+    /// `SOCKET_API_CONCURRENCY` and a tight `RLIMIT_NOFILE` lower it as in
+    /// every other patch-API window (#614).
     proxy_batch_slots: Arc<tokio::sync::Semaphore>,
     /// The vendor loop's download plan, while one is attached
     /// ([`Self::prefetch_vendor_packages`]): [`Self::fetch_vendor_package`]
@@ -321,11 +324,6 @@ enum Sent {
     },
 }
 
-/// Most requests the public proxy's batch path keeps in flight per client:
-/// the legacy per-package fallback's fan-out (one call runs its PURLs in
-/// groups of this size), and the cap all concurrent calls share.
-const PROXY_BATCH_PATH_CONCURRENCY: usize = 10;
-
 /// Retry policy for the vendoring service's package-reference POST and
 /// archive GET: `attempts` tries in total, exponential delays from `base`
 /// with ±25% jitter, each capped at `max_delay` (a `Retry-After`, in
@@ -455,6 +453,12 @@ impl ApiClient {
             ApiRoute::Proxy => None,
         };
 
+        // Only the proxy route uses the batch slots: an org client never
+        // reads the environment for them.
+        let proxy_batch_cap = match options.route {
+            ApiRoute::Proxy => crate::utils::concurrent::api_concurrency(true),
+            ApiRoute::Org { .. } => crate::utils::concurrent::PROXY_API_CONCURRENCY,
+        };
         Self {
             client: api_client(api_token.as_deref(), &timeouts),
             plain: plain_client(&timeouts),
@@ -464,7 +468,7 @@ impl ApiClient {
             org_unresolved: None,
             vendor_retry: VendorRetryPolicy::default(),
             vendor_outage: Arc::new(AtomicU32::new(0)),
-            proxy_batch_slots: Arc::new(tokio::sync::Semaphore::new(PROXY_BATCH_PATH_CONCURRENCY)),
+            proxy_batch_slots: Arc::new(tokio::sync::Semaphore::new(proxy_batch_cap)),
             vendor_prefetch: Arc::new(std::sync::Mutex::new(None)),
             api_retry: ApiRetry::from_env(),
         }
@@ -814,8 +818,10 @@ impl ApiClient {
     /// For authenticated API, uses the POST `/v0/orgs/{slug}/patches/batch`
     /// endpoint. For the public proxy, POSTs `/patch/batch` (served
     /// `Cache-Control: no-store` — POST bodies are not CDN-cacheable) and
-    /// only degrades to individual GET requests per PURL (concurrency 10)
-    /// when the deployed proxy predates the batch endpoint.
+    /// only degrades to individual GET requests per PURL (at most the
+    /// proxy's in-flight cap, [`crate::utils::concurrent::api_concurrency`])
+    /// when the deployed proxy predates the batch endpoint or rejects the
+    /// chunk's validation.
     ///
     /// Maximum 500 PURLs per request.
     ///
@@ -1000,10 +1006,10 @@ impl ApiClient {
     /// proxy gained `POST /patch/batch`, this is the legacy path for
     /// deployments that predate it.
     ///
-    /// Processes PURLs in batches of `PROXY_BATCH_PATH_CONCURRENCY` to
-    /// avoid overwhelming the server while remaining efficient; each GET
-    /// also takes a [`Self::proxy_batch_slots`] slot, so concurrent calls
-    /// on one client share that cap instead of multiplying it.
+    /// Every GET takes a [`Self::proxy_batch_slots`] slot, so this call and
+    /// any concurrent ones on the same client keep at most the proxy's
+    /// in-flight cap between them, and a slow or throttled GET holds one
+    /// slot rather than stalling a whole group.
     async fn search_patches_batch_via_individual_queries(
         &self,
         purls: &[String],
@@ -1016,41 +1022,36 @@ impl ApiClient {
         // would drop the package from the scan without a word.
         let mut throttled: Option<(usize, ApiError)> = None;
 
-        for (chunk_idx, chunk) in purls.chunks(PROXY_BATCH_PATH_CONCURRENCY).enumerate() {
-            // Use tokio::JoinSet for concurrent execution within each chunk
-            let mut join_set = tokio::task::JoinSet::new();
-
-            for (offset, purl) in chunk.iter().enumerate() {
-                let index = chunk_idx * PROXY_BATCH_PATH_CONCURRENCY + offset;
-                let purl = purl.clone();
-                let client = self.clone();
-                join_set.spawn(async move {
-                    let slot = client.proxy_batch_slot().await;
-                    let resp = client.search_patches_by_package(&purl).await;
-                    drop(slot);
-                    match resp {
-                        Ok(r) => (index, purl, Ok(Some(r))),
-                        Err(e) if is_throttle_error(&e) => (index, purl, Err(e)),
-                        Err(e) => {
-                            debug_log(&format!("Error fetching patches for {}: {}", purl, e));
-                            (index, purl, Ok(None))
-                        }
+        let mut join_set = tokio::task::JoinSet::new();
+        for (index, purl) in purls.iter().enumerate() {
+            let purl = purl.clone();
+            let client = self.clone();
+            join_set.spawn(async move {
+                let slot = client.proxy_batch_slot().await;
+                let resp = client.search_patches_by_package(&purl).await;
+                drop(slot);
+                match resp {
+                    Ok(r) => (index, purl, Ok(Some(r))),
+                    Err(e) if is_throttle_error(&e) => (index, purl, Err(e)),
+                    Err(e) => {
+                        debug_log(&format!("Error fetching patches for {}: {}", purl, e));
+                        (index, purl, Ok(None))
                     }
-                });
-            }
+                }
+            });
+        }
 
-            while let Some(result) = join_set.join_next().await {
-                match result {
-                    Ok((_, purl, Ok(resp))) => all_results.push((purl, resp)),
-                    Ok((index, _, Err(e))) => {
-                        if throttled.as_ref().is_none_or(|(first, _)| index < *first) {
-                            throttled = Some((index, e));
-                        }
-                    }
-                    Err(e) => {
-                        debug_log(&format!("Task join error: {}", e));
+        while let Some(result) = join_set.join_next().await {
+            match result {
+                Ok((_, purl, Ok(resp))) => all_results.push((purl, resp)),
+                Ok((index, _, Err(e))) => {
+                    if throttled.as_ref().is_none_or(|(first, _)| index < *first) {
+                        throttled = Some((index, e));
                     }
                 }
+                Err(e) => {
+                    debug_log(&format!("Task join error: {}", e));
+                }
             }
         }
         if let Some((_, e)) = throttled {
@@ -6209,9 +6210,13 @@ mod proxy_batch_path_cap_tests {
     //! Concurrent `search_patches_batch` calls on one public-proxy client
     //! (scan's batch windows) must share the batch path's in-flight cap,
     //! not multiply it: when every chunk degrades to the legacy per-package
-    //! GETs, the proxy sees at most `PROXY_BATCH_PATH_CONCURRENCY` of them
-    //! at once — what the serial batch loop peaked at.
+    //! GETs, the proxy sees at most `api_concurrency(true)` of them at
+    //! once: the one patch-API pacing policy (#614).
+    //!
+    //! Serial with the other tests that set `SOCKET_API_CONCURRENCY`: the
+    //! client reads it when it is built.
     use super::*;
+    use crate::utils::concurrent::API_CONCURRENCY_ENV;
     use std::sync::Mutex;
     use std::time::Instant;
     use wiremock::matchers::{method, path, path_regex};
@@ -6256,8 +6261,38 @@ mod proxy_batch_path_cap_tests {
             .unwrap_or(0)
     }
 
-    #[tokio::test]
-    async fn concurrent_batches_share_the_legacy_fallback_cap() {
+    /// A public-proxy client built while [`API_CONCURRENCY_ENV`] is
+    /// `value` (`None` unsets it). The client reads the variable only when
+    /// it is built, so it is restored right away rather than left set
+    /// while the requests run.
+    fn proxy_client_under(uri: String, value: Option<&str>) -> ApiClient {
+        let orig = std::env::var(API_CONCURRENCY_ENV).ok();
+        match value {
+            Some(v) => std::env::set_var(API_CONCURRENCY_ENV, v),
+            None => std::env::remove_var(API_CONCURRENCY_ENV),
+        }
+        let client = ApiClient::new(ApiClientOptions {
+            api_url: uri,
+            api_token: None,
+            route: ApiRoute::Proxy,
+        });
+        match orig {
+            Some(v) => std::env::set_var(API_CONCURRENCY_ENV, v),
+            None => std::env::remove_var(API_CONCURRENCY_ENV),
+        }
+        client
+    }
+
+    /// `windows` concurrent [`ApiClient::search_patches_batch`] calls of
+    /// `per_window` PURLs on one public-proxy client built under
+    /// [`API_CONCURRENCY_ENV`] = `value`, every chunk rejected by a
+    /// validation 400 and degraded to per-package GETs. Returns the peak
+    /// number of GETs in flight.
+    async fn fallback_peak_in_flight(
+        value: Option<&str>,
+        windows: usize,
+        per_window: usize,
+    ) -> usize {
         let server = MockServer::start().await;
         // A validation 400 for every chunk: each degrades to per-package
         // GETs (one exotic PURL per chunk is enough in the wild).
@@ -6273,15 +6308,10 @@ mod proxy_batch_path_cap_tests {
             .mount(&server)
             .await;
 
-        let client = ApiClient::new(ApiClientOptions {
-            api_url: server.uri(),
-            api_token: None,
-            route: ApiRoute::Proxy,
-        });
-        // Four windows of 10 PURLs, as scan's proxy batch windows run them.
-        let chunks: Vec<Vec<String>> = (0..4)
+        let client = proxy_client_under(server.uri(), value);
+        let chunks: Vec<Vec<String>> = (0..windows)
             .map(|c| {
-                (0..PROXY_BATCH_PATH_CONCURRENCY)
+                (0..per_window)
                     .map(|i| format!("pkg:npm/cap-{c}-{i}@1.0.0"))
                     .collect()
             })
@@ -6298,14 +6328,29 @@ mod proxy_batch_path_cap_tests {
         }
 
         let arrivals = arrivals.lock().unwrap();
-        assert_eq!(arrivals.len(), 40, "one GET per PURL");
-        let peak = peak_in_flight(&arrivals);
+        assert_eq!(arrivals.len(), windows * per_window, "one GET per PURL");
+        peak_in_flight(&arrivals)
+    }
+
+    #[tokio::test]
+    #[serial_test::serial]
+    async fn concurrent_batches_share_the_legacy_fallback_cap() {
+        let cap = crate::utils::concurrent::PROXY_API_CONCURRENCY;
+        // Four windows of 10 PURLs, as scan's proxy batch windows run them.
+        let peak = fallback_peak_in_flight(None, 4, 10).await;
         assert!(
-            peak <= PROXY_BATCH_PATH_CONCURRENCY,
-            "{peak} by-package GETs in flight at once; the cap is \
-             {PROXY_BATCH_PATH_CONCURRENCY}"
+            peak <= cap,
+            "{peak} by-package GETs in flight at once; the cap is {cap}"
         );
         // And the cap is reached, not undershot: the calls still overlap.
-        assert_eq!(peak, PROXY_BATCH_PATH_CONCURRENCY);
+        assert_eq!(peak, cap);
+    }
+
+    /// `SOCKET_API_CONCURRENCY=1` makes the per-package fallback serial,
+    /// like every other patch-API window (#614).
+    #[tokio::test]
+    #[serial_test::serial]
+    async fn the_legacy_fallback_honors_the_concurrency_override() {
+        assert_eq!(fallback_peak_in_flight(Some("1"), 2, 3).await, 1);
     }
 }