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 rayon::in_place_scope;
11use rayon::prelude::*;
12use rustc_hash::FxHashMap;
13use tokio::io::{AsyncRead, AsyncSeekExt, ReadBuf};
14use tokio::sync::Semaphore;
15use tokio_util::compat::FuturesAsyncReadCompatExt;
16use tracing::{Instrument, info_span, instrument, warn};
17use url::Url;
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};
24use uv_configuration::initialize_rayon_once;
25use uv_distribution_filename::WheelFilename;
26use uv_distribution_types::{
27 ArchiveHashPolicy, BuildInfo, BuildableSource, BuiltDist, Dist, DistRef, HashCollection,
28 HashValidation, Hashed, IndexUrl, InstalledDist, MetadataHashPolicy, Name, SourceDist,
29 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_platform_tags::Tags;
36use uv_preview::PreviewFeature;
37use uv_pypi_types::{HashDigest, HashDigests, PyProjectToml};
38use uv_python::PythonVariant;
39use uv_redacted::DisplaySafeUrl;
40use uv_types::{BuildContext, BuildStack};
41
42use crate::archive::Archive;
43use crate::error::PythonVersion;
44use crate::extracted_wheel::{ExtractedWheel, HashedWheel, WheelExtractor};
45use crate::hash::http_hash_algorithms;
46use crate::metadata::{ArchiveMetadata, Metadata};
47use crate::source::SourceDistributionBuilder;
48use crate::{Error, LocalWheel, Reporter, RequiresDist};
49
50pub struct DistributionDatabase<'a, Context: BuildContext> {
63 build_context: &'a Context,
64 builder: SourceDistributionBuilder<'a, Context>,
65 client: ManagedClient<'a>,
66 reporter: Option<Arc<dyn Reporter>>,
67 content_addressed_cache: bool,
68}
69
70impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
71 pub fn new(
72 client: &'a RegistryClient,
73 build_context: &'a Context,
74 downloads_semaphore: Arc<Semaphore>,
75 ) -> Self {
76 let content_addressed_cache = uv_preview::is_enabled(PreviewFeature::ContentAddressedCache)
80 && !uv_extract::insecure_no_validate();
81 Self {
82 build_context,
83 builder: SourceDistributionBuilder::new(build_context),
84 client: ManagedClient::new(client, downloads_semaphore),
85 reporter: None,
86 content_addressed_cache,
87 }
88 }
89
90 #[must_use]
92 pub fn with_build_stack(self, build_stack: &'a BuildStack) -> Self {
93 Self {
94 builder: self.builder.with_build_stack(build_stack),
95 ..self
96 }
97 }
98
99 #[must_use]
101 pub fn with_reporter(self, reporter: Arc<dyn Reporter>) -> Self {
102 Self {
103 builder: self.builder.with_reporter(reporter.clone()),
104 reporter: Some(reporter),
105 ..self
106 }
107 }
108
109 fn handle_response_errors(&self, err: reqwest::Error) -> io::Error {
111 if err.is_timeout() {
112 io::Error::new(
114 io::ErrorKind::TimedOut,
115 format!(
116 "Failed to download distribution due to network timeout. Try increasing UV_HTTP_TIMEOUT (current value: {}s).",
117 self.client.unmanaged.read_timeout().as_secs()
118 ),
119 )
120 } else {
121 io::Error::other(err)
122 }
123 }
124
125 async fn lock_wheel(
134 wheel_entry: &CacheEntry,
135 filename: &WheelFilename,
136 ) -> Result<LockedFile, Error> {
137 #[cfg(windows)]
141 let lock_key = filename.stem();
142 #[cfg(not(windows))]
144 let lock_key = filename.cache_key();
145
146 let lock_entry = wheel_entry.with_file(format!("{lock_key}.lock"));
147 lock_entry.lock().await.map_err(Error::CacheLock)
148 }
149
150 #[instrument(skip_all, fields(%dist))]
157 pub async fn get_or_build_wheel(
158 &self,
159 dist: &Dist,
160 tags: &Tags,
161 hashes: ArchiveHashPolicy<'_>,
162 ) -> Result<LocalWheel, Error> {
163 match dist {
164 Dist::Built(built) => self.get_wheel(built, hashes).await,
165 Dist::Source(source) => self.build_wheel(source, tags, hashes).await,
166 }
167 }
168
169 #[instrument(skip_all, fields(%dist))]
175 pub async fn get_installed_metadata(
176 &self,
177 dist: &InstalledDist,
178 ) -> Result<ArchiveMetadata, Error> {
179 if let Some(metadata) = self
181 .build_context
182 .dependency_metadata()
183 .get(dist.name(), Some(dist.version()))
184 {
185 return Ok(ArchiveMetadata::from_metadata23(metadata));
186 }
187
188 let metadata = dist
189 .read_metadata()
190 .map_err(|err| Error::ReadInstalled(Box::new(dist.clone()), err))?;
191
192 Ok(ArchiveMetadata::from_metadata23(metadata.clone()))
193 }
194
195 #[instrument(skip_all, fields(%dist))]
200 pub async fn get_or_build_wheel_metadata(
201 &self,
202 dist: &Dist,
203 hashes: MetadataHashPolicy<'_>,
204 ) -> Result<ArchiveMetadata, Error> {
205 match dist {
206 Dist::Built(built) => self.get_wheel_metadata(built, hashes).await,
207 Dist::Source(source) => {
208 self.build_wheel_metadata(&BuildableSource::Dist(source), hashes)
209 .await
210 }
211 }
212 }
213
214 async fn get_wheel(
219 &self,
220 dist: &BuiltDist,
221 hashes: ArchiveHashPolicy<'_>,
222 ) -> Result<LocalWheel, Error> {
223 match dist {
224 BuiltDist::Registry(wheels) => {
225 let wheel = wheels.best_wheel();
226 let url = wheel.file.url.to_url()?;
227 let size = wheel.file.size;
228
229 let wheel_entry = self.build_context.cache().entry(
231 CacheBucket::Wheels,
232 WheelCache::Index(&wheel.index).wheel_dir(wheel.name().as_ref()),
233 wheel.filename.cache_key(),
234 );
235
236 if url.scheme() == "file" {
238 let path = url
239 .to_file_path()
240 .map_err(|()| Error::NonFileUrl(url.clone()))?;
241 return self
242 .load_wheel(&path, &wheel.filename, wheel_entry, dist, hashes)
243 .await;
244 }
245
246 match self
248 .stream_wheel(
249 url.clone(),
250 dist.index(),
251 &wheel.filename,
252 size,
253 &wheel_entry,
254 dist,
255 hashes,
256 )
257 .await
258 {
259 Ok(archive) => Ok(LocalWheel {
260 dist: Dist::Built(dist.clone()),
261 archive: self
262 .build_context
263 .cache()
264 .archive(&archive.id)
265 .into_boxed_path(),
266 hashes: archive.hashes,
267 filename: wheel.filename.clone(),
268 cache: CacheInfo::default(),
269 build: None,
270 }),
271 Err(Error::Extract(name, err)) => {
272 if err.is_http_streaming_unsupported() {
273 warn!(
274 "Streaming unsupported for {dist}; downloading wheel to disk ({err})"
275 );
276 } else if err.is_http_streaming_failed() {
277 warn!("Streaming failed for {dist}; downloading wheel to disk ({err})");
278 } else {
279 return Err(Error::Extract(name, err));
280 }
281
282 let archive = self
285 .download_wheel(
286 url,
287 dist.index(),
288 &wheel.filename,
289 size,
290 &wheel_entry,
291 dist,
292 hashes,
293 )
294 .await?;
295
296 Ok(LocalWheel {
297 dist: Dist::Built(dist.clone()),
298 archive: self
299 .build_context
300 .cache()
301 .archive(&archive.id)
302 .into_boxed_path(),
303 hashes: archive.hashes,
304 filename: wheel.filename.clone(),
305 cache: CacheInfo::default(),
306 build: None,
307 })
308 }
309 Err(err) => Err(err),
310 }
311 }
312
313 BuiltDist::DirectUrl(wheel) => {
314 let wheel_entry = self.build_context.cache().entry(
316 CacheBucket::Wheels,
317 WheelCache::Url(&wheel.url).wheel_dir(wheel.name().as_ref()),
318 wheel.filename.cache_key(),
319 );
320
321 match self
323 .stream_wheel(
324 wheel.url.raw().clone(),
325 None,
326 &wheel.filename,
327 wheel.size,
328 &wheel_entry,
329 dist,
330 hashes,
331 )
332 .await
333 {
334 Ok(archive) => Ok(LocalWheel {
335 dist: Dist::Built(dist.clone()),
336 archive: self
337 .build_context
338 .cache()
339 .archive(&archive.id)
340 .into_boxed_path(),
341 hashes: archive.hashes,
342 filename: wheel.filename.clone(),
343 cache: CacheInfo::default(),
344 build: None,
345 }),
346 Err(Error::Extract(name, err)) => {
347 if err.is_http_streaming_unsupported() {
348 warn!(
349 "Streaming unsupported for {dist}; downloading wheel to disk ({err})"
350 );
351 } else if err.is_http_streaming_failed() {
352 warn!("Streaming failed for {dist}; downloading wheel to disk ({err})");
353 } else {
354 return Err(Error::Extract(name, err));
355 }
356
357 let archive = self
360 .download_wheel(
361 wheel.url.raw().clone(),
362 None,
363 &wheel.filename,
364 wheel.size,
365 &wheel_entry,
366 dist,
367 hashes,
368 )
369 .await?;
370 Ok(LocalWheel {
371 dist: Dist::Built(dist.clone()),
372 archive: self
373 .build_context
374 .cache()
375 .archive(&archive.id)
376 .into_boxed_path(),
377 hashes: archive.hashes,
378 filename: wheel.filename.clone(),
379 cache: CacheInfo::default(),
380 build: None,
381 })
382 }
383 Err(err) => Err(err),
384 }
385 }
386
387 BuiltDist::GitPath(wheel) => {
388 let fetch = self
390 .build_context
391 .git()
392 .fetch(
393 &wheel.git,
394 self.client.unmanaged.git_http_settings(wheel.git.url()),
395 self.build_context.cache().bucket(CacheBucket::Git),
396 self.reporter.clone().map(<dyn Reporter>::into_git_reporter),
397 )
398 .await?;
399
400 if wheel.git.lfs().enabled() && !fetch.lfs_ready() {
401 if GIT_LFS.is_err() {
402 return Err(Error::MissingWheelGitLfsArtifacts(
403 wheel.url.to_url(),
404 GitError::GitLfsNotFound,
405 ));
406 }
407 return Err(Error::MissingWheelGitLfsArtifacts(
408 wheel.url.to_url(),
409 GitError::GitLfsNotConfigured,
410 ));
411 }
412
413 let git_sha = fetch.git().precise().expect("Exact commit after checkout");
414 let cache_entry = self.build_context.cache().entry(
415 CacheBucket::Wheels,
416 WheelCache::Git(&wheel.url, git_sha.as_short_str()).root(),
417 wheel.filename.stem(),
418 );
419
420 let install_path = fetch.path().join(&wheel.install_path);
421
422 self.load_wheel(&install_path, &wheel.filename, cache_entry, dist, hashes)
423 .await
424 }
425
426 BuiltDist::Path(wheel) => {
427 let cache_entry = self.build_context.cache().entry(
428 CacheBucket::Wheels,
429 WheelCache::Url(&wheel.url).wheel_dir(wheel.name().as_ref()),
430 wheel.filename.cache_key(),
431 );
432
433 self.load_wheel(
434 &wheel.install_path,
435 &wheel.filename,
436 cache_entry,
437 dist,
438 hashes,
439 )
440 .await
441 }
442 }
443 }
444
445 async fn build_wheel(
451 &self,
452 dist: &SourceDist,
453 tags: &Tags,
454 hashes: ArchiveHashPolicy<'_>,
455 ) -> Result<LocalWheel, Error> {
456 let built_wheel = self
457 .builder
458 .download_and_build(&BuildableSource::Dist(dist), tags, hashes, &self.client)
459 .boxed_local()
460 .await?;
461
462 if !built_wheel.filename.is_compatible(tags) {
468 return if tags.is_cross() {
469 Err(Error::BuiltWheelIncompatibleTargetPlatform {
470 filename: built_wheel.filename,
471 python_platform: tags.python_platform().clone(),
472 python_version: PythonVersion {
473 version: tags.python_version(),
474 variant: if tags.is_freethreaded() {
475 PythonVariant::Freethreaded
476 } else {
477 PythonVariant::Default
478 },
479 },
480 })
481 } else {
482 Err(Error::BuiltWheelIncompatibleHostPlatform {
483 filename: built_wheel.filename,
484 python_platform: tags.python_platform().clone(),
485 python_version: PythonVersion {
486 version: tags.python_version(),
487 variant: if tags.is_freethreaded() {
488 PythonVariant::Freethreaded
489 } else {
490 PythonVariant::Default
491 },
492 },
493 })
494 };
495 }
496
497 let wheel_entry = CacheEntry::from_path(built_wheel.target.as_ref());
499 let _lock = Self::lock_wheel(&wheel_entry, &built_wheel.filename).await?;
500
501 match self.build_context.cache().resolve_link(&built_wheel.target) {
504 Ok(archive) => {
505 return Ok(LocalWheel {
506 dist: Dist::Source(dist.clone()),
507 archive: archive.into_boxed_path(),
508 filename: built_wheel.filename,
509 hashes: built_wheel.hashes,
510 cache: built_wheel.cache_info,
511 build: Some(built_wheel.build_info),
512 });
513 }
514 Err(err) if err.kind() == io::ErrorKind::NotFound => {}
515 Err(err) => return Err(Error::CacheRead(err)),
516 }
517
518 let id = self
520 .unzip_wheel(
521 &built_wheel.path,
522 &built_wheel.target,
523 DistRef::Source(dist),
524 )
525 .await?;
526
527 Ok(LocalWheel {
528 dist: Dist::Source(dist.clone()),
529 archive: self.build_context.cache().archive(&id).into_boxed_path(),
530 hashes: built_wheel.hashes,
531 filename: built_wheel.filename,
532 cache: built_wheel.cache_info,
533 build: Some(built_wheel.build_info),
534 })
535 }
536
537 async fn get_wheel_metadata(
542 &self,
543 dist: &BuiltDist,
544 hashes: MetadataHashPolicy<'_>,
545 ) -> Result<ArchiveMetadata, Error> {
546 let hash_policy = match hashes.validation {
547 HashValidation::None => {
548 let compute_hashes = match hashes.collection {
549 HashCollection::None => false,
550 HashCollection::Url | HashCollection::All => match dist {
551 BuiltDist::Registry(dist) => {
552 hashes.collection == HashCollection::All
555 && dist.best_wheel().file.hashes.is_empty()
556 }
557 BuiltDist::DirectUrl(dist) => parse_url_hashes(&dist.url).is_none(),
558 BuiltDist::Path(_) | BuiltDist::GitPath(_) => true,
559 },
560 };
561 if compute_hashes {
562 ArchiveHashPolicy::Generate
563 } else {
564 ArchiveHashPolicy::None
565 }
566 }
567 HashValidation::Any(_) | HashValidation::All(_) => hashes.validation.into(),
568 };
569
570 if hash_policy == ArchiveHashPolicy::Generate {
573 let wheel = self.get_wheel(dist, hash_policy).await?;
574 let metadata = if let Some(metadata) = self
576 .build_context
577 .dependency_metadata()
578 .get(dist.name(), Some(dist.version()))
579 {
580 metadata
581 } else {
582 wheel.metadata()?
583 };
584 let hashes = wheel.hashes;
585 return Ok(ArchiveMetadata {
586 metadata: Metadata::from_metadata23(metadata),
587 hashes,
588 });
589 }
590
591 if let Some(metadata) = self
593 .build_context
594 .dependency_metadata()
595 .get(dist.name(), Some(dist.version()))
596 {
597 return Ok(ArchiveMetadata::from_metadata23(metadata));
598 }
599
600 let result = self
601 .client
602 .managed(|client| {
603 client
604 .wheel_metadata(
605 dist,
606 self.build_context.git(),
607 self.build_context.capabilities(),
608 self.reporter.clone().map(<dyn Reporter>::into_git_reporter),
609 )
610 .boxed_local()
611 })
612 .await;
613
614 match result {
615 Ok(metadata) => {
616 Ok(ArchiveMetadata::from_metadata23(metadata))
618 }
619 Err(err) if err.is_http_streaming_unsupported() => {
620 warn!(
621 "Streaming unsupported when fetching metadata for {dist}; downloading wheel directly ({err})"
622 );
623
624 let wheel = self.get_wheel(dist, hash_policy).await?;
627 let metadata = wheel.metadata()?;
628 let hashes = wheel.hashes;
629 Ok(ArchiveMetadata {
630 metadata: Metadata::from_metadata23(metadata),
631 hashes,
632 })
633 }
634 Err(err) => Err(err.into()),
635 }
636 }
637
638 pub async fn build_wheel_metadata(
643 &self,
644 source: &BuildableSource<'_>,
645 hashes: MetadataHashPolicy<'_>,
646 ) -> Result<ArchiveMetadata, Error> {
647 if let Some(dist) = source.as_dist() {
649 if let Some(metadata) = self
650 .build_context
651 .dependency_metadata()
652 .get(dist.name(), dist.version())
653 {
654 self.builder.resolve_revision(source, &self.client).await?;
657
658 return Ok(ArchiveMetadata::from_metadata23(metadata));
659 }
660 }
661
662 let url_hashes = if let BuildableSource::Dist(SourceDist::DirectUrl(dist)) = source {
663 parse_url_hashes(&dist.url)
664 } else if let BuildableSource::Url(SourceUrl::Direct(url)) = source {
665 parse_url_hashes(url.url)
666 } else {
667 None
668 };
669
670 let build_hash_policy = match hashes.validation {
671 HashValidation::None => match hashes.collection {
672 HashCollection::None => ArchiveHashPolicy::None,
673 HashCollection::Url | HashCollection::All => match url_hashes.as_ref() {
676 Some(digests) => ArchiveHashPolicy::All(digests.as_slice()),
677 None => ArchiveHashPolicy::Generate,
678 },
679 },
680 HashValidation::Any(_) | HashValidation::All(_) => hashes.validation.into(),
681 };
682 let metadata = self
683 .builder
684 .download_and_build_metadata(source, build_hash_policy, &self.client)
685 .boxed_local()
686 .await?;
687
688 Ok(metadata)
689 }
690
691 pub async fn requires_dist(
693 &self,
694 path: &Path,
695 pyproject_toml: &PyProjectToml,
696 ) -> Result<Option<RequiresDist>, Error> {
697 self.builder
698 .source_tree_requires_dist(
699 path,
700 pyproject_toml,
701 self.client.unmanaged.credentials_cache(),
702 )
703 .await
704 }
705
706 async fn stream_wheel(
708 &self,
709 url: DisplaySafeUrl,
710 index: Option<&IndexUrl>,
711 filename: &WheelFilename,
712 size: Option<u64>,
713 wheel_entry: &CacheEntry,
714 dist: &BuiltDist,
715 hashes: ArchiveHashPolicy<'_>,
716 ) -> Result<Archive, Error> {
717 let expected_size = match dist {
718 BuiltDist::Registry(dist) if dist.best_wheel().size_is_authoritative => size,
719 BuiltDist::DirectUrl(_) => size,
720 _ => None,
721 };
722
723 let _lock = Self::lock_wheel(wheel_entry, filename).await?;
725
726 let http_entry = wheel_entry.with_file(format!("{}.http", filename.cache_key()));
728
729 let download = |response: reqwest::Response| {
730 async {
731 let progress_size = size.or_else(|| content_length(&response));
732
733 let progress = self.reporter.as_ref().map(|reporter| {
734 (
735 reporter,
736 reporter.on_download_start(dist.name(), progress_size),
737 )
738 });
739
740 let reader = response
741 .bytes_stream()
742 .map_err(|err| self.handle_response_errors(err))
743 .into_async_read();
744
745 let algorithms = http_hash_algorithms(hashes);
747 let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
748 let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers);
749
750 let extractor = WheelExtractor::new(
752 self.build_context.cache().root(),
753 self.content_addressed_cache,
754 )
755 .map_err(Error::CacheWrite)?;
756
757 let mut extracted = match progress {
758 Some((reporter, progress)) => {
759 let mut reader = ProgressReader::new(&mut hasher, progress, &**reporter);
760 extractor
761 .extract_streaming(&mut reader)
762 .await
763 .map_err(|err| Error::Extract(filename.to_string(), err))?
764 }
765 None => extractor
766 .extract_streaming(&mut hasher)
767 .await
768 .map_err(|err| Error::Extract(filename.to_string(), err))?,
769 };
770 hasher.finish().await.map_err(Error::HashExhaustion)?;
772 let actual_size = hasher.bytes_read();
773 if let Some(expected) = expected_size
774 && actual_size != expected
775 {
776 return Err(Error::MismatchedSize {
777 distribution: dist.to_string(),
778 expected,
779 actual: actual_size,
780 });
781 }
782
783 extracted.validate_and_heal_record(dist)?;
786
787 let id = self
789 .persist_extracted_wheel(extracted, wheel_entry.path())
790 .await?;
791
792 if let Some((reporter, progress)) = progress {
793 reporter.on_download_complete(dist.name(), progress);
794 }
795
796 Ok(Archive::new(
797 id,
798 hashers.into_iter().map(HashDigest::from).collect(),
799 filename.clone(),
800 Some(actual_size),
801 ))
802 }
803 .instrument(info_span!("wheel", wheel = %dist))
804 };
805
806 let req = self.request(url.clone())?;
808
809 let cache_control = match self.client.unmanaged.connectivity() {
811 Connectivity::Online
812 if let Some(header) = index.and_then(|index| {
813 self.build_context
814 .locations()
815 .artifact_cache_control_for(index)
816 }) =>
817 {
818 CacheControl::Override(header)
819 }
820 Connectivity::Online => CacheControl::from(
821 self.build_context
822 .cache()
823 .freshness(&http_entry, Some(&filename.name), None)
824 .map_err(Error::CacheRead)?,
825 ),
826 Connectivity::Offline => CacheControl::AllowStale,
827 };
828
829 let archive = self
830 .client
831 .managed(|client| {
832 client.cached_client().get_serde_with_retry(
833 req,
834 &http_entry,
835 cache_control.clone(),
836 download,
837 )
838 })
839 .await
840 .map_err(|err| match err {
841 CachedClientError::Callback { err, .. } => err,
842 CachedClientError::Client(err) => Error::Client(err),
843 })?;
844
845 if let (Some(expected), Some(actual)) = (expected_size, archive.size)
846 && expected != actual
847 {
848 return Err(Error::MismatchedSize {
849 distribution: dist.to_string(),
850 expected,
851 actual,
852 });
853 }
854
855 let archive = Some(archive)
857 .filter(|archive| archive.has_digests(hashes))
858 .filter(|archive| archive.exists(self.build_context.cache()))
859 .filter(|archive| expected_size.is_none() || archive.size.is_some());
860
861 let archive = if let Some(archive) = archive {
862 archive
863 } else {
864 self.client
865 .managed(async |client| {
866 client
867 .cached_client()
868 .skip_cache_with_retry(
869 self.request(url)?,
870 &http_entry,
871 cache_control,
872 download,
873 )
874 .await
875 .map_err(|err| match err {
876 CachedClientError::Callback { err, .. } => err,
877 CachedClientError::Client(err) => Error::Client(err),
878 })
879 })
880 .await?
881 };
882
883 Ok(archive)
884 }
885
886 async fn download_wheel(
888 &self,
889 url: DisplaySafeUrl,
890 index: Option<&IndexUrl>,
891 filename: &WheelFilename,
892 size: Option<u64>,
893 wheel_entry: &CacheEntry,
894 dist: &BuiltDist,
895 hashes: ArchiveHashPolicy<'_>,
896 ) -> Result<Archive, Error> {
897 let expected_size = match dist {
898 BuiltDist::Registry(dist) if dist.best_wheel().size_is_authoritative => size,
899 BuiltDist::DirectUrl(_) => size,
900 _ => None,
901 };
902
903 let content_addressed_cache = self.content_addressed_cache;
904
905 let _lock = Self::lock_wheel(wheel_entry, filename).await?;
907
908 let http_entry = wheel_entry.with_file(format!("{}.http", filename.cache_key()));
910
911 let download = |response: reqwest::Response| {
912 async {
913 let progress_size = size.or_else(|| content_length(&response));
914
915 let progress = self.reporter.as_ref().map(|reporter| {
916 (
917 reporter,
918 reporter.on_download_start(dist.name(), progress_size),
919 )
920 });
921
922 let reader = response
923 .bytes_stream()
924 .map_err(|err| self.handle_response_errors(err))
925 .into_async_read();
926 let algorithms = http_hash_algorithms(hashes);
927 let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
928 let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers);
929
930 let temp_file = tempfile::tempfile_in(self.build_context.cache().root())
932 .map_err(Error::CacheWrite)?;
933 let mut writer = tokio::io::BufWriter::new(fs_err::tokio::File::from_std(
934 fs_err::File::from_parts(temp_file, self.build_context.cache().root()),
936 ));
937
938 match progress {
939 Some((reporter, progress)) => {
940 let mut reader = ProgressReader::new(&mut hasher, progress, &**reporter);
943
944 tokio::io::copy(&mut reader, &mut writer)
945 .await
946 .map_err(Error::CacheWrite)?;
947 }
948 None => {
949 tokio::io::copy(&mut hasher, &mut writer)
950 .await
951 .map_err(Error::CacheWrite)?;
952 }
953 }
954
955 if let Some(expected) = expected_size
956 && hasher.bytes_read() != expected
957 {
958 return Err(Error::MismatchedSize {
959 distribution: dist.to_string(),
960 expected,
961 actual: hasher.bytes_read(),
962 });
963 }
964
965 let actual_size = hasher.bytes_read();
966
967 let extractor =
969 WheelExtractor::new(self.build_context.cache().root(), content_addressed_cache)
970 .map_err(Error::CacheWrite)?;
971 let mut file = writer.into_inner();
972 file.seek(io::SeekFrom::Start(0))
973 .await
974 .map_err(Error::CacheWrite)?;
975
976 let file = file.into_std().await;
977 let mut extracted =
978 tokio::task::spawn_blocking(move || extractor.extract_seekable(file))
979 .await?
980 .map_err(|err| Error::Extract(filename.to_string(), err))?;
981 let hashes = hashers.into_iter().map(HashDigest::from).collect();
982
983 extracted.validate_and_heal_record(dist)?;
986
987 let id = self
989 .persist_extracted_wheel(extracted, wheel_entry.path())
990 .await?;
991
992 if let Some((reporter, progress)) = progress {
993 reporter.on_download_complete(dist.name(), progress);
994 }
995
996 Ok(Archive::new(
997 id,
998 hashes,
999 filename.clone(),
1000 Some(actual_size),
1001 ))
1002 }
1003 .instrument(info_span!("wheel", wheel = %dist))
1004 };
1005
1006 let req = self.request(url.clone())?;
1008
1009 let cache_control = match self.client.unmanaged.connectivity() {
1011 Connectivity::Online
1012 if let Some(header) = index.and_then(|index| {
1013 self.build_context
1014 .locations()
1015 .artifact_cache_control_for(index)
1016 }) =>
1017 {
1018 CacheControl::Override(header)
1019 }
1020 Connectivity::Online => CacheControl::from(
1021 self.build_context
1022 .cache()
1023 .freshness(&http_entry, Some(&filename.name), None)
1024 .map_err(Error::CacheRead)?,
1025 ),
1026 Connectivity::Offline => CacheControl::AllowStale,
1027 };
1028
1029 let archive = self
1030 .client
1031 .managed(|client| {
1032 client.cached_client().get_serde_with_retry(
1033 req,
1034 &http_entry,
1035 cache_control.clone(),
1036 download,
1037 )
1038 })
1039 .await
1040 .map_err(|err| match err {
1041 CachedClientError::Callback { err, .. } => err,
1042 CachedClientError::Client(err) => Error::Client(err),
1043 })?;
1044
1045 if let (Some(expected), Some(actual)) = (expected_size, archive.size)
1046 && expected != actual
1047 {
1048 return Err(Error::MismatchedSize {
1049 distribution: dist.to_string(),
1050 expected,
1051 actual,
1052 });
1053 }
1054
1055 let archive = Some(archive)
1057 .filter(|archive| archive.has_digests(hashes))
1058 .filter(|archive| archive.exists(self.build_context.cache()))
1059 .filter(|archive| expected_size.is_none() || archive.size.is_some());
1060
1061 let archive = if let Some(archive) = archive {
1062 archive
1063 } else {
1064 self.client
1065 .managed(async |client| {
1066 client
1067 .cached_client()
1068 .skip_cache_with_retry(
1069 self.request(url)?,
1070 &http_entry,
1071 cache_control,
1072 download,
1073 )
1074 .await
1075 .map_err(|err| match err {
1076 CachedClientError::Callback { err, .. } => err,
1077 CachedClientError::Client(err) => Error::Client(err),
1078 })
1079 })
1080 .await?
1081 };
1082
1083 Ok(archive)
1084 }
1085
1086 async fn load_wheel(
1088 &self,
1089 path: &Path,
1090 filename: &WheelFilename,
1091 wheel_entry: CacheEntry,
1092 dist: &BuiltDist,
1093 hashes: ArchiveHashPolicy<'_>,
1094 ) -> Result<LocalWheel, Error> {
1095 let _lock = Self::lock_wheel(&wheel_entry, filename).await?;
1097
1098 let modified = Timestamp::from_path(path).map_err(Error::CacheRead)?;
1100
1101 let pointer_entry = wheel_entry.with_file(format!("{}.rev", filename.cache_key()));
1103 let pointer = PathArchivePointer::read_from(&pointer_entry)?;
1104
1105 let archive = pointer
1107 .filter(|pointer| pointer.is_up_to_date(modified))
1108 .map(PathArchivePointer::into_archive)
1109 .filter(|archive| archive.has_digests(hashes));
1110
1111 if let Some(archive) = archive {
1113 Ok(LocalWheel {
1114 dist: Dist::Built(dist.clone()),
1115 archive: self
1116 .build_context
1117 .cache()
1118 .archive(&archive.id)
1119 .into_boxed_path(),
1120 hashes: archive.hashes,
1121 filename: filename.clone(),
1122 cache: CacheInfo::from_timestamp(modified),
1123 build: None,
1124 })
1125 } else if hashes.is_none() {
1126 let archive = Archive::new(
1128 self.unzip_wheel(path, wheel_entry.path(), DistRef::Built(dist))
1129 .await?,
1130 HashDigests::empty(),
1131 filename.clone(),
1132 None,
1133 );
1134
1135 let pointer = PathArchivePointer {
1137 timestamp: modified,
1138 archive: archive.clone(),
1139 };
1140 pointer.write_to(&pointer_entry).await?;
1141
1142 Ok(LocalWheel {
1143 dist: Dist::Built(dist.clone()),
1144 archive: self
1145 .build_context
1146 .cache()
1147 .archive(&archive.id)
1148 .into_boxed_path(),
1149 hashes: archive.hashes,
1150 filename: filename.clone(),
1151 cache: CacheInfo::from_timestamp(modified),
1152 build: None,
1153 })
1154 } else {
1155 let file = fs_err::tokio::File::open(path)
1157 .await
1158 .map_err(Error::CacheRead)?;
1159 let extractor = WheelExtractor::new(
1160 self.build_context.cache().root(),
1161 self.content_addressed_cache,
1162 )
1163 .map_err(Error::CacheWrite)?;
1164
1165 let algorithms = hashes.algorithms();
1167 let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
1168 let mut hasher = uv_extract::hash::HashReader::new(file, &mut hashers);
1169
1170 let mut extracted = extractor
1172 .extract_streaming(&mut hasher)
1173 .await
1174 .map_err(|err| Error::Extract(filename.to_string(), err))?;
1175
1176 hasher.finish().await.map_err(Error::HashExhaustion)?;
1178
1179 let hashes = hashers.into_iter().map(HashDigest::from).collect();
1180
1181 extracted.validate_and_heal_record(dist)?;
1184
1185 let id = self
1187 .persist_extracted_wheel(extracted, wheel_entry.path())
1188 .await?;
1189
1190 let archive = Archive::new(id, hashes, filename.clone(), None);
1192
1193 let pointer = PathArchivePointer {
1195 timestamp: modified,
1196 archive: archive.clone(),
1197 };
1198 pointer.write_to(&pointer_entry).await?;
1199
1200 Ok(LocalWheel {
1201 dist: Dist::Built(dist.clone()),
1202 archive: self
1203 .build_context
1204 .cache()
1205 .archive(&archive.id)
1206 .into_boxed_path(),
1207 hashes: archive.hashes,
1208 filename: filename.clone(),
1209 cache: CacheInfo::from_timestamp(modified),
1210 build: None,
1211 })
1212 }
1213 }
1214
1215 async fn unzip_wheel(
1217 &self,
1218 path: &Path,
1219 target: &Path,
1220 dist: DistRef<'_>,
1221 ) -> Result<ArchiveId, Error> {
1222 let content_addressed_cache = self.content_addressed_cache;
1223
1224 let mut extracted = tokio::task::spawn_blocking({
1225 let path = path.to_owned();
1226 let root = self.build_context.cache().root().to_path_buf();
1227 move || -> Result<_, Error> {
1228 let extractor = WheelExtractor::new(&root, content_addressed_cache)
1230 .map_err(Error::CacheWrite)?;
1231 let reader = fs_err::File::open(&path).map_err(Error::CacheWrite)?;
1232 extractor
1233 .extract_seekable(reader)
1234 .map_err(|err| Error::Extract(path.to_string_lossy().into_owned(), err))
1235 }
1236 })
1237 .await??;
1238
1239 extracted.validate_and_heal_record(dist)?;
1241
1242 let id = self.persist_extracted_wheel(extracted, target).await?;
1244
1245 Ok(id)
1246 }
1247
1248 async fn persist_extracted_wheel(
1253 &self,
1254 extracted: ExtractedWheel,
1255 target: &Path,
1256 ) -> Result<ArchiveId, Error> {
1257 let (temp_dir, hashed_wheel) = extracted.into_parts();
1258 let cache = self.build_context.cache();
1259 let (temp_dir, id) = if let Some(HashedWheel { files, tree }) = hashed_wheel {
1260 let digest = DirectoryDigest::from(tree.hash());
1261 let id = ArchiveId::from_digest(digest.into());
1262 let cache = cache.clone();
1263 let temp_dir = tokio::task::spawn_blocking(move || {
1264 persist_archive_files(&cache, temp_dir.path(), &files)
1265 .map_err(Error::CacheWrite)?;
1266 Ok::<_, Error>(temp_dir)
1267 })
1268 .await??;
1269 (temp_dir, id)
1270 } else {
1271 (temp_dir, ArchiveId::default())
1272 };
1273
1274 cache
1275 .persist_with_id(temp_dir, target, id)
1276 .await
1277 .map_err(Error::CacheWrite)
1278 }
1279
1280 fn request(&self, url: DisplaySafeUrl) -> Result<reqwest::Request, reqwest::Error> {
1282 self.client
1283 .unmanaged
1284 .uncached_client(&url)
1285 .get(Url::from(url))
1286 .header(
1287 "accept-encoding",
1291 reqwest::header::HeaderValue::from_static("identity"),
1292 )
1293 .build()
1294 }
1295
1296 pub fn client(&self) -> &ManagedClient<'a> {
1298 &self.client
1299 }
1300}
1301
1302fn persist_archive_files(cache: &Cache, archive: &Path, files: &[HashedFile]) -> io::Result<()> {
1304 initialize_rayon_once();
1305 let targets = files
1306 .par_iter()
1307 .filter(|file| !file.path().ends_with("RECORD"))
1309 .map(|file| {
1310 let id = ArchiveFileId::from_digest(&file.object_digest_hex());
1311 (archive.join(file.path()), cache.archive_file(&id))
1312 })
1313 .collect::<Vec<_>>();
1314
1315 let mut shards: FxHashMap<&Path, Vec<_>> = FxHashMap::default();
1318 for (source, target) in &targets {
1319 let Some(parent) = target.parent() else {
1320 return Err(io::Error::new(
1321 io::ErrorKind::InvalidInput,
1322 "archive file path must have a parent directory",
1323 ));
1324 };
1325 shards.entry(parent).or_default().push((source, target));
1326 }
1327
1328 let mut shards = shards
1329 .into_iter()
1330 .map(|(parent, files)| (parent, files, Ok(())))
1331 .collect::<Vec<_>>();
1332 shards.sort_unstable_by_key(|(_, files, _)| Reverse(files.len()));
1334
1335 in_place_scope(|scope| -> io::Result<()> {
1338 for (parent, files, result) in &mut shards {
1339 fs_err::create_dir_all(parent)?;
1340 scope.spawn(move |_| {
1341 *result = files
1342 .iter()
1343 .try_for_each(|(source, target)| persist_archive_file(source, target));
1344 });
1345 }
1346 Ok(())
1347 })?;
1348
1349 shards.into_iter().try_for_each(|(_, _, result)| result)
1350}
1351
1352fn persist_archive_file(src: &Path, dst: &Path) -> io::Result<()> {
1354 match fs_err::hard_link(src, dst) {
1357 Ok(()) => return Ok(()),
1358 Err(_) if dst.try_exists()? => {}
1359 Err(_) => return uv_fs::copy_atomic_sync(src, dst),
1360 }
1361
1362 if let Err(err) = fs_err::remove_file(src)
1364 && err.kind() != io::ErrorKind::NotFound
1365 {
1366 return Err(err);
1367 }
1368
1369 fs_err::hard_link(dst, src).or_else(|_| uv_fs::copy_atomic_sync(dst, src))
1370}
1371
1372pub struct ManagedClient<'a> {
1374 pub unmanaged: &'a RegistryClient,
1375 control: Arc<Semaphore>,
1376}
1377
1378impl<'a> ManagedClient<'a> {
1379 fn new(client: &'a RegistryClient, control: Arc<Semaphore>) -> Self {
1381 ManagedClient {
1382 unmanaged: client,
1383 control,
1384 }
1385 }
1386
1387 pub async fn managed<F, T>(&self, f: impl FnOnce(&'a RegistryClient) -> F) -> T
1392 where
1393 F: Future<Output = T>,
1394 {
1395 let _permit = self.control.acquire().await.unwrap();
1396 f(self.unmanaged).await
1397 }
1398
1399 pub async fn manual<F, T>(&'a self, f: impl FnOnce(&'a RegistryClient, &'a Semaphore) -> F) -> T
1407 where
1408 F: Future<Output = T>,
1409 {
1410 f(self.unmanaged, &self.control).await
1411 }
1412}
1413
1414fn content_length(response: &reqwest::Response) -> Option<u64> {
1416 response
1417 .headers()
1418 .get(reqwest::header::CONTENT_LENGTH)
1419 .and_then(|val| val.to_str().ok())
1420 .and_then(|val| val.parse::<u64>().ok())
1421}
1422
1423struct ProgressReader<'a, R> {
1425 reader: R,
1426 index: usize,
1427 reporter: &'a dyn Reporter,
1428}
1429
1430impl<'a, R> ProgressReader<'a, R> {
1431 fn new(reader: R, index: usize, reporter: &'a dyn Reporter) -> Self {
1433 Self {
1434 reader,
1435 index,
1436 reporter,
1437 }
1438 }
1439}
1440
1441impl<R> AsyncRead for ProgressReader<'_, R>
1442where
1443 R: AsyncRead + Unpin,
1444{
1445 fn poll_read(
1446 mut self: Pin<&mut Self>,
1447 cx: &mut Context<'_>,
1448 buf: &mut ReadBuf<'_>,
1449 ) -> Poll<io::Result<()>> {
1450 Pin::new(&mut self.as_mut().reader)
1451 .poll_read(cx, buf)
1452 .map_ok(|()| {
1453 self.reporter
1454 .on_download_progress(self.index, buf.filled().len() as u64);
1455 })
1456 }
1457}
1458
1459#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1463pub struct HttpArchivePointer {
1464 archive: Archive,
1465}
1466
1467impl HttpArchivePointer {
1468 pub fn read_from(path: impl AsRef<Path>) -> Result<Option<Self>, Error> {
1470 match fs_err::File::open(path.as_ref()) {
1471 Ok(file) => {
1472 let data = DataWithCachePolicy::from_reader(file)?.data;
1473 let archive = rmp_serde::from_slice::<Archive>(&data)?;
1474 Ok(Some(Self { archive }))
1475 }
1476 Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None),
1477 Err(err) => Err(Error::CacheRead(err)),
1478 }
1479 }
1480
1481 pub fn into_archive(self) -> Archive {
1483 self.archive
1484 }
1485
1486 pub fn to_cache_info(&self) -> CacheInfo {
1488 CacheInfo::default()
1489 }
1490
1491 pub fn to_build_info(&self) -> Option<BuildInfo> {
1493 None
1494 }
1495}
1496
1497#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1501pub struct PathArchivePointer {
1502 timestamp: Timestamp,
1503 archive: Archive,
1504}
1505
1506impl PathArchivePointer {
1507 pub fn read_from(path: impl AsRef<Path>) -> Result<Option<Self>, Error> {
1509 match fs_err::read(path) {
1510 Ok(cached) => Ok(Some(rmp_serde::from_slice::<Self>(&cached)?)),
1511 Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None),
1512 Err(err) => Err(Error::CacheRead(err)),
1513 }
1514 }
1515
1516 async fn write_to(&self, entry: &CacheEntry) -> Result<(), Error> {
1518 write_atomic(entry.path(), rmp_serde::to_vec(&self)?)
1519 .await
1520 .map_err(Error::CacheWrite)
1521 }
1522
1523 pub fn is_up_to_date(&self, modified: Timestamp) -> bool {
1525 self.timestamp == modified
1526 }
1527
1528 pub fn into_archive(self) -> Archive {
1530 self.archive
1531 }
1532
1533 pub fn to_cache_info(&self) -> CacheInfo {
1535 CacheInfo::from_timestamp(self.timestamp)
1536 }
1537
1538 pub fn to_build_info(&self) -> Option<BuildInfo> {
1540 None
1541 }
1542}