66use std:: collections:: HashMap ;
77
88use serde_json:: Value ;
9- use tokio:: sync:: Mutex ;
9+ use tokio:: sync:: { Mutex , Semaphore , SemaphorePermit } ;
1010
1111use 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).
167172pub ( 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
180188impl 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 \n content-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) ]
703808mod 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" ) ;
0 commit comments