Skip to main content

uv_distribution/
distribution_database.rs

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