Skip to main content

uv_distribution/
distribution_database.rs

1use std::cmp::Reverse;
2use std::future::Future;
3use std::io;
4use std::path::Path;
5use std::pin::Pin;
6use std::sync::Arc;
7use std::task::{Context, Poll};
8
9use futures::{FutureExt, TryStreamExt};
10use rayon::in_place_scope;
11use rayon::prelude::*;
12use rustc_hash::FxHashMap;
13use tokio::io::{AsyncRead, AsyncSeekExt, ReadBuf};
14use tokio::sync::Semaphore;
15use tokio_util::compat::FuturesAsyncReadCompatExt;
16use tracing::{Instrument, info_span, instrument, warn};
17use url::Url;
18
19use uv_cache::{ArchiveFileId, ArchiveId, Cache, CacheBucket, CacheEntry, WheelCache};
20use uv_cache_info::{CacheInfo, Timestamp};
21use uv_client::{
22    CacheControl, CachedClientError, Connectivity, DataWithCachePolicy, RegistryClient,
23};
24use uv_configuration::initialize_rayon_once;
25use uv_distribution_filename::WheelFilename;
26use uv_distribution_types::{
27    ArchiveHashPolicy, BuildInfo, BuildableSource, BuiltDist, Dist, DistRef, HashCollection,
28    HashValidation, Hashed, IndexUrl, InstalledDist, MetadataHashPolicy, Name, SourceDist,
29    SourceUrl, parse_url_hashes,
30};
31use uv_extract::dirhash::{DirectoryDigest, HashedFile};
32use uv_extract::hash::Hasher;
33use uv_fs::{LockedFile, write_atomic};
34use uv_git::{GIT_LFS, GitError};
35use uv_platform_tags::Tags;
36use uv_preview::PreviewFeature;
37use uv_pypi_types::{HashDigest, HashDigests, PyProjectToml};
38use uv_python::PythonVariant;
39use uv_redacted::DisplaySafeUrl;
40use uv_types::{BuildContext, BuildStack};
41
42use crate::archive::Archive;
43use crate::error::PythonVersion;
44use crate::extracted_wheel::{ExtractedWheel, HashedWheel, WheelExtractor};
45use crate::hash::http_hash_algorithms;
46use crate::metadata::{ArchiveMetadata, Metadata};
47use crate::source::SourceDistributionBuilder;
48use crate::{Error, LocalWheel, Reporter, RequiresDist};
49
50/// A cached high-level interface to convert distributions (a requirement resolved to a location)
51/// to a wheel or wheel metadata.
52///
53/// For wheel metadata, this happens by either fetching the metadata from the remote wheel or by
54/// building the source distribution. For wheel files, either the wheel is downloaded or a source
55/// distribution is downloaded, built and the new wheel gets returned.
56///
57/// All kinds of wheel sources (index, URL, path) and source distribution source (index, URL, path,
58/// Git) are supported.
59///
60/// This struct also has the task of acquiring locks around source dist builds in general and git
61/// operation especially, as well as respecting concurrency limits.
62pub struct DistributionDatabase<'a, Context: BuildContext> {
63    build_context: &'a Context,
64    builder: SourceDistributionBuilder<'a, Context>,
65    client: ManagedClient<'a>,
66    reporter: Option<Arc<dyn Reporter>>,
67    content_addressed_cache: bool,
68}
69
70impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
71    pub fn new(
72        client: &'a RegistryClient,
73        build_context: &'a Context,
74        downloads_semaphore: Arc<Semaphore>,
75    ) -> Self {
76        // When ZIP validation is disabled, the extracted tree can contain files that aren't
77        // represented in the central directory and therefore aren't included in its digest.
78        // Avoid using an incomplete digest as a content-addressed archive ID.
79        let content_addressed_cache = uv_preview::is_enabled(PreviewFeature::ContentAddressedCache)
80            && !uv_extract::insecure_no_validate();
81        Self {
82            build_context,
83            builder: SourceDistributionBuilder::new(build_context),
84            client: ManagedClient::new(client, downloads_semaphore),
85            reporter: None,
86            content_addressed_cache,
87        }
88    }
89
90    /// Set the build stack to use for the [`DistributionDatabase`].
91    #[must_use]
92    pub fn with_build_stack(self, build_stack: &'a BuildStack) -> Self {
93        Self {
94            builder: self.builder.with_build_stack(build_stack),
95            ..self
96        }
97    }
98
99    /// Set the [`Reporter`] to use for the [`DistributionDatabase`].
100    #[must_use]
101    pub fn with_reporter(self, reporter: Arc<dyn Reporter>) -> Self {
102        Self {
103            builder: self.builder.with_reporter(reporter.clone()),
104            reporter: Some(reporter),
105            ..self
106        }
107    }
108
109    /// Handle a specific `reqwest` error, and convert it to [`io::Error`].
110    fn handle_response_errors(&self, err: reqwest::Error) -> io::Error {
111        if err.is_timeout() {
112            // Assumption: The connect timeout with the 10s default is not the culprit.
113            io::Error::new(
114                io::ErrorKind::TimedOut,
115                format!(
116                    "Failed to download distribution due to network timeout. Try increasing UV_HTTP_TIMEOUT (current value: {}s).",
117                    self.client.unmanaged.read_timeout().as_secs()
118                ),
119            )
120        } else {
121            io::Error::other(err)
122        }
123    }
124
125    /// Acquire an advisory lock for a wheel cache entry.
126    ///
127    /// A remote wheel's content hash is not always available until after the download, so
128    /// concurrent cache fills coordinate on the wheel cache entry instead. The entry is already
129    /// scoped to the distribution's source and wheel filename.
130    ///
131    /// Callers hold the returned lock across cache lookup, download or extraction, and publication.
132    /// A process that waited for another cache fill therefore rechecks and reuses the completed entry.
133    async fn lock_wheel(
134        wheel_entry: &CacheEntry,
135        filename: &WheelFilename,
136    ) -> Result<LockedFile, Error> {
137        // For backwards compatibility, we use the full wheel stem on Windows. Local wheel
138        // extraction and older uv versions use the same key, so changing it would prevent them
139        // from coordinating through a shared cache.
140        #[cfg(windows)]
141        let lock_key = filename.stem();
142        // On other platforms, we use the bounded cache key to avoid filesystem filename limits.
143        #[cfg(not(windows))]
144        let lock_key = filename.cache_key();
145
146        let lock_entry = wheel_entry.with_file(format!("{lock_key}.lock"));
147        lock_entry.lock().await.map_err(Error::CacheLock)
148    }
149
150    /// Either fetch the wheel or fetch and build the source distribution
151    ///
152    /// Returns a wheel that's compliant with the given platform tags.
153    ///
154    /// While hashes will be generated in some cases, hash-checking is only enforced for source
155    /// distributions, and should be enforced by the caller for wheels.
156    #[instrument(skip_all, fields(%dist))]
157    pub async fn get_or_build_wheel(
158        &self,
159        dist: &Dist,
160        tags: &Tags,
161        hashes: ArchiveHashPolicy<'_>,
162    ) -> Result<LocalWheel, Error> {
163        match dist {
164            Dist::Built(built) => self.get_wheel(built, hashes).await,
165            Dist::Source(source) => self.build_wheel(source, tags, hashes).await,
166        }
167    }
168
169    /// Either fetch the only wheel metadata (directly from the index or with range requests) or
170    /// fetch and build the source distribution.
171    ///
172    /// While hashes will be generated in some cases, hash-checking is only enforced for source
173    /// distributions, and should be enforced by the caller for wheels.
174    #[instrument(skip_all, fields(%dist))]
175    pub async fn get_installed_metadata(
176        &self,
177        dist: &InstalledDist,
178    ) -> Result<ArchiveMetadata, Error> {
179        // If the metadata was provided by the user directly, prefer it.
180        if let Some(metadata) = self
181            .build_context
182            .dependency_metadata()
183            .get(dist.name(), Some(dist.version()))
184        {
185            return Ok(ArchiveMetadata::from_metadata23(metadata));
186        }
187
188        let metadata = dist
189            .read_metadata()
190            .map_err(|err| Error::ReadInstalled(Box::new(dist.clone()), err))?;
191
192        Ok(ArchiveMetadata::from_metadata23(metadata.clone()))
193    }
194
195    /// Retrieve distribution metadata and any hashes requested for resolution.
196    ///
197    /// Source archive hashes are validated before executing build backends; wheel archive hash
198    /// validation is deferred to installation.
199    #[instrument(skip_all, fields(%dist))]
200    pub async fn get_or_build_wheel_metadata(
201        &self,
202        dist: &Dist,
203        hashes: MetadataHashPolicy<'_>,
204    ) -> Result<ArchiveMetadata, Error> {
205        match dist {
206            Dist::Built(built) => self.get_wheel_metadata(built, hashes).await,
207            Dist::Source(source) => {
208                self.build_wheel_metadata(&BuildableSource::Dist(source), hashes)
209                    .await
210            }
211        }
212    }
213
214    /// Fetch a wheel from the cache or download it from the index.
215    ///
216    /// While hashes will be generated in all cases, hash-checking is _not_ enforced and should
217    /// instead be enforced by the caller.
218    async fn get_wheel(
219        &self,
220        dist: &BuiltDist,
221        hashes: ArchiveHashPolicy<'_>,
222    ) -> Result<LocalWheel, Error> {
223        match dist {
224            BuiltDist::Registry(wheels) => {
225                let wheel = wheels.best_wheel();
226                let url = wheel.file.url.to_url()?;
227                let size = wheel.file.size;
228
229                // Create a cache entry for the wheel.
230                let wheel_entry = self.build_context.cache().entry(
231                    CacheBucket::Wheels,
232                    WheelCache::Index(&wheel.index).wheel_dir(wheel.name().as_ref()),
233                    wheel.filename.cache_key(),
234                );
235
236                // If the URL is a file URL, load the wheel directly.
237                if url.scheme() == "file" {
238                    let path = url
239                        .to_file_path()
240                        .map_err(|()| Error::NonFileUrl(url.clone()))?;
241                    return self
242                        .load_wheel(&path, &wheel.filename, wheel_entry, dist, hashes)
243                        .await;
244                }
245
246                // Download and unzip.
247                match self
248                    .stream_wheel(
249                        url.clone(),
250                        dist.index(),
251                        &wheel.filename,
252                        size,
253                        &wheel_entry,
254                        dist,
255                        hashes,
256                    )
257                    .await
258                {
259                    Ok(archive) => Ok(LocalWheel {
260                        dist: Dist::Built(dist.clone()),
261                        archive: self
262                            .build_context
263                            .cache()
264                            .archive(&archive.id)
265                            .into_boxed_path(),
266                        hashes: archive.hashes,
267                        filename: wheel.filename.clone(),
268                        cache: CacheInfo::default(),
269                        build: None,
270                    }),
271                    Err(Error::Extract(name, err)) => {
272                        if err.is_http_streaming_unsupported() {
273                            warn!(
274                                "Streaming unsupported for {dist}; downloading wheel to disk ({err})"
275                            );
276                        } else if err.is_http_streaming_failed() {
277                            warn!("Streaming failed for {dist}; downloading wheel to disk ({err})");
278                        } else {
279                            return Err(Error::Extract(name, err));
280                        }
281
282                        // If the request failed because streaming was unsupported or failed,
283                        // download the wheel directly.
284                        let archive = self
285                            .download_wheel(
286                                url,
287                                dist.index(),
288                                &wheel.filename,
289                                size,
290                                &wheel_entry,
291                                dist,
292                                hashes,
293                            )
294                            .await?;
295
296                        Ok(LocalWheel {
297                            dist: Dist::Built(dist.clone()),
298                            archive: self
299                                .build_context
300                                .cache()
301                                .archive(&archive.id)
302                                .into_boxed_path(),
303                            hashes: archive.hashes,
304                            filename: wheel.filename.clone(),
305                            cache: CacheInfo::default(),
306                            build: None,
307                        })
308                    }
309                    Err(err) => Err(err),
310                }
311            }
312
313            BuiltDist::DirectUrl(wheel) => {
314                // Create a cache entry for the wheel.
315                let wheel_entry = self.build_context.cache().entry(
316                    CacheBucket::Wheels,
317                    WheelCache::Url(&wheel.url).wheel_dir(wheel.name().as_ref()),
318                    wheel.filename.cache_key(),
319                );
320
321                // Download and unzip.
322                match self
323                    .stream_wheel(
324                        wheel.url.raw().clone(),
325                        None,
326                        &wheel.filename,
327                        wheel.size,
328                        &wheel_entry,
329                        dist,
330                        hashes,
331                    )
332                    .await
333                {
334                    Ok(archive) => Ok(LocalWheel {
335                        dist: Dist::Built(dist.clone()),
336                        archive: self
337                            .build_context
338                            .cache()
339                            .archive(&archive.id)
340                            .into_boxed_path(),
341                        hashes: archive.hashes,
342                        filename: wheel.filename.clone(),
343                        cache: CacheInfo::default(),
344                        build: None,
345                    }),
346                    Err(Error::Extract(name, err)) => {
347                        if err.is_http_streaming_unsupported() {
348                            warn!(
349                                "Streaming unsupported for {dist}; downloading wheel to disk ({err})"
350                            );
351                        } else if err.is_http_streaming_failed() {
352                            warn!("Streaming failed for {dist}; downloading wheel to disk ({err})");
353                        } else {
354                            return Err(Error::Extract(name, err));
355                        }
356
357                        // If the request failed because streaming was unsupported or failed,
358                        // download the wheel directly.
359                        let archive = self
360                            .download_wheel(
361                                wheel.url.raw().clone(),
362                                None,
363                                &wheel.filename,
364                                wheel.size,
365                                &wheel_entry,
366                                dist,
367                                hashes,
368                            )
369                            .await?;
370                        Ok(LocalWheel {
371                            dist: Dist::Built(dist.clone()),
372                            archive: self
373                                .build_context
374                                .cache()
375                                .archive(&archive.id)
376                                .into_boxed_path(),
377                            hashes: archive.hashes,
378                            filename: wheel.filename.clone(),
379                            cache: CacheInfo::default(),
380                            build: None,
381                        })
382                    }
383                    Err(err) => Err(err),
384                }
385            }
386
387            BuiltDist::GitPath(wheel) => {
388                // Fetch the Git repository.
389                let fetch = self
390                    .build_context
391                    .git()
392                    .fetch(
393                        &wheel.git,
394                        self.client.unmanaged.git_http_settings(wheel.git.url()),
395                        self.build_context.cache().bucket(CacheBucket::Git),
396                        self.reporter.clone().map(<dyn Reporter>::into_git_reporter),
397                    )
398                    .await?;
399
400                if wheel.git.lfs().enabled() && !fetch.lfs_ready() {
401                    if GIT_LFS.is_err() {
402                        return Err(Error::MissingWheelGitLfsArtifacts(
403                            wheel.url.to_url(),
404                            GitError::GitLfsNotFound,
405                        ));
406                    }
407                    return Err(Error::MissingWheelGitLfsArtifacts(
408                        wheel.url.to_url(),
409                        GitError::GitLfsNotConfigured,
410                    ));
411                }
412
413                let git_sha = fetch.git().precise().expect("Exact commit after checkout");
414                let cache_entry = self.build_context.cache().entry(
415                    CacheBucket::Wheels,
416                    WheelCache::Git(&wheel.url, git_sha.as_short_str()).root(),
417                    wheel.filename.stem(),
418                );
419
420                let install_path = fetch.path().join(&wheel.install_path);
421
422                self.load_wheel(&install_path, &wheel.filename, cache_entry, dist, hashes)
423                    .await
424            }
425
426            BuiltDist::Path(wheel) => {
427                let cache_entry = self.build_context.cache().entry(
428                    CacheBucket::Wheels,
429                    WheelCache::Url(&wheel.url).wheel_dir(wheel.name().as_ref()),
430                    wheel.filename.cache_key(),
431                );
432
433                self.load_wheel(
434                    &wheel.install_path,
435                    &wheel.filename,
436                    cache_entry,
437                    dist,
438                    hashes,
439                )
440                .await
441            }
442        }
443    }
444
445    /// Convert a source distribution into a wheel, fetching it from the cache or building it if
446    /// necessary.
447    ///
448    /// The returned wheel is guaranteed to come from a distribution with a matching hash, and
449    /// no build processes will be executed for distributions with mismatched hashes.
450    async fn build_wheel(
451        &self,
452        dist: &SourceDist,
453        tags: &Tags,
454        hashes: ArchiveHashPolicy<'_>,
455    ) -> Result<LocalWheel, Error> {
456        let built_wheel = self
457            .builder
458            .download_and_build(&BuildableSource::Dist(dist), tags, hashes, &self.client)
459            .boxed_local()
460            .await?;
461
462        // Check that the wheel is compatible with its install target.
463        //
464        // When building a build dependency for a cross-install, the build dependency needs
465        // to install and run on the host instead of the target. In this case the `tags` are already
466        // for the host instead of the target, so this check passes.
467        if !built_wheel.filename.is_compatible(tags) {
468            return if tags.is_cross() {
469                Err(Error::BuiltWheelIncompatibleTargetPlatform {
470                    filename: built_wheel.filename,
471                    python_platform: tags.python_platform().clone(),
472                    python_version: PythonVersion {
473                        version: tags.python_version(),
474                        variant: if tags.is_freethreaded() {
475                            PythonVariant::Freethreaded
476                        } else {
477                            PythonVariant::Default
478                        },
479                    },
480                })
481            } else {
482                Err(Error::BuiltWheelIncompatibleHostPlatform {
483                    filename: built_wheel.filename,
484                    python_platform: tags.python_platform().clone(),
485                    python_version: PythonVersion {
486                        version: tags.python_version(),
487                        variant: if tags.is_freethreaded() {
488                            PythonVariant::Freethreaded
489                        } else {
490                            PythonVariant::Default
491                        },
492                    },
493                })
494            };
495        }
496
497        // Acquire the advisory lock.
498        let wheel_entry = CacheEntry::from_path(built_wheel.target.as_ref());
499        let _lock = Self::lock_wheel(&wheel_entry, &built_wheel.filename).await?;
500
501        // If the wheel was unzipped previously, respect it. Source distributions are
502        // cached under a unique revision ID, so unzipped directories are never stale.
503        match self.build_context.cache().resolve_link(&built_wheel.target) {
504            Ok(archive) => {
505                return Ok(LocalWheel {
506                    dist: Dist::Source(dist.clone()),
507                    archive: archive.into_boxed_path(),
508                    filename: built_wheel.filename,
509                    hashes: built_wheel.hashes,
510                    cache: built_wheel.cache_info,
511                    build: Some(built_wheel.build_info),
512                });
513            }
514            Err(err) if err.kind() == io::ErrorKind::NotFound => {}
515            Err(err) => return Err(Error::CacheRead(err)),
516        }
517
518        // Otherwise, unzip the wheel.
519        let id = self
520            .unzip_wheel(
521                &built_wheel.path,
522                &built_wheel.target,
523                DistRef::Source(dist),
524            )
525            .await?;
526
527        Ok(LocalWheel {
528            dist: Dist::Source(dist.clone()),
529            archive: self.build_context.cache().archive(&id).into_boxed_path(),
530            hashes: built_wheel.hashes,
531            filename: built_wheel.filename,
532            cache: built_wheel.cache_info,
533            build: Some(built_wheel.build_info),
534        })
535    }
536
537    /// Fetch wheel metadata, along with any hashes requested for resolution.
538    ///
539    /// Wheel archive hash validation is deferred to installation. Metadata sidecar hashes are
540    /// checked by the client when downloading sidecars with index-provided hashes.
541    async fn get_wheel_metadata(
542        &self,
543        dist: &BuiltDist,
544        hashes: MetadataHashPolicy<'_>,
545    ) -> Result<ArchiveMetadata, Error> {
546        let hash_policy = match hashes.validation {
547            HashValidation::None => {
548                let compute_hashes = match hashes.collection {
549                    HashCollection::None => false,
550                    HashCollection::Url | HashCollection::All => match dist {
551                        BuiltDist::Registry(dist) => {
552                            // Generate missing hashes for indexes and `--find-links` without hashes.
553                            // This hashes only the selected wheel, not every distribution for the version.
554                            hashes.collection == HashCollection::All
555                                && dist.best_wheel().file.hashes.is_empty()
556                        }
557                        BuiltDist::DirectUrl(dist) => parse_url_hashes(&dist.url).is_none(),
558                        BuiltDist::Path(_) | BuiltDist::GitPath(_) => true,
559                    },
560                };
561                if compute_hashes {
562                    ArchiveHashPolicy::Generate
563                } else {
564                    ArchiveHashPolicy::None
565                }
566            }
567            HashValidation::Any(_) | HashValidation::All(_) => hashes.validation.into(),
568        };
569
570        // Fetch the entire wheel only when we need to compute a hash for resolution.
571        // TODO(charlie): Request the hashes via a separate method, to reduce the coupling in this API.
572        if hash_policy == ArchiveHashPolicy::Generate {
573            let wheel = self.get_wheel(dist, hash_policy).await?;
574            // If the metadata was provided by the user directly, prefer it.
575            let metadata = if let Some(metadata) = self
576                .build_context
577                .dependency_metadata()
578                .get(dist.name(), Some(dist.version()))
579            {
580                metadata
581            } else {
582                wheel.metadata()?
583            };
584            let hashes = wheel.hashes;
585            return Ok(ArchiveMetadata {
586                metadata: Metadata::from_metadata23(metadata),
587                hashes,
588            });
589        }
590
591        // If the metadata was provided by the user directly, prefer it.
592        if let Some(metadata) = self
593            .build_context
594            .dependency_metadata()
595            .get(dist.name(), Some(dist.version()))
596        {
597            return Ok(ArchiveMetadata::from_metadata23(metadata));
598        }
599
600        let result = self
601            .client
602            .managed(|client| {
603                client
604                    .wheel_metadata(
605                        dist,
606                        self.build_context.git(),
607                        self.build_context.capabilities(),
608                        self.reporter.clone().map(<dyn Reporter>::into_git_reporter),
609                    )
610                    .boxed_local()
611            })
612            .await;
613
614        match result {
615            Ok(metadata) => {
616                // Validate that the metadata is consistent with the distribution.
617                Ok(ArchiveMetadata::from_metadata23(metadata))
618            }
619            Err(err) if err.is_http_streaming_unsupported() => {
620                warn!(
621                    "Streaming unsupported when fetching metadata for {dist}; downloading wheel directly ({err})"
622                );
623
624                // If the request failed due to an error that could be resolved by
625                // downloading the wheel directly, try that.
626                let wheel = self.get_wheel(dist, hash_policy).await?;
627                let metadata = wheel.metadata()?;
628                let hashes = wheel.hashes;
629                Ok(ArchiveMetadata {
630                    metadata: Metadata::from_metadata23(metadata),
631                    hashes,
632                })
633            }
634            Err(err) => Err(err.into()),
635        }
636    }
637
638    /// Build the wheel metadata for a source distribution, or fetch it from the cache if possible.
639    ///
640    /// Requested hashes are validated before executing a build backend. User-provided metadata
641    /// can avoid downloading or validating the source archive.
642    pub async fn build_wheel_metadata(
643        &self,
644        source: &BuildableSource<'_>,
645        hashes: MetadataHashPolicy<'_>,
646    ) -> Result<ArchiveMetadata, Error> {
647        // If the metadata was provided by the user directly, prefer it.
648        if let Some(dist) = source.as_dist() {
649            if let Some(metadata) = self
650                .build_context
651                .dependency_metadata()
652                .get(dist.name(), dist.version())
653            {
654                // If we skipped the build, we should still resolve any Git dependencies to precise
655                // commits.
656                self.builder.resolve_revision(source, &self.client).await?;
657
658                return Ok(ArchiveMetadata::from_metadata23(metadata));
659            }
660        }
661
662        let url_hashes = if let BuildableSource::Dist(SourceDist::DirectUrl(dist)) = source {
663            parse_url_hashes(&dist.url)
664        } else if let BuildableSource::Url(SourceUrl::Direct(url)) = source {
665            parse_url_hashes(url.url)
666        } else {
667            None
668        };
669
670        let build_hash_policy = match hashes.validation {
671            HashValidation::None => match hashes.collection {
672                HashCollection::None => ArchiveHashPolicy::None,
673                // If resolving metadata requires a build, validate any URL hash before executing
674                // the backend, even when the caller only requested hash collection.
675                HashCollection::Url | HashCollection::All => match url_hashes.as_ref() {
676                    Some(digests) => ArchiveHashPolicy::All(digests.as_slice()),
677                    None => ArchiveHashPolicy::Generate,
678                },
679            },
680            HashValidation::Any(_) | HashValidation::All(_) => hashes.validation.into(),
681        };
682        let metadata = self
683            .builder
684            .download_and_build_metadata(source, build_hash_policy, &self.client)
685            .boxed_local()
686            .await?;
687
688        Ok(metadata)
689    }
690
691    /// Return the [`RequiresDist`] from a `pyproject.toml`, if it can be statically extracted.
692    pub async fn requires_dist(
693        &self,
694        path: &Path,
695        pyproject_toml: &PyProjectToml,
696    ) -> Result<Option<RequiresDist>, Error> {
697        self.builder
698            .source_tree_requires_dist(
699                path,
700                pyproject_toml,
701                self.client.unmanaged.credentials_cache(),
702            )
703            .await
704    }
705
706    /// Stream a wheel from a URL, unzipping it into the cache as it's downloaded.
707    async fn stream_wheel(
708        &self,
709        url: DisplaySafeUrl,
710        index: Option<&IndexUrl>,
711        filename: &WheelFilename,
712        size: Option<u64>,
713        wheel_entry: &CacheEntry,
714        dist: &BuiltDist,
715        hashes: ArchiveHashPolicy<'_>,
716    ) -> Result<Archive, Error> {
717        let expected_size = match dist {
718            BuiltDist::Registry(dist) if dist.best_wheel().size_is_authoritative => size,
719            BuiltDist::DirectUrl(_) => size,
720            _ => None,
721        };
722
723        // Acquire an advisory lock, to guard against concurrent writes.
724        let _lock = Self::lock_wheel(wheel_entry, filename).await?;
725
726        // Create an entry for the HTTP cache.
727        let http_entry = wheel_entry.with_file(format!("{}.http", filename.cache_key()));
728
729        let download = |response: reqwest::Response| {
730            async {
731                let progress_size = size.or_else(|| content_length(&response));
732
733                let progress = self.reporter.as_ref().map(|reporter| {
734                    (
735                        reporter,
736                        reporter.on_download_start(dist.name(), progress_size),
737                    )
738                });
739
740                let reader = response
741                    .bytes_stream()
742                    .map_err(|err| self.handle_response_errors(err))
743                    .into_async_read();
744
745                // Create a hasher for each hash algorithm.
746                let algorithms = http_hash_algorithms(hashes);
747                let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
748                let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers);
749
750                // Download and unzip the wheel to a temporary directory.
751                let extractor = WheelExtractor::new(
752                    self.build_context.cache().root(),
753                    self.content_addressed_cache,
754                )
755                .map_err(Error::CacheWrite)?;
756
757                let mut extracted = match progress {
758                    Some((reporter, progress)) => {
759                        let mut reader = ProgressReader::new(&mut hasher, progress, &**reporter);
760                        extractor
761                            .extract_streaming(&mut reader)
762                            .await
763                            .map_err(|err| Error::Extract(filename.to_string(), err))?
764                    }
765                    None => extractor
766                        .extract_streaming(&mut hasher)
767                        .await
768                        .map_err(|err| Error::Extract(filename.to_string(), err))?,
769                };
770                // Exhaust the reader to compute the hashes.
771                hasher.finish().await.map_err(Error::HashExhaustion)?;
772                let actual_size = hasher.bytes_read();
773                if let Some(expected) = expected_size
774                    && actual_size != expected
775                {
776                    return Err(Error::MismatchedSize {
777                        distribution: dist.to_string(),
778                        expected,
779                        actual: actual_size,
780                    });
781                }
782
783                // Before we make the wheel accessible by persisting it, ensure that the RECORD is
784                // valid.
785                extracted.validate_and_heal_record(dist)?;
786
787                // Persist the temporary directory to the directory store.
788                let id = self
789                    .persist_extracted_wheel(extracted, wheel_entry.path())
790                    .await?;
791
792                if let Some((reporter, progress)) = progress {
793                    reporter.on_download_complete(dist.name(), progress);
794                }
795
796                Ok(Archive::new(
797                    id,
798                    hashers.into_iter().map(HashDigest::from).collect(),
799                    filename.clone(),
800                    Some(actual_size),
801                ))
802            }
803            .instrument(info_span!("wheel", wheel = %dist))
804        };
805
806        // Fetch the archive from the cache, or download it if necessary.
807        let req = self.request(url.clone())?;
808
809        // Determine the cache control policy for the URL.
810        let cache_control = match self.client.unmanaged.connectivity() {
811            Connectivity::Online
812                if let Some(header) = index.and_then(|index| {
813                    self.build_context
814                        .locations()
815                        .artifact_cache_control_for(index)
816                }) =>
817            {
818                CacheControl::Override(header)
819            }
820            Connectivity::Online => CacheControl::from(
821                self.build_context
822                    .cache()
823                    .freshness(&http_entry, Some(&filename.name), None)
824                    .map_err(Error::CacheRead)?,
825            ),
826            Connectivity::Offline => CacheControl::AllowStale,
827        };
828
829        let archive = self
830            .client
831            .managed(|client| {
832                client.cached_client().get_serde_with_retry(
833                    req,
834                    &http_entry,
835                    cache_control.clone(),
836                    download,
837                )
838            })
839            .await
840            .map_err(|err| match err {
841                CachedClientError::Callback { err, .. } => err,
842                CachedClientError::Client(err) => Error::Client(err),
843            })?;
844
845        if let (Some(expected), Some(actual)) = (expected_size, archive.size)
846            && expected != actual
847        {
848            return Err(Error::MismatchedSize {
849                distribution: dist.to_string(),
850                expected,
851                actual,
852            });
853        }
854
855        // If the archive is missing the required hashes or size, or has since been removed, force a refresh.
856        let archive = Some(archive)
857            .filter(|archive| archive.has_digests(hashes))
858            .filter(|archive| archive.exists(self.build_context.cache()))
859            .filter(|archive| expected_size.is_none() || archive.size.is_some());
860
861        let archive = if let Some(archive) = archive {
862            archive
863        } else {
864            self.client
865                .managed(async |client| {
866                    client
867                        .cached_client()
868                        .skip_cache_with_retry(
869                            self.request(url)?,
870                            &http_entry,
871                            cache_control,
872                            download,
873                        )
874                        .await
875                        .map_err(|err| match err {
876                            CachedClientError::Callback { err, .. } => err,
877                            CachedClientError::Client(err) => Error::Client(err),
878                        })
879                })
880                .await?
881        };
882
883        Ok(archive)
884    }
885
886    /// Download a wheel from a URL, then unzip it into the cache.
887    async fn download_wheel(
888        &self,
889        url: DisplaySafeUrl,
890        index: Option<&IndexUrl>,
891        filename: &WheelFilename,
892        size: Option<u64>,
893        wheel_entry: &CacheEntry,
894        dist: &BuiltDist,
895        hashes: ArchiveHashPolicy<'_>,
896    ) -> Result<Archive, Error> {
897        let expected_size = match dist {
898            BuiltDist::Registry(dist) if dist.best_wheel().size_is_authoritative => size,
899            BuiltDist::DirectUrl(_) => size,
900            _ => None,
901        };
902
903        let content_addressed_cache = self.content_addressed_cache;
904
905        // Acquire an advisory lock, to guard against concurrent writes.
906        let _lock = Self::lock_wheel(wheel_entry, filename).await?;
907
908        // Create an entry for the HTTP cache.
909        let http_entry = wheel_entry.with_file(format!("{}.http", filename.cache_key()));
910
911        let download = |response: reqwest::Response| {
912            async {
913                let progress_size = size.or_else(|| content_length(&response));
914
915                let progress = self.reporter.as_ref().map(|reporter| {
916                    (
917                        reporter,
918                        reporter.on_download_start(dist.name(), progress_size),
919                    )
920                });
921
922                let reader = response
923                    .bytes_stream()
924                    .map_err(|err| self.handle_response_errors(err))
925                    .into_async_read();
926                let algorithms = http_hash_algorithms(hashes);
927                let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
928                let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers);
929
930                // Download the wheel to a temporary file.
931                let temp_file = tempfile::tempfile_in(self.build_context.cache().root())
932                    .map_err(Error::CacheWrite)?;
933                let mut writer = tokio::io::BufWriter::new(fs_err::tokio::File::from_std(
934                    // It's an unnamed file on Linux so that's the best approximation.
935                    fs_err::File::from_parts(temp_file, self.build_context.cache().root()),
936                ));
937
938                match progress {
939                    Some((reporter, progress)) => {
940                        // Wrap the reader in a progress reporter. This will report 100% progress once
941                        // the download is complete, before the wheel is unzipped.
942                        let mut reader = ProgressReader::new(&mut hasher, progress, &**reporter);
943
944                        tokio::io::copy(&mut reader, &mut writer)
945                            .await
946                            .map_err(Error::CacheWrite)?;
947                    }
948                    None => {
949                        tokio::io::copy(&mut hasher, &mut writer)
950                            .await
951                            .map_err(Error::CacheWrite)?;
952                    }
953                }
954
955                if let Some(expected) = expected_size
956                    && hasher.bytes_read() != expected
957                {
958                    return Err(Error::MismatchedSize {
959                        distribution: dist.to_string(),
960                        expected,
961                        actual: hasher.bytes_read(),
962                    });
963                }
964
965                let actual_size = hasher.bytes_read();
966
967                // Unzip the wheel to a temporary directory.
968                let extractor =
969                    WheelExtractor::new(self.build_context.cache().root(), content_addressed_cache)
970                        .map_err(Error::CacheWrite)?;
971                let mut file = writer.into_inner();
972                file.seek(io::SeekFrom::Start(0))
973                    .await
974                    .map_err(Error::CacheWrite)?;
975
976                let file = file.into_std().await;
977                let mut extracted =
978                    tokio::task::spawn_blocking(move || extractor.extract_seekable(file))
979                        .await?
980                        .map_err(|err| Error::Extract(filename.to_string(), err))?;
981                let hashes = hashers.into_iter().map(HashDigest::from).collect();
982
983                // Before we make the wheel accessible by persisting it, ensure that the RECORD is
984                // valid.
985                extracted.validate_and_heal_record(dist)?;
986
987                // Persist the temporary directory to the directory store.
988                let id = self
989                    .persist_extracted_wheel(extracted, wheel_entry.path())
990                    .await?;
991
992                if let Some((reporter, progress)) = progress {
993                    reporter.on_download_complete(dist.name(), progress);
994                }
995
996                Ok(Archive::new(
997                    id,
998                    hashes,
999                    filename.clone(),
1000                    Some(actual_size),
1001                ))
1002            }
1003            .instrument(info_span!("wheel", wheel = %dist))
1004        };
1005
1006        // Fetch the archive from the cache, or download it if necessary.
1007        let req = self.request(url.clone())?;
1008
1009        // Determine the cache control policy for the URL.
1010        let cache_control = match self.client.unmanaged.connectivity() {
1011            Connectivity::Online
1012                if let Some(header) = index.and_then(|index| {
1013                    self.build_context
1014                        .locations()
1015                        .artifact_cache_control_for(index)
1016                }) =>
1017            {
1018                CacheControl::Override(header)
1019            }
1020            Connectivity::Online => CacheControl::from(
1021                self.build_context
1022                    .cache()
1023                    .freshness(&http_entry, Some(&filename.name), None)
1024                    .map_err(Error::CacheRead)?,
1025            ),
1026            Connectivity::Offline => CacheControl::AllowStale,
1027        };
1028
1029        let archive = self
1030            .client
1031            .managed(|client| {
1032                client.cached_client().get_serde_with_retry(
1033                    req,
1034                    &http_entry,
1035                    cache_control.clone(),
1036                    download,
1037                )
1038            })
1039            .await
1040            .map_err(|err| match err {
1041                CachedClientError::Callback { err, .. } => err,
1042                CachedClientError::Client(err) => Error::Client(err),
1043            })?;
1044
1045        if let (Some(expected), Some(actual)) = (expected_size, archive.size)
1046            && expected != actual
1047        {
1048            return Err(Error::MismatchedSize {
1049                distribution: dist.to_string(),
1050                expected,
1051                actual,
1052            });
1053        }
1054
1055        // If the archive is missing the required hashes or size, or has since been removed, force a refresh.
1056        let archive = Some(archive)
1057            .filter(|archive| archive.has_digests(hashes))
1058            .filter(|archive| archive.exists(self.build_context.cache()))
1059            .filter(|archive| expected_size.is_none() || archive.size.is_some());
1060
1061        let archive = if let Some(archive) = archive {
1062            archive
1063        } else {
1064            self.client
1065                .managed(async |client| {
1066                    client
1067                        .cached_client()
1068                        .skip_cache_with_retry(
1069                            self.request(url)?,
1070                            &http_entry,
1071                            cache_control,
1072                            download,
1073                        )
1074                        .await
1075                        .map_err(|err| match err {
1076                            CachedClientError::Callback { err, .. } => err,
1077                            CachedClientError::Client(err) => Error::Client(err),
1078                        })
1079                })
1080                .await?
1081        };
1082
1083        Ok(archive)
1084    }
1085
1086    /// Load a wheel from a local path.
1087    async fn load_wheel(
1088        &self,
1089        path: &Path,
1090        filename: &WheelFilename,
1091        wheel_entry: CacheEntry,
1092        dist: &BuiltDist,
1093        hashes: ArchiveHashPolicy<'_>,
1094    ) -> Result<LocalWheel, Error> {
1095        // Acquire an advisory lock, to guard against concurrent writes.
1096        let _lock = Self::lock_wheel(&wheel_entry, filename).await?;
1097
1098        // Determine the last-modified time of the wheel.
1099        let modified = Timestamp::from_path(path).map_err(Error::CacheRead)?;
1100
1101        // Attempt to read the archive pointer from the cache.
1102        let pointer_entry = wheel_entry.with_file(format!("{}.rev", filename.cache_key()));
1103        let pointer = PathArchivePointer::read_from(&pointer_entry)?;
1104
1105        // Extract the archive from the pointer.
1106        let archive = pointer
1107            .filter(|pointer| pointer.is_up_to_date(modified))
1108            .map(PathArchivePointer::into_archive)
1109            .filter(|archive| archive.has_digests(hashes));
1110
1111        // If the file is already unzipped, and the cache is up-to-date, return it.
1112        if let Some(archive) = archive {
1113            Ok(LocalWheel {
1114                dist: Dist::Built(dist.clone()),
1115                archive: self
1116                    .build_context
1117                    .cache()
1118                    .archive(&archive.id)
1119                    .into_boxed_path(),
1120                hashes: archive.hashes,
1121                filename: filename.clone(),
1122                cache: CacheInfo::from_timestamp(modified),
1123                build: None,
1124            })
1125        } else if hashes.is_none() {
1126            // Otherwise, unzip the wheel.
1127            let archive = Archive::new(
1128                self.unzip_wheel(path, wheel_entry.path(), DistRef::Built(dist))
1129                    .await?,
1130                HashDigests::empty(),
1131                filename.clone(),
1132                None,
1133            );
1134
1135            // Write the archive pointer to the cache.
1136            let pointer = PathArchivePointer {
1137                timestamp: modified,
1138                archive: archive.clone(),
1139            };
1140            pointer.write_to(&pointer_entry).await?;
1141
1142            Ok(LocalWheel {
1143                dist: Dist::Built(dist.clone()),
1144                archive: self
1145                    .build_context
1146                    .cache()
1147                    .archive(&archive.id)
1148                    .into_boxed_path(),
1149                hashes: archive.hashes,
1150                filename: filename.clone(),
1151                cache: CacheInfo::from_timestamp(modified),
1152                build: None,
1153            })
1154        } else {
1155            // If necessary, compute the hashes of the wheel.
1156            let file = fs_err::tokio::File::open(path)
1157                .await
1158                .map_err(Error::CacheRead)?;
1159            let extractor = WheelExtractor::new(
1160                self.build_context.cache().root(),
1161                self.content_addressed_cache,
1162            )
1163            .map_err(Error::CacheWrite)?;
1164
1165            // Create a hasher for each hash algorithm.
1166            let algorithms = hashes.algorithms();
1167            let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
1168            let mut hasher = uv_extract::hash::HashReader::new(file, &mut hashers);
1169
1170            // Unzip the wheel to a temporary directory.
1171            let mut extracted = extractor
1172                .extract_streaming(&mut hasher)
1173                .await
1174                .map_err(|err| Error::Extract(filename.to_string(), err))?;
1175
1176            // Exhaust the reader to compute the hash.
1177            hasher.finish().await.map_err(Error::HashExhaustion)?;
1178
1179            let hashes = hashers.into_iter().map(HashDigest::from).collect();
1180
1181            // Before we make the wheel accessible by persisting it, ensure that the RECORD is
1182            // valid.
1183            extracted.validate_and_heal_record(dist)?;
1184
1185            // Persist the temporary directory to the directory store.
1186            let id = self
1187                .persist_extracted_wheel(extracted, wheel_entry.path())
1188                .await?;
1189
1190            // Create an archive.
1191            let archive = Archive::new(id, hashes, filename.clone(), None);
1192
1193            // Write the archive pointer to the cache.
1194            let pointer = PathArchivePointer {
1195                timestamp: modified,
1196                archive: archive.clone(),
1197            };
1198            pointer.write_to(&pointer_entry).await?;
1199
1200            Ok(LocalWheel {
1201                dist: Dist::Built(dist.clone()),
1202                archive: self
1203                    .build_context
1204                    .cache()
1205                    .archive(&archive.id)
1206                    .into_boxed_path(),
1207                hashes: archive.hashes,
1208                filename: filename.clone(),
1209                cache: CacheInfo::from_timestamp(modified),
1210                build: None,
1211            })
1212        }
1213    }
1214
1215    /// Unzip a wheel into the cache, returning the path to the unzipped directory.
1216    async fn unzip_wheel(
1217        &self,
1218        path: &Path,
1219        target: &Path,
1220        dist: DistRef<'_>,
1221    ) -> Result<ArchiveId, Error> {
1222        let content_addressed_cache = self.content_addressed_cache;
1223
1224        let mut extracted = tokio::task::spawn_blocking({
1225            let path = path.to_owned();
1226            let root = self.build_context.cache().root().to_path_buf();
1227            move || -> Result<_, Error> {
1228                // Unzip the wheel into a temporary directory.
1229                let extractor = WheelExtractor::new(&root, content_addressed_cache)
1230                    .map_err(Error::CacheWrite)?;
1231                let reader = fs_err::File::open(&path).map_err(Error::CacheWrite)?;
1232                extractor
1233                    .extract_seekable(reader)
1234                    .map_err(|err| Error::Extract(path.to_string_lossy().into_owned(), err))
1235            }
1236        })
1237        .await??;
1238
1239        // Before we make the wheel accessible by persisting it, ensure that the RECORD is valid.
1240        extracted.validate_and_heal_record(dist)?;
1241
1242        // Persist the temporary directory to the directory store.
1243        let id = self.persist_extracted_wheel(extracted, target).await?;
1244
1245        Ok(id)
1246    }
1247
1248    /// Persist an extracted wheel into the archive store.
1249    ///
1250    /// A hash tree makes identical extracted trees converge on one archive entry. Without one,
1251    /// persistence retains the existing behavior of assigning a unique archive ID.
1252    async fn persist_extracted_wheel(
1253        &self,
1254        extracted: ExtractedWheel,
1255        target: &Path,
1256    ) -> Result<ArchiveId, Error> {
1257        let (temp_dir, hashed_wheel) = extracted.into_parts();
1258        let cache = self.build_context.cache();
1259        let (temp_dir, id) = if let Some(HashedWheel { files, tree }) = hashed_wheel {
1260            let digest = DirectoryDigest::from(tree.hash());
1261            let id = ArchiveId::from_digest(digest.into());
1262            let cache = cache.clone();
1263            let temp_dir = tokio::task::spawn_blocking(move || {
1264                persist_archive_files(&cache, temp_dir.path(), &files)
1265                    .map_err(Error::CacheWrite)?;
1266                Ok::<_, Error>(temp_dir)
1267            })
1268            .await??;
1269            (temp_dir, id)
1270        } else {
1271            (temp_dir, ArchiveId::default())
1272        };
1273
1274        cache
1275            .persist_with_id(temp_dir, target, id)
1276            .await
1277            .map_err(Error::CacheWrite)
1278    }
1279
1280    /// Returns a GET [`reqwest::Request`] for the given URL.
1281    fn request(&self, url: DisplaySafeUrl) -> Result<reqwest::Request, reqwest::Error> {
1282        self.client
1283            .unmanaged
1284            .uncached_client(&url)
1285            .get(Url::from(url))
1286            .header(
1287                // `reqwest` defaults to accepting compressed responses.
1288                // Specify identity encoding to get consistent .whl downloading
1289                // behavior from servers. ref: https://github.com/pypa/pip/pull/1688
1290                "accept-encoding",
1291                reqwest::header::HeaderValue::from_static("identity"),
1292            )
1293            .build()
1294    }
1295
1296    /// Return the [`ManagedClient`] used by this resolver.
1297    pub fn client(&self) -> &ManagedClient<'a> {
1298        &self.client
1299    }
1300}
1301
1302/// Share extracted files other than `RECORD` while keeping the unpublished archive complete.
1303fn persist_archive_files(cache: &Cache, archive: &Path, files: &[HashedFile]) -> io::Result<()> {
1304    initialize_rayon_once();
1305    let targets = files
1306        .par_iter()
1307        // Keep RECORD private, since it may have been healed after hashing.
1308        .filter(|file| !file.path().ends_with("RECORD"))
1309        .map(|file| {
1310            let id = ArchiveFileId::from_digest(&file.object_digest_hex());
1311            (archive.join(file.path()), cache.archive_file(&id))
1312        })
1313        .collect::<Vec<_>>();
1314
1315    // Group files by shard so its directory is created once and its files are linked by the
1316    // same worker, avoiding contention between workers on each shard directory.
1317    let mut shards: FxHashMap<&Path, Vec<_>> = FxHashMap::default();
1318    for (source, target) in &targets {
1319        let Some(parent) = target.parent() else {
1320            return Err(io::Error::new(
1321                io::ErrorKind::InvalidInput,
1322                "archive file path must have a parent directory",
1323            ));
1324        };
1325        shards.entry(parent).or_default().push((source, target));
1326    }
1327
1328    let mut shards = shards
1329        .into_iter()
1330        .map(|(parent, files)| (parent, files, Ok(())))
1331        .collect::<Vec<_>>();
1332    // Start larger shards first so their work can overlap the remaining directory creation.
1333    shards.sort_unstable_by_key(|(_, files, _)| Reverse(files.len()));
1334
1335    // Creating shards concurrently contends on their shared parent. Keep creation on this
1336    // thread, while workers link files in the shards that are already available.
1337    in_place_scope(|scope| -> io::Result<()> {
1338        for (parent, files, result) in &mut shards {
1339            fs_err::create_dir_all(parent)?;
1340            scope.spawn(move |_| {
1341                *result = files
1342                    .iter()
1343                    .try_for_each(|(source, target)| persist_archive_file(source, target));
1344            });
1345        }
1346        Ok(())
1347    })?;
1348
1349    shards.into_iter().try_for_each(|(_, _, result)| result)
1350}
1351
1352/// Publish a shared object and retain a hardlink in the archive, with a copy fallback.
1353fn persist_archive_file(src: &Path, dst: &Path) -> io::Result<()> {
1354    // The shard already exists, and most objects are new, so try linking before checking for an
1355    // existing object. This avoids an extra filesystem lookup for every new object.
1356    match fs_err::hard_link(src, dst) {
1357        Ok(()) => return Ok(()),
1358        Err(_) if dst.try_exists()? => {}
1359        Err(_) => return uv_fs::copy_atomic_sync(src, dst),
1360    }
1361
1362    // This archive is still private, so it is safe to replace its extracted copy before publication.
1363    if let Err(err) = fs_err::remove_file(src)
1364        && err.kind() != io::ErrorKind::NotFound
1365    {
1366        return Err(err);
1367    }
1368
1369    fs_err::hard_link(dst, src).or_else(|_| uv_fs::copy_atomic_sync(dst, src))
1370}
1371
1372/// A wrapper around `RegistryClient` that manages a concurrency limit.
1373pub struct ManagedClient<'a> {
1374    pub unmanaged: &'a RegistryClient,
1375    control: Arc<Semaphore>,
1376}
1377
1378impl<'a> ManagedClient<'a> {
1379    /// Create a new `ManagedClient` using the given client and concurrency semaphore.
1380    fn new(client: &'a RegistryClient, control: Arc<Semaphore>) -> Self {
1381        ManagedClient {
1382            unmanaged: client,
1383            control,
1384        }
1385    }
1386
1387    /// Perform a request using the client, respecting the concurrency limit.
1388    ///
1389    /// If the concurrency limit has been reached, this method will wait until a pending
1390    /// operation completes before executing the closure.
1391    pub async fn managed<F, T>(&self, f: impl FnOnce(&'a RegistryClient) -> F) -> T
1392    where
1393        F: Future<Output = T>,
1394    {
1395        let _permit = self.control.acquire().await.unwrap();
1396        f(self.unmanaged).await
1397    }
1398
1399    /// Perform a request using a client that internally manages the concurrency limit.
1400    ///
1401    /// The callback is passed the client and a semaphore. It must acquire the semaphore before
1402    /// any request through the client and drop it after.
1403    ///
1404    /// This method serves as an escape hatch for functions that may want to send multiple requests
1405    /// in parallel.
1406    pub async fn manual<F, T>(&'a self, f: impl FnOnce(&'a RegistryClient, &'a Semaphore) -> F) -> T
1407    where
1408        F: Future<Output = T>,
1409    {
1410        f(self.unmanaged, &self.control).await
1411    }
1412}
1413
1414/// Returns the value of the `Content-Length` header from the [`reqwest::Response`], if present.
1415fn content_length(response: &reqwest::Response) -> Option<u64> {
1416    response
1417        .headers()
1418        .get(reqwest::header::CONTENT_LENGTH)
1419        .and_then(|val| val.to_str().ok())
1420        .and_then(|val| val.parse::<u64>().ok())
1421}
1422
1423/// An asynchronous reader that reports progress as bytes are read.
1424struct ProgressReader<'a, R> {
1425    reader: R,
1426    index: usize,
1427    reporter: &'a dyn Reporter,
1428}
1429
1430impl<'a, R> ProgressReader<'a, R> {
1431    /// Create a new [`ProgressReader`] that wraps another reader.
1432    fn new(reader: R, index: usize, reporter: &'a dyn Reporter) -> Self {
1433        Self {
1434            reader,
1435            index,
1436            reporter,
1437        }
1438    }
1439}
1440
1441impl<R> AsyncRead for ProgressReader<'_, R>
1442where
1443    R: AsyncRead + Unpin,
1444{
1445    fn poll_read(
1446        mut self: Pin<&mut Self>,
1447        cx: &mut Context<'_>,
1448        buf: &mut ReadBuf<'_>,
1449    ) -> Poll<io::Result<()>> {
1450        Pin::new(&mut self.as_mut().reader)
1451            .poll_read(cx, buf)
1452            .map_ok(|()| {
1453                self.reporter
1454                    .on_download_progress(self.index, buf.filled().len() as u64);
1455            })
1456    }
1457}
1458
1459/// A pointer to an archive in the cache, fetched from an HTTP archive.
1460///
1461/// Encoded with `MsgPack`, and represented on disk by a `.http` file.
1462#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1463pub struct HttpArchivePointer {
1464    archive: Archive,
1465}
1466
1467impl HttpArchivePointer {
1468    /// Read an [`HttpArchivePointer`] from the cache.
1469    pub fn read_from(path: impl AsRef<Path>) -> Result<Option<Self>, Error> {
1470        match fs_err::File::open(path.as_ref()) {
1471            Ok(file) => {
1472                let data = DataWithCachePolicy::from_reader(file)?.data;
1473                let archive = rmp_serde::from_slice::<Archive>(&data)?;
1474                Ok(Some(Self { archive }))
1475            }
1476            Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None),
1477            Err(err) => Err(Error::CacheRead(err)),
1478        }
1479    }
1480
1481    /// Return the [`Archive`] from the pointer.
1482    pub fn into_archive(self) -> Archive {
1483        self.archive
1484    }
1485
1486    /// Return the [`CacheInfo`] from the pointer.
1487    pub fn to_cache_info(&self) -> CacheInfo {
1488        CacheInfo::default()
1489    }
1490
1491    /// Return the [`BuildInfo`] from the pointer.
1492    pub fn to_build_info(&self) -> Option<BuildInfo> {
1493        None
1494    }
1495}
1496
1497/// A pointer to an archive in the cache, fetched from a local path.
1498///
1499/// Encoded with `MsgPack`, and represented on disk by a `.rev` file.
1500#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1501pub struct PathArchivePointer {
1502    timestamp: Timestamp,
1503    archive: Archive,
1504}
1505
1506impl PathArchivePointer {
1507    /// Read an [`PathArchivePointer`] from the cache.
1508    pub fn read_from(path: impl AsRef<Path>) -> Result<Option<Self>, Error> {
1509        match fs_err::read(path) {
1510            Ok(cached) => Ok(Some(rmp_serde::from_slice::<Self>(&cached)?)),
1511            Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None),
1512            Err(err) => Err(Error::CacheRead(err)),
1513        }
1514    }
1515
1516    /// Write an [`PathArchivePointer`] to the cache.
1517    async fn write_to(&self, entry: &CacheEntry) -> Result<(), Error> {
1518        write_atomic(entry.path(), rmp_serde::to_vec(&self)?)
1519            .await
1520            .map_err(Error::CacheWrite)
1521    }
1522
1523    /// Returns `true` if the archive is up-to-date with the given modified timestamp.
1524    pub fn is_up_to_date(&self, modified: Timestamp) -> bool {
1525        self.timestamp == modified
1526    }
1527
1528    /// Return the [`Archive`] from the pointer.
1529    pub fn into_archive(self) -> Archive {
1530        self.archive
1531    }
1532
1533    /// Return the [`CacheInfo`] from the pointer.
1534    pub fn to_cache_info(&self) -> CacheInfo {
1535        CacheInfo::from_timestamp(self.timestamp)
1536    }
1537
1538    /// Return the [`BuildInfo`] from the pointer.
1539    pub fn to_build_info(&self) -> Option<BuildInfo> {
1540        None
1541    }
1542}