Skip to content

Commit a7223e8

Browse files
wan9chiclaude
andauthored
refactor(cache): let remote cache requests carry auth headers
`vt_remote_cache::Client::new` now takes an `Auth`, which supplies the headers for each request by operation (fetch, download, or store). It can use the client's HTTP client to get credentials, and when it fails, the request isn't sent and the operation fails with `Error::Auth`. The plan resolves how requests authenticate into `ResolvedRemoteCacheConfig::auth`, which holds everything needed to build the credentials, so nothing reads envs after planning. `vt` builds the `Auth` for it with `build_auth` and caches clients by endpoint and auth. The only kind is `anonymous`, which adds no headers, so requests are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent d4615fb commit a7223e8

16 files changed

Lines changed: 356 additions & 105 deletions

‎crates/vt/src/session/cache/mod.rs‎

Lines changed: 22 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -355,16 +355,18 @@ impl ExecutionCache {
355355
clippy::manual_let_else,
356356
reason = "naming every access mode makes adding one a compile error here"
357357
)]
358-
let url = match &cache_metadata.remote_cache {
359-
Some(ResolvedRemoteCacheConfig {
360-
access: RemoteCacheAccess::Read | RemoteCacheAccess::ReadWrite,
361-
url,
362-
}) => url,
358+
let remote_config = match &cache_metadata.remote_cache {
359+
Some(
360+
remote_config @ ResolvedRemoteCacheConfig {
361+
access: RemoteCacheAccess::Read | RemoteCacheAccess::ReadWrite,
362+
..
363+
},
364+
) => remote_config,
363365
None => return Ok(Err(local_miss)),
364366
};
365367
let remote_miss = match self
366368
.try_hit_remote(
367-
url,
369+
remote_config,
368370
cache_metadata,
369371
&cache_key,
370372
globbed_inputs,
@@ -421,15 +423,16 @@ impl ExecutionCache {
421423
Ok(Err(CacheMiss::NotFound))
422424
}
423425

424-
/// Fetch the entry from the remote cache at `endpoint`. An exact entry
425-
/// that passes validation is a hit once its output archive is downloaded.
426+
/// Fetch the entry from the remote cache that `remote_config` configures.
427+
/// An exact entry that passes validation is a hit once its output archive
428+
/// is downloaded.
426429
/// A fallback entry, a failed validation, or a failed read is a miss. An
427430
/// error while validating counts as a failed read, so the remote entry
428431
/// never fails the task.
429432
#[expect(clippy::too_many_arguments, reason = "forwarded from `try_hit`")]
430433
async fn try_hit_remote(
431434
&self,
432-
endpoint: &Arc<str>,
435+
remote_config: &ResolvedRemoteCacheConfig,
433436
cache_metadata: &CacheMetadata,
434437
cache_key: &CacheEntryKey,
435438
globbed_inputs: &BTreeMap<RelativePathBuf, u64>,
@@ -439,7 +442,7 @@ impl ExecutionCache {
439442
) -> anyhow::Result<Result<CacheEntryValue, CacheMiss>> {
440443
let fetched = self
441444
.remote_clients
442-
.fetch(endpoint, cache_key, &cache_metadata.execution_cache_key, cancel_token)
445+
.fetch(remote_config, cache_key, &cache_metadata.execution_cache_key, cancel_token)
443446
.await;
444447
let validate = |cache_value: &CacheEntryValue| {
445448
cache_value.validate(&cache_metadata.unfiltered_envs, globbed_inputs, workspace_root)
@@ -454,7 +457,7 @@ impl ExecutionCache {
454457
Some(blob_id) => {
455458
match self
456459
.remote_clients
457-
.download_archive(endpoint, &blob_id, cache_dir, cancel_token)
460+
.download_archive(remote_config, &blob_id, cache_dir, cancel_token)
458461
.await
459462
{
460463
Ok(archive_name) => Some(archive_name),
@@ -517,14 +520,19 @@ impl ExecutionCache {
517520
self.record(&cache_key, execution_cache_key, &cache_value, cache_dir).await?;
518521

519522
let upload_error = Arc::new(OnceLock::new());
520-
let url = match &cache_metadata.remote_cache {
521-
Some(ResolvedRemoteCacheConfig { access: RemoteCacheAccess::ReadWrite, url }) => url,
523+
let remote_config = match &cache_metadata.remote_cache {
524+
Some(
525+
remote_config @ ResolvedRemoteCacheConfig {
526+
access: RemoteCacheAccess::ReadWrite,
527+
..
528+
},
529+
) => remote_config,
522530
Some(ResolvedRemoteCacheConfig { access: RemoteCacheAccess::Read, .. }) | None => {
523531
return Ok(upload_error);
524532
}
525533
};
526534
match self.remote_clients.prepare_upload(
527-
url,
535+
remote_config,
528536
&cache_key,
529537
execution_cache_key,
530538
&cache_value,

‎crates/vt/src/session/cache/remote.rs‎

Lines changed: 92 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,14 @@ use rustc_hash::FxHashMap;
2424
use tokio::sync::mpsc;
2525
use tokio_util::{sync::CancellationToken, task::TaskTracker};
2626
use vt_path::AbsolutePath;
27-
use vt_plan::cache_metadata::ExecutionCacheKey;
28-
use vt_remote_cache::{Client, Download, Fetched};
27+
use vt_plan::{
28+
cache_metadata::ExecutionCacheKey,
29+
remote_cache::{RemoteCacheAuth, ResolvedRemoteCacheConfig},
30+
};
31+
use vt_remote_cache::{
32+
Client, Download, Fetched,
33+
auth::{Anonymous, Auth},
34+
};
2935
use vt_str::Str;
3036
use wincode::{
3137
SchemaWrite,
@@ -130,35 +136,49 @@ pub(super) fn resolve(
130136
}
131137
}
132138

133-
/// Remote cache clients, each created when its endpoint is first used.
139+
/// An endpoint and the auth that its requests use. Each has its own client.
140+
#[derive(Debug, PartialEq, Eq, Hash)]
141+
struct ClientKey {
142+
url: Arc<str>,
143+
auth: RemoteCacheAuth,
144+
}
145+
146+
/// Remote cache clients, each created when its endpoint is first used with
147+
/// its auth.
134148
#[derive(Debug, Default)]
135149
pub struct RemoteClients {
136-
clients: Mutex<FxHashMap<Arc<str>, Arc<Client>>>,
150+
clients: Mutex<FxHashMap<ClientKey, Arc<Client>>>,
137151
}
138152

139153
impl RemoteClients {
140-
fn client(&self, endpoint: &Arc<str>) -> Result<Arc<Client>, vt_remote_cache::Error> {
154+
fn client(
155+
&self,
156+
remote_config: &ResolvedRemoteCacheConfig,
157+
) -> Result<Arc<Client>, vt_remote_cache::Error> {
158+
let key =
159+
ClientKey { url: Arc::clone(&remote_config.url), auth: remote_config.auth.clone() };
141160
let mut clients = self.clients.lock().unwrap_or_else(PoisonError::into_inner);
142-
if let Some(client) = clients.get(endpoint) {
161+
if let Some(client) = clients.get(&key) {
143162
return Ok(Arc::clone(client));
144163
}
145-
let client = Arc::new(Client::new(endpoint)?);
146-
clients.insert(Arc::clone(endpoint), Arc::clone(&client));
164+
let client = Arc::new(Client::new(&remote_config.url, build_auth(&remote_config.auth))?);
165+
clients.insert(key, Arc::clone(&client));
147166
drop(clients);
148167
Ok(client)
149168
}
150169

151-
/// Fetch the entry stored under `cache_key`, falling back to the entry
152-
/// last stored for `execution_cache_key`. Returns `None` if neither key
153-
/// matched. Stops when `cancel_token` is cancelled.
170+
/// Fetch the entry stored under `cache_key` in the remote cache that
171+
/// `remote_config` configures, falling back to the entry last stored for
172+
/// `execution_cache_key`. Returns `None` if neither key matched. Stops
173+
/// when `cancel_token` is cancelled.
154174
pub(super) async fn fetch(
155175
&self,
156-
endpoint: &Arc<str>,
176+
remote_config: &ResolvedRemoteCacheConfig,
157177
cache_key: &CacheEntryKey,
158178
execution_cache_key: &ExecutionCacheKey,
159179
cancel_token: &CancellationToken,
160180
) -> Result<Option<Fetched>, ReadError> {
161-
let client = self.client(endpoint).map_err(ReadError::Fetch)?;
181+
let client = self.client(remote_config).map_err(ReadError::Fetch)?;
162182
let key = encode_key(cache_key)?;
163183
let secondary_key = encode_key(execution_cache_key)?;
164184
cancel_token
@@ -168,18 +188,19 @@ impl RemoteClients {
168188
.map_err(ReadError::Fetch)
169189
}
170190

171-
/// Download the blob `blob_id` into `cache_dir`, checking that it decodes
172-
/// as an output archive as it arrives. It's downloaded to a `.tmp` file,
173-
/// which is renamed once the check passes and removed otherwise, such as
174-
/// when `cancel_token` is cancelled. Returns the archive's file name.
191+
/// Download the blob `blob_id` from the remote cache that `remote_config`
192+
/// configures into `cache_dir`, checking that it decodes as an output
193+
/// archive as it arrives. It's downloaded to a `.tmp` file, which is
194+
/// renamed once the check passes and removed otherwise, such as when
195+
/// `cancel_token` is cancelled. Returns the archive's file name.
175196
pub(super) async fn download_archive(
176197
&self,
177-
endpoint: &Arc<str>,
198+
remote_config: &ResolvedRemoteCacheConfig,
178199
blob_id: &str,
179200
cache_dir: &AbsolutePath,
180201
cancel_token: &CancellationToken,
181202
) -> Result<Str, ReadError> {
182-
let client = self.client(endpoint).map_err(ReadError::Download)?;
203+
let client = self.client(remote_config).map_err(ReadError::Download)?;
183204
let archive_name = vt_str::format!("{}.tar.zst", uuid::Uuid::new_v4());
184205
let archive_path = cache_dir.join(archive_name.as_str());
185206
let temp_path = cache_dir.join(vt_str::format!("{archive_name}.tmp").as_str());
@@ -196,19 +217,20 @@ impl RemoteClients {
196217
}
197218

198219
/// Prepare to upload an entry that was just recorded locally, along with
199-
/// its output archive in `cache_dir`: get the endpoint's client and encode
200-
/// the entry. Nothing is sent until the returned future is polled, so an
201-
/// invalid endpoint or an entry that doesn't encode fails here. The future
202-
/// owns everything the upload needs, so it can run in a spawned task.
220+
/// its output archive in `cache_dir`, to the remote cache that
221+
/// `remote_config` configures: get its client and encode the entry.
222+
/// Nothing is sent until the returned future is polled, so an invalid
223+
/// endpoint or an entry that doesn't encode fails here. The future owns
224+
/// everything the upload needs, so it can run in a spawned task.
203225
pub(super) fn prepare_upload(
204226
&self,
205-
endpoint: &Arc<str>,
227+
remote_config: &ResolvedRemoteCacheConfig,
206228
cache_key: &CacheEntryKey,
207229
execution_cache_key: &ExecutionCacheKey,
208230
cache_value: &CacheEntryValue,
209231
cache_dir: &AbsolutePath,
210232
) -> Result<impl Future<Output = Result<(), UploadError>> + Send + use<>, UploadError> {
211-
let client = self.client(endpoint)?;
233+
let client = self.client(remote_config)?;
212234
let key = encode_key(cache_key)?;
213235
let secondary_key = encode_key(execution_cache_key)?;
214236
let value = serialize_cache(cache_value)?;
@@ -220,6 +242,13 @@ impl RemoteClients {
220242
}
221243
}
222244

245+
/// The credentials that requests carry for `auth`.
246+
fn build_auth(auth: &RemoteCacheAuth) -> Arc<dyn Auth> {
247+
match auth {
248+
RemoteCacheAuth::Anonymous => Arc::new(Anonymous),
249+
}
250+
}
251+
223252
/// Uploads running in the background. Each keeps running after its task
224253
/// finishes, until [`Self::wait`] waits for all of them.
225254
#[derive(Debug, Default)]
@@ -395,7 +424,10 @@ mod tests {
395424
use tokio::sync::oneshot;
396425
use vt_graph::config::ResolvedGlobConfig;
397426
use vt_path::{AbsolutePathBuf, RelativePathBuf};
398-
use vt_plan::cache_metadata::{EnvValueHash, SpawnFingerprint};
427+
use vt_plan::{
428+
cache_metadata::{EnvValueHash, SpawnFingerprint},
429+
remote_cache::RemoteCacheAccess,
430+
};
399431

400432
use super::*;
401433
use crate::session::{
@@ -604,14 +636,25 @@ mod tests {
604636
}
605637
}
606638

639+
/// A `read-write` remote cache at `url` without credentials.
640+
fn anonymous(url: &str) -> ResolvedRemoteCacheConfig {
641+
ResolvedRemoteCacheConfig {
642+
access: RemoteCacheAccess::ReadWrite,
643+
url: Arc::from(url),
644+
auth: RemoteCacheAuth::Anonymous,
645+
}
646+
}
647+
607648
/// Serve one request on a loopback endpoint: once the request head
608649
/// arrives, write `response`, then send nothing more and keep the
609-
/// connection open until the client closes it. Returns the endpoint and a
610-
/// receiver that resolves once `response` is written.
611-
fn serve_stalled(response: &'static [u8]) -> (Arc<str>, oneshot::Receiver<()>) {
650+
/// connection open until the client closes it. Returns the remote cache
651+
/// there and a receiver that resolves once `response` is written.
652+
fn serve_stalled(
653+
response: &'static [u8],
654+
) -> (ResolvedRemoteCacheConfig, oneshot::Receiver<()>) {
612655
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
613-
let endpoint =
614-
Arc::from(vt_str::format!("http://{}/projects/test", listener.local_addr().unwrap()));
656+
let remote_config =
657+
anonymous(&vt_str::format!("http://{}/projects/test", listener.local_addr().unwrap()));
615658
let (responded_sender, responded) = oneshot::channel();
616659
std::thread::spawn(move || {
617660
let (mut stream, _) = listener.accept().unwrap();
@@ -626,20 +669,20 @@ mod tests {
626669
let _ = responded_sender.send(());
627670
while stream.read(&mut buf).is_ok_and(|n| n > 0) {}
628671
});
629-
(endpoint, responded)
672+
(remote_config, responded)
630673
}
631674

632675
#[tokio::test]
633676
async fn failed_archive_check_stops_the_download() {
634677
// The response announces more than it sends, so only the failed check
635678
// can end the download.
636-
let (endpoint, _) =
679+
let (remote_config, _) =
637680
serve_stalled(b"HTTP/1.1 200 OK\r\ncontent-length: 1000\r\n\r\nnot an archive");
638681
let dir = tempfile::tempdir().unwrap();
639682
let cache_dir = AbsolutePathBuf::new(dir.path().to_path_buf()).unwrap();
640683

641684
let error = RemoteClients::default()
642-
.download_archive(&endpoint, "1", &cache_dir, &CancellationToken::new())
685+
.download_archive(&remote_config, "1", &cache_dir, &CancellationToken::new())
643686
.await
644687
.unwrap_err();
645688
assert!(matches!(error, ReadError::CorruptArchive(_)), "{error:?}");
@@ -648,13 +691,13 @@ mod tests {
648691

649692
#[tokio::test]
650693
async fn cancelling_stops_a_fetch() {
651-
let (endpoint, requested) = serve_stalled(b"");
694+
let (remote_config, requested) = serve_stalled(b"");
652695
let cancel_token = CancellationToken::new();
653696
let key = cache_key(ResolvedGlobConfig::default_auto());
654697
let execution_key = ExecutionCacheKey::ExecAPI(Arc::from([]));
655698

656699
let clients = RemoteClients::default();
657-
let fetch = clients.fetch(&endpoint, &key, &execution_key, &cancel_token);
700+
let fetch = clients.fetch(&remote_config, &key, &execution_key, &cancel_token);
658701
let (fetched, ()) = tokio::join!(fetch, async {
659702
requested.await.unwrap();
660703
cancel_token.cancel();
@@ -666,14 +709,14 @@ mod tests {
666709
async fn cancelling_stops_a_download_and_removes_it() {
667710
// The response announces a body that never arrives, so only
668711
// cancelling can end the download.
669-
let (endpoint, responded) =
712+
let (remote_config, responded) =
670713
serve_stalled(b"HTTP/1.1 200 OK\r\ncontent-length: 1000\r\n\r\n");
671714
let cancel_token = CancellationToken::new();
672715
let dir = tempfile::tempdir().unwrap();
673716
let cache_dir = AbsolutePathBuf::new(dir.path().to_path_buf()).unwrap();
674717

675718
let clients = RemoteClients::default();
676-
let download = clients.download_archive(&endpoint, "1", &cache_dir, &cancel_token);
719+
let download = clients.download_archive(&remote_config, "1", &cache_dir, &cancel_token);
677720
let (downloaded, ()) = tokio::join!(download, async {
678721
responded.await.unwrap();
679722
// Cancel once the `.tmp` file exists, so the download has started
@@ -687,16 +730,17 @@ mod tests {
687730
assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 0);
688731
}
689732

690-
/// Prepare to upload an entry without an output archive to `endpoint`.
733+
/// Prepare to upload an entry without an output archive to the remote
734+
/// cache that `remote_config` configures.
691735
fn prepare_upload(
692736
clients: &RemoteClients,
693-
endpoint: &Arc<str>,
737+
remote_config: &ResolvedRemoteCacheConfig,
694738
) -> Result<impl Future<Output = Result<(), UploadError>> + use<>, UploadError> {
695739
let key = cache_key(ResolvedGlobConfig::default_auto());
696740
let execution_key = ExecutionCacheKey::ExecAPI(Arc::from([]));
697741
let value = CacheEntryValue { output_archive: None, ..cache_value() };
698742
let cache_dir = vt_path::current_dir().unwrap();
699-
clients.prepare_upload(endpoint, &key, &execution_key, &value, &cache_dir)
743+
clients.prepare_upload(remote_config, &key, &execution_key, &value, &cache_dir)
700744
}
701745

702746
fn upload_error(error: &OnceLock<UploadError>) -> Option<Str> {
@@ -705,8 +749,8 @@ mod tests {
705749

706750
#[test]
707751
fn upload_to_an_invalid_endpoint_fails_before_it_starts() {
708-
let endpoint = Arc::from("cache.example/projects/test");
709-
let Err(error) = prepare_upload(&RemoteClients::default(), &endpoint) else {
752+
let remote_config = anonymous("cache.example/projects/test");
753+
let Err(error) = prepare_upload(&RemoteClients::default(), &remote_config) else {
710754
panic!("an invalid endpoint should fail");
711755
};
712756
assert!(
@@ -720,15 +764,15 @@ mod tests {
720764
let (error_status, _) =
721765
serve_stalled(b"HTTP/1.1 500 Internal Server Error\r\ncontent-length: 0\r\n\r\n");
722766
// Nothing can listen on port 0.
723-
let unreachable = Arc::from("http://127.0.0.1:0/projects/test");
767+
let unreachable = anonymous("http://127.0.0.1:0/projects/test");
724768
let clients = RemoteClients::default();
725769
let uploads = RemoteUploads::default();
726770

727-
for (endpoint, message) in
771+
for (remote_config, message) in
728772
[(error_status, "HTTP status 500"), (unreachable, "network error")]
729773
{
730774
let error = Arc::new(OnceLock::new());
731-
uploads.spawn(prepare_upload(&clients, &endpoint).unwrap(), Arc::clone(&error));
775+
uploads.spawn(prepare_upload(&clients, &remote_config).unwrap(), Arc::clone(&error));
732776
uploads.wait(&CancellationToken::new()).await;
733777
assert_eq!(uploads.pending(), 0);
734778
assert_eq!(upload_error(&error).as_deref(), Some(message));
@@ -737,11 +781,11 @@ mod tests {
737781

738782
#[tokio::test]
739783
async fn interrupting_the_wait_cancels_the_uploads() {
740-
let (endpoint, requested) = serve_stalled(b"");
784+
let (remote_config, requested) = serve_stalled(b"");
741785
let clients = RemoteClients::default();
742786
let uploads = RemoteUploads::default();
743787
let error = Arc::new(OnceLock::new());
744-
uploads.spawn(prepare_upload(&clients, &endpoint).unwrap(), Arc::clone(&error));
788+
uploads.spawn(prepare_upload(&clients, &remote_config).unwrap(), Arc::clone(&error));
745789
requested.await.unwrap();
746790
assert_eq!(uploads.pending(), 1);
747791

0 commit comments

Comments
 (0)