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