Skip to content

Commit d0efe1e

Browse files
charliermarshcharliermarsh
authored
Apply Cache-Control overrides to response, not request headers (astral-sh#14736)
## Summary This was just an oversight on my part in the initial implementation. Closes astral-sh#14719. ## Test Plan With: ```toml [project] name = "foo" version = "0.1.0" description = "Add your description here" readme = "README.md" requires-python = ">=3.13.2" dependencies = [ ] [[tool.uv.index]] url = "https://download.pytorch.org/whl/cpu" cache-control = { api = "max-age=600" } ``` Ran `cargo run lock -vvv` and verified that the PyTorch index response was cached (whereas it typically returns `cache-control: no-cache,no-store,must-revalidate`).
1 parent 574aa1e commit d0efe1e

File tree

4 files changed

+201
-50
lines changed

4 files changed

+201
-50
lines changed

‎crates/uv-client/src/cached_client.rs‎

Lines changed: 60 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -304,7 +304,7 @@ impl CachedClient {
304304
.await?
305305
} else {
306306
debug!("No cache entry for: {}", req.url());
307-
let (response, cache_policy) = self.fresh_request(req).await?;
307+
let (response, cache_policy) = self.fresh_request(req, cache_control).await?;
308308
CachedResponse::ModifiedOrNew {
309309
response,
310310
cache_policy,
@@ -318,8 +318,13 @@ impl CachedClient {
318318
"Broken fresh cache entry (for payload) at {}, removing: {err}",
319319
cache_entry.path().display()
320320
);
321-
self.resend_and_heal_cache(fresh_req, cache_entry, response_callback)
322-
.await
321+
self.resend_and_heal_cache(
322+
fresh_req,
323+
cache_entry,
324+
cache_control,
325+
response_callback,
326+
)
327+
.await
323328
}
324329
},
325330
CachedResponse::NotModified { cached, new_policy } => {
@@ -339,8 +344,13 @@ impl CachedClient {
339344
(for payload) at {}, removing: {err}",
340345
cache_entry.path().display()
341346
);
342-
self.resend_and_heal_cache(fresh_req, cache_entry, response_callback)
343-
.await
347+
self.resend_and_heal_cache(
348+
fresh_req,
349+
cache_entry,
350+
cache_control,
351+
response_callback,
352+
)
353+
.await
344354
}
345355
}
346356
}
@@ -355,8 +365,13 @@ impl CachedClient {
355365
// ETag didn't match). We need to make a fresh request.
356366
if response.status() == http::StatusCode::NOT_MODIFIED {
357367
warn!("Server returned unusable 304 for: {}", fresh_req.url());
358-
self.resend_and_heal_cache(fresh_req, cache_entry, response_callback)
359-
.await
368+
self.resend_and_heal_cache(
369+
fresh_req,
370+
cache_entry,
371+
cache_control,
372+
response_callback,
373+
)
374+
.await
360375
} else {
361376
self.run_response_callback(
362377
cache_entry,
@@ -379,9 +394,10 @@ impl CachedClient {
379394
&self,
380395
req: Request,
381396
cache_entry: &CacheEntry,
397+
cache_control: CacheControl<'_>,
382398
response_callback: Callback,
383399
) -> Result<Payload, CachedClientError<CallBackError>> {
384-
let (response, cache_policy) = self.fresh_request(req).await?;
400+
let (response, cache_policy) = self.fresh_request(req, cache_control).await?;
385401

386402
let payload = self
387403
.run_response_callback(cache_entry, cache_policy, response, async |resp| {
@@ -401,10 +417,11 @@ impl CachedClient {
401417
&self,
402418
req: Request,
403419
cache_entry: &CacheEntry,
420+
cache_control: CacheControl<'_>,
404421
response_callback: Callback,
405422
) -> Result<Payload::Target, CachedClientError<CallBackError>> {
406423
let _ = fs_err::tokio::remove_file(&cache_entry.path()).await;
407-
let (response, cache_policy) = self.fresh_request(req).await?;
424+
let (response, cache_policy) = self.fresh_request(req, cache_control).await?;
408425
self.run_response_callback(cache_entry, cache_policy, response, response_callback)
409426
.await
410427
}
@@ -476,20 +493,13 @@ impl CachedClient {
476493
) -> Result<CachedResponse, Error> {
477494
// Apply the cache control header, if necessary.
478495
match cache_control {
479-
CacheControl::None | CacheControl::AllowStale => {}
496+
CacheControl::None | CacheControl::AllowStale | CacheControl::Override(..) => {}
480497
CacheControl::MustRevalidate => {
481498
req.headers_mut().insert(
482499
http::header::CACHE_CONTROL,
483500
http::HeaderValue::from_static("no-cache"),
484501
);
485502
}
486-
CacheControl::Override(value) => {
487-
req.headers_mut().insert(
488-
http::header::CACHE_CONTROL,
489-
http::HeaderValue::from_str(value)
490-
.map_err(|_| ErrorKind::InvalidCacheControl(value.to_string()))?,
491-
);
492-
}
493503
}
494504
Ok(match cached.cache_policy.before_request(&mut req) {
495505
BeforeRequest::Fresh => {
@@ -499,8 +509,13 @@ impl CachedClient {
499509
BeforeRequest::Stale(new_cache_policy_builder) => match cache_control {
500510
CacheControl::None | CacheControl::MustRevalidate | CacheControl::Override(_) => {
501511
debug!("Found stale response for: {}", req.url());
502-
self.send_cached_handle_stale(req, cached, new_cache_policy_builder)
503-
.await?
512+
self.send_cached_handle_stale(
513+
req,
514+
cache_control,
515+
cached,
516+
new_cache_policy_builder,
517+
)
518+
.await?
504519
}
505520
CacheControl::AllowStale => {
506521
debug!("Found stale (but allowed) response for: {}", req.url());
@@ -513,7 +528,7 @@ impl CachedClient {
513528
"Cached request doesn't match current request for: {}",
514529
req.url()
515530
);
516-
let (response, cache_policy) = self.fresh_request(req).await?;
531+
let (response, cache_policy) = self.fresh_request(req, cache_control).await?;
517532
CachedResponse::ModifiedOrNew {
518533
response,
519534
cache_policy,
@@ -525,19 +540,30 @@ impl CachedClient {
525540
async fn send_cached_handle_stale(
526541
&self,
527542
req: Request,
543+
cache_control: CacheControl<'_>,
528544
cached: DataWithCachePolicy,
529545
new_cache_policy_builder: CachePolicyBuilder,
530546
) -> Result<CachedResponse, Error> {
531547
let url = DisplaySafeUrl::from(req.url().clone());
532548
debug!("Sending revalidation request for: {url}");
533-
let response = self
549+
let mut response = self
534550
.0
535551
.execute(req)
536552
.instrument(info_span!("revalidation_request", url = url.as_str()))
537553
.await
538554
.map_err(|err| ErrorKind::from_reqwest_middleware(url.clone(), err))?
539555
.error_for_status()
540556
.map_err(|err| ErrorKind::from_reqwest(url.clone(), err))?;
557+
558+
// If the user set a custom `Cache-Control` header, override it.
559+
if let CacheControl::Override(header) = cache_control {
560+
response.headers_mut().insert(
561+
http::header::CACHE_CONTROL,
562+
http::HeaderValue::from_str(header)
563+
.expect("Cache-Control header must be valid UTF-8"),
564+
);
565+
}
566+
541567
match cached
542568
.cache_policy
543569
.after_response(new_cache_policy_builder, &response)
@@ -566,16 +592,26 @@ impl CachedClient {
566592
async fn fresh_request(
567593
&self,
568594
req: Request,
595+
cache_control: CacheControl<'_>,
569596
) -> Result<(Response, Option<Box<CachePolicy>>), Error> {
570597
let url = DisplaySafeUrl::from(req.url().clone());
571598
trace!("Sending fresh {} request for {}", req.method(), url);
572599
let cache_policy_builder = CachePolicyBuilder::new(&req);
573-
let response = self
600+
let mut response = self
574601
.0
575602
.execute(req)
576603
.await
577604
.map_err(|err| ErrorKind::from_reqwest_middleware(url.clone(), err))?;
578605

606+
// If the user set a custom `Cache-Control` header, override it.
607+
if let CacheControl::Override(header) = cache_control {
608+
response.headers_mut().insert(
609+
http::header::CACHE_CONTROL,
610+
http::HeaderValue::from_str(header)
611+
.expect("Cache-Control header must be valid UTF-8"),
612+
);
613+
}
614+
579615
let retry_count = response
580616
.extensions()
581617
.get::<reqwest_retry::RetryCount>()
@@ -690,6 +726,7 @@ impl CachedClient {
690726
&self,
691727
req: Request,
692728
cache_entry: &CacheEntry,
729+
cache_control: CacheControl<'_>,
693730
response_callback: Callback,
694731
) -> Result<Payload, CachedClientError<CallBackError>> {
695732
let mut past_retries = 0;
@@ -698,7 +735,7 @@ impl CachedClient {
698735
loop {
699736
let fresh_req = req.try_clone().expect("HTTP request must be cloneable");
700737
let result = self
701-
.skip_cache(fresh_req, cache_entry, &response_callback)
738+
.skip_cache(fresh_req, cache_entry, cache_control, &response_callback)
702739
.await;
703740

704741
// Check if the middleware already performed retries

‎crates/uv-distribution-types/src/index_url.rs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -441,6 +441,26 @@ impl<'a> IndexLocations {
441441
}
442442
}
443443
}
444+
445+
/// Return the Simple API cache control header for an [`IndexUrl`], if configured.
446+
pub fn simple_api_cache_control_for(&self, url: &IndexUrl) -> Option<&str> {
447+
for index in &self.indexes {
448+
if index.url() == url {
449+
return index.cache_control.as_ref()?.api.as_deref();
450+
}
451+
}
452+
None
453+
}
454+
455+
/// Return the artifact cache control header for an [`IndexUrl`], if configured.
456+
pub fn artifact_cache_control_for(&self, url: &IndexUrl) -> Option<&str> {
457+
for index in &self.indexes {
458+
if index.url() == url {
459+
return index.cache_control.as_ref()?.files.as_deref();
460+
}
461+
}
462+
None
463+
}
444464
}
445465

446466
impl From<&IndexLocations> for uv_auth::Indexes {

‎crates/uv-distribution/src/distribution_database.rs‎

Lines changed: 53 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use uv_client::{
2020
};
2121
use uv_distribution_filename::WheelFilename;
2222
use uv_distribution_types::{
23-
BuildableSource, BuiltDist, Dist, HashPolicy, Hashed, InstalledDist, Name, SourceDist,
23+
BuildableSource, BuiltDist, Dist, HashPolicy, Hashed, IndexUrl, InstalledDist, Name, SourceDist,
2424
};
2525
use uv_extract::hash::Hasher;
2626
use uv_fs::write_atomic;
@@ -201,6 +201,7 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
201201
match self
202202
.stream_wheel(
203203
url.clone(),
204+
dist.index(),
204205
&wheel.filename,
205206
wheel.file.size,
206207
&wheel_entry,
@@ -236,6 +237,7 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
236237
let archive = self
237238
.download_wheel(
238239
url,
240+
dist.index(),
239241
&wheel.filename,
240242
wheel.file.size,
241243
&wheel_entry,
@@ -272,6 +274,7 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
272274
match self
273275
.stream_wheel(
274276
wheel.url.raw().clone(),
277+
None,
275278
&wheel.filename,
276279
None,
277280
&wheel_entry,
@@ -301,6 +304,7 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
301304
let archive = self
302305
.download_wheel(
303306
wheel.url.raw().clone(),
307+
None,
304308
&wheel.filename,
305309
None,
306310
&wheel_entry,
@@ -534,6 +538,7 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
534538
async fn stream_wheel(
535539
&self,
536540
url: DisplaySafeUrl,
541+
index: Option<&IndexUrl>,
537542
filename: &WheelFilename,
538543
size: Option<u64>,
539544
wheel_entry: &CacheEntry,
@@ -616,13 +621,24 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
616621
// Fetch the archive from the cache, or download it if necessary.
617622
let req = self.request(url.clone())?;
618623

624+
// Determine the cache control policy for the URL.
619625
let cache_control = match self.client.unmanaged.connectivity() {
620-
Connectivity::Online => CacheControl::from(
621-
self.build_context
622-
.cache()
623-
.freshness(&http_entry, Some(&filename.name), None)
624-
.map_err(Error::CacheRead)?,
625-
),
626+
Connectivity::Online => {
627+
if let Some(header) = index.and_then(|index| {
628+
self.build_context
629+
.locations()
630+
.artifact_cache_control_for(index)
631+
}) {
632+
CacheControl::Override(header)
633+
} else {
634+
CacheControl::from(
635+
self.build_context
636+
.cache()
637+
.freshness(&http_entry, Some(&filename.name), None)
638+
.map_err(Error::CacheRead)?,
639+
)
640+
}
641+
}
626642
Connectivity::Offline => CacheControl::AllowStale,
627643
};
628644

@@ -654,7 +670,12 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
654670
.managed(async |client| {
655671
client
656672
.cached_client()
657-
.skip_cache_with_retry(self.request(url)?, &http_entry, download)
673+
.skip_cache_with_retry(
674+
self.request(url)?,
675+
&http_entry,
676+
cache_control,
677+
download,
678+
)
658679
.await
659680
.map_err(|err| match err {
660681
CachedClientError::Callback { err, .. } => err,
@@ -671,6 +692,7 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
671692
async fn download_wheel(
672693
&self,
673694
url: DisplaySafeUrl,
695+
index: Option<&IndexUrl>,
674696
filename: &WheelFilename,
675697
size: Option<u64>,
676698
wheel_entry: &CacheEntry,
@@ -783,13 +805,24 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
783805
// Fetch the archive from the cache, or download it if necessary.
784806
let req = self.request(url.clone())?;
785807

808+
// Determine the cache control policy for the URL.
786809
let cache_control = match self.client.unmanaged.connectivity() {
787-
Connectivity::Online => CacheControl::from(
788-
self.build_context
789-
.cache()
790-
.freshness(&http_entry, Some(&filename.name), None)
791-
.map_err(Error::CacheRead)?,
792-
),
810+
Connectivity::Online => {
811+
if let Some(header) = index.and_then(|index| {
812+
self.build_context
813+
.locations()
814+
.artifact_cache_control_for(index)
815+
}) {
816+
CacheControl::Override(header)
817+
} else {
818+
CacheControl::from(
819+
self.build_context
820+
.cache()
821+
.freshness(&http_entry, Some(&filename.name), None)
822+
.map_err(Error::CacheRead)?,
823+
)
824+
}
825+
}
793826
Connectivity::Offline => CacheControl::AllowStale,
794827
};
795828

@@ -821,7 +854,12 @@ impl<'a, Context: BuildContext> DistributionDatabase<'a, Context> {
821854
.managed(async |client| {
822855
client
823856
.cached_client()
824-
.skip_cache_with_retry(self.request(url)?, &http_entry, download)
857+
.skip_cache_with_retry(
858+
self.request(url)?,
859+
&http_entry,
860+
cache_control,
861+
download,
862+
)
825863
.await
826864
.map_err(|err| match err {
827865
CachedClientError::Callback { err, .. } => err,

0 commit comments

Comments
 (0)