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