diff --git a/crates/socket-patch-core/src/api/client.rs b/crates/socket-patch-core/src/api/client.rs index aca4ca2a6..0ed3470e3 100644 --- a/crates/socket-patch-core/src/api/client.rs +++ b/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, /// 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> = (0..4) + let client = proxy_client_under(server.uri(), value); + let chunks: Vec> = (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); } }