Skip to content

Commit 188501e

Browse files
committed
Cap registry lookups in the upstream restore
A rollback, remove or vendored takeover of many hosted pins sent every public-registry lookup at once (40 pins, 40 requests in flight), with no 429 handling and no regard for a tight file descriptor limit. registry_concurrency() existed for exactly this but had no caller. UpstreamClient now holds a semaphore sized by registry_concurrency() (4, or 1 under a tight RLIMIT_NOFILE), and every request it sends, JSON or download, holds a slot from send to the end of its body. The restorers' join_all folds are unchanged, so output order and contents stay the same. Fixes #1220 Assisted-by: Claude Code:claude-opus-5-5
1 parent 68d8e8b commit 188501e

2 files changed

Lines changed: 175 additions & 9 deletions

File tree

‎crates/socket-patch-core/src/patch/redirect/upstream/client.rs‎

Lines changed: 139 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
use std::collections::HashMap;
77

88
use serde_json::Value;
9-
use tokio::sync::Mutex;
9+
use tokio::sync::{Mutex, Semaphore, SemaphorePermit};
1010

1111
use crate::vendor::registry_fetch::{
1212
build_registry_client, npm_registry_base, pypi_json_api_base, RegistryClient,
@@ -164,8 +164,16 @@ type RegistryCache<T> = Mutex<HashMap<(String, Option<String>, String, String),
164164

165165
/// One client per restore run; every lookup is cached (success and
166166
/// failure alike) so a pin wired in several files costs one request.
167+
///
168+
/// The restorers start every pin's lookup at once (`join_all`), so the cap
169+
/// on requests in flight lives here, at the one place they are sent: each
170+
/// request holds a [`Self::permits`] slot from send to the end of its body
171+
/// (#1220).
167172
pub(crate) struct UpstreamClient {
168173
http: RegistryClient,
174+
/// At most [`crate::utils::concurrent::registry_concurrency`] requests
175+
/// to the public registries in flight at once.
176+
permits: Semaphore,
169177
offline: bool,
170178
npm: RegistryCache<NpmDist>,
171179
npm_berry: Cache<String>,
@@ -179,8 +187,14 @@ pub(crate) struct UpstreamClient {
179187

180188
impl UpstreamClient {
181189
pub(crate) fn new(offline: bool) -> Self {
190+
Self::with_concurrency(offline, crate::utils::concurrent::registry_concurrency())
191+
}
192+
193+
/// [`Self::new`] with at most `limit` requests in flight (0 counts as 1).
194+
pub(crate) fn with_concurrency(offline: bool, limit: usize) -> Self {
182195
UpstreamClient {
183196
http: build_registry_client(),
197+
permits: Semaphore::new(limit.max(1)),
184198
offline,
185199
npm: Mutex::default(),
186200
npm_berry: Mutex::default(),
@@ -193,6 +207,21 @@ impl UpstreamClient {
193207
}
194208
}
195209

210+
/// A slot to send one request in. The semaphore is never closed, so
211+
/// the wait always ends in a permit.
212+
async fn permit(&self) -> SemaphorePermit<'_> {
213+
self.permits
214+
.acquire()
215+
.await
216+
.expect("the upstream request semaphore is never closed")
217+
}
218+
219+
/// `url`'s bytes, sent under a [`Self::permit`].
220+
async fn download(&self, url: &str) -> Result<Vec<u8>, String> {
221+
let _permit = self.permit().await;
222+
crate::vendor::registry_fetch::download(&self.http, url).await
223+
}
224+
196225
async fn get_json(&self, url: &str) -> Result<Value, String> {
197226
self.get_json_authorized(url, None).await
198227
}
@@ -205,6 +234,7 @@ impl UpstreamClient {
205234
url: &str,
206235
authorization: Option<&str>,
207236
) -> Result<Value, String> {
237+
let _permit = self.permit().await;
208238
let mut request = self.http.get(url).header("accept", "application/json");
209239
if let Some(authorization) = authorization {
210240
let mut value =
@@ -230,7 +260,7 @@ impl UpstreamClient {
230260
}
231261

232262
async fn get_text(&self, url: &str) -> Result<String, String> {
233-
let bytes = crate::vendor::registry_fetch::download(&self.http, url).await?;
263+
let bytes = self.download(url).await?;
234264
String::from_utf8(bytes).map_err(|_| format!("{url} is not UTF-8"))
235265
}
236266

@@ -352,8 +382,7 @@ impl UpstreamClient {
352382
.as_deref()
353383
.is_some_and(|s| s.split_whitespace().any(|v| v == integrity))
354384
{
355-
let bytes =
356-
crate::vendor::registry_fetch::download(&self.http, &dist.tarball).await?;
385+
let bytes = self.download(&dist.tarball).await?;
357386
crate::vendor::registry_fetch::verify_sri(&bytes, integrity)?;
358387
if let Some(sri) = dist.integrity.as_deref() {
359388
crate::vendor::registry_fetch::verify_sri(&bytes, sri)?;
@@ -474,11 +503,9 @@ impl UpstreamClient {
474503
let escaped = crate::crawlers::go_crawler::encode_module_path(module);
475504
let escaped_version = crate::crawlers::go_crawler::encode_module_path(version);
476505
let base = format!("{proxy}/{escaped}/@v/{escaped_version}");
477-
let zip =
478-
crate::vendor::registry_fetch::download(&self.http, &format!("{base}.zip")).await?;
506+
let zip = self.download(&format!("{base}.zip")).await?;
479507
let zip_h1 = crate::vendor::registry_fetch::go_h1_of_zip(&zip)?;
480-
let go_mod =
481-
crate::vendor::registry_fetch::download(&self.http, &format!("{base}.mod")).await?;
508+
let go_mod = self.download(&format!("{base}.mod")).await?;
482509
Ok(GoSums {
483510
zip_h1,
484511
mod_h1: go_mod_h1(&go_mod),
@@ -587,7 +614,7 @@ impl UpstreamClient {
587614
"{}/v3-flatcontainer/{id_seg}/{version_seg}/{id_seg}.{version_seg}.nupkg",
588615
nuget_api_base(),
589616
);
590-
let bytes = crate::vendor::registry_fetch::download(&self.http, &url).await?;
617+
let bytes = self.download(&url).await?;
591618
crate::formats::nuget::package::package_content_hash(&bytes)
592619
.map_err(|why| format!("{url}: {why}"))
593620
}
@@ -699,10 +726,113 @@ impl UpstreamClient {
699726
}
700727
}
701728

729+
/// A local registry that counts how many requests it holds open at once
730+
/// (#1220). Each request is answered after a fixed delay with the same
731+
/// JSON body, over HTTP/1.1 keep-alive, so a client that pipelines several
732+
/// requests per connection is counted as accurately as one that opens a
733+
/// connection per request.
734+
#[cfg(test)]
735+
pub(crate) mod peak_server {
736+
use std::sync::atomic::{AtomicUsize, Ordering};
737+
use std::sync::Arc;
738+
use std::time::Duration;
739+
740+
use tokio::io::{AsyncReadExt, AsyncWriteExt};
741+
742+
pub(crate) struct PeakServer {
743+
pub base: String,
744+
pub requests: Arc<AtomicUsize>,
745+
peak: Arc<AtomicUsize>,
746+
}
747+
748+
impl PeakServer {
749+
pub(crate) async fn start(body: &'static str, delay: Duration) -> Self {
750+
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
751+
let base = format!("http://{}", listener.local_addr().unwrap());
752+
let live = Arc::new(AtomicUsize::new(0));
753+
let peak = Arc::new(AtomicUsize::new(0));
754+
let requests = Arc::new(AtomicUsize::new(0));
755+
let (l, p, r) = (live.clone(), peak.clone(), requests.clone());
756+
tokio::spawn(async move {
757+
loop {
758+
let Ok((mut sock, _)) = listener.accept().await else {
759+
return;
760+
};
761+
let (live, peak, requests) = (l.clone(), p.clone(), r.clone());
762+
tokio::spawn(async move {
763+
let mut buf = Vec::new();
764+
let mut chunk = [0u8; 4096];
765+
loop {
766+
// One request: a GET has no body, so its head
767+
// ends the request.
768+
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
769+
match sock.read(&mut chunk).await {
770+
Ok(0) | Err(_) => return,
771+
Ok(n) => buf.extend_from_slice(&chunk[..n]),
772+
}
773+
}
774+
let end = buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap();
775+
buf.drain(..end + 4);
776+
requests.fetch_add(1, Ordering::SeqCst);
777+
let now = live.fetch_add(1, Ordering::SeqCst) + 1;
778+
peak.fetch_max(now, Ordering::SeqCst);
779+
tokio::time::sleep(delay).await;
780+
live.fetch_sub(1, Ordering::SeqCst);
781+
let response = format!(
782+
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\n\
783+
content-length: {}\r\n\r\n{body}",
784+
body.len()
785+
);
786+
if sock.write_all(response.as_bytes()).await.is_err() {
787+
return;
788+
}
789+
}
790+
});
791+
}
792+
});
793+
PeakServer {
794+
base,
795+
requests,
796+
peak,
797+
}
798+
}
799+
800+
/// The most requests the server ever held open at once.
801+
pub(crate) fn peak(&self) -> usize {
802+
self.peak.load(Ordering::SeqCst)
803+
}
804+
}
805+
}
806+
702807
#[cfg(test)]
703808
mod tests {
704809
use super::*;
705810

811+
/// #1220: a client built with a cap of 1 (what a tight
812+
/// `RLIMIT_NOFILE` gives) sends its lookups strictly one at a time,
813+
/// however many the caller starts at once.
814+
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
815+
async fn a_cap_of_one_sends_lookups_one_at_a_time() {
816+
let server = peak_server::PeakServer::start(
817+
r#"{"dist":{"tarball":"http://registry.test/x.tgz"}}"#,
818+
std::time::Duration::from_millis(50),
819+
)
820+
.await;
821+
let client = UpstreamClient::with_concurrency(false, 1);
822+
let names: Vec<String> = (0..10).map(|i| format!("pkg{i}")).collect();
823+
let lookups = names
824+
.iter()
825+
.map(|name| client.npm_dist_on(&server.base, name, "1.0.0"));
826+
for found in futures_util::future::join_all(lookups).await {
827+
found.unwrap();
828+
}
829+
assert_eq!(
830+
server.requests.load(std::sync::atomic::Ordering::SeqCst),
831+
10
832+
);
833+
assert_eq!(server.peak(), 1);
834+
}
835+
706836
#[test]
707837
fn crates_index_paths_follow_cargo_layout() {
708838
assert_eq!(crates_index_path("a"), "1/a");

‎crates/socket-patch-core/src/patch/redirect/upstream/npm.rs‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2745,6 +2745,42 @@ mod tests {
27452745
};
27462746
use crate::patch::redirect::upstream::client::NpmDist;
27472747

2748+
/// #1220: a restore with many npm pins never has more lookups in
2749+
/// flight than `registry_concurrency()`, and still uses that many (the
2750+
/// restore stays concurrent, not serial).
2751+
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
2752+
async fn restore_lookups_never_exceed_the_registry_cap() {
2753+
use crate::patch::redirect::upstream::client::{peak_server::PeakServer, UpstreamClient};
2754+
use crate::patch::redirect::upstream::{Ctx, FormatResult};
2755+
let server = PeakServer::start(
2756+
r#"{"dist":{"tarball":"http://registry.test/x.tgz"}}"#,
2757+
std::time::Duration::from_millis(50),
2758+
)
2759+
.await;
2760+
let client = UpstreamClient::new(false);
2761+
let ctx = Ctx {
2762+
client: &client,
2763+
origins: &[],
2764+
bun_lockb: false,
2765+
};
2766+
let wanted: std::collections::BTreeSet<(String, String, String)> = (0..40)
2767+
.map(|i| (format!("uuid-{i}"), format!("pkg{i}"), "1.0.0".to_string()))
2768+
.collect();
2769+
let mut result = FormatResult::default();
2770+
let base = server.base.clone();
2771+
let dists = super::fetch_dists_on(&wanted, |_| Some(base.clone()), &ctx, &mut result).await;
2772+
assert!(result.refused.is_empty(), "{:?}", result.refused);
2773+
assert_eq!(dists.len(), 40);
2774+
assert!(dists.values().all(|d| d.from_project));
2775+
assert_eq!(
2776+
server.requests.load(std::sync::atomic::Ordering::SeqCst),
2777+
40
2778+
);
2779+
let cap = crate::utils::concurrent::registry_concurrency();
2780+
assert!(server.peak() <= cap, "peak {} > cap {cap}", server.peak());
2781+
assert_eq!(server.peak(), cap, "the restore should still fill the cap");
2782+
}
2783+
27482784
/// #1435: what still declares a descriptor the hosted restore keys
27492785
/// the registry entry by.
27502786
#[test]

0 commit comments

Comments
 (0)