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