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