1use std::fmt::Display;
2use std::future::Future;
3use std::io;
4use std::path::{Path, PathBuf};
5use std::pin::Pin;
6use std::sync::Arc;
7use std::task::{Context, Poll};
8
9use futures::{FutureExt, TryStreamExt};
10use tokio::io::{AsyncRead, AsyncSeekExt, ReadBuf};
11use tokio::sync::Semaphore;
12use tokio_util::compat::FuturesAsyncReadCompatExt;
13use tracing::{Instrument, info_span, instrument, warn};
14use url::Url;
15
16use uv_cache::{ArchiveId, CacheBucket, CacheEntry, WheelCache};
17use uv_cache_info::{CacheInfo, Timestamp};
18use uv_client::{
19 CacheControl, CachedClientError, Connectivity, DataWithCachePolicy, RegistryClient,
20};
21use uv_distribution_filename::WheelFilename;
22use uv_distribution_types::{
23 BuildInfo, BuildableSource, BuiltDist, Dist, DistRef, HashPolicy, Hashed, IndexUrl,
24 InstalledDist, Name, SourceDist,
25};
26use uv_extract::dirhash::{DirectoryDigest, DirhashTree, dirhash_path};
27use uv_extract::hash::Hasher;
28use uv_fs::{PortablePath, write_atomic};
29use uv_git::{GIT_LFS, GitError};
30use uv_install_wheel::validate_and_heal_record;
31use uv_platform_tags::Tags;
32use uv_preview::PreviewFeature;
33use uv_pypi_types::{HashDigest, HashDigests, PyProjectToml};
34use uv_python::PythonVariant;
35use uv_redacted::DisplaySafeUrl;
36use uv_types::{BuildContext, BuildStack};
37
38use crate::archive::Archive;
39use crate::error::PythonVersion;
40use crate::hash::http_hash_algorithms;
41use crate::metadata::{ArchiveMetadata, Metadata};
42use crate::source::SourceDistributionBuilder;
43use crate::{Error, LocalWheel, Reporter, RequiresDist};
44
45pub struct DistributionDatabase<'a, Context: BuildContext> {
58 build_context: &'a Context,
59 builder: SourceDistributionBuilder<'a, Context>,
60 client: ManagedClient<'a>,
61 reporter: Option<Arc<dyn Reporter>>,
62 content_addressed_cache: bool,
63}
64
65impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
66 pub fn new(
67 client: &'a RegistryClient,
68 build_context: &'a Context,
69 downloads_semaphore: Arc<Semaphore>,
70 ) -> Self {
71 let content_addressed_cache = uv_preview::is_enabled(PreviewFeature::ContentAddressedCache)
75 && !uv_extract::insecure_no_validate();
76 Self {
77 build_context,
78 builder: SourceDistributionBuilder::new(build_context),
79 client: ManagedClient::new(client, downloads_semaphore),
80 reporter: None,
81 content_addressed_cache,
82 }
83 }
84
85 #[must_use]
87 pub fn with_build_stack(self, build_stack: &'a BuildStack) -> Self {
88 Self {
89 builder: self.builder.with_build_stack(build_stack),
90 ..self
91 }
92 }
93
94 #[must_use]
96 pub fn with_reporter(self, reporter: Arc<dyn Reporter>) -> Self {
97 Self {
98 builder: self.builder.with_reporter(reporter.clone()),
99 reporter: Some(reporter),
100 ..self
101 }
102 }
103
104 fn handle_response_errors(&self, err: reqwest::Error) -> io::Error {
106 if err.is_timeout() {
107 io::Error::new(
109 io::ErrorKind::TimedOut,
110 format!(
111 "Failed to download distribution due to network timeout. Try increasing UV_HTTP_TIMEOUT (current value: {}s).",
112 self.client.unmanaged.read_timeout().as_secs()
113 ),
114 )
115 } else {
116 io::Error::other(err)
117 }
118 }
119
120 #[instrument(skip_all, fields(%dist))]
127 pub async fn get_or_build_wheel(
128 &self,
129 dist: &Dist,
130 tags: &Tags,
131 hashes: HashPolicy<'_>,
132 ) -> Result<LocalWheel, Error> {
133 match dist {
134 Dist::Built(built) => self.get_wheel(built, hashes).await,
135 Dist::Source(source) => self.build_wheel(source, tags, hashes).await,
136 }
137 }
138
139 #[instrument(skip_all, fields(%dist))]
145 pub async fn get_installed_metadata(
146 &self,
147 dist: &InstalledDist,
148 ) -> Result<ArchiveMetadata, Error> {
149 if let Some(metadata) = self
151 .build_context
152 .dependency_metadata()
153 .get(dist.name(), Some(dist.version()))
154 {
155 return Ok(ArchiveMetadata::from_metadata23(metadata));
156 }
157
158 let metadata = dist
159 .read_metadata()
160 .map_err(|err| Error::ReadInstalled(Box::new(dist.clone()), err))?;
161
162 Ok(ArchiveMetadata::from_metadata23(metadata.clone()))
163 }
164
165 #[instrument(skip_all, fields(%dist))]
171 pub async fn get_or_build_wheel_metadata(
172 &self,
173 dist: &Dist,
174 hashes: HashPolicy<'_>,
175 ) -> Result<ArchiveMetadata, Error> {
176 match dist {
177 Dist::Built(built) => self.get_wheel_metadata(built, hashes).await,
178 Dist::Source(source) => {
179 self.build_wheel_metadata(&BuildableSource::Dist(source), hashes)
180 .await
181 }
182 }
183 }
184
185 async fn get_wheel(
190 &self,
191 dist: &BuiltDist,
192 hashes: HashPolicy<'_>,
193 ) -> Result<LocalWheel, Error> {
194 match dist {
195 BuiltDist::Registry(wheels) => {
196 let wheel = wheels.best_wheel();
197 let url = wheel.file.url.to_url()?;
198 let size = wheel.file.size;
199
200 let wheel_entry = self.build_context.cache().entry(
202 CacheBucket::Wheels,
203 WheelCache::Index(&wheel.index).wheel_dir(wheel.name().as_ref()),
204 wheel.filename.cache_key(),
205 );
206
207 if url.scheme() == "file" {
209 let path = url
210 .to_file_path()
211 .map_err(|()| Error::NonFileUrl(url.clone()))?;
212 return self
213 .load_wheel(&path, &wheel.filename, wheel_entry, dist, hashes)
214 .await;
215 }
216
217 match self
219 .stream_wheel(
220 url.clone(),
221 dist.index(),
222 &wheel.filename,
223 size,
224 &wheel_entry,
225 dist,
226 hashes,
227 )
228 .await
229 {
230 Ok(archive) => Ok(LocalWheel {
231 dist: Dist::Built(dist.clone()),
232 archive: self
233 .build_context
234 .cache()
235 .archive(&archive.id)
236 .into_boxed_path(),
237 hashes: archive.hashes,
238 filename: wheel.filename.clone(),
239 cache: CacheInfo::default(),
240 build: None,
241 }),
242 Err(Error::Extract(name, err)) => {
243 if err.is_http_streaming_unsupported() {
244 warn!(
245 "Streaming unsupported for {dist}; downloading wheel to disk ({err})"
246 );
247 } else if err.is_http_streaming_failed() {
248 warn!("Streaming failed for {dist}; downloading wheel to disk ({err})");
249 } else {
250 return Err(Error::Extract(name, err));
251 }
252
253 let archive = self
256 .download_wheel(
257 url,
258 dist.index(),
259 &wheel.filename,
260 size,
261 &wheel_entry,
262 dist,
263 hashes,
264 )
265 .await?;
266
267 Ok(LocalWheel {
268 dist: Dist::Built(dist.clone()),
269 archive: self
270 .build_context
271 .cache()
272 .archive(&archive.id)
273 .into_boxed_path(),
274 hashes: archive.hashes,
275 filename: wheel.filename.clone(),
276 cache: CacheInfo::default(),
277 build: None,
278 })
279 }
280 Err(err) => Err(err),
281 }
282 }
283
284 BuiltDist::DirectUrl(wheel) => {
285 let wheel_entry = self.build_context.cache().entry(
287 CacheBucket::Wheels,
288 WheelCache::Url(&wheel.url).wheel_dir(wheel.name().as_ref()),
289 wheel.filename.cache_key(),
290 );
291
292 match self
294 .stream_wheel(
295 wheel.url.raw().clone(),
296 None,
297 &wheel.filename,
298 wheel.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 let archive = self
331 .download_wheel(
332 wheel.url.raw().clone(),
333 None,
334 &wheel.filename,
335 wheel.size,
336 &wheel_entry,
337 dist,
338 hashes,
339 )
340 .await?;
341 Ok(LocalWheel {
342 dist: Dist::Built(dist.clone()),
343 archive: self
344 .build_context
345 .cache()
346 .archive(&archive.id)
347 .into_boxed_path(),
348 hashes: archive.hashes,
349 filename: wheel.filename.clone(),
350 cache: CacheInfo::default(),
351 build: None,
352 })
353 }
354 Err(err) => Err(err),
355 }
356 }
357
358 BuiltDist::GitPath(wheel) => {
359 let fetch = self
361 .build_context
362 .git()
363 .fetch(
364 &wheel.git,
365 self.client.unmanaged.git_http_settings(wheel.git.url()),
366 self.build_context.cache().bucket(CacheBucket::Git),
367 self.reporter.clone().map(<dyn Reporter>::into_git_reporter),
368 )
369 .await?;
370
371 if wheel.git.lfs().enabled() && !fetch.lfs_ready() {
372 if GIT_LFS.is_err() {
373 return Err(Error::MissingWheelGitLfsArtifacts(
374 wheel.url.to_url(),
375 GitError::GitLfsNotFound,
376 ));
377 }
378 return Err(Error::MissingWheelGitLfsArtifacts(
379 wheel.url.to_url(),
380 GitError::GitLfsNotConfigured,
381 ));
382 }
383
384 let git_sha = fetch.git().precise().expect("Exact commit after checkout");
385 let cache_entry = self.build_context.cache().entry(
386 CacheBucket::Wheels,
387 WheelCache::Git(&wheel.url, git_sha.as_short_str()).root(),
388 wheel.filename.stem(),
389 );
390
391 let install_path = fetch.path().join(&wheel.install_path);
392
393 self.load_wheel(&install_path, &wheel.filename, cache_entry, dist, hashes)
394 .await
395 }
396
397 BuiltDist::Path(wheel) => {
398 let cache_entry = self.build_context.cache().entry(
399 CacheBucket::Wheels,
400 WheelCache::Url(&wheel.url).wheel_dir(wheel.name().as_ref()),
401 wheel.filename.cache_key(),
402 );
403
404 self.load_wheel(
405 &wheel.install_path,
406 &wheel.filename,
407 cache_entry,
408 dist,
409 hashes,
410 )
411 .await
412 }
413 }
414 }
415
416 async fn build_wheel(
422 &self,
423 dist: &SourceDist,
424 tags: &Tags,
425 hashes: HashPolicy<'_>,
426 ) -> Result<LocalWheel, Error> {
427 let built_wheel = self
428 .builder
429 .download_and_build(&BuildableSource::Dist(dist), tags, hashes, &self.client)
430 .boxed_local()
431 .await?;
432
433 if !built_wheel.filename.is_compatible(tags) {
439 return if tags.is_cross() {
440 Err(Error::BuiltWheelIncompatibleTargetPlatform {
441 filename: built_wheel.filename,
442 python_platform: tags.python_platform().clone(),
443 python_version: PythonVersion {
444 version: tags.python_version(),
445 variant: if tags.is_freethreaded() {
446 PythonVariant::Freethreaded
447 } else {
448 PythonVariant::Default
449 },
450 },
451 })
452 } else {
453 Err(Error::BuiltWheelIncompatibleHostPlatform {
454 filename: built_wheel.filename,
455 python_platform: tags.python_platform().clone(),
456 python_version: PythonVersion {
457 version: tags.python_version(),
458 variant: if tags.is_freethreaded() {
459 PythonVariant::Freethreaded
460 } else {
461 PythonVariant::Default
462 },
463 },
464 })
465 };
466 }
467
468 #[cfg(windows)]
470 let _lock = {
471 let lock_entry = CacheEntry::new(
472 built_wheel.target.parent().unwrap(),
473 format!(
474 "{}.lock",
475 built_wheel.target.file_name().unwrap().to_str().unwrap()
476 ),
477 );
478 lock_entry.lock().await.map_err(Error::CacheLock)?
479 };
480
481 match self.build_context.cache().resolve_link(&built_wheel.target) {
484 Ok(archive) => {
485 return Ok(LocalWheel {
486 dist: Dist::Source(dist.clone()),
487 archive: archive.into_boxed_path(),
488 filename: built_wheel.filename,
489 hashes: built_wheel.hashes,
490 cache: built_wheel.cache_info,
491 build: Some(built_wheel.build_info),
492 });
493 }
494 Err(err) if err.kind() == io::ErrorKind::NotFound => {}
495 Err(err) => return Err(Error::CacheRead(err)),
496 }
497
498 let id = self
500 .unzip_wheel(
501 &built_wheel.path,
502 &built_wheel.target,
503 DistRef::Source(dist),
504 )
505 .await?;
506
507 Ok(LocalWheel {
508 dist: Dist::Source(dist.clone()),
509 archive: self.build_context.cache().archive(&id).into_boxed_path(),
510 hashes: built_wheel.hashes,
511 filename: built_wheel.filename,
512 cache: built_wheel.cache_info,
513 build: Some(built_wheel.build_info),
514 })
515 }
516
517 async fn get_wheel_metadata(
522 &self,
523 dist: &BuiltDist,
524 hashes: HashPolicy<'_>,
525 ) -> Result<ArchiveMetadata, Error> {
526 if hashes.is_generate(dist) {
541 let wheel = self.get_wheel(dist, hashes).await?;
542 let metadata = if let Some(metadata) = self
544 .build_context
545 .dependency_metadata()
546 .get(dist.name(), Some(dist.version()))
547 {
548 metadata
549 } else {
550 wheel.metadata()?
551 };
552 let hashes = wheel.hashes;
553 return Ok(ArchiveMetadata {
554 metadata: Metadata::from_metadata23(metadata),
555 hashes,
556 });
557 }
558
559 if let Some(metadata) = self
561 .build_context
562 .dependency_metadata()
563 .get(dist.name(), Some(dist.version()))
564 {
565 return Ok(ArchiveMetadata::from_metadata23(metadata));
566 }
567
568 let result = self
569 .client
570 .managed(|client| {
571 client
572 .wheel_metadata(
573 dist,
574 self.build_context.git(),
575 self.build_context.capabilities(),
576 self.reporter.clone().map(<dyn Reporter>::into_git_reporter),
577 )
578 .boxed_local()
579 })
580 .await;
581
582 match result {
583 Ok(metadata) => {
584 Ok(ArchiveMetadata::from_metadata23(metadata))
586 }
587 Err(err) if err.is_http_streaming_unsupported() => {
588 warn!(
589 "Streaming unsupported when fetching metadata for {dist}; downloading wheel directly ({err})"
590 );
591
592 let wheel = self.get_wheel(dist, hashes).await?;
595 let metadata = wheel.metadata()?;
596 let hashes = wheel.hashes;
597 Ok(ArchiveMetadata {
598 metadata: Metadata::from_metadata23(metadata),
599 hashes,
600 })
601 }
602 Err(err) => Err(err.into()),
603 }
604 }
605
606 pub async fn build_wheel_metadata(
611 &self,
612 source: &BuildableSource<'_>,
613 hashes: HashPolicy<'_>,
614 ) -> Result<ArchiveMetadata, Error> {
615 if let Some(dist) = source.as_dist() {
617 if let Some(metadata) = self
618 .build_context
619 .dependency_metadata()
620 .get(dist.name(), dist.version())
621 {
622 self.builder.resolve_revision(source, &self.client).await?;
625
626 return Ok(ArchiveMetadata::from_metadata23(metadata));
627 }
628 }
629
630 let metadata = self
631 .builder
632 .download_and_build_metadata(source, hashes, &self.client)
633 .boxed_local()
634 .await?;
635
636 Ok(metadata)
637 }
638
639 pub async fn requires_dist(
641 &self,
642 path: &Path,
643 pyproject_toml: &PyProjectToml,
644 ) -> Result<Option<RequiresDist>, Error> {
645 self.builder
646 .source_tree_requires_dist(
647 path,
648 pyproject_toml,
649 self.client.unmanaged.credentials_cache(),
650 )
651 .await
652 }
653
654 async fn stream_wheel(
656 &self,
657 url: DisplaySafeUrl,
658 index: Option<&IndexUrl>,
659 filename: &WheelFilename,
660 size: Option<u64>,
661 wheel_entry: &CacheEntry,
662 dist: &BuiltDist,
663 hashes: HashPolicy<'_>,
664 ) -> Result<Archive, Error> {
665 let expected_size = match dist {
666 BuiltDist::Registry(dist) if dist.best_wheel().size_is_authoritative => size,
667 BuiltDist::DirectUrl(_) => size,
668 _ => None,
669 };
670
671 #[cfg(windows)]
673 let _lock = {
674 let lock_entry = wheel_entry.with_file(format!("{}.lock", filename.stem()));
675 lock_entry.lock().await.map_err(Error::CacheLock)?
676 };
677
678 let http_entry = wheel_entry.with_file(format!("{}.http", filename.cache_key()));
680
681 let download = |response: reqwest::Response| {
682 async {
683 let progress_size = size.or_else(|| content_length(&response));
684
685 let progress = self.reporter.as_ref().map(|reporter| {
686 (
687 reporter,
688 reporter.on_download_start(dist.name(), progress_size),
689 )
690 });
691
692 let reader = response
693 .bytes_stream()
694 .map_err(|err| self.handle_response_errors(err))
695 .into_async_read();
696
697 let algorithms = http_hash_algorithms(hashes);
699 let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
700 let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers);
701
702 let temp_dir = tempfile::tempdir_in(self.build_context.cache().root())
704 .map_err(Error::CacheWrite)?;
705
706 let mut extracted = match progress {
707 Some((reporter, progress)) => {
708 let mut reader = ProgressReader::new(&mut hasher, progress, &**reporter);
709 ExtractedWheelManifest::extract_streaming(
710 &mut reader,
711 temp_dir.path(),
712 self.content_addressed_cache,
713 )
714 .await
715 .map_err(|err| Error::Extract(filename.to_string(), err))?
716 }
717 None => ExtractedWheelManifest::extract_streaming(
718 &mut hasher,
719 temp_dir.path(),
720 self.content_addressed_cache,
721 )
722 .await
723 .map_err(|err| Error::Extract(filename.to_string(), err))?,
724 };
725 hasher.finish().await.map_err(Error::HashExhaustion)?;
727 let actual_size = hasher.bytes_read();
728 if let Some(expected) = expected_size
729 && actual_size != expected
730 {
731 return Err(Error::MismatchedSize {
732 distribution: dist.to_string(),
733 expected,
734 actual: actual_size,
735 });
736 }
737
738 extracted.validate_and_heal_record(temp_dir.path(), dist)?;
741
742 let id = self
744 .persist_extracted_wheel(temp_dir, wheel_entry.path(), extracted.tree)
745 .await?;
746
747 if let Some((reporter, progress)) = progress {
748 reporter.on_download_complete(dist.name(), progress);
749 }
750
751 Ok(Archive::new(
752 id,
753 hashers.into_iter().map(HashDigest::from).collect(),
754 filename.clone(),
755 Some(actual_size),
756 ))
757 }
758 .instrument(info_span!("wheel", wheel = %dist))
759 };
760
761 let req = self.request(url.clone())?;
763
764 let cache_control = match self.client.unmanaged.connectivity() {
766 Connectivity::Online
767 if let Some(header) = index.and_then(|index| {
768 self.build_context
769 .locations()
770 .artifact_cache_control_for(index)
771 }) =>
772 {
773 CacheControl::Override(header)
774 }
775 Connectivity::Online => CacheControl::from(
776 self.build_context
777 .cache()
778 .freshness(&http_entry, Some(&filename.name), None)
779 .map_err(Error::CacheRead)?,
780 ),
781 Connectivity::Offline => CacheControl::AllowStale,
782 };
783
784 let archive = self
785 .client
786 .managed(|client| {
787 client.cached_client().get_serde_with_retry(
788 req,
789 &http_entry,
790 cache_control.clone(),
791 download,
792 )
793 })
794 .await
795 .map_err(|err| match err {
796 CachedClientError::Callback { err, .. } => err,
797 CachedClientError::Client(err) => Error::Client(err),
798 })?;
799
800 if let (Some(expected), Some(actual)) = (expected_size, archive.size)
801 && expected != actual
802 {
803 return Err(Error::MismatchedSize {
804 distribution: dist.to_string(),
805 expected,
806 actual,
807 });
808 }
809
810 let archive = Some(archive)
812 .filter(|archive| archive.has_digests(hashes))
813 .filter(|archive| archive.exists(self.build_context.cache()))
814 .filter(|archive| expected_size.is_none() || archive.size.is_some());
815
816 let archive = if let Some(archive) = archive {
817 archive
818 } else {
819 self.client
820 .managed(async |client| {
821 client
822 .cached_client()
823 .skip_cache_with_retry(
824 self.request(url)?,
825 &http_entry,
826 cache_control,
827 download,
828 )
829 .await
830 .map_err(|err| match err {
831 CachedClientError::Callback { err, .. } => err,
832 CachedClientError::Client(err) => Error::Client(err),
833 })
834 })
835 .await?
836 };
837
838 Ok(archive)
839 }
840
841 async fn download_wheel(
843 &self,
844 url: DisplaySafeUrl,
845 index: Option<&IndexUrl>,
846 filename: &WheelFilename,
847 size: Option<u64>,
848 wheel_entry: &CacheEntry,
849 dist: &BuiltDist,
850 hashes: HashPolicy<'_>,
851 ) -> Result<Archive, Error> {
852 let expected_size = match dist {
853 BuiltDist::Registry(dist) if dist.best_wheel().size_is_authoritative => size,
854 BuiltDist::DirectUrl(_) => size,
855 _ => None,
856 };
857
858 let content_addressed_cache = self.content_addressed_cache;
859
860 #[cfg(windows)]
862 let _lock = {
863 let lock_entry = wheel_entry.with_file(format!("{}.lock", filename.stem()));
864 lock_entry.lock().await.map_err(Error::CacheLock)?
865 };
866
867 let http_entry = wheel_entry.with_file(format!("{}.http", filename.cache_key()));
869
870 let download = |response: reqwest::Response| {
871 async {
872 let progress_size = size.or_else(|| content_length(&response));
873
874 let progress = self.reporter.as_ref().map(|reporter| {
875 (
876 reporter,
877 reporter.on_download_start(dist.name(), progress_size),
878 )
879 });
880
881 let reader = response
882 .bytes_stream()
883 .map_err(|err| self.handle_response_errors(err))
884 .into_async_read();
885 let algorithms = http_hash_algorithms(hashes);
886 let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
887 let mut hasher = uv_extract::hash::HashReader::new(reader.compat(), &mut hashers);
888
889 let temp_file = tempfile::tempfile_in(self.build_context.cache().root())
891 .map_err(Error::CacheWrite)?;
892 let mut writer = tokio::io::BufWriter::new(fs_err::tokio::File::from_std(
893 fs_err::File::from_parts(temp_file, self.build_context.cache().root()),
895 ));
896
897 match progress {
898 Some((reporter, progress)) => {
899 let mut reader = ProgressReader::new(&mut hasher, progress, &**reporter);
902
903 tokio::io::copy(&mut reader, &mut writer)
904 .await
905 .map_err(Error::CacheWrite)?;
906 }
907 None => {
908 tokio::io::copy(&mut hasher, &mut writer)
909 .await
910 .map_err(Error::CacheWrite)?;
911 }
912 }
913
914 if let Some(expected) = expected_size
915 && hasher.bytes_read() != expected
916 {
917 return Err(Error::MismatchedSize {
918 distribution: dist.to_string(),
919 expected,
920 actual: hasher.bytes_read(),
921 });
922 }
923
924 let actual_size = hasher.bytes_read();
925
926 let temp_dir = tempfile::tempdir_in(self.build_context.cache().root())
928 .map_err(Error::CacheWrite)?;
929 let mut file = writer.into_inner();
930 file.seek(io::SeekFrom::Start(0))
931 .await
932 .map_err(Error::CacheWrite)?;
933
934 let target = temp_dir.path().to_owned();
935 let file = file.into_std().await;
936 let mut extracted = tokio::task::spawn_blocking(move || {
937 ExtractedWheelManifest::extract_seekable(file, &target, content_addressed_cache)
938 })
939 .await?
940 .map_err(|err| Error::Extract(filename.to_string(), err))?;
941 let hashes = hashers.into_iter().map(HashDigest::from).collect();
942
943 extracted.validate_and_heal_record(temp_dir.path(), dist)?;
946
947 let id = self
949 .persist_extracted_wheel(temp_dir, wheel_entry.path(), extracted.tree)
950 .await?;
951
952 if let Some((reporter, progress)) = progress {
953 reporter.on_download_complete(dist.name(), progress);
954 }
955
956 Ok(Archive::new(
957 id,
958 hashes,
959 filename.clone(),
960 Some(actual_size),
961 ))
962 }
963 .instrument(info_span!("wheel", wheel = %dist))
964 };
965
966 let req = self.request(url.clone())?;
968
969 let cache_control = match self.client.unmanaged.connectivity() {
971 Connectivity::Online
972 if let Some(header) = index.and_then(|index| {
973 self.build_context
974 .locations()
975 .artifact_cache_control_for(index)
976 }) =>
977 {
978 CacheControl::Override(header)
979 }
980 Connectivity::Online => CacheControl::from(
981 self.build_context
982 .cache()
983 .freshness(&http_entry, Some(&filename.name), None)
984 .map_err(Error::CacheRead)?,
985 ),
986 Connectivity::Offline => CacheControl::AllowStale,
987 };
988
989 let archive = self
990 .client
991 .managed(|client| {
992 client.cached_client().get_serde_with_retry(
993 req,
994 &http_entry,
995 cache_control.clone(),
996 download,
997 )
998 })
999 .await
1000 .map_err(|err| match err {
1001 CachedClientError::Callback { err, .. } => err,
1002 CachedClientError::Client(err) => Error::Client(err),
1003 })?;
1004
1005 if let (Some(expected), Some(actual)) = (expected_size, archive.size)
1006 && expected != actual
1007 {
1008 return Err(Error::MismatchedSize {
1009 distribution: dist.to_string(),
1010 expected,
1011 actual,
1012 });
1013 }
1014
1015 let archive = Some(archive)
1017 .filter(|archive| archive.has_digests(hashes))
1018 .filter(|archive| archive.exists(self.build_context.cache()))
1019 .filter(|archive| expected_size.is_none() || archive.size.is_some());
1020
1021 let archive = if let Some(archive) = archive {
1022 archive
1023 } else {
1024 self.client
1025 .managed(async |client| {
1026 client
1027 .cached_client()
1028 .skip_cache_with_retry(
1029 self.request(url)?,
1030 &http_entry,
1031 cache_control,
1032 download,
1033 )
1034 .await
1035 .map_err(|err| match err {
1036 CachedClientError::Callback { err, .. } => err,
1037 CachedClientError::Client(err) => Error::Client(err),
1038 })
1039 })
1040 .await?
1041 };
1042
1043 Ok(archive)
1044 }
1045
1046 async fn load_wheel(
1048 &self,
1049 path: &Path,
1050 filename: &WheelFilename,
1051 wheel_entry: CacheEntry,
1052 dist: &BuiltDist,
1053 hashes: HashPolicy<'_>,
1054 ) -> Result<LocalWheel, Error> {
1055 #[cfg(windows)]
1056 let _lock = {
1057 let lock_entry = wheel_entry.with_file(format!("{}.lock", filename.stem()));
1058 lock_entry.lock().await.map_err(Error::CacheLock)?
1059 };
1060
1061 let modified = Timestamp::from_path(path).map_err(Error::CacheRead)?;
1063
1064 let pointer_entry = wheel_entry.with_file(format!("{}.rev", filename.cache_key()));
1066 let pointer = PathArchivePointer::read_from(&pointer_entry)?;
1067
1068 let archive = pointer
1070 .filter(|pointer| pointer.is_up_to_date(modified))
1071 .map(PathArchivePointer::into_archive)
1072 .filter(|archive| archive.has_digests(hashes));
1073
1074 if let Some(archive) = archive {
1076 Ok(LocalWheel {
1077 dist: Dist::Built(dist.clone()),
1078 archive: self
1079 .build_context
1080 .cache()
1081 .archive(&archive.id)
1082 .into_boxed_path(),
1083 hashes: archive.hashes,
1084 filename: filename.clone(),
1085 cache: CacheInfo::from_timestamp(modified),
1086 build: None,
1087 })
1088 } else if hashes.is_none() {
1089 let archive = Archive::new(
1091 self.unzip_wheel(path, wheel_entry.path(), DistRef::Built(dist))
1092 .await?,
1093 HashDigests::empty(),
1094 filename.clone(),
1095 None,
1096 );
1097
1098 let pointer = PathArchivePointer {
1100 timestamp: modified,
1101 archive: archive.clone(),
1102 };
1103 pointer.write_to(&pointer_entry).await?;
1104
1105 Ok(LocalWheel {
1106 dist: Dist::Built(dist.clone()),
1107 archive: self
1108 .build_context
1109 .cache()
1110 .archive(&archive.id)
1111 .into_boxed_path(),
1112 hashes: archive.hashes,
1113 filename: filename.clone(),
1114 cache: CacheInfo::from_timestamp(modified),
1115 build: None,
1116 })
1117 } else {
1118 let file = fs_err::tokio::File::open(path)
1120 .await
1121 .map_err(Error::CacheRead)?;
1122 let temp_dir = tempfile::tempdir_in(self.build_context.cache().root())
1123 .map_err(Error::CacheWrite)?;
1124
1125 let algorithms = hashes.algorithms();
1127 let mut hashers = algorithms.into_iter().map(Hasher::from).collect::<Vec<_>>();
1128 let mut hasher = uv_extract::hash::HashReader::new(file, &mut hashers);
1129
1130 let mut extracted = ExtractedWheelManifest::extract_streaming(
1132 &mut hasher,
1133 temp_dir.path(),
1134 self.content_addressed_cache,
1135 )
1136 .await
1137 .map_err(|err| Error::Extract(filename.to_string(), err))?;
1138
1139 hasher.finish().await.map_err(Error::HashExhaustion)?;
1141
1142 let hashes = hashers.into_iter().map(HashDigest::from).collect();
1143
1144 extracted.validate_and_heal_record(temp_dir.path(), dist)?;
1147
1148 let id = self
1150 .persist_extracted_wheel(temp_dir, wheel_entry.path(), extracted.tree)
1151 .await?;
1152
1153 let archive = Archive::new(id, hashes, filename.clone(), None);
1155
1156 let pointer = PathArchivePointer {
1158 timestamp: modified,
1159 archive: archive.clone(),
1160 };
1161 pointer.write_to(&pointer_entry).await?;
1162
1163 Ok(LocalWheel {
1164 dist: Dist::Built(dist.clone()),
1165 archive: self
1166 .build_context
1167 .cache()
1168 .archive(&archive.id)
1169 .into_boxed_path(),
1170 hashes: archive.hashes,
1171 filename: filename.clone(),
1172 cache: CacheInfo::from_timestamp(modified),
1173 build: None,
1174 })
1175 }
1176 }
1177
1178 async fn unzip_wheel(
1180 &self,
1181 path: &Path,
1182 target: &Path,
1183 dist: DistRef<'_>,
1184 ) -> Result<ArchiveId, Error> {
1185 let content_addressed_cache = self.content_addressed_cache;
1186
1187 let (temp_dir, mut extracted) = tokio::task::spawn_blocking({
1188 let path = path.to_owned();
1189 let root = self.build_context.cache().root().to_path_buf();
1190 move || -> Result<_, Error> {
1191 let temp_dir = tempfile::tempdir_in(root).map_err(Error::CacheWrite)?;
1193 let reader = fs_err::File::open(&path).map_err(Error::CacheWrite)?;
1194 let extracted = ExtractedWheelManifest::extract_seekable(
1195 reader,
1196 temp_dir.path(),
1197 content_addressed_cache,
1198 )
1199 .map_err(|err| Error::Extract(path.to_string_lossy().into_owned(), err))?;
1200 Ok((temp_dir, extracted))
1201 }
1202 })
1203 .await??;
1204
1205 extracted.validate_and_heal_record(temp_dir.path(), dist)?;
1207
1208 let id = self
1210 .persist_extracted_wheel(temp_dir, target, extracted.tree)
1211 .await?;
1212
1213 Ok(id)
1214 }
1215
1216 async fn persist_extracted_wheel(
1221 &self,
1222 temp_dir: tempfile::TempDir,
1223 target: &Path,
1224 tree: Option<DirhashTree>,
1225 ) -> Result<ArchiveId, Error> {
1226 let cache = self.build_context.cache();
1227 match tree {
1228 Some(tree) => {
1229 let digest = DirectoryDigest::from(tree.hash());
1230 let id = ArchiveId::from_digest(digest.into());
1231 cache.persist_with_id(temp_dir, target, id).await
1232 }
1233 None => cache.persist(temp_dir.keep(), target).await,
1234 }
1235 .map_err(Error::CacheWrite)
1236 }
1237
1238 fn request(&self, url: DisplaySafeUrl) -> Result<reqwest::Request, reqwest::Error> {
1240 self.client
1241 .unmanaged
1242 .uncached_client(&url)
1243 .get(Url::from(url))
1244 .header(
1245 "accept-encoding",
1249 reqwest::header::HeaderValue::from_static("identity"),
1250 )
1251 .build()
1252 }
1253
1254 pub fn client(&self) -> &ManagedClient<'a> {
1256 &self.client
1257 }
1258}
1259
1260struct ExtractedWheelManifest {
1262 files: Vec<(PathBuf, u64)>,
1263 tree: Option<DirhashTree>,
1264}
1265
1266impl ExtractedWheelManifest {
1267 async fn extract_streaming<R>(
1269 reader: R,
1270 target: &Path,
1271 content_addressed: bool,
1272 ) -> Result<Self, uv_extract::Error>
1273 where
1274 R: AsyncRead + Unpin,
1275 {
1276 if content_addressed {
1277 let (files, tree) = uv_extract::stream::unzip_and_hash(reader, target).await?;
1278 Ok(Self {
1279 files,
1280 tree: Some(tree),
1281 })
1282 } else {
1283 let files = uv_extract::stream::unzip(reader, target).await?;
1284 Ok(Self::without_tree(files))
1285 }
1286 }
1287
1288 fn extract_seekable(
1290 reader: fs_err::File,
1291 target: &Path,
1292 content_addressed: bool,
1293 ) -> Result<Self, uv_extract::Error> {
1294 if content_addressed {
1295 let (files, tree) = uv_extract::unzip_and_hash(reader, target)?;
1296 Ok(Self {
1297 files,
1298 tree: Some(tree),
1299 })
1300 } else {
1301 let files = uv_extract::unzip(reader, target)?;
1302 Ok(Self::without_tree(files))
1303 }
1304 }
1305
1306 fn without_tree(files: Vec<(PathBuf, u64)>) -> Self {
1307 Self { files, tree: None }
1308 }
1309
1310 fn validate_and_heal_record(&mut self, root: &Path, dist: impl Display) -> Result<(), Error> {
1312 let Some(record_path) = validate_and_heal_record(root, self.files.iter(), dist)
1313 .map_err(Error::InstallWheelError)?
1314 else {
1315 return Ok(());
1316 };
1317 let Some(tree) = self.tree.as_mut() else {
1318 return Ok(());
1319 };
1320
1321 let hash = dirhash_path(&root.join(&record_path)).map_err(|err| {
1322 Error::Extract(
1323 record_path.display().to_string(),
1324 uv_extract::Error::from(err),
1325 )
1326 })?;
1327 let record_path = PortablePath::from(record_path.as_path()).to_string();
1328 tree.update_file(&record_path, hash)
1329 .map_err(|err| Error::Extract(record_path, uv_extract::Error::from(err)))
1330 }
1331}
1332
1333pub struct ManagedClient<'a> {
1335 pub unmanaged: &'a RegistryClient,
1336 control: Arc<Semaphore>,
1337}
1338
1339impl<'a> ManagedClient<'a> {
1340 fn new(client: &'a RegistryClient, control: Arc<Semaphore>) -> Self {
1342 ManagedClient {
1343 unmanaged: client,
1344 control,
1345 }
1346 }
1347
1348 pub async fn managed<F, T>(&self, f: impl FnOnce(&'a RegistryClient) -> F) -> T
1353 where
1354 F: Future<Output = T>,
1355 {
1356 let _permit = self.control.acquire().await.unwrap();
1357 f(self.unmanaged).await
1358 }
1359
1360 pub async fn manual<F, T>(&'a self, f: impl FnOnce(&'a RegistryClient, &'a Semaphore) -> F) -> T
1368 where
1369 F: Future<Output = T>,
1370 {
1371 f(self.unmanaged, &self.control).await
1372 }
1373}
1374
1375fn content_length(response: &reqwest::Response) -> Option<u64> {
1377 response
1378 .headers()
1379 .get(reqwest::header::CONTENT_LENGTH)
1380 .and_then(|val| val.to_str().ok())
1381 .and_then(|val| val.parse::<u64>().ok())
1382}
1383
1384struct ProgressReader<'a, R> {
1386 reader: R,
1387 index: usize,
1388 reporter: &'a dyn Reporter,
1389}
1390
1391impl<'a, R> ProgressReader<'a, R> {
1392 fn new(reader: R, index: usize, reporter: &'a dyn Reporter) -> Self {
1394 Self {
1395 reader,
1396 index,
1397 reporter,
1398 }
1399 }
1400}
1401
1402impl<R> AsyncRead for ProgressReader<'_, R>
1403where
1404 R: AsyncRead + Unpin,
1405{
1406 fn poll_read(
1407 mut self: Pin<&mut Self>,
1408 cx: &mut Context<'_>,
1409 buf: &mut ReadBuf<'_>,
1410 ) -> Poll<io::Result<()>> {
1411 Pin::new(&mut self.as_mut().reader)
1412 .poll_read(cx, buf)
1413 .map_ok(|()| {
1414 self.reporter
1415 .on_download_progress(self.index, buf.filled().len() as u64);
1416 })
1417 }
1418}
1419
1420#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1424pub struct HttpArchivePointer {
1425 archive: Archive,
1426}
1427
1428impl HttpArchivePointer {
1429 pub fn read_from(path: impl AsRef<Path>) -> Result<Option<Self>, Error> {
1431 match fs_err::File::open(path.as_ref()) {
1432 Ok(file) => {
1433 let data = DataWithCachePolicy::from_reader(file)?.data;
1434 let archive = rmp_serde::from_slice::<Archive>(&data)?;
1435 Ok(Some(Self { archive }))
1436 }
1437 Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None),
1438 Err(err) => Err(Error::CacheRead(err)),
1439 }
1440 }
1441
1442 pub fn into_archive(self) -> Archive {
1444 self.archive
1445 }
1446
1447 pub fn to_cache_info(&self) -> CacheInfo {
1449 CacheInfo::default()
1450 }
1451
1452 pub fn to_build_info(&self) -> Option<BuildInfo> {
1454 None
1455 }
1456}
1457
1458#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1462pub struct PathArchivePointer {
1463 timestamp: Timestamp,
1464 archive: Archive,
1465}
1466
1467impl PathArchivePointer {
1468 pub fn read_from(path: impl AsRef<Path>) -> Result<Option<Self>, Error> {
1470 match fs_err::read(path) {
1471 Ok(cached) => Ok(Some(rmp_serde::from_slice::<Self>(&cached)?)),
1472 Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(None),
1473 Err(err) => Err(Error::CacheRead(err)),
1474 }
1475 }
1476
1477 async fn write_to(&self, entry: &CacheEntry) -> Result<(), Error> {
1479 write_atomic(entry.path(), rmp_serde::to_vec(&self)?)
1480 .await
1481 .map_err(Error::CacheWrite)
1482 }
1483
1484 pub fn is_up_to_date(&self, modified: Timestamp) -> bool {
1486 self.timestamp == modified
1487 }
1488
1489 pub fn into_archive(self) -> Archive {
1491 self.archive
1492 }
1493
1494 pub fn to_cache_info(&self) -> CacheInfo {
1496 CacheInfo::from_timestamp(self.timestamp)
1497 }
1498
1499 pub fn to_build_info(&self) -> Option<BuildInfo> {
1501 None
1502 }
1503}