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 http_content_range::{ContentRange, ContentRangeBytes, ContentRangeUnbound};
11use rayon::in_place_scope;
12use rayon::prelude::*;
13use rustc_hash::FxHashMap;
14use tokio::io::{AsyncRead, AsyncSeekExt, AsyncWriteExt, ReadBuf};
15use tokio::sync::Semaphore;
16use tokio_util::compat::FuturesAsyncReadCompatExt;
17use tracing::{Instrument, debug, info_span, instrument, warn};
18use url::Url;
19
20use uv_cache::{ArchiveFileId, ArchiveId, Cache, CacheBucket, CacheEntry, WheelCache};
21use uv_cache_info::{CacheInfo, Timestamp};
22use uv_client::{
23    CacheControl, CachedClientError, Connectivity, DataWithCachePolicy, RegistryClient,
24    RequestBuilder, RetryState,
25};
26use uv_distribution_filename::WheelFilename;
27use uv_distribution_types::{
28    ArchiveHashPolicy, BuildInfo, BuildableSource, BuiltDist, Dist, DistRef, HashCollection,
29    HashValidation, Hashed, IndexUrl, InstalledDist, MetadataHashPolicy, Name, ResolutionRecorder,
30    SourceDist, SourceUrl, parse_url_hashes,
31};
32use uv_extract::dirhash::{DirectoryDigest, HashedFile};
33use uv_extract::hash::Hasher;
34use uv_fs::{LockedFile, write_atomic};
35use uv_git::{GIT_LFS, GitError};
36use uv_normalize::PackageName;
37use uv_pep440::Version;
38use uv_platform_tags::Tags;
39use uv_preview::PreviewFeature;
40use uv_pypi_types::{HashDigest, HashDigests, PyProjectToml, ResolutionMetadata};
41use uv_python::PythonVariant;
42use uv_redacted::DisplaySafeUrl;
43use uv_threads::initialize_rayon_once;
44use uv_types::{BuildContext, BuildStack};
45
46use crate::archive::Archive;
47use crate::error::PythonVersion;
48use crate::extracted_wheel::{ExtractedWheel, HashedWheel, WheelExtractor};
49use crate::hash::http_hash_algorithms;
50use crate::metadata::{ArchiveMetadata, Metadata};
51use crate::source::SourceDistributionBuilder;
52use crate::{Error, FirstPartyPackages, LocalWheel, Reporter, RequiresDist};
53
54/// A cached high-level interface to convert distributions (a requirement resolved to a location)
55/// to a wheel or wheel metadata.
56///
57/// For wheel metadata, this happens by either fetching the metadata from the remote wheel or by
58/// building the source distribution. For wheel files, either the wheel is downloaded or a source
59/// distribution is downloaded, built and the new wheel gets returned.
60///
61/// All kinds of wheel sources (index, URL, path) and source distribution source (index, URL, path,
62/// Git) are supported.
63///
64/// This struct also has the task of acquiring locks around source dist builds in general and git
65/// operation especially, as well as respecting concurrency limits.
66pub struct DistributionDatabase<'a, Context: BuildContext> {
67    build_context: &'a Context,
68    recorder: Option<ResolutionRecorder>,
69    builder: SourceDistributionBuilder<'a, Context>,
70    client: ManagedClient<'a>,
71    reporter: Option<Arc<dyn Reporter>>,
72    content_addressed_cache: bool,
73    first_party_packages: Option<&'a FirstPartyPackages>,
74}
75
76impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
77    pub fn new(
78        client: &'a RegistryClient,
79        build_context: &'a Context,
80        downloads_semaphore: Arc<Semaphore>,
81    ) -> Self {
82        // When ZIP validation is disabled, the extracted tree can contain files that aren't
83        // represented in the central directory and therefore aren't included in its digest.
84        // Avoid using an incomplete digest as a content-addressed archive ID.
85        let content_addressed_cache = uv_preview::is_enabled(PreviewFeature::ContentAddressedCache)
86            && !uv_extract::insecure_no_validate();
87        Self {
88            recorder: None,
89            build_context,
90            builder: SourceDistributionBuilder::new(build_context),
91            client: ManagedClient::new(client, downloads_semaphore),
92            reporter: None,
93            content_addressed_cache,
94            first_party_packages: None,
95        }
96    }
97
98    /// Allow metadata builds for the given first-party workspace source trees.
99    #[must_use]
100    pub fn with_first_party_packages(
101        mut self,
102        first_party_packages: &'a FirstPartyPackages,
103    ) -> Self {
104        self.first_party_packages = Some(first_party_packages);
105        self
106    }
107
108    /// Return whether the name and path identify an eligible workspace member.
109    pub fn is_first_party(&self, name: &PackageName, path: &Path) -> bool {
110        self.first_party_packages
111            .is_some_and(|packages| packages.contains(name, path))
112    }
113
114    /// Record which static metadata entries are consulted while resolving runtime dependencies.
115    #[must_use]
116    pub fn with_recorder(mut self, recorder: Option<ResolutionRecorder>) -> Self {
117        self.recorder = recorder;
118        self
119    }
120
121    /// Record a metadata lookup before reading an in-memory cache.
122    pub fn record_metadata(&self, dist: &Dist) {
123        if let Some(recorder) = &self.recorder {
124            recorder.dependency_metadata(dist.name());
125        }
126    }
127
128    /// Look up user-provided metadata, recording misses as well as matches.
129    pub fn dependency_metadata(
130        &self,
131        name: &PackageName,
132        version: Option<&Version>,
133    ) -> Option<ResolutionMetadata> {
134        if let Some(recorder) = &self.recorder {
135            recorder.dependency_metadata(name);
136        }
137        self.build_context.dependency_metadata().get(name, version)
138    }
139
140    /// Set the build stack to use for the [`DistributionDatabase`].
141    #[must_use]
142    pub fn with_build_stack(self, build_stack: &'a BuildStack) -> Self {
143        Self {
144            builder: self.builder.with_build_stack(build_stack),
145            ..self
146        }
147    }
148
149    /// Set the [`Reporter`] to use for the [`DistributionDatabase`].
150    #[must_use]
151    pub fn with_reporter(self, reporter: Arc<dyn Reporter>) -> Self {
152        Self {
153            builder: self.builder.with_reporter(reporter.clone()),
154            reporter: Some(reporter),
155            ..self
156        }
157    }
158
159    /// Handle a specific `reqwest` error, and convert it to [`io::Error`].
160    fn handle_response_errors(&self, err: reqwest::Error) -> io::Error {
161        if err.is_timeout() {
162            // Assumption: The connect timeout with the 10s default is not the culprit.
163            io::Error::new(
164                io::ErrorKind::TimedOut,
165                format!(
166                    "Failed to download distribution due to network timeout. Try increasing UV_HTTP_TIMEOUT (current value: {}s).",
167                    self.client.unmanaged.read_timeout().as_secs()
168                ),
169            )
170        } else {
171            io::Error::other(err)
172        }
173    }
174
175    /// Acquire an advisory lock for a wheel cache entry.
176    ///
177    /// A remote wheel's content hash is not always available until after the download, so
178    /// concurrent cache fills coordinate on the wheel cache entry instead. The entry is already
179    /// scoped to the distribution's source and wheel filename.
180    ///
181    /// Callers hold the returned lock across cache lookup, download or extraction, and publication.
182    /// A process that waited for another cache fill therefore rechecks and reuses the completed entry.
183    async fn lock_wheel(
184        wheel_entry: &CacheEntry,
185        filename: &WheelFilename,
186    ) -> Result<LockedFile, Error> {
187        // For backwards compatibility, we use the full wheel stem on Windows. Local wheel
188        // extraction and older uv versions use the same key, so changing it would prevent them
189        // from coordinating through a shared cache.
190        #[cfg(windows)]
191        let lock_key = filename.stem();
192        // On other platforms, we use the bounded cache key to avoid filesystem filename limits.
193        #[cfg(not(windows))]
194        let lock_key = filename.cache_key();
195
196        let lock_entry = wheel_entry.with_file(format!("{lock_key}.lock"));
197        lock_entry.lock().await.map_err(Error::CacheLock)
198    }
199
200    /// Validate dist hashes, and return them.
201    fn validate_hashes(
202        dist: &BuiltDist,
203        hashes: ArchiveHashPolicy<'_>,
204        hashers: Vec<Hasher>,
205    ) -> Result<HashDigests, Error> {
206        let computed_hashes: HashDigests = hashers.into_iter().map(HashDigest::from).collect();
207        if hashes.requires_validation() && !hashes.matches(computed_hashes.as_slice()) {
208            return Err(Error::hash_mismatch(
209                dist.to_string(),
210                hashes.digests(),
211                computed_hashes.as_slice(),
212            ));
213        }
214        Ok(computed_hashes)
215    }
216
217    /// Either fetch the wheel or fetch and build the source distribution
218    ///
219    /// Returns a wheel that's compliant with the given platform tags.
220    ///
221    /// Applicable hash checks are enforced before newly fetched archives are published to the cache.
222    /// Callers must enforce their hash policy when reusing cached wheels.
223    #[instrument(skip_all, fields(%dist))]
224    pub async fn get_or_build_wheel(
225        &self,
226        dist: &Dist,
227        tags: &Tags,
228        hashes: ArchiveHashPolicy<'_>,
229    ) -> Result<LocalWheel, Error> {
230        match dist {
231            Dist::Built(built) => self.get_wheel(built, hashes).await,
232            Dist::Source(source) => self.build_wheel(source, tags, hashes).await,
233        }
234    }
235
236    /// Either fetch the only wheel metadata (directly from the index or with range requests) or
237    /// fetch and build the source distribution.
238    ///
239    /// While hashes will be generated in some cases, hash-checking is only enforced for source
240    /// distributions, and should be enforced by the caller for wheels.
241    #[instrument(skip_all, fields(%dist))]
242    pub async fn get_installed_metadata(
243        &self,
244        dist: &InstalledDist,
245    ) -> Result<ArchiveMetadata, Error> {
246        // If the metadata was provided by the user directly, prefer it.
247        if let Some(metadata) = self.dependency_metadata(dist.name(), Some(dist.version())) {
248            return Ok(Metadata::from_dependency_metadata(metadata).into());
249        }
250
251        let metadata = dist
252            .read_metadata()
253            .map_err(|err| Error::ReadInstalled(Box::new(dist.clone()), err))?;
254
255        Ok(ArchiveMetadata::from_metadata23(metadata.clone()))
256    }
257
258    /// Retrieve distribution metadata and any hashes requested for resolution.
259    ///
260    /// Source archive hashes are validated before executing build backends; wheel archive hash
261    /// validation is deferred to installation.
262    #[instrument(skip_all, fields(%dist))]
263    pub async fn get_or_build_wheel_metadata(
264        &self,
265        dist: &Dist,
266        hashes: MetadataHashPolicy<'_>,
267    ) -> Result<ArchiveMetadata, Error> {
268        match dist {
269            Dist::Built(built) => self.get_wheel_metadata(built, hashes).await,
270            Dist::Source(source) => {
271                self.build_wheel_metadata(&BuildableSource::Dist(source), hashes)
272                    .await
273            }
274        }
275    }
276
277    /// Fetch a wheel from the cache or download it from the index.
278    ///
279    /// Applicable hash checks are enforced before newly fetched wheels are published to the cache.
280    /// Registry hashes are used when available and no explicit verification policy is provided.
281    /// Callers must enforce their hash policy when reusing cached wheels.
282    async fn get_wheel(
283        &self,
284        dist: &BuiltDist,
285        hashes: ArchiveHashPolicy<'_>,
286    ) -> Result<LocalWheel, Error> {
287        match dist {
288            BuiltDist::Registry(wheels) => {
289                let wheel = wheels.best_wheel();
290                let url = wheel.file.url.to_url()?;
291                let size = wheel.file.size;
292
293                // Create a cache entry for the wheel.
294                let wheel_entry = self.build_context.cache().entry(
295                    CacheBucket::Wheels,
296                    WheelCache::Index(&wheel.index).wheel_dir(wheel.name().as_ref()),
297                    wheel.filename.cache_key(),
298                );
299
300                // If the URL is a file URL, load the wheel directly.
301                if url.scheme() == "file" {
302                    let path = url
303                        .to_file_path()
304                        .map_err(|()| Error::NonFileUrl(url.clone()))?;
305                    return self
306                        .load_wheel(&path, &wheel.filename, wheel_entry, dist, hashes)
307                        .await;
308                }
309
310                // Download and unzip.
311                match self
312                    .stream_wheel(
313                        url.clone(),
314                        dist.index(),
315                        &wheel.filename,
316                        size,
317                        &wheel_entry,
318                        dist,
319                        hashes,
320                    )
321                    .await
322                {
323                    Ok(archive) => Ok(LocalWheel {
324                        dist: Dist::Built(dist.clone()),
325                        archive: self
326                            .build_context
327                            .cache()
328                            .archive(&archive.id)
329                            .into_boxed_path(),
330                        hashes: archive.hashes,
331                        filename: wheel.filename.clone(),
332                        cache: CacheInfo::default(),
333                        build: None,
334                    }),
335                    Err(Error::Extract(name, err)) => {
336                        if err.is_http_streaming_unsupported() {
337                            warn!(
338                                "Streaming unsupported for {dist}; downloading wheel to disk ({err})"
339                            );
340                        } else if err.is_http_streaming_failed() {
341                            warn!("Streaming failed for {dist}; downloading wheel to disk ({err})");
342                        } else {
343                            return Err(Error::Extract(name, err));
344                        }
345
346                        // If the request failed because streaming was unsupported or failed,
347                        // download the wheel directly.
348                        let archive = self
349                            .download_wheel(
350                                url,
351                                dist.index(),
352                                &wheel.filename,
353                                size,
354                                &wheel_entry,
355                                dist,
356                                hashes,
357                            )
358                            .await?;
359
360                        Ok(LocalWheel {
361                            dist: Dist::Built(dist.clone()),
362                            archive: self
363                                .build_context
364                                .cache()
365                                .archive(&archive.id)
366                                .into_boxed_path(),
367                            hashes: archive.hashes,
368                            filename: wheel.filename.clone(),
369                            cache: CacheInfo::default(),
370                            build: None,
371                        })
372                    }
373                    Err(err) => Err(err),
374                }
375            }
376
377            BuiltDist::DirectUrl(wheel) => {
378                // Create a cache entry for the wheel.
379                let wheel_entry = self.build_context.cache().entry(
380                    CacheBucket::Wheels,
381                    WheelCache::Url(&wheel.url).wheel_dir(wheel.name().as_ref()),
382                    wheel.filename.cache_key(),
383                );
384
385                // Download and unzip.
386                match self
387                    .stream_wheel(
388                        wheel.url.raw().clone(),
389                        None,
390                        &wheel.filename,
391                        wheel.size,
392                        &wheel_entry,
393                        dist,
394                        hashes,
395                    )
396                    .await
397                {
398                    Ok(archive) => Ok(LocalWheel {
399                        dist: Dist::Built(dist.clone()),
400                        archive: self
401                            .build_context
402                            .cache()
403                            .archive(&archive.id)
404                            .into_boxed_path(),
405                        hashes: archive.hashes,
406                        filename: wheel.filename.clone(),
407                        cache: CacheInfo::default(),
408                        build: None,
409                    }),
410                    Err(Error::Extract(name, err)) => {
411                        if err.is_http_streaming_unsupported() {
412                            warn!(
413                                "Streaming unsupported for {dist}; downloading wheel to disk ({err})"
414                            );
415                        } else if err.is_http_streaming_failed() {
416                            warn!("Streaming failed for {dist}; downloading wheel to disk ({err})");
417                        } else {
418                            return Err(Error::Extract(name, err));
419                        }
420
421                        // If the request failed because streaming was unsupported or failed,
422                        // download the wheel directly.
423                        let archive = self
424                            .download_wheel(
425                                wheel.url.raw().clone(),
426                                None,
427                                &wheel.filename,
428                                wheel.size,
429                                &wheel_entry,
430                                dist,
431                                hashes,
432                            )
433                            .await?;
434                        Ok(LocalWheel {
435                            dist: Dist::Built(dist.clone()),
436                            archive: self
437                                .build_context
438                                .cache()
439                                .archive(&archive.id)
440                                .into_boxed_path(),
441                            hashes: archive.hashes,
442                            filename: wheel.filename.clone(),
443                            cache: CacheInfo::default(),
444                            build: None,
445                        })
446                    }
447                    Err(err) => Err(err),
448                }
449            }
450
451            BuiltDist::GitPath(wheel) => {
452                // Fetch the Git repository.
453                let fetch = self
454                    .build_context
455                    .git()
456                    .fetch(
457                        &wheel.git,
458                        self.client.unmanaged.git_http_settings(wheel.git.url()),
459                        self.build_context.cache().bucket(CacheBucket::Git),
460                        self.reporter.clone().map(<dyn Reporter>::into_git_reporter),
461                    )
462                    .await?;
463
464                if wheel.git.lfs().enabled() && !fetch.lfs_ready() {
465                    if GIT_LFS.is_err() {
466                        return Err(Error::MissingWheelGitLfsArtifacts(
467                            wheel.url.to_url(),
468                            GitError::GitLfsNotFound,
469                        ));
470                    }
471                    return Err(Error::MissingWheelGitLfsArtifacts(
472                        wheel.url.to_url(),
473                        GitError::GitLfsNotConfigured,
474                    ));
475                }
476
477                let git_sha = fetch.git().precise().expect("Exact commit after checkout");
478                let cache_entry = self.build_context.cache().entry(
479                    CacheBucket::Wheels,
480                    WheelCache::Git(&wheel.url, git_sha.as_short_str()).root(),
481                    wheel.filename.stem(),
482                );
483
484                let install_path = fetch.path().join(&wheel.install_path);
485
486                self.load_wheel(&install_path, &wheel.filename, cache_entry, dist, hashes)
487                    .await
488            }
489
490            BuiltDist::Path(wheel) => {
491                let cache_entry = self.build_context.cache().entry(
492                    CacheBucket::Wheels,
493                    WheelCache::Url(&wheel.url).wheel_dir(wheel.name().as_ref()),
494                    wheel.filename.cache_key(),
495                );
496
497                self.load_wheel(
498                    &wheel.install_path,
499                    &wheel.filename,
500                    cache_entry,
501                    dist,
502                    hashes,
503                )
504                .await
505            }
506        }
507    }
508
509    /// Convert a source distribution into a wheel, fetching it from the cache or building it if
510    /// necessary.
511    ///
512    /// The returned wheel is guaranteed to come from a distribution with a matching hash, and
513    /// no build processes will be executed for distributions with mismatched hashes.
514    async fn build_wheel(
515        &self,
516        dist: &SourceDist,
517        tags: &Tags,
518        hashes: ArchiveHashPolicy<'_>,
519    ) -> Result<LocalWheel, Error> {
520        let built_wheel = self
521            .builder
522            .download_and_build(&BuildableSource::Dist(dist), tags, hashes, &self.client)
523            .boxed_local()
524            .await?;
525
526        // Check that the wheel is compatible with its install target.
527        //
528        // When building a build dependency for a cross-install, the build dependency needs
529        // to install and run on the host instead of the target. In this case the `tags` are already
530        // for the host instead of the target, so this check passes.
531        if !built_wheel.filename.is_compatible(tags) {
532            return if tags.is_cross() {
533                Err(Error::BuiltWheelIncompatibleTargetPlatform {
534                    filename: built_wheel.filename,
535                    python_platform: tags.python_platform().clone(),
536                    python_version: PythonVersion {
537                        version: tags.python_version(),
538                        variant: if tags.is_freethreaded() {
539                            PythonVariant::Freethreaded
540                        } else {
541                            PythonVariant::Default
542                        },
543                    },
544                })
545            } else {
546                Err(Error::BuiltWheelIncompatibleHostPlatform {
547                    filename: built_wheel.filename,
548                    python_platform: tags.python_platform().clone(),
549                    python_version: PythonVersion {
550                        version: tags.python_version(),
551                        variant: if tags.is_freethreaded() {
552                            PythonVariant::Freethreaded
553                        } else {
554                            PythonVariant::Default
555                        },
556                    },
557                })
558            };
559        }
560
561        // Acquire the advisory lock.
562        let wheel_entry = CacheEntry::from_path(built_wheel.target.as_ref());
563        let _lock = Self::lock_wheel(&wheel_entry, &built_wheel.filename).await?;
564
565        // If the wheel was unzipped previously, respect it. Source distributions are
566        // cached under a unique revision ID, so unzipped directories are never stale.
567        match self.build_context.cache().resolve_link(&built_wheel.target) {
568            Ok(archive) => {
569                return Ok(LocalWheel {
570                    dist: Dist::Source(dist.clone()),
571                    archive: archive.into_boxed_path(),
572                    filename: built_wheel.filename,
573                    hashes: built_wheel.hashes,
574                    cache: built_wheel.cache_info,
575                    build: Some(built_wheel.build_info),
576                });
577            }
578            Err(err) if err.kind() == io::ErrorKind::NotFound => {}
579            Err(err) => return Err(Error::CacheRead(err)),
580        }
581
582        // Otherwise, unzip the wheel.
583        let id = self
584            .unzip_wheel(
585                &built_wheel.path,
586                &built_wheel.target,
587                DistRef::Source(dist),
588            )
589            .await?;
590
591        Ok(LocalWheel {
592            dist: Dist::Source(dist.clone()),
593            archive: self.build_context.cache().archive(&id).into_boxed_path(),
594            hashes: built_wheel.hashes,
595            filename: built_wheel.filename,
596            cache: built_wheel.cache_info,
597            build: Some(built_wheel.build_info),
598        })
599    }
600
601    /// Fetch wheel metadata, along with any hashes requested for resolution.
602    ///
603    /// Wheel archive hash validation is deferred to installation. Metadata sidecar hashes are
604    /// checked by the client when downloading sidecars with index-provided hashes.
605    async fn get_wheel_metadata(
606        &self,
607        dist: &BuiltDist,
608        hashes: MetadataHashPolicy<'_>,
609    ) -> Result<ArchiveMetadata, Error> {
610        let hash_policy = match hashes.validation {
611            HashValidation::None => {
612                let compute_hashes = match hashes.collection {
613                    HashCollection::None => false,
614                    HashCollection::Url | HashCollection::All => match dist {
615                        BuiltDist::Registry(dist) => {
616                            // Generate missing hashes for indexes and `--find-links` without hashes.
617                            // This hashes only the selected wheel, not every distribution for the version.
618                            hashes.collection == HashCollection::All
619                                && dist.best_wheel().file.hashes.is_empty()
620                        }
621                        BuiltDist::DirectUrl(dist) => parse_url_hashes(&dist.url).is_none(),
622                        BuiltDist::Path(_) | BuiltDist::GitPath(_) => true,
623                    },
624                };
625                if compute_hashes {
626                    ArchiveHashPolicy::Generate
627                } else {
628                    ArchiveHashPolicy::None
629                }
630            }
631            HashValidation::Any(_) | HashValidation::All(_) => hashes.validation.into(),
632        };
633
634        // Fetch the entire wheel only when we need to compute a hash for resolution.
635        // TODO(charlie): Request the hashes via a separate method, to reduce the coupling in this API.
636        if hash_policy == ArchiveHashPolicy::Generate {
637            let wheel = self.get_wheel(dist, hash_policy).await?;
638            // If the metadata was provided by the user directly, prefer it.
639            let metadata = if let Some(metadata) =
640                self.dependency_metadata(dist.name(), Some(dist.version()))
641            {
642                Metadata::from_dependency_metadata(metadata)
643            } else {
644                Metadata::from_metadata23(wheel.metadata()?)
645            };
646            let hashes = wheel.hashes;
647            return Ok(ArchiveMetadata { metadata, hashes });
648        }
649
650        // If the metadata was provided by the user directly, prefer it.
651        if let Some(metadata) = self.dependency_metadata(dist.name(), Some(dist.version())) {
652            return Ok(Metadata::from_dependency_metadata(metadata).into());
653        }
654
655        let result = self
656            .client
657            .managed(|client| {
658                client
659                    .wheel_metadata(
660                        dist,
661                        self.build_context.git(),
662                        self.build_context.capabilities(),
663                        self.reporter.clone().map(<dyn Reporter>::into_git_reporter),
664                    )
665                    .boxed_local()
666            })
667            .await;
668
669        match result {
670            Ok(metadata) => {
671                // Validate that the metadata is consistent with the distribution.
672                Ok(ArchiveMetadata::from_metadata23(metadata))
673            }
674            Err(err) if err.is_http_streaming_unsupported() => {
675                warn!(
676                    "Streaming unsupported when fetching metadata for {dist}; downloading wheel directly ({err})"
677                );
678
679                // If the request failed due to an error that could be resolved by
680                // downloading the wheel directly, try that.
681                let wheel = self.get_wheel(dist, hash_policy).await?;
682                let metadata = wheel.metadata()?;
683                let hashes = wheel.hashes;
684                Ok(ArchiveMetadata {
685                    metadata: Metadata::from_metadata23(metadata),
686                    hashes,
687                })
688            }
689            Err(err) => Err(err.into()),
690        }
691    }
692
693    /// Build the wheel metadata for a source distribution, or fetch it from the cache if possible.
694    ///
695    /// Requested hashes are validated before executing a build backend. User-provided metadata
696    /// can avoid downloading or validating the source archive.
697    pub async fn build_wheel_metadata(
698        &self,
699        source: &BuildableSource<'_>,
700        hashes: MetadataHashPolicy<'_>,
701    ) -> Result<ArchiveMetadata, Error> {
702        // If the metadata was provided by the user directly, prefer it.
703        if let Some(dist) = source.as_dist() {
704            if let Some(metadata) = self.dependency_metadata(dist.name(), dist.version()) {
705                // If we skipped the build, we should still resolve any Git dependencies to precise
706                // commits.
707                self.builder.resolve_revision(source, &self.client).await?;
708
709                return Ok(Metadata::from_dependency_metadata(metadata).into());
710            }
711        }
712
713        let url_hashes = if let BuildableSource::Dist(SourceDist::DirectUrl(dist)) = source {
714            parse_url_hashes(&dist.url)
715        } else if let BuildableSource::Url(SourceUrl::Direct(url)) = source {
716            parse_url_hashes(url.url)
717        } else {
718            None
719        };
720
721        let build_hash_policy = match hashes.validation {
722            HashValidation::None => match hashes.collection {
723                HashCollection::None => ArchiveHashPolicy::None,
724                // If resolving metadata requires a build, validate any URL hash before executing
725                // the backend, even when the caller only requested hash collection.
726                HashCollection::Url | HashCollection::All => match url_hashes.as_ref() {
727                    Some(digests) => ArchiveHashPolicy::All(digests.as_slice()),
728                    None => ArchiveHashPolicy::Generate,
729                },
730            },
731            HashValidation::Any(_) | HashValidation::All(_) => hashes.validation.into(),
732        };
733        let ArchiveMetadata { metadata, hashes } = self
734            .builder
735            .for_metadata(self.first_party_packages)
736            .download_and_build_metadata(source, build_hash_policy, &self.client)
737            .boxed_local()
738            .await?;
739
740        // This also handles external dependencies, so force relative paths by default.
741        // For explicit inputs and current project/workspace metadata, this is later overridden with
742        // `with_force_relative(false)` to respect the user's relative/absolute path preference.
743        Ok(ArchiveMetadata {
744            metadata: metadata.with_force_relative(true),
745            hashes,
746        })
747    }
748
749    /// Return the [`RequiresDist`] from a `pyproject.toml`, if it can be statically extracted.
750    pub async fn requires_dist(
751        &self,
752        path: &Path,
753        pyproject_toml: &PyProjectToml,
754    ) -> Result<Option<RequiresDist>, Error> {
755        self.builder
756            .source_tree_requires_dist(
757                path,
758                pyproject_toml,
759                self.client.unmanaged.credentials_cache(),
760            )
761            .await
762    }
763
764    /// Stream a wheel from a URL, unzipping it into the cache as it's downloaded.
765    ///
766    /// Note that `progress_size_hint` is a size hint for the progress reporter,
767    /// and is not guaranteed to be the true size of the wheel if/when fetched from the
768    /// origin.
769    async fn stream_wheel(
770        &self,
771        url: DisplaySafeUrl,
772        index: Option<&IndexUrl>,
773        filename: &WheelFilename,
774        progress_size_hint: Option<u64>,
775        wheel_entry: &CacheEntry,
776        dist: &BuiltDist,
777        hashes: ArchiveHashPolicy<'_>,
778    ) -> Result<Archive, Error> {
779        // Sometimes we can promote the size hint to a trusted effective size.
780        let expected_size = match dist {
781            BuiltDist::Registry(dist) if dist.best_wheel().size_is_authoritative => {
782                progress_size_hint
783            }
784            BuiltDist::DirectUrl(_) => progress_size_hint,
785            _ => None,
786        };
787
788        // Acquire an advisory lock, to guard against concurrent writes.
789        let _lock = Self::lock_wheel(wheel_entry, filename).await?;
790
791        // Create an entry for the HTTP cache.
792        let http_entry = wheel_entry.with_file(format!("{}.http", filename.cache_key()));
793
794        let download = |response: reqwest::Response, _: &mut RetryState| {
795            async {
796                let hashes = if let BuiltDist::Registry(wheels) = dist {
797                    hashes.with_index_hashes(wheels.best_wheel().file.hashes.as_slice())
798                } else {
799                    hashes
800                };
801
802                let progress_size_hint = progress_size_hint.or_else(|| content_length(&response));
803
804                let progress = self.reporter.as_ref().map(|reporter| {
805                    (
806                        reporter,
807                        reporter.on_download_start(dist.name(), progress_size_hint),
808                    )
809                });
810
811                let reader = response
812                    .bytes_stream()
813                    .map_err(|err| self.handle_response_errors(err))
814                    .into_async_read();
815
816                // Create a hasher for each hash algorithm.
817                let algorithms = http_hash_algorithms(hashes);
818                let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
819                let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers);
820
821                // Download and unzip the wheel to a temporary directory.
822                let extractor = WheelExtractor::new(
823                    self.build_context.cache().root(),
824                    self.content_addressed_cache,
825                )
826                .map_err(Error::CacheWrite)?;
827
828                let mut extracted = match progress {
829                    Some((reporter, progress)) => {
830                        let mut reader = ProgressReader::new(&mut hasher, progress, &**reporter);
831                        extractor
832                            .extract_streaming(&mut reader)
833                            .await
834                            .map_err(|err| Error::Extract(filename.to_string(), err))?
835                    }
836                    None => extractor
837                        .extract_streaming(&mut hasher)
838                        .await
839                        .map_err(|err| Error::Extract(filename.to_string(), err))?,
840                };
841                // Exhaust the reader to compute the hashes.
842                hasher.finish().await.map_err(Error::HashExhaustion)?;
843                let actual_size = hasher.bytes_read();
844                if let Some(expected) = expected_size
845                    && actual_size != expected
846                {
847                    return Err(Error::MismatchedSize {
848                        distribution: dist.to_string(),
849                        expected,
850                        actual: actual_size,
851                    });
852                }
853
854                let computed_hashes = Self::validate_hashes(dist, hashes, hashers)?;
855
856                // Before we make the wheel accessible by persisting it, ensure that the RECORD is
857                // valid.
858                extracted.validate_and_heal_record(dist)?;
859
860                // Persist the temporary directory to the directory store.
861                let id = self
862                    .persist_extracted_wheel(extracted, wheel_entry.path())
863                    .await?;
864
865                if let Some((reporter, progress)) = progress {
866                    reporter.on_download_complete(dist.name(), progress);
867                }
868
869                Ok(Archive::new(
870                    id,
871                    computed_hashes,
872                    filename.clone(),
873                    Some(actual_size),
874                ))
875            }
876            .instrument(info_span!("wheel", wheel = %dist))
877        };
878
879        // Fetch the archive from the cache, or download it if necessary.
880        let req = self.request(url.clone())?;
881
882        // Determine the cache control policy for the URL.
883        let cache_control = match self.client.unmanaged.connectivity() {
884            Connectivity::Online
885                if let Some(header) = index.and_then(|index| {
886                    self.build_context
887                        .locations()
888                        .artifact_cache_control_for(index)
889                }) =>
890            {
891                CacheControl::Override(header)
892            }
893            Connectivity::Online => CacheControl::from(
894                self.build_context
895                    .cache()
896                    .freshness(&http_entry, Some(&filename.name), None)
897                    .map_err(Error::CacheRead)?,
898            ),
899            Connectivity::Offline => CacheControl::AllowStale,
900        };
901
902        let archive = self
903            .client
904            .managed(|client| {
905                client.cached_client().get_serde_with_retry(
906                    req,
907                    &http_entry,
908                    cache_control.clone(),
909                    download,
910                )
911            })
912            .await
913            .map_err(|err| match err {
914                CachedClientError::Callback { err, .. } => err,
915                CachedClientError::Client(err) => Error::Client(err),
916            })?;
917
918        if let (Some(expected), Some(actual)) = (expected_size, archive.size)
919            && expected != actual
920        {
921            return Err(Error::MismatchedSize {
922                distribution: dist.to_string(),
923                expected,
924                actual,
925            });
926        }
927
928        // If the archive is missing the required hashes or size, or has since been removed, force a refresh.
929        let archive = Some(archive)
930            .filter(|archive| archive.has_digests(hashes))
931            .filter(|archive| archive.exists(self.build_context.cache()))
932            .filter(|archive| expected_size.is_none() || archive.size.is_some());
933
934        let archive = if let Some(archive) = archive {
935            archive
936        } else {
937            self.client
938                .managed(async |client| {
939                    client
940                        .cached_client()
941                        .skip_cache_with_retry(
942                            self.request(url)?,
943                            &http_entry,
944                            cache_control,
945                            download,
946                        )
947                        .await
948                        .map_err(|err| match err {
949                            CachedClientError::Callback { err, .. } => err,
950                            CachedClientError::Client(err) => Error::Client(err),
951                        })
952                })
953                .await?
954        };
955
956        Ok(archive)
957    }
958
959    /// Download a wheel from a URL, then unzip it into the cache.
960    ///
961    /// Note that, like [`Self::stream_wheel`], `progress_size_hint` is
962    /// a size hint for the progress reporter, and is not guaranteed to be the size
963    /// of the wheel if/when fetched from the origin.
964    async fn download_wheel(
965        &self,
966        url: DisplaySafeUrl,
967        index: Option<&IndexUrl>,
968        filename: &WheelFilename,
969        progress_size_hint: Option<u64>,
970        wheel_entry: &CacheEntry,
971        dist: &BuiltDist,
972        hashes: ArchiveHashPolicy<'_>,
973    ) -> Result<Archive, Error> {
974        // Sometimes we can promote the size hint to a trusted effective size.
975        let expected_size = match dist {
976            BuiltDist::Registry(dist) if dist.best_wheel().size_is_authoritative => {
977                progress_size_hint
978            }
979            BuiltDist::DirectUrl(_) => progress_size_hint,
980            _ => None,
981        };
982
983        // Acquire an advisory lock, to guard against concurrent writes.
984        let _lock = Self::lock_wheel(wheel_entry, filename).await?;
985
986        // Create an entry for the HTTP cache.
987        let http_entry = wheel_entry.with_file(format!("{}.http", filename.cache_key()));
988
989        let download_url = url.clone();
990
991        let download = async |response, retry_state: &mut RetryState| {
992            self.download_wheel_response(
993                response,
994                &download_url,
995                retry_state,
996                filename,
997                progress_size_hint,
998                expected_size,
999                wheel_entry,
1000                dist,
1001                hashes,
1002            )
1003            .instrument(info_span!("wheel", wheel = %dist))
1004            .await
1005        };
1006
1007        // Fetch the archive from the cache, or download it if necessary.
1008        let req = self.request(url.clone())?;
1009
1010        // Determine the cache control policy for the URL.
1011        let cache_control = match self.client.unmanaged.connectivity() {
1012            Connectivity::Online
1013                if let Some(header) = index.and_then(|index| {
1014                    self.build_context
1015                        .locations()
1016                        .artifact_cache_control_for(index)
1017                }) =>
1018            {
1019                CacheControl::Override(header)
1020            }
1021            Connectivity::Online => CacheControl::from(
1022                self.build_context
1023                    .cache()
1024                    .freshness(&http_entry, Some(&filename.name), None)
1025                    .map_err(Error::CacheRead)?,
1026            ),
1027            Connectivity::Offline => CacheControl::AllowStale,
1028        };
1029
1030        let archive = self
1031            .client
1032            .managed(|client| {
1033                client.cached_client().get_serde_with_retry(
1034                    req,
1035                    &http_entry,
1036                    cache_control.clone(),
1037                    download,
1038                )
1039            })
1040            .await
1041            .map_err(|err| match err {
1042                CachedClientError::Callback { err, .. } => err,
1043                CachedClientError::Client(err) => Error::Client(err),
1044            })?;
1045
1046        if let (Some(expected), Some(actual)) = (expected_size, archive.size)
1047            && expected != actual
1048        {
1049            return Err(Error::MismatchedSize {
1050                distribution: dist.to_string(),
1051                expected,
1052                actual,
1053            });
1054        }
1055
1056        // If the archive is missing the required hashes or size, or has since been removed, force a refresh.
1057        let archive = Some(archive)
1058            .filter(|archive| archive.has_digests(hashes))
1059            .filter(|archive| archive.exists(self.build_context.cache()))
1060            .filter(|archive| expected_size.is_none() || archive.size.is_some());
1061
1062        let archive = if let Some(archive) = archive {
1063            archive
1064        } else {
1065            self.client
1066                .managed(async |client| {
1067                    client
1068                        .cached_client()
1069                        .skip_cache_with_retry(
1070                            self.request(url)?,
1071                            &http_entry,
1072                            cache_control,
1073                            download,
1074                        )
1075                        .await
1076                        .map_err(|err| match err {
1077                            CachedClientError::Callback { err, .. } => err,
1078                            CachedClientError::Client(err) => Error::Client(err),
1079                        })
1080                })
1081                .await?
1082        };
1083
1084        Ok(archive)
1085    }
1086
1087    /// Download and extract a wheel from an HTTP response.
1088    ///
1089    /// This helper is called by [`Self::download_wheel`] and handles partial content/resumptions
1090    /// of interrupted downloads if the origin supports it. Hard failures are propagated
1091    /// back to the caller and thus to our general retry machinery.
1092    ///
1093    /// The caller is responsible for obtaining a lock on the wheel cache.
1094    async fn download_wheel_response(
1095        &self,
1096        mut response: reqwest::Response,
1097        url: &DisplaySafeUrl,
1098        retry_state: &mut RetryState,
1099        filename: &WheelFilename,
1100        progress_size_hint: Option<u64>,
1101        expected_size: Option<u64>,
1102        wheel_entry: &CacheEntry,
1103        dist: &BuiltDist,
1104        hashes: ArchiveHashPolicy<'_>,
1105    ) -> Result<Archive, Error> {
1106        let hashes = if let BuiltDist::Registry(wheels) = dist {
1107            hashes.with_index_hashes(wheels.best_wheel().file.hashes.as_slice())
1108        } else {
1109            hashes
1110        };
1111
1112        let progress_size_hint = progress_size_hint.or_else(|| content_length(&response));
1113        let mut download_size = content_length(&response).or(expected_size);
1114
1115        let progress = self.reporter.as_ref().map(|reporter| {
1116            (
1117                reporter,
1118                reporter.on_download_start(dist.name(), progress_size_hint),
1119            )
1120        });
1121
1122        let algorithms = http_hash_algorithms(hashes);
1123
1124        // Download the wheel to a temporary file.
1125        let temp_file =
1126            tempfile::tempfile_in(self.build_context.cache().root()).map_err(Error::CacheWrite)?;
1127        let mut writer = tokio::io::BufWriter::new(fs_err::tokio::File::from_std(
1128            // It's an unnamed file on Linux so that's the best approximation.
1129            fs_err::File::from_parts(temp_file, self.build_context.cache().root()),
1130        ));
1131
1132        // States for the download loop below.
1133        let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
1134        // The total number of bytes retrieved, accumulated over individual requests.
1135        // Note that this is *not* the same as the number of bytes actually written
1136        // to the temporary file, since a copy from the response to the file can fail.
1137        let mut bytes_retrieved = 0;
1138        // The most recent range request's starting offset.
1139        let mut resumed_at = None;
1140        // The `Content-Range` that the download was last resumed at.
1141        let mut resumed_range: Option<ContentRangeBytes> = None;
1142
1143        // This loop is where we handle resumption of interrupted downloads, as well as
1144        // range requests (if the origin supports them).
1145        //
1146        // Errors returned from this loop reach the outer retry classifier, which may restart
1147        // the full download.
1148        loop {
1149            // Reject conflicting full-response lengths before reading the body. A range response's
1150            // Content-Length describes only that range, so it cannot be compared to the wheel size.
1151            if response.status() == reqwest::StatusCode::OK
1152                && let (Some(expected), Some(actual)) = (expected_size, content_length(&response))
1153                && expected != actual
1154            {
1155                return Err(Error::MismatchedContentLength {
1156                    distribution: dist.to_string(),
1157                    expected,
1158                    actual,
1159                });
1160            }
1161
1162            // Check whether the response indicates range request support. A `206 Partial Content`
1163            // implies range support while an `Accept-Ranges: bytes` header explicitly advertises it.
1164            let supports_range_requests = response.status() == reqwest::StatusCode::PARTIAL_CONTENT
1165                || response
1166                    .headers()
1167                    .get(reqwest::header::ACCEPT_RANGES)
1168                    .is_some_and(|value| value == "bytes");
1169
1170            // A server can advertise range requests but ignore one. In that case, the
1171            // response is a complete download and must replace the partial bytes.
1172            let replaces_partial_download =
1173                resumed_at.is_some() && response.status() != reqwest::StatusCode::PARTIAL_CONTENT;
1174            if replaces_partial_download {
1175                writer
1176                    .get_mut()
1177                    .set_len(0)
1178                    .await
1179                    .map_err(Error::CacheWrite)?;
1180                writer
1181                    .seek(io::SeekFrom::Start(0))
1182                    .await
1183                    .map_err(Error::CacheWrite)?;
1184                hashers = http_hash_algorithms(hashes)
1185                    .into_iter()
1186                    .map(Hasher::from)
1187                    .collect();
1188                bytes_retrieved = 0;
1189            }
1190
1191            let reader = response
1192                .bytes_stream()
1193                .map_err(|err| self.handle_response_errors(err))
1194                .into_async_read();
1195            let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers);
1196
1197            // Drain the response. This could be a partial response or a full one.
1198            // Note that the partial response here can take several forms: it can be an interrupted
1199            // request *or* it can be a `206 Partial Content`.
1200            let copy_result = match progress {
1201                Some((reporter, progress)) => {
1202                    // Wrap the reader in a progress reporter. This will report 100%
1203                    // progress once the download is complete, before the wheel is unzipped.
1204                    let mut reader = ProgressReader::new(&mut hasher, progress, &**reporter);
1205
1206                    tokio::io::copy(&mut reader, &mut writer)
1207                        .await
1208                        .map_err(Error::CacheWrite)
1209                }
1210                None => tokio::io::copy(&mut hasher, &mut writer)
1211                    .await
1212                    .map_err(Error::CacheWrite),
1213            };
1214
1215            bytes_retrieved += hasher.bytes_read();
1216
1217            let interrupted = copy_result.is_err();
1218            let err = match copy_result {
1219                // Draining the response succeeded. However, the response could be a `206 Partial Content`,
1220                // so we can't assume that we're done.
1221                Ok(_) => {
1222                    // No `resumed_range` means this was a normal full response, so there's
1223                    // no resumption to do. We can leave the loop.
1224                    let Some(range) = resumed_range else {
1225                        break;
1226                    };
1227
1228                    // We should be in sync with the origin. Failure here means that the
1229                    // origin actually sent us an under- or over-length response. Treat this
1230                    // as a non-retryable error, since the HTTP body was received successfully.
1231                    //
1232                    // Byte ranges are inclusive, not exclusive.
1233                    if bytes_retrieved != range.last_byte + 1 {
1234                        return Err(Error::MismatchedRangeSize {
1235                            distribution: dist.to_string(),
1236                            expected: range.last_byte - range.first_byte + 1,
1237                            actual: hasher.bytes_read(),
1238                        });
1239                    }
1240
1241                    // We've successfully performed the download over one or more
1242                    // range requests. We can leave the loop.
1243                    if bytes_retrieved == range.complete_length {
1244                        break;
1245                    }
1246
1247                    // A successful range response may cover only part of the requested
1248                    // bytes. Keep requesting the remainder before extracting the wheel.
1249                    //
1250                    // Observe that we don't return this error; we bind it for handling
1251                    // below.
1252                    Error::CacheWrite(io::Error::new(
1253                        io::ErrorKind::UnexpectedEof,
1254                        "Range response did not complete the download",
1255                    ))
1256                }
1257                // Draining the response failed, potentially because the connection was interrupted.
1258                // Like above, we don't return this error until we know how we want to treat it.
1259                Err(err) => err,
1260            };
1261
1262            // At this point we have *some* error state: either a synthetic error (from a
1263            // partial response that did not complete the download) or a real error (e.g.
1264            // from an interrupted download). We can only resume if certain conditions
1265            // below are met.
1266
1267            // If the response replaces an already in-flight partial download or
1268            // indicates a lack of range-request support, then we can't resume.
1269            // Return the error back to the retry stack.
1270            if replaces_partial_download || !supports_range_requests {
1271                return Err(err);
1272            }
1273
1274            // Some sanity checks: we should have made *some* progress as part of
1275            // partial response, plus our writer's offset should agree with the number
1276            // of bytes retrieved, *and* we should be making forward progress (i.e. not
1277            // resuming behind where we already are).
1278            writer.flush().await.map_err(Error::CacheWrite)?;
1279            let offset = writer
1280                .get_mut()
1281                .stream_position()
1282                .await
1283                .map_err(Error::CacheWrite)?;
1284            if offset == 0
1285                || offset != bytes_retrieved
1286                || resumed_at.is_some_and(|previous| offset <= previous)
1287            {
1288                return Err(err);
1289            }
1290
1291            // Recovering from a failed body consumes the same budget as a full restart.
1292            // A successfully completed range needs no retry, even if more bytes remain.
1293            if interrupted {
1294                let Some(backoff) = retry_state.should_retry(&err, 0) else {
1295                    return Err(err);
1296                };
1297                retry_state.sleep_backoff(backoff).await;
1298            }
1299
1300            // Finally our resumption, which is a range request.
1301            debug!("Resuming download of {url} at byte {offset}");
1302            let resumed_response = retry_state
1303                .send(self.request_with_offset(url.clone(), offset))
1304                .await?;
1305
1306            // A chunked response can fail after all wheel bytes arrive, leaving no satisfiable
1307            // range. Return the original error so the outer retry policy can restart in full.
1308            if resumed_response.status() == reqwest::StatusCode::RANGE_NOT_SATISFIABLE {
1309                debug!("Range not satisfiable while resuming {url}; abandoning resumed download");
1310                return Err(err);
1311            }
1312            resumed_response.error_for_status_ref()?;
1313
1314            resumed_range = if resumed_response.status() == reqwest::StatusCode::PARTIAL_CONTENT {
1315                // The origin honored our range request, so we update `resumed_range`
1316                // for the next iteration.
1317                let Some(range) = content_range(&resumed_response, offset, download_size) else {
1318                    warn!(
1319                        "Invalid range request response from server that declares HTTP range \
1320                         request support, abandoning resumed download: {url}"
1321                    );
1322                    return Err(err);
1323                };
1324                download_size = Some(range.complete_length);
1325                Some(range)
1326            } else {
1327                // The origin is allowed to ignore our range request and send a full response instead.
1328                // That means we have no `resumed_range` to honor on the next iteration.
1329                None
1330            };
1331
1332            response = resumed_response;
1333            resumed_at = Some(offset);
1334        }
1335
1336        // We've left the resumption loop.
1337        // Sanity check: we should have written as many bytes as we expected.
1338        if let Some(expected) = expected_size
1339            && bytes_retrieved != expected
1340        {
1341            return Err(Error::MismatchedSize {
1342                distribution: dist.to_string(),
1343                expected,
1344                actual: bytes_retrieved,
1345            });
1346        }
1347
1348        // Unzip the wheel to a temporary directory.
1349        let extractor = WheelExtractor::new(
1350            self.build_context.cache().root(),
1351            self.content_addressed_cache,
1352        )
1353        .map_err(Error::CacheWrite)?;
1354        let mut file = writer.into_inner();
1355        file.seek(io::SeekFrom::Start(0))
1356            .await
1357            .map_err(Error::CacheWrite)?;
1358
1359        let file = file.into_std().await;
1360        let mut extracted = tokio::task::spawn_blocking(move || extractor.extract_seekable(file))
1361            .await?
1362            .map_err(|err| Error::Extract(filename.to_string(), err))?;
1363        let computed_hashes = Self::validate_hashes(dist, hashes, hashers)?;
1364
1365        // Before we make the wheel accessible by persisting it, ensure that the RECORD is
1366        // valid.
1367        extracted.validate_and_heal_record(dist)?;
1368
1369        // Persist the temporary directory to the directory store.
1370        let id = self
1371            .persist_extracted_wheel(extracted, wheel_entry.path())
1372            .await?;
1373
1374        if let Some((reporter, progress)) = progress {
1375            reporter.on_download_complete(dist.name(), progress);
1376        }
1377
1378        Ok(Archive::new(
1379            id,
1380            computed_hashes,
1381            filename.clone(),
1382            Some(bytes_retrieved),
1383        ))
1384    }
1385
1386    /// Load a wheel from a local path.
1387    async fn load_wheel(
1388        &self,
1389        path: &Path,
1390        filename: &WheelFilename,
1391        wheel_entry: CacheEntry,
1392        dist: &BuiltDist,
1393        hashes: ArchiveHashPolicy<'_>,
1394    ) -> Result<LocalWheel, Error> {
1395        // Acquire an advisory lock, to guard against concurrent writes.
1396        let _lock = Self::lock_wheel(&wheel_entry, filename).await?;
1397
1398        // Determine the last-modified time of the wheel.
1399        let modified = Timestamp::from_path(path).map_err(Error::CacheRead)?;
1400
1401        // Attempt to read the archive pointer from the cache.
1402        let pointer_entry = wheel_entry.with_file(format!("{}.rev", filename.cache_key()));
1403        let pointer = PathArchivePointer::read_from(&pointer_entry)?;
1404
1405        // Extract the archive from the pointer.
1406        let archive = pointer
1407            .filter(|pointer| pointer.is_up_to_date(modified))
1408            .map(PathArchivePointer::into_archive)
1409            .filter(|archive| archive.has_digests(hashes));
1410
1411        // Index hashes may replace `Generate`, but the caller still needs its SHA-256.
1412        let caller_hashes = hashes;
1413        let hashes = if let BuiltDist::Registry(wheels) = dist {
1414            hashes.with_index_hashes(wheels.best_wheel().file.hashes.as_slice())
1415        } else {
1416            hashes
1417        };
1418
1419        // If the file is already unzipped, and the cache is up-to-date, return it.
1420        if let Some(archive) = archive {
1421            Ok(LocalWheel {
1422                dist: Dist::Built(dist.clone()),
1423                archive: self
1424                    .build_context
1425                    .cache()
1426                    .archive(&archive.id)
1427                    .into_boxed_path(),
1428                hashes: archive.hashes,
1429                filename: filename.clone(),
1430                cache: CacheInfo::from_timestamp(modified),
1431                build: None,
1432            })
1433        } else if hashes.is_none() {
1434            // Otherwise, unzip the wheel.
1435            let archive = Archive::new(
1436                self.unzip_wheel(path, wheel_entry.path(), DistRef::Built(dist))
1437                    .await?,
1438                HashDigests::empty(),
1439                filename.clone(),
1440                None,
1441            );
1442
1443            // Write the archive pointer to the cache.
1444            let pointer = PathArchivePointer {
1445                timestamp: modified,
1446                archive: archive.clone(),
1447            };
1448            pointer.write_to(&pointer_entry).await?;
1449
1450            Ok(LocalWheel {
1451                dist: Dist::Built(dist.clone()),
1452                archive: self
1453                    .build_context
1454                    .cache()
1455                    .archive(&archive.id)
1456                    .into_boxed_path(),
1457                hashes: archive.hashes,
1458                filename: filename.clone(),
1459                cache: CacheInfo::from_timestamp(modified),
1460                build: None,
1461            })
1462        } else {
1463            // If necessary, compute the hashes of the wheel.
1464            let file = fs_err::tokio::File::open(path)
1465                .await
1466                .map_err(Error::CacheRead)?;
1467            let extractor = WheelExtractor::new(
1468                self.build_context.cache().root(),
1469                self.content_addressed_cache,
1470            )
1471            .map_err(Error::CacheWrite)?;
1472
1473            // Include the caller's algorithms alongside any index algorithms selected for validation.
1474            let mut algorithms = caller_hashes.algorithms();
1475            algorithms.extend(hashes.algorithms());
1476            algorithms.sort();
1477            algorithms.dedup();
1478            let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
1479            let mut hasher = uv_extract::hash::HashReader::new(file, &mut hashers);
1480
1481            // Unzip the wheel to a temporary directory.
1482            let mut extracted = extractor
1483                .extract_streaming(&mut hasher)
1484                .await
1485                .map_err(|err| Error::Extract(filename.to_string(), err))?;
1486
1487            // Exhaust the reader to compute the hash.
1488            hasher.finish().await.map_err(Error::HashExhaustion)?;
1489
1490            let computed_hashes = Self::validate_hashes(dist, hashes, hashers)?;
1491
1492            // Before we make the wheel accessible by persisting it, ensure that the RECORD is
1493            // valid.
1494            extracted.validate_and_heal_record(dist)?;
1495
1496            // Persist the temporary directory to the directory store.
1497            let id = self
1498                .persist_extracted_wheel(extracted, wheel_entry.path())
1499                .await?;
1500
1501            // Create an archive.
1502            let archive = Archive::new(id, computed_hashes, filename.clone(), None);
1503
1504            // Write the archive pointer to the cache.
1505            let pointer = PathArchivePointer {
1506                timestamp: modified,
1507                archive: archive.clone(),
1508            };
1509            pointer.write_to(&pointer_entry).await?;
1510
1511            Ok(LocalWheel {
1512                dist: Dist::Built(dist.clone()),
1513                archive: self
1514                    .build_context
1515                    .cache()
1516                    .archive(&archive.id)
1517                    .into_boxed_path(),
1518                hashes: archive.hashes,
1519                filename: filename.clone(),
1520                cache: CacheInfo::from_timestamp(modified),
1521                build: None,
1522            })
1523        }
1524    }
1525
1526    /// Unzip a wheel into the cache, returning the path to the unzipped directory.
1527    async fn unzip_wheel(
1528        &self,
1529        path: &Path,
1530        target: &Path,
1531        dist: DistRef<'_>,
1532    ) -> Result<ArchiveId, Error> {
1533        let content_addressed_cache = self.content_addressed_cache;
1534
1535        let mut extracted = tokio::task::spawn_blocking({
1536            let path = path.to_owned();
1537            let root = self.build_context.cache().root().to_path_buf();
1538            move || -> Result<_, Error> {
1539                // Unzip the wheel into a temporary directory.
1540                let extractor = WheelExtractor::new(&root, content_addressed_cache)
1541                    .map_err(Error::CacheWrite)?;
1542                let reader = fs_err::File::open(&path).map_err(Error::CacheWrite)?;
1543                extractor
1544                    .extract_seekable(reader)
1545                    .map_err(|err| Error::Extract(path.to_string_lossy().into_owned(), err))
1546            }
1547        })
1548        .await??;
1549
1550        // Before we make the wheel accessible by persisting it, ensure that the RECORD is valid.
1551        extracted.validate_and_heal_record(dist)?;
1552
1553        // Persist the temporary directory to the directory store.
1554        let id = self.persist_extracted_wheel(extracted, target).await?;
1555
1556        Ok(id)
1557    }
1558
1559    /// Persist an extracted wheel into the archive store.
1560    ///
1561    /// A hash tree makes identical extracted trees converge on one archive entry. Without one,
1562    /// persistence retains the existing behavior of assigning a unique archive ID.
1563    async fn persist_extracted_wheel(
1564        &self,
1565        extracted: ExtractedWheel,
1566        target: &Path,
1567    ) -> Result<ArchiveId, Error> {
1568        let (temp_dir, hashed_wheel) = extracted.into_parts();
1569        let cache = self.build_context.cache();
1570        let (temp_dir, id) = if let Some(HashedWheel { files, tree }) = hashed_wheel {
1571            let digest = DirectoryDigest::from(tree.hash());
1572            let id = ArchiveId::from_digest(digest.into());
1573            let cache = cache.clone();
1574            let temp_dir = tokio::task::spawn_blocking(move || {
1575                persist_archive_files(&cache, temp_dir.path(), &files)
1576                    .map_err(Error::CacheWrite)?;
1577                Ok::<_, Error>(temp_dir)
1578            })
1579            .await??;
1580            (temp_dir, id)
1581        } else {
1582            (temp_dir, ArchiveId::default())
1583        };
1584
1585        cache
1586            .persist_with_id(temp_dir, target, id)
1587            .await
1588            .map_err(Error::CacheWrite)
1589    }
1590
1591    /// Returns a GET [`reqwest::Request`] for the given URL.
1592    fn request(&self, url: DisplaySafeUrl) -> Result<reqwest::Request, reqwest::Error> {
1593        self.client
1594            .unmanaged
1595            .uncached_client(&url)
1596            .get(Url::from(url))
1597            .header(
1598                // `reqwest` defaults to accepting compressed responses.
1599                // Specify identity encoding to get consistent .whl downloading
1600                // behavior from servers. ref: https://github.com/pypa/pip/pull/1688
1601                "accept-encoding",
1602                reqwest::header::HeaderValue::from_static("identity"),
1603            )
1604            .build()
1605    }
1606
1607    /// Build a GET request with a `Range: bytes=<offset>-` header.
1608    ///
1609    /// Used to resume an interrupted download from `offset` bytes into the file.
1610    fn request_with_offset(&self, url: DisplaySafeUrl, offset: u64) -> RequestBuilder<'_> {
1611        self.client
1612            .unmanaged
1613            .uncached_client(&url)
1614            .get(Url::from(url))
1615            .header(
1616                "accept-encoding",
1617                reqwest::header::HeaderValue::from_static("identity"),
1618            )
1619            .header(reqwest::header::RANGE, format!("bytes={offset}-"))
1620    }
1621
1622    /// Return the [`ManagedClient`] used by this resolver.
1623    pub fn client(&self) -> &ManagedClient<'a> {
1624        &self.client
1625    }
1626}
1627
1628/// Share extracted files other than `RECORD` while keeping the unpublished archive complete.
1629fn persist_archive_files(cache: &Cache, archive: &Path, files: &[HashedFile]) -> io::Result<()> {
1630    initialize_rayon_once();
1631    let targets = files
1632        .par_iter()
1633        // Keep RECORD private, since it may have been healed after hashing.
1634        .filter(|file| !file.path().ends_with("RECORD"))
1635        .map(|file| {
1636            let id = ArchiveFileId::from_digest(&file.object_digest_hex());
1637            (archive.join(file.path()), cache.archive_file(&id))
1638        })
1639        .collect::<Vec<_>>();
1640
1641    // Group files by shard so its directory is created once and its files are linked by the
1642    // same worker, avoiding contention between workers on each shard directory.
1643    let mut shards: FxHashMap<&Path, Vec<_>> = FxHashMap::default();
1644    for (source, target) in &targets {
1645        let Some(parent) = target.parent() else {
1646            return Err(io::Error::new(
1647                io::ErrorKind::InvalidInput,
1648                "archive file path must have a parent directory",
1649            ));
1650        };
1651        shards.entry(parent).or_default().push((source, target));
1652    }
1653
1654    let mut shards = shards
1655        .into_iter()
1656        .map(|(parent, files)| (parent, files, Ok(())))
1657        .collect::<Vec<_>>();
1658    // Start larger shards first so their work can overlap the remaining directory creation.
1659    shards.sort_unstable_by_key(|(_, files, _)| Reverse(files.len()));
1660
1661    // Creating shards concurrently contends on their shared parent. Keep creation on this
1662    // thread, while workers link files in the shards that are already available.
1663    in_place_scope(|scope| -> io::Result<()> {
1664        for (parent, files, result) in &mut shards {
1665            fs_err::create_dir_all(parent)?;
1666            scope.spawn(move |_| {
1667                *result = files
1668                    .iter()
1669                    .try_for_each(|(source, target)| persist_archive_file(source, target));
1670            });
1671        }
1672        Ok(())
1673    })?;
1674
1675    shards.into_iter().try_for_each(|(_, _, result)| result)
1676}
1677
1678/// Publish a shared object and retain a hardlink in the archive, with a copy fallback.
1679fn persist_archive_file(src: &Path, dst: &Path) -> io::Result<()> {
1680    // The shard already exists, and most objects are new, so try linking before checking for an
1681    // existing object. This avoids an extra filesystem lookup for every new object.
1682    match fs_err::hard_link(src, dst) {
1683        Ok(()) => return Ok(()),
1684        Err(_) if dst.try_exists()? => {}
1685        Err(_) => return uv_fs::copy_atomic_sync(src, dst),
1686    }
1687
1688    // This archive is still private, so it is safe to replace its extracted copy before publication.
1689    if let Err(err) = fs_err::remove_file(src)
1690        && err.kind() != io::ErrorKind::NotFound
1691    {
1692        return Err(err);
1693    }
1694
1695    fs_err::hard_link(dst, src).or_else(|_| uv_fs::copy_atomic_sync(dst, src))
1696}
1697
1698/// A wrapper around `RegistryClient` that manages a concurrency limit.
1699pub struct ManagedClient<'a> {
1700    pub unmanaged: &'a RegistryClient,
1701    control: Arc<Semaphore>,
1702}
1703
1704impl<'a> ManagedClient<'a> {
1705    /// Create a new `ManagedClient` using the given client and concurrency semaphore.
1706    fn new(client: &'a RegistryClient, control: Arc<Semaphore>) -> Self {
1707        ManagedClient {
1708            unmanaged: client,
1709            control,
1710        }
1711    }
1712
1713    /// Perform a request using the client, respecting the concurrency limit.
1714    ///
1715    /// If the concurrency limit has been reached, this method will wait until a pending
1716    /// operation completes before executing the closure.
1717    pub async fn managed<F, T>(&self, f: impl FnOnce(&'a RegistryClient) -> F) -> T
1718    where
1719        F: Future<Output = T>,
1720    {
1721        let _permit = self.control.acquire().await.unwrap();
1722        f(self.unmanaged).await
1723    }
1724
1725    /// Perform a request using a client that internally manages the concurrency limit.
1726    ///
1727    /// The callback is passed the client and a semaphore. It must acquire the semaphore before
1728    /// any request through the client and drop it after.
1729    ///
1730    /// This method serves as an escape hatch for functions that may want to send multiple requests
1731    /// in parallel.
1732    pub async fn manual<F, T>(&'a self, f: impl FnOnce(&'a RegistryClient, &'a Semaphore) -> F) -> T
1733    where
1734        F: Future<Output = T>,
1735    {
1736        f(self.unmanaged, &self.control).await
1737    }
1738}
1739
1740/// Returns the value of the `Content-Length` header from the [`reqwest::Response`], if present.
1741fn content_length(response: &reqwest::Response) -> Option<u64> {
1742    response
1743        .headers()
1744        .get(reqwest::header::CONTENT_LENGTH)
1745        .and_then(|val| val.to_str().ok())
1746        .and_then(|val| val.parse::<u64>().ok())
1747}
1748
1749/// Return the bounds of a range response starting at `offset` with a known complete length.
1750fn content_range(
1751    response: &reqwest::Response,
1752    offset: u64,
1753    download_size: Option<u64>,
1754) -> Option<ContentRangeBytes> {
1755    let range = response
1756        .headers()
1757        .get(reqwest::header::CONTENT_RANGE)
1758        .and_then(|value| value.to_str().ok())
1759        .and_then(ContentRange::parse)?;
1760
1761    let range = match range {
1762        ContentRange::Bytes(range) => range,
1763        ContentRange::UnboundBytes(ContentRangeUnbound {
1764            first_byte,
1765            last_byte,
1766        }) => ContentRangeBytes {
1767            first_byte,
1768            last_byte,
1769            complete_length: download_size?,
1770        },
1771        ContentRange::Unsatisfied(_) => return None,
1772    };
1773
1774    if range.first_byte != offset
1775        || range.last_byte >= range.complete_length
1776        || download_size.is_some_and(|size| size != range.complete_length)
1777    {
1778        return None;
1779    }
1780
1781    Some(range)
1782}
1783
1784/// An asynchronous reader that reports progress as bytes are read.
1785struct ProgressReader<'a, R> {
1786    reader: R,
1787    index: usize,
1788    reporter: &'a dyn Reporter,
1789}
1790
1791impl<'a, R> ProgressReader<'a, R> {
1792    /// Create a new [`ProgressReader`] that wraps another reader.
1793    fn new(reader: R, index: usize, reporter: &'a dyn Reporter) -> Self {
1794        Self {
1795            reader,
1796            index,
1797            reporter,
1798        }
1799    }
1800}
1801
1802impl<R> AsyncRead for ProgressReader<'_, R>
1803where
1804    R: AsyncRead + Unpin,
1805{
1806    fn poll_read(
1807        mut self: Pin<&mut Self>,
1808        cx: &mut Context<'_>,
1809        buf: &mut ReadBuf<'_>,
1810    ) -> Poll<io::Result<()>> {
1811        Pin::new(&mut self.as_mut().reader)
1812            .poll_read(cx, buf)
1813            .map_ok(|()| {
1814                self.reporter
1815                    .on_download_progress(self.index, buf.filled().len() as u64);
1816            })
1817    }
1818}
1819
1820/// A pointer to an archive in the cache, fetched from an HTTP archive.
1821///
1822/// Encoded with `MsgPack`, and represented on disk by a `.http` file.
1823#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1824pub struct HttpArchivePointer {
1825    archive: Archive,
1826}
1827
1828impl HttpArchivePointer {
1829    /// Read an [`HttpArchivePointer`] from the cache.
1830    pub fn read_from(path: impl AsRef<Path>) -> Result<Option<Self>, Error> {
1831        match fs_err::File::open(path.as_ref()) {
1832            Ok(file) => {
1833                let data = DataWithCachePolicy::from_reader(file)?.data;
1834                let archive = rmp_serde::from_slice::<Archive>(&data)?;
1835                Ok(Some(Self { archive }))
1836            }
1837            Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None),
1838            Err(err) => Err(Error::CacheRead(err)),
1839        }
1840    }
1841
1842    /// Return the [`Archive`] from the pointer.
1843    pub fn into_archive(self) -> Archive {
1844        self.archive
1845    }
1846
1847    /// Return the [`CacheInfo`] from the pointer.
1848    pub fn to_cache_info(&self) -> CacheInfo {
1849        CacheInfo::default()
1850    }
1851
1852    /// Return the [`BuildInfo`] from the pointer.
1853    pub fn to_build_info(&self) -> Option<BuildInfo> {
1854        None
1855    }
1856}
1857
1858/// A pointer to an archive in the cache, fetched from a local path.
1859///
1860/// Encoded with `MsgPack`, and represented on disk by a `.rev` file.
1861#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1862pub struct PathArchivePointer {
1863    timestamp: Timestamp,
1864    archive: Archive,
1865}
1866
1867impl PathArchivePointer {
1868    /// Read an [`PathArchivePointer`] from the cache.
1869    pub fn read_from(path: impl AsRef<Path>) -> Result<Option<Self>, Error> {
1870        match fs_err::read(path) {
1871            Ok(cached) => Ok(Some(rmp_serde::from_slice::<Self>(&cached)?)),
1872            Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None),
1873            Err(err) => Err(Error::CacheRead(err)),
1874        }
1875    }
1876
1877    /// Write an [`PathArchivePointer`] to the cache.
1878    async fn write_to(&self, entry: &CacheEntry) -> Result<(), Error> {
1879        write_atomic(entry.path(), rmp_serde::to_vec(&self)?)
1880            .await
1881            .map_err(Error::CacheWrite)
1882    }
1883
1884    /// Returns `true` if the archive is up-to-date with the given modified timestamp.
1885    pub fn is_up_to_date(&self, modified: Timestamp) -> bool {
1886        self.timestamp == modified
1887    }
1888
1889    /// Return the [`Archive`] from the pointer.
1890    pub fn into_archive(self) -> Archive {
1891        self.archive
1892    }
1893
1894    /// Return the [`CacheInfo`] from the pointer.
1895    pub fn to_cache_info(&self) -> CacheInfo {
1896        CacheInfo::from_timestamp(self.timestamp)
1897    }
1898
1899    /// Return the [`BuildInfo`] from the pointer.
1900    pub fn to_build_info(&self) -> Option<BuildInfo> {
1901        None
1902    }
1903}