Skip to main content

mbx_cache_core/
agent.rs

1use crate::uploads::{ConnectionUploads, UploadQueue, UploadSink};
2use crate::{
3    ActionPrediction, BlobPackLimits, CacheDigest, CacheDirectory, LocalActionCache, LocalCas,
4    MAX_STAGED_BLOB_PACK_BYTES, MAX_STAGED_BLOB_PACK_ITEMS, ManifestPutOutcome, RemoteActionResult,
5    RemoteCacheClient, RemoteCacheMode, RustcMetadata, TaskActionManifest, blob_pack_chunk,
6    canonical_json,
7};
8use eyre::{Context, Result, bail};
9use futures_util::{FutureExt, StreamExt, future::BoxFuture, stream};
10use log::warn;
11use serde::{Deserialize, Serialize};
12use std::collections::{BTreeMap, BTreeSet};
13use std::fs;
14use std::path::{Path, PathBuf};
15use std::sync::atomic::{AtomicU64, Ordering};
16use std::sync::{Arc, Mutex, Weak};
17use std::time::{Duration, Instant, SystemTime};
18use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
19
20const MAX_EXECUTABLE_IDENTITIES: usize = 64;
21const MAX_EXECUTABLE_IDENTITY_SIZE: usize = 64 * 1024;
22const MAX_EXECUTABLE_IDENTITY_BYTES: usize = 256 * 1024;
23const TASK_ACTION_MANIFEST_VERSION: u8 = 1;
24const MAX_TASK_ACTION_PREDICTIONS: usize = 16 * 1024;
25/// Longest shim diagnostic the agent accepts, in bytes.
26const MAX_WARNING_BYTES: usize = 4 * 1024;
27/// Most distinct shim diagnostics one session surfaces before going quiet.
28///
29/// A failure mode that repeats across a build tends to repeat with the same
30/// message, which deduplication already collapses; the cap only guards
31/// against a message that embeds something unique per compilation.
32const MAX_WARNINGS: usize = 128;
33const MAX_REMOTE_TRANSFERS: usize = 64;
34const MAX_PREFETCH_TRANSFERS: usize = 48;
35const MAX_PREFETCH_ACTION_BATCH: usize = 256;
36/// Batched action lookups issued at once, well inside the transfer budget: each
37/// asks about hundreds of actions, so a handful covers a large workspace.
38const MAX_PREFETCH_BATCH_LOOKUPS: usize = 4;
39const PREFETCH_ACTION_BATCH_DELAY: Duration = Duration::from_millis(5);
40const MAX_PREFETCH_DIRECTORY_OBJECTS: usize = 100_000;
41const MAX_PREFETCH_OBJECTS_PER_WAVE: usize = 100_000;
42const DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES: u64 = 5 * 1024 * 1024 * 1024;
43
44/// Remote action-cache access owned by one task session.
45pub struct AgentRemoteCache {
46    /// Remote protocol client used by the agent.
47    pub client: RemoteCacheClient,
48    /// Permitted remote read/write operations.
49    pub mode: RemoteCacheMode,
50    /// Directory used for verified downloads before CAS ingestion.
51    pub staging_dir: PathBuf,
52}
53
54/// Wire protocol version used between an in-process cache agent and its shims.
55pub const AGENT_PROTOCOL_VERSION: u8 = 4;
56/// Largest single protocol request the agent will read.
57///
58/// Requests are small JSON objects; the largest legitimate ones carry an output
59/// tree or a batch of digests, which stay far below this.
60const MAX_REQUEST_BYTES: usize = 16 * 1024 * 1024;
61
62/// A request accepted by the task-scoped cache agent.
63#[derive(Debug, Clone, Serialize, Deserialize)]
64#[serde(tag = "type", rename_all = "snake_case")]
65pub enum AgentRequest {
66    /// Negotiate protocol and application versions.
67    Hello {
68        /// Agent protocol version understood by the caller.
69        protocol: u8,
70        /// Human-readable mbx client version.
71        client_version: String,
72    },
73    /// Begin a prediction-manifest run for a stable Cargo or task identity.
74    BeginTask {
75        /// Stable 64-character lowercase hexadecimal task identity.
76        task: String,
77    },
78    /// Commit predictions collected by an earlier [`Self::BeginTask`].
79    CommitTask {
80        /// Opaque run identifier returned by the agent.
81        run: String,
82    },
83    /// Resolve a blob to a session-verified local CAS path.
84    FindBlob {
85        /// Blob to resolve.
86        digest: CacheDigest,
87    },
88    /// Resolve blobs to session-verified local CAS paths.
89    FindBlobs {
90        /// Blobs to resolve, preserving request order in the response.
91        digests: Vec<CacheDigest>,
92    },
93    /// Import a file into the local content-addressed store.
94    StoreBlob {
95        /// Digest the source must match.
96        digest: CacheDigest,
97        /// File to verify and import.
98        source: PathBuf,
99    },
100    /// Look up an action-result record.
101    FindActionResult {
102        /// Action digest to resolve.
103        action: CacheDigest,
104    },
105    /// Account for a successfully restored cache hit.
106    RecordActionHit {
107        /// Action that supplied the outputs.
108        action: CacheDigest,
109        /// Restoration work performed by the adapter.
110        restore: RestoreStats,
111        /// Compiler crate name, when the invocation supplied one.
112        #[serde(default)]
113        crate_name: Option<String>,
114    },
115    /// A compilation the adapter declined to cache, grouped by reason.
116    RecordBypass {
117        /// Stable, low-cardinality bypass-reason name.
118        kind: String,
119    },
120    /// A compilation the adapter could not look up, having no key to look up
121    /// with. Distinct from a bypass: these are cached once compiled.
122    RecordUnconsulted,
123    /// Account for one real compiler invocation performed by the adapter.
124    RecordCompilerInvocation {
125        /// Stable outcome category such as `miss`, `unconsulted`, `bypass`, or
126        /// `incremental` for a compilation the adapter deliberately ran with
127        /// incremental state instead of publishing.
128        outcome: String,
129        /// Compiler crate name, when the invocation supplied one.
130        crate_name: Option<String>,
131        /// Wall time spent running the compiler.
132        duration_ns: u64,
133    },
134    /// Account for a cache hit that was rebuilt for correctness verification.
135    RecordActionVerification {
136        /// Whether rebuilt and cached outputs matched.
137        matched: bool,
138        /// Restoration work performed before rebuilding.
139        restore: RestoreStats,
140    },
141    /// Store an action-result record locally and enqueue remote publication.
142    StoreActionResult {
143        /// Action-result record to store.
144        result: RemoteActionResult,
145    },
146    /// Find an earlier input prediction for a task and invocation.
147    FindActionPrediction {
148        /// Stable task identity.
149        task: String,
150        /// Digest of the compiler invocation without discovered inputs.
151        invocation: CacheDigest,
152    },
153    /// Record an input prediction after a successful compilation.
154    RecordActionPrediction {
155        /// Stable task identity.
156        task: String,
157        /// Adapter-owned prediction record.
158        prediction: ActionPrediction,
159    },
160    /// Find cached identity output for an executable and environment.
161    FindExecutableIdentity {
162        /// Executable whose identity command would run.
163        executable: PathBuf,
164        /// Environment variables affecting identity output.
165        environment: BTreeMap<String, Option<String>>,
166    },
167    /// Cache identity output for an executable and environment.
168    StoreExecutableIdentity {
169        /// Executable whose identity command ran.
170        executable: PathBuf,
171        /// Environment variables affecting identity output.
172        environment: BTreeMap<String, Option<String>>,
173        /// Captured identity-command standard output.
174        stdout: Vec<u8>,
175    },
176    /// Surface a shim diagnostic through the session that owns the build.
177    ///
178    /// A shim must not print diagnostics itself: its stderr belongs to the
179    /// compiler it stands in for, and build scripts read that stream as part
180    /// of the compiler's answer -- cc-rs treats any stderr output from a flag
181    /// probe as "unsupported", which changes the flags of every compilation
182    /// that follows and, with them, every action key the build produces.
183    ///
184    /// Appended rather than grouped with the other `Record*` requests it
185    /// belongs beside: these variants carry no `repr`, so inserting one moves
186    /// the discriminant of every variant after it, and a break nobody asked
187    /// for is worth less than the grouping. New variants go here.
188    RecordWarning {
189        /// Human-readable single-line diagnostic.
190        message: String,
191    },
192}
193
194/// Local output restoration work performed by one action-cache adapter hit.
195#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
196#[serde(deny_unknown_fields)]
197pub struct RestoreStats {
198    /// Cumulative time spent materializing and validating output files.
199    pub duration_ns: u64,
200    /// Compiler wall time recorded when this action was originally produced.
201    /// Zero means no timing hint was available.
202    pub avoided_compiler_duration_ns: u64,
203    /// Number of compiler output files restored.
204    pub output_files: u64,
205    /// Declared size of compiler output files restored.
206    pub output_bytes: u64,
207    /// Number of restored output files that share data blocks with the CAS.
208    pub reflinked_output_files: u64,
209    /// Declared size of restored outputs that share data blocks with the CAS.
210    pub reflinked_output_bytes: u64,
211    /// Number of restored output files that required a byte-for-byte copy.
212    pub copied_output_files: u64,
213    /// Declared size of restored outputs that required a byte-for-byte copy.
214    pub copied_output_bytes: u64,
215}
216
217/// One accounted cache decision, as it happens.
218///
219/// The agent already folds every one of these into [`AgentStats`]; an observer
220/// sees the same decisions individually, before that summing loses the crate
221/// they belong to. Delivered synchronously from the request handler, so an
222/// observer that blocks slows the build it is watching.
223#[derive(Debug, Clone)]
224#[non_exhaustive]
225pub enum AgentEvent {
226    /// An action's outputs were restored from cache.
227    ActionHit {
228        /// Compiler crate name, when the invocation supplied one.
229        crate_name: Option<String>,
230        /// Restoration work performed by the adapter.
231        restore: RestoreStats,
232    },
233    /// A compilation the adapter declined to cache.
234    Bypass {
235        /// Stable, low-cardinality bypass-reason name.
236        kind: String,
237    },
238    /// A compilation no lookup was possible for.
239    Unconsulted,
240    /// A real compiler invocation ran.
241    CompilerInvocation {
242        /// Stable outcome category such as `miss`, `unconsulted`, or `bypass`.
243        outcome: String,
244        /// Compiler crate name, when the invocation supplied one.
245        crate_name: Option<String>,
246        /// Wall time spent running the compiler.
247        duration_ns: u64,
248    },
249    /// A hit was rebuilt to verify it.
250    Verification {
251        /// Whether rebuilt and cached outputs matched.
252        matched: bool,
253        /// Restoration work performed before rebuilding.
254        restore: RestoreStats,
255    },
256    /// A shim reported a diagnostic for the session to surface.
257    Warning {
258        /// Human-readable single-line diagnostic.
259        message: String,
260    },
261}
262
263/// A sink for [`AgentEvent`]s observed during one session.
264pub trait AgentEventObserver: Send + Sync {
265    /// Handle one event. Must not panic, and should not block.
266    fn event(&self, event: AgentEvent);
267}
268
269/// A response returned by the task-scoped cache agent.
270#[derive(Debug, Clone, Serialize, Deserialize)]
271#[serde(tag = "type", rename_all = "snake_case")]
272pub enum AgentResponse {
273    /// Successful protocol negotiation.
274    Hello {
275        /// Agent protocol version.
276        protocol: u8,
277        /// Human-readable agent version.
278        agent_version: String,
279    },
280    /// A task prediction run was begun.
281    TaskBegun {
282        /// Opaque run identifier to pass to compiler shims and commit later.
283        run: String,
284    },
285    /// A task prediction run was committed.
286    TaskCommitted,
287    /// A local CAS path already verified against the requested digest.
288    Blob {
289        /// Verified local path, or `None` on a cache miss.
290        path: Option<PathBuf>,
291    },
292    /// Local CAS paths already verified against the requested digests.
293    Blobs {
294        /// Verified local paths or misses, in request order.
295        paths: Vec<Option<PathBuf>>,
296    },
297    /// A blob was stored locally.
298    Stored {
299        /// Path of the stored object in the local CAS.
300        path: PathBuf,
301    },
302    /// Result of an action lookup.
303    ActionResult {
304        /// Validated action result, or `None` on a cache miss.
305        result: Option<RemoteActionResult>,
306    },
307    /// Hit statistics were updated.
308    ActionHitRecorded,
309    /// Verification statistics were updated.
310    ActionVerificationRecorded,
311    /// Bypass statistics were updated.
312    BypassRecorded,
313    /// Unconsulted-compilation statistics were updated.
314    UnconsultedRecorded,
315    /// Compiler invocation accounting was recorded.
316    CompilerInvocationRecorded,
317    /// An action result was stored.
318    ActionStored {
319        /// Path of the stored local action-result record.
320        path: PathBuf,
321    },
322    /// Result of an input-prediction lookup.
323    ActionPrediction {
324        /// Matching prediction, or `None` when none is known.
325        prediction: Option<ActionPrediction>,
326    },
327    /// An input prediction was recorded.
328    ActionPredictionRecorded,
329    /// Result of an executable-identity lookup.
330    ExecutableIdentity {
331        /// Captured output, or `None` when no identity is cached.
332        stdout: Option<Vec<u8>>,
333    },
334    /// The request failed without terminating the agent connection.
335    Error {
336        /// Human-readable failure description.
337        message: String,
338    },
339    /// A shim diagnostic was accepted for the session to surface.
340    ///
341    /// Sits past `Error` for the reason [`AgentRequest::RecordWarning`] sits
342    /// last: anywhere earlier renumbers the variants below it. New variants go
343    /// here.
344    WarningRecorded,
345}
346
347/// Aggregate cache activity for one task session.
348///
349/// The agent produces these; nothing outside this crate has cause to build one.
350/// Saying so keeps a new counter from being a breaking change, which is what
351/// this type exists to accumulate -- reach for [`AgentStats::default`] and
352/// assign the fields a test needs.
353#[derive(Debug, Clone, Default, PartialEq, Eq)]
354#[non_exhaustive]
355pub struct AgentStats {
356    /// End-to-end lifetime of the task-scoped cache session.
357    pub session_duration_ns: u64,
358    /// Number of action-result lookups.
359    pub lookups: u64,
360    /// Compilations no action-result lookup was possible for, because no usable
361    /// action key was available.
362    ///
363    /// Counted separately from a miss, which is a lookup that found nothing.
364    /// Both compile, but only a miss says a lookup happened.
365    pub unconsulted: u64,
366    /// Number of lookups that found a valid local action result.
367    pub hits: u64,
368    /// Number of newly stored content-addressed objects.
369    pub stores: u64,
370    /// Total size of newly stored objects.
371    pub stored_bytes: u64,
372    /// Number of cache hits compiled again for qualification.
373    pub verifications: u64,
374    /// Number of qualification builds that diverged from the cached result.
375    pub divergences: u64,
376    /// CAS payload bytes downloaded from the remote cache.
377    pub downloaded_bytes: u64,
378    /// CAS payload bytes uploaded to the remote cache.
379    pub uploaded_bytes: u64,
380    /// Objects published to the remote cache after the build asked for them.
381    ///
382    /// A store request returns once the object is in the local CAS, so the
383    /// upload it implies happens off the build's critical path. This counts the
384    /// blobs and action results that publication actually completed for.
385    pub background_uploads: u64,
386    /// Queued uploads that did not publish, having been reported and recovered
387    /// from.
388    pub background_upload_failures: u64,
389    /// Framed requests that published several blobs at once.
390    pub remote_blob_pack_uploads: u64,
391    /// Blobs published through those framed requests.
392    pub remote_blob_pack_upload_blobs: u64,
393    /// Time the session spent waiting for queued uploads once the build ended.
394    pub upload_drain_duration_ns: u64,
395    /// Complete actions staged before an adapter requested them.
396    pub prefetched_actions: u64,
397    /// Predictions carried by the task manifest this session loaded.
398    ///
399    /// Zero means no earlier build left a manifest behind -- a genuinely cold
400    /// start. Together with `lookups`, this is what tells a session that loaded
401    /// hundreds of predictions and matched none of them apart from one that had
402    /// nothing to match against; the first means the invocations changed (a
403    /// compiler update does this to every one of them at once), the second that
404    /// the store was empty.
405    pub predictions_loaded: u64,
406    /// Compilations that were not cacheable, counted by reason.
407    pub bypasses: BTreeMap<String, u64>,
408    /// Estimated compiler time avoided by restored action hits.
409    pub avoided_compiler_duration_ns: u64,
410    /// Real compiler work performed in this session, grouped by outcome.
411    pub compiler: BTreeMap<String, CompilerStats>,
412    /// Cumulative real compiler time by crate name.
413    pub slow_compilations: BTreeMap<String, u64>,
414    /// Remote cache operations that failed and were degraded to a local result.
415    ///
416    /// A remote cache that cannot be reached, or that answers in a way this
417    /// client refuses, costs hit rate rather than correctness, so every one of
418    /// these is recovered from rather than raised. Counting them is what keeps a
419    /// remote that is failing every request from reading as one that merely had
420    /// nothing to offer.
421    pub remote_failures: u64,
422    /// Number of task manifest requests made to the remote cache.
423    pub remote_manifest_lookups: u64,
424    /// Cumulative time spent requesting remote task manifests.
425    pub remote_manifest_lookup_duration_ns: u64,
426    /// Number of action-result requests made to the remote cache.
427    pub remote_action_lookups: u64,
428    /// Cumulative time spent requesting remote action results.
429    pub remote_action_lookup_duration_ns: u64,
430    /// Number of blob requests made to the remote cache.
431    pub remote_blob_requests: u64,
432    /// Number of packed blob requests made to the remote cache.
433    pub remote_blob_pack_requests: u64,
434    /// Number of verified blobs received through packed responses.
435    pub remote_blob_pack_blobs: u64,
436    /// Cumulative time spent downloading and verifying remote blobs.
437    pub remote_blob_transfer_duration_ns: u64,
438    /// Cumulative time spent ingesting downloaded blobs into the local CAS.
439    pub local_cas_write_duration_ns: u64,
440    /// Number of speculative prefetch runs started for task manifests.
441    pub prefetch_runs: u64,
442    /// Cumulative wall time of speculative task-manifest prefetch runs.
443    pub prefetch_duration_ns: u64,
444    /// Cumulative time spent staging or materializing and validating cached outputs.
445    pub materialization_duration_ns: u64,
446    /// Number of compiler output files restored from action hits.
447    pub restored_output_files: u64,
448    /// Declared size of compiler output files restored from action hits.
449    pub restored_output_bytes: u64,
450    /// Number of restored output files materialized with filesystem reflinks.
451    pub reflinked_output_files: u64,
452    /// Declared size of outputs materialized with filesystem reflinks.
453    pub reflinked_output_bytes: u64,
454    /// Number of restored output files materialized by copying their bytes.
455    pub copied_output_files: u64,
456    /// Declared size of outputs materialized by copying their bytes.
457    pub copied_output_bytes: u64,
458}
459
460/// Count and cumulative wall time for one compiler-invocation outcome.
461///
462/// Non-exhaustive for the same reason as [`AgentStats`], which holds these:
463/// leaving the outer bag open does not help if describing an outcome in more
464/// detail still breaks the type inside it.
465#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
466#[non_exhaustive]
467pub struct CompilerStats {
468    /// Number of compiler invocations observed.
469    pub invocations: u64,
470    /// Cumulative wall time spent in those invocations.
471    pub duration_ns: u64,
472}
473
474impl CompilerStats {
475    /// The counts observed for one outcome.
476    pub fn new(invocations: u64, duration_ns: u64) -> Self {
477        Self {
478            invocations,
479            duration_ns,
480        }
481    }
482}
483
484#[derive(Default)]
485struct AtomicAgentStats {
486    lookups: AtomicU64,
487    unconsulted: AtomicU64,
488    hits: AtomicU64,
489    stores: AtomicU64,
490    stored_bytes: AtomicU64,
491    verifications: AtomicU64,
492    divergences: AtomicU64,
493    downloaded_bytes: AtomicU64,
494    uploaded_bytes: AtomicU64,
495    background_uploads: AtomicU64,
496    background_upload_failures: AtomicU64,
497    remote_blob_pack_uploads: AtomicU64,
498    remote_blob_pack_upload_blobs: AtomicU64,
499    upload_drain_duration_ns: AtomicU64,
500    prefetched_actions: AtomicU64,
501    predictions_loaded: AtomicU64,
502    remote_failures: AtomicU64,
503    remote_manifest_lookups: AtomicU64,
504    remote_manifest_lookup_duration_ns: AtomicU64,
505    remote_action_lookups: AtomicU64,
506    remote_action_lookup_duration_ns: AtomicU64,
507    remote_blob_requests: AtomicU64,
508    remote_blob_pack_requests: AtomicU64,
509    remote_blob_pack_blobs: AtomicU64,
510    remote_blob_transfer_duration_ns: AtomicU64,
511    local_cas_write_duration_ns: AtomicU64,
512    prefetch_runs: AtomicU64,
513    prefetch_duration_ns: AtomicU64,
514    materialization_duration_ns: AtomicU64,
515    bypasses: Mutex<BTreeMap<String, u64>>,
516    avoided_compiler_duration_ns: AtomicU64,
517    compiler: Mutex<BTreeMap<String, CompilerStats>>,
518    slow_compilations: Mutex<BTreeMap<String, u64>>,
519    restored_output_files: AtomicU64,
520    restored_output_bytes: AtomicU64,
521    reflinked_output_files: AtomicU64,
522    reflinked_output_bytes: AtomicU64,
523    copied_output_files: AtomicU64,
524    copied_output_bytes: AtomicU64,
525}
526
527struct AtomicDurationTimer<'a> {
528    started: Instant,
529    target: &'a AtomicU64,
530}
531
532impl<'a> AtomicDurationTimer<'a> {
533    fn start(target: &'a AtomicU64) -> Self {
534        Self {
535            started: Instant::now(),
536            target,
537        }
538    }
539}
540
541impl Drop for AtomicDurationTimer<'_> {
542    fn drop(&mut self) {
543        atomic_saturating_add(self.target, duration_ns(self.started));
544    }
545}
546
547fn duration_ns(started: Instant) -> u64 {
548    started.elapsed().as_nanos().try_into().unwrap_or(u64::MAX)
549}
550
551fn validate_crate_name(crate_name: Option<&str>) -> Result<()> {
552    if let Some(crate_name) = crate_name
553        && (crate_name.len() > 256 || crate_name.contains(['\0', '\n', '\r']))
554    {
555        bail!("invalid compiler crate name");
556    }
557    Ok(())
558}
559
560fn atomic_saturating_add(target: &AtomicU64, value: u64) {
561    let _ = target.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
562        Some(current.saturating_add(value))
563    });
564}
565
566fn queue_prefetch_digest(
567    verified: &BTreeMap<CacheDigest, PathBuf>,
568    pending: &mut BTreeMap<CacheDigest, ()>,
569    digest: CacheDigest,
570) {
571    if verified.contains_key(&digest) || pending.contains_key(&digest) {
572        return;
573    }
574    pending.insert(digest, ());
575}
576
577fn queue_prefetch_directory(
578    seen: &BTreeMap<CacheDigest, ()>,
579    pending: &mut BTreeMap<CacheDigest, ()>,
580    digest: CacheDigest,
581    limit: usize,
582) -> bool {
583    if seen.contains_key(&digest) || pending.contains_key(&digest) {
584        return true;
585    }
586    if seen.len().saturating_add(pending.len()) >= limit {
587        return false;
588    }
589    pending.insert(digest, ());
590    true
591}
592
593/// Shared state for an agent hosted by the process that owns a build session.
594///
595/// Transport listeners deliberately live in the embedder so the session
596/// lifecycle owns them. This type only contains ecosystem-independent CAS and
597/// protocol logic.
598#[derive(Clone)]
599pub struct CacheAgent {
600    cas: LocalCas,
601    actions: LocalActionCache,
602    verified_blobs: Arc<Mutex<BTreeMap<CacheDigest, VerifiedBlob>>>,
603    version: Arc<str>,
604    write_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
605    action_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
606    stats: Arc<AtomicAgentStats>,
607    observer: Option<Arc<dyn AgentEventObserver>>,
608    executable_identities: Arc<Mutex<BTreeMap<ExecutableIdentityKey, Vec<u8>>>>,
609    manifest_dir: Arc<PathBuf>,
610    task_actions: Arc<Mutex<BTreeMap<String, TaskActionState>>>,
611    next_task_run: Arc<AtomicU64>,
612    manifest_write_lock: Arc<Mutex<()>>,
613    remote: Option<Arc<RemoteCacheClient>>,
614    remote_mode: RemoteCacheMode,
615    remote_staging_dir: Arc<PathBuf>,
616    remote_download_limit: u64,
617    remote_download_bytes: Arc<AtomicU64>,
618    pending_remote_actions: Arc<Mutex<BTreeMap<CacheDigest, RemoteActionResult>>>,
619    remote_transfers: Arc<tokio::sync::Semaphore>,
620    prefetch_transfers: Arc<tokio::sync::Semaphore>,
621    prefetch_tasks: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
622    /// Distinct shim diagnostics already surfaced, so a warning that fires
623    /// once per compilation is printed once per session.
624    warnings: Arc<Mutex<BTreeSet<String>>>,
625    /// Deferred remote publication, present only when the session may write.
626    uploads: Option<UploadQueue>,
627}
628
629/// Records background upload activity against a session's statistics.
630struct AgentUploadSink {
631    stats: Arc<AtomicAgentStats>,
632}
633
634impl UploadSink for AgentUploadSink {
635    fn record_blob_uploaded(&self, bytes: u64) {
636        self.stats
637            .background_uploads
638            .fetch_add(1, Ordering::Relaxed);
639        self.stats
640            .uploaded_bytes
641            .fetch_add(bytes, Ordering::Relaxed);
642    }
643
644    fn record_action_uploaded(&self) {
645        self.stats
646            .background_uploads
647            .fetch_add(1, Ordering::Relaxed);
648    }
649
650    fn record_blob_pack_uploaded(&self, blobs: u64) {
651        self.stats
652            .remote_blob_pack_uploads
653            .fetch_add(1, Ordering::Relaxed);
654        self.stats
655            .remote_blob_pack_upload_blobs
656            .fetch_add(blobs, Ordering::Relaxed);
657    }
658
659    fn record_upload_failure(&self) {
660        self.stats
661            .background_upload_failures
662            .fetch_add(1, Ordering::Relaxed);
663        self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
664    }
665}
666
667/// A CAS blob whose contents this session hashed, and the file identity that
668/// says the hash still describes what is on disk.
669///
670/// Length alone would miss an overwrite that keeps the size; the modification
671/// time is what makes a rewrite visible without reading the bytes again.
672#[derive(Debug, Clone)]
673struct VerifiedBlob {
674    path: PathBuf,
675    len: u64,
676    modified: SystemTime,
677}
678
679impl VerifiedBlob {
680    /// Describe a file that was just verified, or nothing when the filesystem
681    /// will not report a modification time to compare against later.
682    fn describe(path: &Path) -> Option<Self> {
683        let metadata = std::fs::metadata(path).ok()?;
684        Some(Self {
685            path: path.to_path_buf(),
686            len: metadata.len(),
687            modified: metadata.modified().ok()?,
688        })
689    }
690
691    /// Whether the file still has the identity it had when it was verified.
692    fn is_unchanged(&self) -> bool {
693        let Ok(metadata) = std::fs::metadata(&self.path) else {
694            return false;
695        };
696        metadata.len() == self.len && metadata.modified().is_ok_and(|now| now == self.modified)
697    }
698}
699
700#[derive(Debug, Clone, Default)]
701struct TaskActionState {
702    manifest: String,
703    baseline_loaded: bool,
704    predictions: BTreeMap<CacheDigest, ActionPrediction>,
705    pending_predictions: BTreeMap<CacheDigest, ActionPrediction>,
706    remote_etag: Option<String>,
707}
708
709struct PrefetchedAction {
710    adapter: String,
711    result: RemoteActionResult,
712}
713
714struct RemoteDownloadReservation {
715    counter: Arc<AtomicU64>,
716    reserved: u64,
717    committed: bool,
718}
719
720impl RemoteDownloadReservation {
721    fn bytes(&self) -> u64 {
722        self.reserved
723    }
724
725    fn commit(mut self, bytes: u64) {
726        debug_assert!(bytes <= self.reserved);
727        self.counter
728            .fetch_sub(self.reserved.saturating_sub(bytes), Ordering::AcqRel);
729        self.committed = true;
730    }
731}
732
733impl Drop for RemoteDownloadReservation {
734    fn drop(&mut self) {
735        if !self.committed {
736            self.counter.fetch_sub(self.reserved, Ordering::AcqRel);
737        }
738    }
739}
740
741#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
742struct ExecutableIdentityKey {
743    executable: PathBuf,
744    environment: BTreeMap<String, Option<String>>,
745}
746
747impl CacheAgent {
748    /// Create an agent backed by the cache rooted at `cache_dir`.
749    pub fn new(cache_dir: impl Into<PathBuf>, version: impl Into<Arc<str>>) -> Self {
750        Self::build(cache_dir.into(), version.into(), None, 0)
751    }
752
753    /// Create an agent with local-first access to a remote action cache.
754    pub fn new_remote(
755        cache_dir: impl Into<PathBuf>,
756        version: impl Into<Arc<str>>,
757        remote: AgentRemoteCache,
758    ) -> Self {
759        Self::build(
760            cache_dir.into(),
761            version.into(),
762            Some(remote),
763            DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES,
764        )
765    }
766
767    /// Create a remote agent with a cumulative download budget for the session.
768    pub fn new_remote_with_download_limit(
769        cache_dir: impl Into<PathBuf>,
770        version: impl Into<Arc<str>>,
771        remote: AgentRemoteCache,
772        max_remote_download_bytes: u64,
773    ) -> Self {
774        Self::build(
775            cache_dir.into(),
776            version.into(),
777            Some(remote),
778            max_remote_download_bytes,
779        )
780    }
781
782    fn build(
783        cache_dir: PathBuf,
784        version: Arc<str>,
785        remote: Option<AgentRemoteCache>,
786        remote_download_limit: u64,
787    ) -> Self {
788        let remote_mode = remote
789            .as_ref()
790            .map_or(RemoteCacheMode::ReadOnly, |remote| remote.mode);
791        let remote_staging_dir = remote.as_ref().map_or_else(
792            || cache_dir.join("remote"),
793            |remote| remote.staging_dir.clone(),
794        );
795        let remote = remote.map(|remote| Arc::new(remote.client));
796        let stats = Arc::new(AtomicAgentStats::default());
797        let remote_transfers = Arc::new(tokio::sync::Semaphore::new(MAX_REMOTE_TRANSFERS));
798        // A session that cannot publish never queues an upload, so it does not
799        // need somewhere to queue one.
800        let uploads = remote
801            .clone()
802            .filter(|_| remote_mode.writes())
803            .map(|client| {
804                UploadQueue::new(
805                    client,
806                    Arc::new(AgentUploadSink {
807                        stats: stats.clone(),
808                    }),
809                    remote_transfers.clone(),
810                )
811            });
812        Self {
813            cas: LocalCas::new(cache_dir.clone()),
814            actions: LocalActionCache::new(cache_dir.clone()),
815            verified_blobs: Arc::new(Mutex::new(BTreeMap::new())),
816            version,
817            write_locks: Arc::new(Mutex::new(BTreeMap::new())),
818            action_locks: Arc::new(Mutex::new(BTreeMap::new())),
819            stats,
820            observer: None,
821            executable_identities: Arc::new(Mutex::new(BTreeMap::new())),
822            manifest_dir: Arc::new(task_manifest_dir(&cache_dir)),
823            task_actions: Arc::new(Mutex::new(BTreeMap::new())),
824            next_task_run: Arc::new(AtomicU64::new(0)),
825            manifest_write_lock: Arc::new(Mutex::new(())),
826            remote,
827            remote_mode,
828            remote_staging_dir: Arc::new(remote_staging_dir),
829            remote_download_limit,
830            remote_download_bytes: Arc::new(AtomicU64::new(0)),
831            pending_remote_actions: Arc::new(Mutex::new(BTreeMap::new())),
832            remote_transfers,
833            prefetch_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_PREFETCH_TRANSFERS)),
834            prefetch_tasks: Arc::new(Mutex::new(Vec::new())),
835            warnings: Arc::new(Mutex::new(BTreeSet::new())),
836            uploads,
837        }
838    }
839
840    /// Report each accounted cache decision to `observer` as it happens.
841    #[must_use]
842    pub fn with_observer(mut self, observer: Arc<dyn AgentEventObserver>) -> Self {
843        self.observer = Some(observer);
844        self
845    }
846
847    fn emit(&self, event: impl FnOnce() -> AgentEvent) {
848        if let Some(observer) = &self.observer {
849            observer.event(event());
850        }
851    }
852
853    fn reserve_remote_download(&self, bytes: u64) -> Result<RemoteDownloadReservation> {
854        let mut current = self.remote_download_bytes.load(Ordering::Acquire);
855        loop {
856            let next = current
857                .checked_add(bytes)
858                .ok_or_else(|| eyre::eyre!("remote cache download budget overflowed"))?;
859            if next > self.remote_download_limit {
860                bail!(
861                    "remote cache download budget exceeded: {} bytes requested with {} of {} bytes already used",
862                    bytes,
863                    current,
864                    self.remote_download_limit
865                );
866            }
867            match self.remote_download_bytes.compare_exchange_weak(
868                current,
869                next,
870                Ordering::AcqRel,
871                Ordering::Acquire,
872            ) {
873                Ok(_) => {
874                    return Ok(RemoteDownloadReservation {
875                        counter: self.remote_download_bytes.clone(),
876                        reserved: bytes,
877                        committed: false,
878                    });
879                }
880                Err(observed) => current = observed,
881            }
882        }
883    }
884
885    fn reserve_remote_download_up_to(&self, requested: u64) -> Result<RemoteDownloadReservation> {
886        let mut current = self.remote_download_bytes.load(Ordering::Acquire);
887        loop {
888            let reserved = requested.min(self.remote_download_limit.saturating_sub(current));
889            let next = current + reserved;
890            match self.remote_download_bytes.compare_exchange_weak(
891                current,
892                next,
893                Ordering::AcqRel,
894                Ordering::Acquire,
895            ) {
896                Ok(_) => {
897                    return Ok(RemoteDownloadReservation {
898                        counter: self.remote_download_bytes.clone(),
899                        reserved,
900                        committed: false,
901                    });
902                }
903                Err(observed) => current = observed,
904            }
905        }
906    }
907
908    /// Load the last committed action manifest for a task into this session.
909    pub async fn begin_task(&self, task: &str) -> Result<String> {
910        self.begin_task_with_remote_errors(task, false).await
911    }
912
913    /// Load a task and finish its prefetch, surfacing remote lookup failures.
914    pub async fn prefetch_task(&self, task: &str) -> Result<String> {
915        let run = self.begin_task_with_remote_errors(task, true).await?;
916        self.wait_for_prefetches().await;
917        Ok(run)
918    }
919
920    async fn begin_task_with_remote_errors(&self, task: &str, strict: bool) -> Result<String> {
921        validate_task_identity(task)?;
922        let (remote_manifest, mut remote_etag) = if self.remote_mode.reads() {
923            match self.get_remote_task_manifest(task).await {
924                Ok(Some((manifest, etag))) => (Some(manifest), Some(etag)),
925                Ok(None) => (None, None),
926                Err(error) => {
927                    if strict {
928                        return Err(error).wrap_err_with(|| {
929                            format!("remote task action manifest lookup failed for {task}")
930                        });
931                    }
932                    self.note_remote_failure();
933                    warn!("remote task action manifest lookup failed for {task}: {error}");
934                    (None, None)
935                }
936            }
937        } else {
938            (None, None)
939        };
940        let manifest = {
941            let _write_guard = self.manifest_write_lock.lock().unwrap();
942            let _file_guard = self.lock_task_manifest(task)?;
943            let local_manifest = self.load_task_manifest(task)?;
944            let manifest = match (remote_manifest, local_manifest) {
945                (Some(remote), Some(local)) => {
946                    let (manifest, merged) = merge_remote_task_manifest(task, remote, local);
947                    if !merged {
948                        remote_etag = None;
949                    }
950                    Some(manifest)
951                }
952                (Some(remote), None) => Some(remote),
953                (None, local) => local,
954            };
955            if let Some(manifest) = &manifest {
956                self.persist_task_manifest(manifest)?;
957            }
958            manifest
959        };
960        let state = if let Some(manifest) = manifest {
961            TaskActionState {
962                manifest: task.to_string(),
963                baseline_loaded: true,
964                predictions: manifest
965                    .predictions
966                    .into_iter()
967                    .map(|prediction| (prediction.invocation.clone(), prediction))
968                    .collect(),
969                pending_predictions: BTreeMap::new(),
970                remote_etag,
971            }
972        } else {
973            TaskActionState {
974                manifest: task.to_string(),
975                baseline_loaded: true,
976                remote_etag,
977                ..TaskActionState::default()
978            }
979        };
980        let sequence = self.next_task_run.fetch_add(1, Ordering::Relaxed);
981        let run =
982            CacheDigest::blake3(format!("{task}\0{}\0{sequence}", std::process::id()).as_bytes())
983                .hash;
984        // The largest baseline rather than a sum: beginning the same task again
985        // in one session reloads the same manifest, and counting it twice would
986        // overstate what there was to match.
987        self.stats.predictions_loaded.fetch_max(
988            state.predictions.len().try_into().unwrap_or(u64::MAX),
989            Ordering::Relaxed,
990        );
991        let predictions = state.predictions.values().cloned().collect();
992        self.task_actions.lock().unwrap().insert(run.clone(), state);
993        self.spawn_prefetch_predictions(predictions);
994        Ok(run)
995    }
996
997    /// Cancel speculative downloads before the owning session exits.
998    pub async fn cancel_prefetches(&self) {
999        let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1000        for task in &tasks {
1001            task.abort();
1002        }
1003        for task in tasks {
1004            if let Err(error) = task.await
1005                && !error.is_cancelled()
1006            {
1007                warn!("remote action prefetch task failed: {error}");
1008            }
1009        }
1010    }
1011
1012    /// Publish everything a build queued, before the session stops.
1013    ///
1014    /// Store requests return once an object is durable locally, so at this point
1015    /// the remote cache may still be behind the local one. Uploads run on the
1016    /// session's runtime and are abandoned if it goes away, so a session that
1017    /// wants them published has to wait here.
1018    pub async fn wait_for_uploads(&self) {
1019        let Some(uploads) = &self.uploads else {
1020            return;
1021        };
1022        let _timer = AtomicDurationTimer::start(&self.stats.upload_drain_duration_ns);
1023        uploads.drain().await;
1024    }
1025
1026    async fn wait_for_prefetches(&self) {
1027        let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1028        for task in tasks {
1029            if let Err(error) = task.await {
1030                warn!("remote action prefetch task failed: {error}");
1031            }
1032        }
1033    }
1034
1035    /// Atomically publish the completed actions collected by a task run.
1036    pub async fn commit_task(&self, run: &str) -> Result<()> {
1037        validate_task_identity(run)?;
1038        let state = self
1039            .task_actions
1040            .lock()
1041            .unwrap()
1042            .get(run)
1043            .cloned()
1044            .ok_or_else(|| eyre::eyre!("task action manifest baseline was not loaded"))?;
1045        if !state.baseline_loaded {
1046            bail!("task action manifest baseline was not loaded");
1047        }
1048        let task = state.manifest;
1049        validate_task_identity(&task)?;
1050        let (manifest, introduced) = {
1051            let _write_guard = self.manifest_write_lock.lock().unwrap();
1052            let _file_guard = self.lock_task_manifest(&task)?;
1053            let mut predictions = self
1054                .load_task_manifest(&task)?
1055                .map(|manifest| {
1056                    manifest
1057                        .predictions
1058                        .into_iter()
1059                        .map(|prediction| (prediction.invocation.clone(), prediction))
1060                        .collect::<BTreeMap<_, _>>()
1061                })
1062                .unwrap_or_default();
1063            // Which predictions this run adds, as opposed to inherits. Only a
1064            // new one may be withheld for a failed upload: retracting an
1065            // inherited one would un-advertise a result that is plausibly still
1066            // on the server from whichever session put it there.
1067            let introduced: BTreeSet<CacheDigest> = state
1068                .pending_predictions
1069                .keys()
1070                .filter(|invocation| !predictions.contains_key(*invocation))
1071                .cloned()
1072                .collect();
1073            // Only publish predictions recorded by this run. `predictions`
1074            // also contains the baseline loaded by `begin_task`; extending
1075            // with that snapshot would overwrite newer entries committed by
1076            // another agent process after this run began.
1077            predictions.extend(state.pending_predictions);
1078            let manifest = TaskActionManifest {
1079                version: TASK_ACTION_MANIFEST_VERSION,
1080                task: task.clone(),
1081                predictions: predictions.into_values().collect(),
1082            };
1083            validate_task_manifest(&manifest, &task)?;
1084            self.persist_task_manifest(&manifest)?;
1085            (manifest, introduced)
1086        };
1087        self.task_actions.lock().unwrap().remove(run);
1088        if self.remote_mode.writes() {
1089            // A manifest advertises the actions it predicts, so it must not
1090            // reach the remote cache before the results a reader would then go
1091            // looking for -- nor name a result that never got there at all.
1092            let mut manifest = manifest;
1093            if let Some(uploads) = &self.uploads {
1094                let actions: Vec<CacheDigest> = manifest
1095                    .predictions
1096                    .iter()
1097                    .map(|prediction| prediction.action.clone())
1098                    .collect();
1099                let unpublished = uploads.wait_for_actions(&actions).await;
1100                let withheld = manifest
1101                    .predictions
1102                    .iter()
1103                    .filter(|prediction| {
1104                        introduced.contains(&prediction.invocation)
1105                            && unpublished.contains(&prediction.action)
1106                    })
1107                    .count();
1108                if withheld > 0 {
1109                    // The local manifest keeps them: this checkout can still use
1110                    // what it built, and a later session can publish it.
1111                    warn!(
1112                        "{withheld} of {} predicted actions were not published, so the remote task action manifest omits them",
1113                        manifest.predictions.len()
1114                    );
1115                    manifest.predictions.retain(|prediction| {
1116                        !(introduced.contains(&prediction.invocation)
1117                            && unpublished.contains(&prediction.action))
1118                    });
1119                }
1120            }
1121            match self
1122                .put_remote_task_manifest(&task, manifest, state.remote_etag)
1123                .await
1124            {
1125                Ok(remote_manifest) => {
1126                    let _write_guard = self.manifest_write_lock.lock().unwrap();
1127                    let reconciliation = (|| {
1128                        let _file_guard = self.lock_task_manifest(&task)?;
1129                        let manifest = match self.load_task_manifest(&task)? {
1130                            Some(local) => {
1131                                merge_remote_task_manifest(&task, remote_manifest, local).0
1132                            }
1133                            None => remote_manifest,
1134                        };
1135                        self.persist_task_manifest(&manifest)
1136                    })();
1137                    if let Err(error) = reconciliation {
1138                        warn!(
1139                            "remote task action manifest reconciliation failed for {task}: {error}"
1140                        );
1141                    }
1142                }
1143                Err(error) => {
1144                    self.note_remote_failure();
1145                    warn!("remote task action manifest upload failed for {task}: {error}");
1146                }
1147            }
1148        }
1149        Ok(())
1150    }
1151
1152    fn task_manifest_path(&self, task: &str) -> PathBuf {
1153        self.manifest_dir.join(format!("{task}.json"))
1154    }
1155
1156    fn task_manifest_lock_path(&self, task: &str) -> PathBuf {
1157        self.manifest_dir.join("locks").join(format!("{task}.lock"))
1158    }
1159
1160    fn lock_task_manifest(&self, task: &str) -> Result<fslock::LockFile> {
1161        let path = self.task_manifest_lock_path(task);
1162        fs::create_dir_all(path.parent().expect("task manifest lock has a parent"))?;
1163        let mut lock = fslock::LockFile::open(&path)?;
1164        lock.lock()?;
1165        Ok(lock)
1166    }
1167
1168    fn load_task_manifest(&self, task: &str) -> Result<Option<TaskActionManifest>> {
1169        match fs::read(self.task_manifest_path(task)) {
1170            Ok(contents) => Ok(Some(self.parse_task_manifest(task, &contents, false)?)),
1171            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
1172            Err(error) => Err(error.into()),
1173        }
1174    }
1175
1176    fn parse_task_manifest(
1177        &self,
1178        task: &str,
1179        contents: &[u8],
1180        require_canonical: bool,
1181    ) -> Result<TaskActionManifest> {
1182        let manifest: TaskActionManifest = serde_json::from_slice(contents)?;
1183        validate_task_manifest(&manifest, task)?;
1184        if require_canonical && canonical_json(&manifest)? != contents {
1185            bail!("task action manifest is not canonical JSON");
1186        }
1187        Ok(manifest)
1188    }
1189
1190    fn task_manifest_selector(task: &str) -> Result<(Vec<u8>, CacheDigest)> {
1191        TaskActionManifest::selector(task)
1192    }
1193
1194    fn persist_task_manifest(&self, manifest: &TaskActionManifest) -> Result<()> {
1195        let bytes = canonical_json(manifest)?;
1196        fs::create_dir_all(self.manifest_dir.as_path())?;
1197        let mut temporary = tempfile::NamedTempFile::new_in(self.manifest_dir.as_path())?;
1198        std::io::Write::write_all(temporary.as_file_mut(), &bytes)?;
1199        temporary.as_file_mut().sync_all()?;
1200        temporary
1201            .persist(self.task_manifest_path(&manifest.task))
1202            .map_err(|error| error.error)?;
1203        Ok(())
1204    }
1205
1206    async fn get_remote_task_manifest(
1207        &self,
1208        task: &str,
1209    ) -> Result<Option<(TaskActionManifest, String)>> {
1210        let Some(remote) = &self.remote else {
1211            return Ok(None);
1212        };
1213        let (_, selector) = Self::task_manifest_selector(task)?;
1214        let _permit = self.remote_transfers.acquire().await?;
1215        self.stats
1216            .remote_manifest_lookups
1217            .fetch_add(1, Ordering::Relaxed);
1218        let _timer = AtomicDurationTimer::start(&self.stats.remote_manifest_lookup_duration_ns);
1219        let Some(remote_manifest) = remote.get_action_manifest(&selector).await? else {
1220            return Ok(None);
1221        };
1222        let manifest = self.parse_task_manifest(task, &remote_manifest.bytes, true)?;
1223        Ok(Some((manifest, remote_manifest.etag)))
1224    }
1225
1226    async fn put_remote_task_manifest(
1227        &self,
1228        task: &str,
1229        mut manifest: TaskActionManifest,
1230        mut expected_etag: Option<String>,
1231    ) -> Result<TaskActionManifest> {
1232        let Some(remote) = &self.remote else {
1233            return Ok(manifest);
1234        };
1235        let (_, selector) = Self::task_manifest_selector(task)?;
1236        for _ in 0..4 {
1237            let bytes = canonical_json(&manifest)?;
1238            let outcome = {
1239                let _permit = self.remote_transfers.acquire().await?;
1240                remote
1241                    .put_action_manifest(&selector, &bytes, expected_etag.as_deref())
1242                    .await?
1243            };
1244            match outcome {
1245                ManifestPutOutcome::Stored => return Ok(manifest),
1246                ManifestPutOutcome::PreconditionFailed => {
1247                    let Some((remote_manifest, etag)) = self.get_remote_task_manifest(task).await?
1248                    else {
1249                        expected_etag = None;
1250                        continue;
1251                    };
1252                    manifest = merge_task_manifests(task, Some(remote_manifest), manifest)?;
1253                    expected_etag = Some(etag);
1254                }
1255            }
1256        }
1257        bail!("remote task action manifest changed too frequently")
1258    }
1259
1260    /// Return a snapshot of this session's cache activity.
1261    pub fn stats(&self) -> AgentStats {
1262        AgentStats {
1263            session_duration_ns: 0,
1264            lookups: self.stats.lookups.load(Ordering::Relaxed),
1265            unconsulted: self.stats.unconsulted.load(Ordering::Relaxed),
1266            hits: self.stats.hits.load(Ordering::Relaxed),
1267            stores: self.stats.stores.load(Ordering::Relaxed),
1268            stored_bytes: self.stats.stored_bytes.load(Ordering::Relaxed),
1269            verifications: self.stats.verifications.load(Ordering::Relaxed),
1270            divergences: self.stats.divergences.load(Ordering::Relaxed),
1271            downloaded_bytes: self.stats.downloaded_bytes.load(Ordering::Relaxed),
1272            uploaded_bytes: self.stats.uploaded_bytes.load(Ordering::Relaxed),
1273            background_uploads: self.stats.background_uploads.load(Ordering::Relaxed),
1274            background_upload_failures: self
1275                .stats
1276                .background_upload_failures
1277                .load(Ordering::Relaxed),
1278            remote_blob_pack_uploads: self.stats.remote_blob_pack_uploads.load(Ordering::Relaxed),
1279            remote_blob_pack_upload_blobs: self
1280                .stats
1281                .remote_blob_pack_upload_blobs
1282                .load(Ordering::Relaxed),
1283            upload_drain_duration_ns: self.stats.upload_drain_duration_ns.load(Ordering::Relaxed),
1284            prefetched_actions: self.stats.prefetched_actions.load(Ordering::Relaxed),
1285            predictions_loaded: self.stats.predictions_loaded.load(Ordering::Relaxed),
1286            bypasses: self.stats.bypasses.lock().unwrap().clone(),
1287            avoided_compiler_duration_ns: self
1288                .stats
1289                .avoided_compiler_duration_ns
1290                .load(Ordering::Relaxed),
1291            compiler: self.stats.compiler.lock().unwrap().clone(),
1292            slow_compilations: self.stats.slow_compilations.lock().unwrap().clone(),
1293            remote_failures: self.stats.remote_failures.load(Ordering::Relaxed),
1294            remote_manifest_lookups: self.stats.remote_manifest_lookups.load(Ordering::Relaxed),
1295            remote_manifest_lookup_duration_ns: self
1296                .stats
1297                .remote_manifest_lookup_duration_ns
1298                .load(Ordering::Relaxed),
1299            remote_action_lookups: self.stats.remote_action_lookups.load(Ordering::Relaxed),
1300            remote_action_lookup_duration_ns: self
1301                .stats
1302                .remote_action_lookup_duration_ns
1303                .load(Ordering::Relaxed),
1304            remote_blob_requests: self.stats.remote_blob_requests.load(Ordering::Relaxed),
1305            remote_blob_pack_requests: self.stats.remote_blob_pack_requests.load(Ordering::Relaxed),
1306            remote_blob_pack_blobs: self.stats.remote_blob_pack_blobs.load(Ordering::Relaxed),
1307            remote_blob_transfer_duration_ns: self
1308                .stats
1309                .remote_blob_transfer_duration_ns
1310                .load(Ordering::Relaxed),
1311            local_cas_write_duration_ns: self
1312                .stats
1313                .local_cas_write_duration_ns
1314                .load(Ordering::Relaxed),
1315            prefetch_runs: self.stats.prefetch_runs.load(Ordering::Relaxed),
1316            prefetch_duration_ns: self.stats.prefetch_duration_ns.load(Ordering::Relaxed),
1317            materialization_duration_ns: self
1318                .stats
1319                .materialization_duration_ns
1320                .load(Ordering::Relaxed),
1321            restored_output_files: self.stats.restored_output_files.load(Ordering::Relaxed),
1322            restored_output_bytes: self.stats.restored_output_bytes.load(Ordering::Relaxed),
1323            reflinked_output_files: self.stats.reflinked_output_files.load(Ordering::Relaxed),
1324            reflinked_output_bytes: self.stats.reflinked_output_bytes.load(Ordering::Relaxed),
1325            copied_output_files: self.stats.copied_output_files.load(Ordering::Relaxed),
1326            copied_output_bytes: self.stats.copied_output_bytes.load(Ordering::Relaxed),
1327        }
1328    }
1329
1330    /// Handle requests without a transport connection.
1331    ///
1332    /// Persistent wrappers use this entry point when Cargo invokes mbx outside
1333    /// an orchestrated session. It intentionally has the same response
1334    /// semantics as [`Self::handle_connection`], while leaving framing and the
1335    /// version handshake to callers that actually cross a process boundary.
1336    pub async fn handle_requests(
1337        &self,
1338        requests: impl IntoIterator<Item = AgentRequest>,
1339    ) -> Vec<AgentResponse> {
1340        let mut connection = ConnectionUploads::default();
1341        let mut responses = Vec::new();
1342        for request in requests {
1343            responses.push(self.respond_on(request, &mut connection).await);
1344        }
1345        responses
1346    }
1347
1348    fn write_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1349        Self::digest_lock(&self.write_locks, digest)
1350    }
1351
1352    fn action_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1353        Self::digest_lock(&self.action_locks, digest)
1354    }
1355
1356    fn digest_lock(
1357        locks: &Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>,
1358        digest: &CacheDigest,
1359    ) -> Arc<tokio::sync::Mutex<()>> {
1360        let mut locks = locks.lock().unwrap();
1361        locks.retain(|_, lock| lock.strong_count() > 0);
1362        if let Some(lock) = locks.get(digest).and_then(Weak::upgrade) {
1363            return lock;
1364        }
1365        let lock = Arc::new(tokio::sync::Mutex::new(()));
1366        locks.insert(digest.clone(), Arc::downgrade(&lock));
1367        lock
1368    }
1369
1370    fn spawn_prefetch_predictions(&self, predictions: Vec<ActionPrediction>) {
1371        if predictions.is_empty() || !self.remote_mode.reads() || self.remote.is_none() {
1372            return;
1373        }
1374        let agent = self.clone();
1375        let task = tokio::spawn(async move {
1376            agent.prefetch_predictions(predictions.iter()).await;
1377        });
1378        self.prefetch_tasks.lock().unwrap().push(task);
1379    }
1380
1381    async fn prefetch_predictions<'a>(
1382        &self,
1383        predictions: impl Iterator<Item = &'a ActionPrediction>,
1384    ) {
1385        if !self.remote_mode.reads() || self.remote.is_none() {
1386            return;
1387        }
1388        self.stats.prefetch_runs.fetch_add(1, Ordering::Relaxed);
1389        let _timer = AtomicDurationTimer::start(&self.stats.prefetch_duration_ns);
1390        let mut actions = BTreeMap::new();
1391        for prediction in predictions {
1392            actions
1393                .entry(prediction.action.clone())
1394                .or_insert_with(|| prediction.adapter.clone());
1395        }
1396        // One request per predicted action is the bulk of a prefetch's latency on
1397        // a large workspace, so ask for them together where the server allows it.
1398        match self.prefetch_action_batches(&actions).await {
1399            Ok(true) => return,
1400            Ok(false) => {}
1401            Err(error) => {
1402                self.note_remote_failure();
1403                warn!("remote action batch lookup failed: {error}");
1404            }
1405        }
1406        let mut actions = actions.into_iter();
1407        let mut tasks = tokio::task::JoinSet::new();
1408        for _ in 0..MAX_PREFETCH_TRANSFERS {
1409            let Some((action, adapter)) = actions.next() else {
1410                break;
1411            };
1412            let agent = self.clone();
1413            tasks.spawn(async move { agent.resolve_prefetch_action(action, adapter).await });
1414        }
1415        let mut resolved = Vec::new();
1416        while !tasks.is_empty() {
1417            let result = if resolved.is_empty() {
1418                tasks.join_next().await
1419            } else {
1420                match tokio::time::timeout(PREFETCH_ACTION_BATCH_DELAY, tasks.join_next()).await {
1421                    Ok(result) => result,
1422                    Err(_) => {
1423                        self.prefetch_resolved_actions(std::mem::take(&mut resolved))
1424                            .await;
1425                        continue;
1426                    }
1427                }
1428            };
1429            let Some(result) = result else {
1430                break;
1431            };
1432            match result {
1433                Ok(Ok(Some(action))) => resolved.push(action),
1434                Ok(Ok(None)) => {}
1435                Ok(Err(error)) => {
1436                    self.note_remote_failure();
1437                    warn!("remote action prefetch failed: {error}");
1438                }
1439                Err(error) => {
1440                    self.note_remote_failure();
1441                    warn!("remote action prefetch task failed: {error}");
1442                }
1443            }
1444            if let Some((action, adapter)) = actions.next() {
1445                let agent = self.clone();
1446                tasks.spawn(async move { agent.resolve_prefetch_action(action, adapter).await });
1447            }
1448            if resolved.len() == MAX_PREFETCH_ACTION_BATCH {
1449                self.prefetch_resolved_actions(std::mem::take(&mut resolved))
1450                    .await;
1451            }
1452        }
1453        if !resolved.is_empty() {
1454            self.prefetch_resolved_actions(resolved).await;
1455        }
1456    }
1457
1458    /// Resolve predicted actions in batched lookups, staging what comes back.
1459    ///
1460    /// Returns whether the batch extension answered. A server without it, or one
1461    /// that stops answering part way through, leaves the per-action path to
1462    /// resolve whatever is left -- the results already staged here are memoized,
1463    /// so nothing is looked up twice.
1464    async fn prefetch_action_batches(
1465        &self,
1466        actions: &BTreeMap<CacheDigest, String>,
1467    ) -> Result<bool> {
1468        let Some(remote) = self.remote.as_deref() else {
1469            return Ok(false);
1470        };
1471        let Some(limit) = remote.action_batch_limit().await? else {
1472            return Ok(false);
1473        };
1474        let wanted: Vec<CacheDigest> = actions
1475            .keys()
1476            .filter(|action| !self.action_is_staged(action))
1477            .cloned()
1478            .collect();
1479        if wanted.is_empty() {
1480            return Ok(true);
1481        }
1482        let mut chunks: Vec<Vec<CacheDigest>> = Vec::new();
1483        for chunk in wanted.chunks(limit) {
1484            chunks.push(chunk.to_vec());
1485        }
1486        let mut lookups = stream::iter(chunks)
1487            .map(|chunk| async move {
1488                let _prefetch_permit = self.prefetch_transfers.acquire().await?;
1489                let _permit = self.remote_transfers.acquire().await?;
1490                self.stats
1491                    .remote_action_lookups
1492                    .fetch_add(1, Ordering::Relaxed);
1493                let _timer =
1494                    AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
1495                remote.get_action_results(&chunk).await
1496            })
1497            .buffer_unordered(MAX_PREFETCH_BATCH_LOOKUPS);
1498        let mut answered = true;
1499        while let Some(lookup) = lookups.next().await {
1500            let results = match lookup {
1501                Ok(Some(results)) => results,
1502                // Advertised but unavailable. What is left falls back.
1503                Ok(None) => {
1504                    answered = false;
1505                    continue;
1506                }
1507                Err(error) => {
1508                    self.note_remote_failure();
1509                    warn!("remote action batch lookup failed: {error}");
1510                    answered = false;
1511                    continue;
1512                }
1513            };
1514            let mut resolved = Vec::with_capacity(results.len());
1515            for result in results {
1516                let Some(adapter) = actions.get(&result.action).cloned() else {
1517                    continue;
1518                };
1519                let lock = self.action_lock(&result.action);
1520                let _guard = lock.lock().await;
1521                if self.actions.find(&result.action)?.is_some() {
1522                    continue;
1523                }
1524                self.pending_remote_actions
1525                    .lock()
1526                    .unwrap()
1527                    .insert(result.action.clone(), result.clone());
1528                resolved.push(PrefetchedAction { adapter, result });
1529            }
1530            while !resolved.is_empty() {
1531                let wave = resolved
1532                    .drain(..resolved.len().min(MAX_PREFETCH_ACTION_BATCH))
1533                    .collect();
1534                self.prefetch_resolved_actions(wave).await;
1535            }
1536        }
1537        Ok(answered)
1538    }
1539
1540    /// Whether an action's result is already local or already looked up.
1541    fn action_is_staged(&self, action: &CacheDigest) -> bool {
1542        if self
1543            .pending_remote_actions
1544            .lock()
1545            .unwrap()
1546            .contains_key(action)
1547        {
1548            return true;
1549        }
1550        self.actions.find(action).is_ok_and(|found| found.is_some())
1551    }
1552
1553    #[cfg(test)]
1554    async fn prefetch_action(&self, action: CacheDigest, adapter: String) -> Result<()> {
1555        if let Some(action) = self.resolve_prefetch_action(action, adapter).await? {
1556            self.prefetch_resolved_actions(vec![action]).await;
1557        }
1558        Ok(())
1559    }
1560
1561    async fn resolve_prefetch_action(
1562        &self,
1563        action: CacheDigest,
1564        adapter: String,
1565    ) -> Result<Option<PrefetchedAction>> {
1566        let remote = self
1567            .remote
1568            .as_ref()
1569            .ok_or_else(|| eyre::eyre!("remote cache is not configured"))?;
1570        let result = {
1571            let lock = self.action_lock(&action);
1572            let _guard = lock.lock().await;
1573            if self.actions.find(&action)?.is_some() {
1574                return Ok(None);
1575            }
1576            if let Some(result) = self
1577                .pending_remote_actions
1578                .lock()
1579                .unwrap()
1580                .get(&action)
1581                .cloned()
1582            {
1583                result
1584            } else {
1585                let _prefetch_permit = self.prefetch_transfers.acquire().await?;
1586                let result = {
1587                    let _permit = self.remote_transfers.acquire().await?;
1588                    self.get_remote_action_result(remote, &action).await?
1589                };
1590                let Some(result) = result else {
1591                    return Ok(None);
1592                };
1593                self.pending_remote_actions
1594                    .lock()
1595                    .unwrap()
1596                    .insert(action.clone(), result.clone());
1597                result
1598            }
1599        };
1600        Ok(Some(PrefetchedAction { adapter, result }))
1601    }
1602
1603    fn prefetch_resolved_actions(&self, actions: Vec<PrefetchedAction>) -> BoxFuture<'_, ()> {
1604        self.prefetch_resolved_actions_inner(actions).boxed()
1605    }
1606
1607    async fn prefetch_resolved_actions_inner(&self, actions: Vec<PrefetchedAction>) {
1608        let Some(remote) = self.remote.as_deref() else {
1609            return;
1610        };
1611        if actions.is_empty() {
1612            return;
1613        }
1614
1615        let mut top_level = BTreeMap::new();
1616        for action in &actions {
1617            for digest in [
1618                Some(&action.result.action),
1619                action.result.metadata.as_ref(),
1620                action.result.output_root.as_ref(),
1621            ]
1622            .into_iter()
1623            .flatten()
1624            {
1625                top_level.insert(digest.clone(), ());
1626            }
1627        }
1628        let mut verified = self
1629            .fetch_remote_blobs(
1630                remote,
1631                top_level.into_keys().collect(),
1632                Some(&self.prefetch_transfers),
1633            )
1634            .await;
1635
1636        let mut next = BTreeMap::new();
1637        let mut pending_directories = BTreeMap::new();
1638        let mut parsed_directories = BTreeMap::new();
1639        let mut rustc_metadata = BTreeMap::new();
1640        for action in &actions {
1641            if action.adapter == "rustc"
1642                && let Some(metadata_digest) = &action.result.metadata
1643            {
1644                match verified
1645                    .get(metadata_digest)
1646                    .ok_or_else(|| eyre::eyre!("remote rustc action metadata is missing"))
1647                    .and_then(|path| Self::parse_rustc_metadata(path))
1648                {
1649                    Ok(metadata) => {
1650                        queue_prefetch_digest(&verified, &mut next, metadata.stdout.clone());
1651                        queue_prefetch_digest(&verified, &mut next, metadata.stderr.clone());
1652                        rustc_metadata.insert(metadata_digest.clone(), metadata);
1653                    }
1654                    Err(error) => warn!(
1655                        "remote rustc action metadata prefetch failed for {}: {error}",
1656                        action.result.action.hash
1657                    ),
1658                }
1659            }
1660            if let Some(output_root) = &action.result.output_root {
1661                pending_directories.insert(output_root.clone(), ());
1662            }
1663        }
1664
1665        let mut seen_directories = BTreeMap::new();
1666        loop {
1667            let mut following = BTreeMap::new();
1668            let mut directory_limit_exceeded = false;
1669            for digest in pending_directories.into_keys() {
1670                following.remove(&digest);
1671                if seen_directories.insert(digest.clone(), ()).is_some() {
1672                    continue;
1673                }
1674                if seen_directories.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
1675                    warn!("remote action output tree is too large to prefetch");
1676                    following.clear();
1677                    break;
1678                }
1679                match verified
1680                    .get(&digest)
1681                    .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))
1682                    .and_then(|path| Self::parse_cache_directory(path))
1683                {
1684                    Ok(directory) => {
1685                        for file in &directory.files {
1686                            queue_prefetch_digest(&verified, &mut next, file.digest.clone());
1687                            if next.len() >= MAX_PREFETCH_OBJECTS_PER_WAVE {
1688                                self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1689                                    .await;
1690                            }
1691                        }
1692                        for child in &directory.directories {
1693                            if !queue_prefetch_directory(
1694                                &seen_directories,
1695                                &mut following,
1696                                child.digest.clone(),
1697                                MAX_PREFETCH_DIRECTORY_OBJECTS,
1698                            ) {
1699                                warn!("remote action output tree is too large to prefetch");
1700                                directory_limit_exceeded = true;
1701                                break;
1702                            }
1703                            queue_prefetch_digest(&verified, &mut next, child.digest.clone());
1704                            if next.len() >= MAX_PREFETCH_OBJECTS_PER_WAVE {
1705                                self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1706                                    .await;
1707                            }
1708                        }
1709                        parsed_directories.insert(digest, directory);
1710                    }
1711                    Err(error) => warn!(
1712                        "remote action output directory prefetch failed for {}: {error}",
1713                        digest.hash
1714                    ),
1715                }
1716                if directory_limit_exceeded {
1717                    following.clear();
1718                    break;
1719                }
1720            }
1721            self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1722                .await;
1723            if following.is_empty() {
1724                break;
1725            }
1726            pending_directories = following;
1727        }
1728
1729        for action in actions {
1730            match Self::validate_prefetched_action(
1731                &action,
1732                &verified,
1733                &rustc_metadata,
1734                &parsed_directories,
1735            ) {
1736                Ok(()) => {
1737                    if let Err(error) = self.actions.store(&action.result) {
1738                        warn!(
1739                            "remote action prefetch could not publish {}: {error}",
1740                            action.result.action.hash
1741                        );
1742                        continue;
1743                    }
1744                    self.pending_remote_actions
1745                        .lock()
1746                        .unwrap()
1747                        .remove(&action.result.action);
1748                    self.stats
1749                        .prefetched_actions
1750                        .fetch_add(1, Ordering::Relaxed);
1751                }
1752                Err(error) => warn!(
1753                    "remote action prefetch was incomplete for {}: {error}",
1754                    action.result.action.hash
1755                ),
1756            }
1757        }
1758    }
1759
1760    async fn flush_prefetch_digest_batch(
1761        &self,
1762        remote: &RemoteCacheClient,
1763        verified: &mut BTreeMap<CacheDigest, PathBuf>,
1764        pending: &mut BTreeMap<CacheDigest, ()>,
1765    ) {
1766        if pending.is_empty() {
1767            return;
1768        }
1769        let digests = std::mem::take(pending).into_keys().collect();
1770        verified.extend(
1771            self.fetch_remote_blobs(remote, digests, Some(&self.prefetch_transfers))
1772                .await,
1773        );
1774    }
1775
1776    async fn fetch_remote_blobs(
1777        &self,
1778        remote: &RemoteCacheClient,
1779        digests: Vec<CacheDigest>,
1780        prefetch_limit: Option<&tokio::sync::Semaphore>,
1781    ) -> BTreeMap<CacheDigest, PathBuf> {
1782        let mut verified = BTreeMap::new();
1783        let mut missing = BTreeMap::new();
1784        for digest in digests {
1785            match self.find_verified_blob(&digest) {
1786                Ok(Some(path)) => {
1787                    verified.insert(digest, path);
1788                }
1789                Ok(None) => {
1790                    missing.insert(digest, ());
1791                }
1792                Err(error) => warn!(
1793                    "local cache blob lookup failed for {}: {error}",
1794                    digest.hash
1795                ),
1796            }
1797        }
1798        if missing.is_empty() {
1799            return verified;
1800        }
1801
1802        let mut pack_candidates = missing.clone();
1803        while !pack_candidates.is_empty() {
1804            let candidates = match blob_pack_chunk(
1805                &pack_candidates.keys().cloned().collect::<Vec<_>>(),
1806                BlobPackLimits {
1807                    max_items: MAX_STAGED_BLOB_PACK_ITEMS,
1808                    max_bytes: MAX_STAGED_BLOB_PACK_BYTES,
1809                },
1810            ) {
1811                Ok(candidates) if !candidates.is_empty() => candidates,
1812                Ok(_) => break,
1813                Err(error) => {
1814                    warn!("remote cache blob pack skipped: {error}");
1815                    break;
1816                }
1817            };
1818            // A pack and an individual fetch share these per-digest locks. Hold
1819            // them through ingestion so overlapping prefetch and foreground
1820            // requests cannot download or charge the same object twice.
1821            let mut pack_guards = BTreeMap::new();
1822            for digest in candidates {
1823                let guard = self.write_lock(&digest).lock_owned().await;
1824                match self.find_verified_blob(&digest) {
1825                    Ok(Some(path)) => {
1826                        pack_candidates.remove(&digest);
1827                        missing.remove(&digest);
1828                        verified.insert(digest, path);
1829                    }
1830                    Ok(None) => {
1831                        pack_guards.insert(digest, guard);
1832                    }
1833                    Err(error) => {
1834                        warn!(
1835                            "local cache blob lookup failed for {}: {error}",
1836                            digest.hash
1837                        );
1838                        pack_guards.insert(digest, guard);
1839                    }
1840                }
1841            }
1842            let requested = pack_guards.keys().cloned().collect::<Vec<_>>();
1843            if requested.is_empty() {
1844                continue;
1845            }
1846            let requested_bytes = requested
1847                .iter()
1848                .fold(0_u64, |total, digest| total.saturating_add(digest.size));
1849            let pack_reservation = match self
1850                .reserve_remote_download_up_to(requested_bytes.min(MAX_STAGED_BLOB_PACK_BYTES))
1851            {
1852                Ok(reservation) if reservation.bytes() > 0 => reservation,
1853                Ok(_) => break,
1854                Err(error) => {
1855                    warn!("remote cache blob pack skipped: {error}");
1856                    break;
1857                }
1858            };
1859            let (pack, transfer_duration_ns) = {
1860                let _prefetch_permit = match prefetch_limit {
1861                    Some(limit) => match limit.acquire().await {
1862                        Ok(permit) => Some(permit),
1863                        Err(error) => {
1864                            warn!(
1865                                "remote cache blob pack could not acquire prefetch limit: {error}"
1866                            );
1867                            break;
1868                        }
1869                    },
1870                    None => None,
1871                };
1872                let _transfer_permit = match self.remote_transfers.acquire().await {
1873                    Ok(permit) => permit,
1874                    Err(error) => {
1875                        warn!("remote cache blob pack could not acquire transfer limit: {error}");
1876                        break;
1877                    }
1878                };
1879                let transfer_started = Instant::now();
1880                let pack = remote
1881                    .get_blob_pack_with_limit(
1882                        &requested,
1883                        self.remote_staging_dir.as_path(),
1884                        pack_reservation.bytes(),
1885                    )
1886                    .await;
1887                (pack, duration_ns(transfer_started))
1888            };
1889            let pack = match pack {
1890                Ok(Some(pack)) => pack,
1891                Ok(None) => break,
1892                Err(error) => {
1893                    atomic_saturating_add(
1894                        &self.stats.remote_blob_transfer_duration_ns,
1895                        transfer_duration_ns,
1896                    );
1897                    warn!(
1898                        "remote cache blob pack failed; falling back to individual blobs: {error}"
1899                    );
1900                    break;
1901                }
1902            };
1903            pack_reservation.commit(pack.payload_bytes);
1904            atomic_saturating_add(
1905                &self.stats.remote_blob_transfer_duration_ns,
1906                transfer_duration_ns,
1907            );
1908            atomic_saturating_add(&self.stats.remote_blob_pack_requests, pack.requests);
1909            atomic_saturating_add(&self.stats.remote_blob_pack_blobs, pack.blob_count);
1910            atomic_saturating_add(&self.stats.downloaded_bytes, pack.payload_bytes);
1911            if pack.requested.is_empty() {
1912                // The server's negotiated cap can be smaller than the local
1913                // staging cap used to select this locked slice. Fall back for
1914                // this slice, but keep packing later candidates that may fit.
1915                for digest in &requested {
1916                    pack_candidates.remove(digest);
1917                }
1918                continue;
1919            }
1920            for digest in &pack.requested {
1921                pack_candidates.remove(digest);
1922            }
1923            let mut ingests = stream::iter(pack.blobs.into_iter().map(|(digest, source)| {
1924                let digest_for_result = digest.clone();
1925                let guard = pack_guards
1926                    .remove(&digest)
1927                    .expect("requested packed blob has a write lock");
1928                async move {
1929                    (
1930                        digest_for_result,
1931                        self.ingest_packed_blob(digest, source, guard).await,
1932                    )
1933                }
1934            }))
1935            .buffer_unordered(MAX_PREFETCH_TRANSFERS);
1936            while let Some((digest, result)) = ingests.next().await {
1937                match result {
1938                    Ok(path) => {
1939                        missing.remove(&digest);
1940                        verified.insert(digest, path);
1941                    }
1942                    Err(error) => warn!(
1943                        "remote cache packed blob ingest failed for {}: {error}",
1944                        digest.hash
1945                    ),
1946                }
1947            }
1948        }
1949
1950        let mut transfers = stream::iter(missing.into_keys().map(|digest| {
1951            let digest_for_result = digest.clone();
1952            async move {
1953                (
1954                    digest_for_result,
1955                    self.fetch_remote_blob_with_limit(remote, &digest, prefetch_limit)
1956                        .await,
1957                )
1958            }
1959        }))
1960        .buffer_unordered(MAX_PREFETCH_TRANSFERS);
1961        while let Some((digest, result)) = transfers.next().await {
1962            match result {
1963                Ok(path) => {
1964                    verified.insert(digest, path);
1965                }
1966                Err(error) => warn!(
1967                    "remote cache blob prefetch failed for {}: {error}",
1968                    digest.hash
1969                ),
1970            }
1971        }
1972        verified
1973    }
1974
1975    async fn ingest_packed_blob(
1976        &self,
1977        digest: CacheDigest,
1978        source: PathBuf,
1979        _guard: tokio::sync::OwnedMutexGuard<()>,
1980    ) -> Result<PathBuf> {
1981        let digest_size = digest.size;
1982        let agent = self.clone();
1983        let (path, stored, cas_duration_ns) = tokio::task::spawn_blocking(move || {
1984            if let Some(path) = agent.find_verified_blob(&digest)? {
1985                return Ok::<_, eyre::Report>((path, false, 0));
1986            }
1987            let cas_started = Instant::now();
1988            let path = agent.cas.store_verified_file(&digest, &source)?;
1989            let cas_duration_ns = duration_ns(cas_started);
1990            agent.remember_verified_blob(&digest, &path);
1991            Ok((path, true, cas_duration_ns))
1992        })
1993        .await??;
1994        atomic_saturating_add(&self.stats.local_cas_write_duration_ns, cas_duration_ns);
1995        if stored {
1996            self.stats.stores.fetch_add(1, Ordering::Relaxed);
1997            atomic_saturating_add(&self.stats.stored_bytes, digest_size);
1998        }
1999        Ok(path)
2000    }
2001
2002    fn parse_rustc_metadata(path: &Path) -> Result<RustcMetadata> {
2003        let bytes = fs::read(path)?;
2004        let metadata: RustcMetadata = serde_json::from_slice(&bytes)?;
2005        if metadata.version != 1 || metadata.kind != "rustc" || canonical_json(&metadata)? != bytes
2006        {
2007            bail!("remote rustc action metadata is invalid");
2008        }
2009        Ok(metadata)
2010    }
2011
2012    fn parse_cache_directory(path: &Path) -> Result<CacheDirectory> {
2013        let bytes = fs::read(path)?;
2014        let directory: CacheDirectory = serde_json::from_slice(&bytes)?;
2015        if directory.version != 1 || canonical_json(&directory)? != bytes {
2016            bail!("remote action output directory is invalid");
2017        }
2018        Ok(directory)
2019    }
2020
2021    #[cfg(test)]
2022    fn load_cache_directory(&self, digest: &CacheDigest) -> Result<CacheDirectory> {
2023        let path = self
2024            .find_verified_blob(digest)?
2025            .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))?;
2026        Self::parse_cache_directory(&path)
2027    }
2028
2029    fn validate_prefetched_action(
2030        action: &PrefetchedAction,
2031        verified: &BTreeMap<CacheDigest, PathBuf>,
2032        rustc_metadata: &BTreeMap<CacheDigest, RustcMetadata>,
2033        directories: &BTreeMap<CacheDigest, CacheDirectory>,
2034    ) -> Result<()> {
2035        if !verified.contains_key(&action.result.action) {
2036            bail!("remote action descriptor is missing");
2037        }
2038        if let Some(metadata) = &action.result.metadata {
2039            if action.adapter == "rustc" {
2040                let metadata = rustc_metadata
2041                    .get(metadata)
2042                    .ok_or_else(|| eyre::eyre!("remote rustc action metadata is missing"))?;
2043                for digest in [&metadata.stdout, &metadata.stderr] {
2044                    if !verified.contains_key(digest) {
2045                        bail!("remote rustc action diagnostic blob is missing");
2046                    }
2047                }
2048            } else if !verified.contains_key(metadata) {
2049                bail!("remote action metadata is missing");
2050            }
2051        }
2052        let mut pending = action
2053            .result
2054            .output_root
2055            .iter()
2056            .cloned()
2057            .collect::<Vec<_>>();
2058        let mut seen = BTreeMap::new();
2059        while let Some(digest) = pending.pop() {
2060            if seen.insert(digest.clone(), ()).is_some() {
2061                continue;
2062            }
2063            if seen.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
2064                bail!("remote action output tree is too large");
2065            }
2066            let directory = directories
2067                .get(&digest)
2068                .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))?;
2069            for file in &directory.files {
2070                if !verified.contains_key(&file.digest) {
2071                    bail!("remote action output file is missing");
2072                }
2073            }
2074            pending.extend(
2075                directory
2076                    .directories
2077                    .iter()
2078                    .map(|directory| directory.digest.clone()),
2079            );
2080        }
2081        Ok(())
2082    }
2083
2084    #[cfg(test)]
2085    async fn prefetch_output_tree(
2086        &self,
2087        remote: &RemoteCacheClient,
2088        output_root: &CacheDigest,
2089    ) -> Result<()> {
2090        let mut pending = vec![output_root.clone()];
2091        let mut seen = BTreeMap::new();
2092        while let Some(digest) = pending.pop() {
2093            if seen.insert(digest.clone(), ()).is_some() {
2094                continue;
2095            }
2096            if seen.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
2097                bail!("remote action output tree is too large");
2098            }
2099            self.fetch_remote_blob_with_limit(remote, &digest, Some(&self.prefetch_transfers))
2100                .await?;
2101            let directory = self.load_cache_directory(&digest)?;
2102            let mut transfers = stream::iter(directory.files.into_iter().map(|file| async move {
2103                self.fetch_remote_blob_with_limit(
2104                    remote,
2105                    &file.digest,
2106                    Some(&self.prefetch_transfers),
2107                )
2108                .await
2109                .map(|_| ())
2110            }))
2111            .buffer_unordered(MAX_PREFETCH_TRANSFERS);
2112            while let Some(result) = transfers.next().await {
2113                result?;
2114            }
2115            pending.extend(
2116                directory
2117                    .directories
2118                    .into_iter()
2119                    .map(|directory| directory.digest),
2120            );
2121        }
2122        Ok(())
2123    }
2124
2125    async fn fetch_remote_blob(
2126        &self,
2127        remote: &RemoteCacheClient,
2128        digest: &CacheDigest,
2129    ) -> Result<PathBuf> {
2130        self.fetch_remote_blob_with_limit(remote, digest, None)
2131            .await
2132    }
2133
2134    async fn fetch_remote_blob_with_limit(
2135        &self,
2136        remote: &RemoteCacheClient,
2137        digest: &CacheDigest,
2138        prefetch_limit: Option<&tokio::sync::Semaphore>,
2139    ) -> Result<PathBuf> {
2140        let lock = self.write_lock(digest);
2141        let _guard = lock.lock().await;
2142        if let Some(path) = self.find_verified_blob(digest)? {
2143            return Ok(path);
2144        }
2145        let _prefetch_permit = match prefetch_limit {
2146            Some(limit) => Some(limit.acquire().await?),
2147            None => None,
2148        };
2149        let _permit = self.remote_transfers.acquire().await?;
2150        let reservation = self.reserve_remote_download(digest.size)?;
2151        self.stats
2152            .remote_blob_requests
2153            .fetch_add(1, Ordering::Relaxed);
2154        let transfer_timer =
2155            AtomicDurationTimer::start(&self.stats.remote_blob_transfer_duration_ns);
2156        let temporary = remote
2157            .get_blob_file(digest, self.remote_staging_dir.as_path())
2158            .await?;
2159        drop(transfer_timer);
2160        let _cas_timer = AtomicDurationTimer::start(&self.stats.local_cas_write_duration_ns);
2161        let path = self.cas.store_verified_file(digest, temporary.path())?;
2162        reservation.commit(digest.size);
2163        self.remember_verified_blob(digest, &path);
2164        self.stats.stores.fetch_add(1, Ordering::Relaxed);
2165        self.stats
2166            .stored_bytes
2167            .fetch_add(digest.size, Ordering::Relaxed);
2168        self.stats
2169            .downloaded_bytes
2170            .fetch_add(digest.size, Ordering::Relaxed);
2171        Ok(path)
2172    }
2173
2174    /// Answer one request outside any connection.
2175    ///
2176    /// An upload queued this way has no sibling requests to order against, so
2177    /// the queue falls back to treating every blob it holds as a prerequisite.
2178    #[cfg(test)]
2179    async fn respond(&self, request: AgentRequest) -> AgentResponse {
2180        self.respond_on(request, &mut ConnectionUploads::default())
2181            .await
2182    }
2183
2184    async fn respond_on(
2185        &self,
2186        request: AgentRequest,
2187        connection: &mut ConnectionUploads,
2188    ) -> AgentResponse {
2189        let result = match request {
2190            AgentRequest::BeginTask { task } => self
2191                .begin_task(&task)
2192                .await
2193                .map(|run| AgentResponse::TaskBegun { run }),
2194            AgentRequest::CommitTask { run } => self
2195                .commit_task(&run)
2196                .await
2197                .map(|()| AgentResponse::TaskCommitted),
2198            AgentRequest::FindBlob { digest } => self.find_blob(&digest).await,
2199            AgentRequest::FindBlobs { digests } => self.find_blobs(digests).await,
2200            AgentRequest::StoreBlob { digest, source } => {
2201                self.store_blob(&digest, &source, connection).await
2202            }
2203            AgentRequest::FindActionResult { action } => {
2204                self.stats.lookups.fetch_add(1, Ordering::Relaxed);
2205                self.find_action_result(&action).await
2206            }
2207            AgentRequest::RecordActionHit {
2208                action,
2209                restore,
2210                crate_name,
2211            } => self.record_action_hit(&action, restore, crate_name),
2212            AgentRequest::RecordBypass { kind } => {
2213                *self
2214                    .stats
2215                    .bypasses
2216                    .lock()
2217                    .unwrap()
2218                    .entry(kind.clone())
2219                    .or_insert(0) += 1;
2220                self.emit(|| AgentEvent::Bypass { kind });
2221                Ok(AgentResponse::BypassRecorded)
2222            }
2223            AgentRequest::RecordUnconsulted => {
2224                self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
2225                self.emit(|| AgentEvent::Unconsulted);
2226                Ok(AgentResponse::UnconsultedRecorded)
2227            }
2228            AgentRequest::RecordWarning { message } => self.record_warning(message),
2229            AgentRequest::RecordCompilerInvocation {
2230                outcome,
2231                crate_name,
2232                duration_ns,
2233            } => self.record_compiler_invocation(&outcome, crate_name.as_deref(), duration_ns),
2234            AgentRequest::RecordActionVerification { matched, restore } => {
2235                self.record_materialization(restore);
2236                self.stats.verifications.fetch_add(1, Ordering::Relaxed);
2237                if !matched {
2238                    self.stats.divergences.fetch_add(1, Ordering::Relaxed);
2239                }
2240                self.emit(|| AgentEvent::Verification { matched, restore });
2241                Ok(AgentResponse::ActionVerificationRecorded)
2242            }
2243            AgentRequest::StoreActionResult { result } => {
2244                self.store_action_result(&result, connection).await
2245            }
2246            AgentRequest::FindActionPrediction { task, invocation } => {
2247                self.find_action_prediction(&task, &invocation)
2248            }
2249            AgentRequest::RecordActionPrediction { task, prediction } => {
2250                self.record_action_prediction(&task, prediction)
2251            }
2252            AgentRequest::FindExecutableIdentity {
2253                executable,
2254                environment,
2255            } => self.find_executable_identity(executable, environment),
2256            AgentRequest::StoreExecutableIdentity {
2257                executable,
2258                environment,
2259                stdout,
2260            } => self.store_executable_identity(executable, environment, stdout),
2261            AgentRequest::Hello { .. } => {
2262                Err(eyre::eyre!("hello is only valid as the first request"))
2263            }
2264        };
2265        result.unwrap_or_else(|error| AgentResponse::Error {
2266            message: error.to_string(),
2267        })
2268    }
2269
2270    async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
2271        if let Some(path) = self.find_verified_blob(digest)? {
2272            return Ok(AgentResponse::Blob { path: Some(path) });
2273        }
2274        if !self.remote_mode.reads() {
2275            return Ok(AgentResponse::Blob { path: None });
2276        }
2277        let Some(remote) = &self.remote else {
2278            return Ok(AgentResponse::Blob { path: None });
2279        };
2280        match self.fetch_remote_blob(remote, digest).await {
2281            Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
2282            Err(error) => {
2283                warn!(
2284                    "remote cache blob lookup failed for {}: {error}",
2285                    digest.hash
2286                );
2287                Ok(AgentResponse::Blob { path: None })
2288            }
2289        }
2290    }
2291
2292    async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
2293        let mut paths = BTreeMap::new();
2294        let mut missing = Vec::new();
2295        for digest in &digests {
2296            match self.find_verified_blob(digest)? {
2297                Some(path) => {
2298                    paths.insert(digest.clone(), path);
2299                }
2300                None => {
2301                    missing.push(digest.clone());
2302                }
2303            }
2304        }
2305
2306        if !missing.is_empty()
2307            && self.remote_mode.reads()
2308            && let Some(remote) = &self.remote
2309        {
2310            paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
2311        }
2312
2313        Ok(AgentResponse::Blobs {
2314            paths: digests
2315                .into_iter()
2316                .map(|digest| paths.get(&digest).cloned())
2317                .collect(),
2318        })
2319    }
2320
2321    async fn store_blob(
2322        &self,
2323        digest: &CacheDigest,
2324        source: &Path,
2325        connection: &mut ConnectionUploads,
2326    ) -> Result<AgentResponse> {
2327        let path = {
2328            let lock = self.write_lock(digest);
2329            let _guard = lock.lock().await;
2330            if let Some(path) = self.find_verified_blob(digest)? {
2331                path
2332            } else {
2333                let path = self.cas.store_file(digest, source)?;
2334                self.remember_verified_blob(digest, &path);
2335                self.stats.stores.fetch_add(1, Ordering::Relaxed);
2336                self.stats
2337                    .stored_bytes
2338                    .fetch_add(digest.size, Ordering::Relaxed);
2339                path
2340            }
2341        };
2342        // The object is durable locally, so the build has what it needs and the
2343        // remote publication can happen after this request returns.
2344        if let Some(uploads) = &self.uploads {
2345            uploads.queue_blob(digest, path.clone(), connection);
2346        }
2347        Ok(AgentResponse::Stored { path })
2348    }
2349
2350    fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
2351        let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
2352        if let Some(remembered) = remembered {
2353            // The contents behind this digest were hashed in full when the
2354            // session first reached for them. Hashing again on every later
2355            // lookup would re-read each dependency's rlib once per crate that
2356            // links it, so what is rechecked here is that the file has not
2357            // been written since: an overwrite moves the modification time,
2358            // and truncation or eviction changes the length or removes it.
2359            // Only a replacement that reproduces both -- which no writer in
2360            // the store does, since blobs are published by rename under a
2361            // content-derived name -- would go unnoticed until `mbx cache
2362            // verify` reads the store back in full.
2363            if remembered.is_unchanged() {
2364                return Ok(Some(remembered.path));
2365            }
2366            self.verified_blobs.lock().unwrap().remove(digest);
2367        }
2368        let path = self.cas.find(digest)?;
2369        if let Some(path) = &path {
2370            self.remember_verified_blob(digest, path);
2371        }
2372        Ok(path)
2373    }
2374
2375    fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
2376        let Some(verified) = VerifiedBlob::describe(path) else {
2377            // Without an identity to compare against there is nothing to
2378            // shortcut safely, so the digest stays unremembered and every
2379            // lookup re-reads it.
2380            return;
2381        };
2382        self.verified_blobs
2383            .lock()
2384            .unwrap()
2385            .insert(digest.clone(), verified);
2386    }
2387
2388    async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
2389        if let Some(result) = self.actions.find(action)? {
2390            return Ok(AgentResponse::ActionResult {
2391                result: Some(result),
2392            });
2393        }
2394        if !self.remote_mode.reads() {
2395            return Ok(AgentResponse::ActionResult { result: None });
2396        }
2397        let Some(remote) = &self.remote else {
2398            return Ok(AgentResponse::ActionResult { result: None });
2399        };
2400        let lock = self.action_lock(action);
2401        let _guard = lock.lock().await;
2402        if let Some(result) = self.actions.find(action)? {
2403            return Ok(AgentResponse::ActionResult {
2404                result: Some(result),
2405            });
2406        }
2407        if let Some(result) = self
2408            .pending_remote_actions
2409            .lock()
2410            .unwrap()
2411            .get(action)
2412            .cloned()
2413        {
2414            return Ok(AgentResponse::ActionResult {
2415                result: Some(result),
2416            });
2417        }
2418        let _permit = self.remote_transfers.acquire().await?;
2419        match self.get_remote_action_result(remote, action).await {
2420            Ok(Some(result)) => {
2421                self.pending_remote_actions
2422                    .lock()
2423                    .unwrap()
2424                    .insert(action.clone(), result.clone());
2425                Ok(AgentResponse::ActionResult {
2426                    result: Some(result),
2427                })
2428            }
2429            Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
2430            Err(error) => {
2431                self.note_remote_failure();
2432                warn!(
2433                    "remote cache action lookup failed for {}: {error}",
2434                    action.hash
2435                );
2436                Ok(AgentResponse::ActionResult { result: None })
2437            }
2438        }
2439    }
2440
2441    async fn store_action_result(
2442        &self,
2443        result: &RemoteActionResult,
2444        connection: &ConnectionUploads,
2445    ) -> Result<AgentResponse> {
2446        let path = self.actions.store(result)?;
2447        if let Some(uploads) = &self.uploads {
2448            uploads.queue_action_result(result, connection);
2449        }
2450        Ok(AgentResponse::ActionStored { path })
2451    }
2452
2453    /// Record that a remote operation failed and the build carried on without it.
2454    fn note_remote_failure(&self) {
2455        self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
2456    }
2457
2458    async fn get_remote_action_result(
2459        &self,
2460        remote: &RemoteCacheClient,
2461        action: &CacheDigest,
2462    ) -> Result<Option<RemoteActionResult>> {
2463        self.stats
2464            .remote_action_lookups
2465            .fetch_add(1, Ordering::Relaxed);
2466        let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
2467        remote.get_action_result(action).await
2468    }
2469
2470    fn record_action_hit(
2471        &self,
2472        action: &CacheDigest,
2473        restore: RestoreStats,
2474        crate_name: Option<String>,
2475    ) -> Result<AgentResponse> {
2476        validate_crate_name(crate_name.as_deref())?;
2477        if self.actions.find(action)?.is_none() {
2478            let pending = self.pending_remote_actions.lock().unwrap().remove(action);
2479            if let Some(result) = pending {
2480                self.actions.store(&result)?;
2481            } else {
2482                bail!("cannot record a hit for a missing action result");
2483            }
2484        }
2485        self.record_restore(restore);
2486        self.stats.hits.fetch_add(1, Ordering::Relaxed);
2487        self.emit(|| AgentEvent::ActionHit {
2488            crate_name,
2489            restore,
2490        });
2491        Ok(AgentResponse::ActionHitRecorded)
2492    }
2493
2494    fn record_restore(&self, restore: RestoreStats) {
2495        self.record_materialization(restore);
2496        atomic_saturating_add(
2497            &self.stats.avoided_compiler_duration_ns,
2498            restore.avoided_compiler_duration_ns,
2499        );
2500        atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
2501        atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
2502        atomic_saturating_add(
2503            &self.stats.reflinked_output_files,
2504            restore.reflinked_output_files,
2505        );
2506        atomic_saturating_add(
2507            &self.stats.reflinked_output_bytes,
2508            restore.reflinked_output_bytes,
2509        );
2510        atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
2511        atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
2512    }
2513
2514    fn record_compiler_invocation(
2515        &self,
2516        outcome: &str,
2517        crate_name: Option<&str>,
2518        duration_ns: u64,
2519    ) -> Result<AgentResponse> {
2520        if !matches!(
2521            outcome,
2522            "miss" | "unconsulted" | "bypass" | "verification" | "incremental"
2523        ) {
2524            bail!("invalid compiler invocation outcome");
2525        }
2526        validate_crate_name(crate_name)?;
2527        let mut compiler = self.stats.compiler.lock().unwrap();
2528        let stats = compiler.entry(outcome.to_string()).or_default();
2529        stats.invocations = stats.invocations.saturating_add(1);
2530        stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
2531        drop(compiler);
2532        if outcome != "verification"
2533            && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
2534        {
2535            let mut slow = self.stats.slow_compilations.lock().unwrap();
2536            let duration = slow.entry(crate_name.to_string()).or_default();
2537            *duration = duration.saturating_add(duration_ns);
2538        }
2539        self.emit(|| AgentEvent::CompilerInvocation {
2540            outcome: outcome.to_string(),
2541            crate_name: crate_name.map(str::to_string),
2542            duration_ns,
2543        });
2544        Ok(AgentResponse::CompilerInvocationRecorded)
2545    }
2546
2547    fn record_materialization(&self, restore: RestoreStats) {
2548        atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
2549    }
2550
2551    /// Surface a shim diagnostic on this process's stderr.
2552    ///
2553    /// The agent lives in the process that owns the session, so printing here
2554    /// reaches the terminal running the build rather than the stderr of the
2555    /// compilation the shim stands in for -- which build scripts read as part
2556    /// of the compiler's answer. Deduplicated because one cause tends to fire
2557    /// once per compilation, and capped so a message unique per compilation
2558    /// cannot scroll the build away.
2559    fn record_warning(&self, message: String) -> Result<AgentResponse> {
2560        if message.is_empty()
2561            || message.len() > MAX_WARNING_BYTES
2562            || message.contains(['\n', '\r', '\0'])
2563        {
2564            bail!("invalid shim warning");
2565        }
2566        let mut warnings = self.warnings.lock().unwrap();
2567        if !warnings.contains(&message) && warnings.len() < MAX_WARNINGS {
2568            eprintln!("mbx[warning]: {message}");
2569            self.emit(|| AgentEvent::Warning {
2570                message: message.clone(),
2571            });
2572            warnings.insert(message);
2573        }
2574        Ok(AgentResponse::WarningRecorded)
2575    }
2576
2577    fn find_action_prediction(
2578        &self,
2579        task: &str,
2580        invocation: &CacheDigest,
2581    ) -> Result<AgentResponse> {
2582        validate_task_identity(task)?;
2583        invocation.validate()?;
2584        let prediction = self
2585            .task_actions
2586            .lock()
2587            .unwrap()
2588            .get(task)
2589            .and_then(|state| state.predictions.get(invocation))
2590            .cloned();
2591        Ok(AgentResponse::ActionPrediction { prediction })
2592    }
2593
2594    fn record_action_prediction(
2595        &self,
2596        task: &str,
2597        prediction: ActionPrediction,
2598    ) -> Result<AgentResponse> {
2599        validate_task_identity(task)?;
2600        validate_action_prediction(&prediction)?;
2601        let mut tasks = self.task_actions.lock().unwrap();
2602        let state = tasks.entry(task.to_string()).or_default();
2603        if !state.predictions.contains_key(&prediction.invocation)
2604            && state.predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2605        {
2606            bail!("task action manifest contains too many predictions");
2607        }
2608        state
2609            .predictions
2610            .insert(prediction.invocation.clone(), prediction.clone());
2611        state
2612            .pending_predictions
2613            .insert(prediction.invocation.clone(), prediction);
2614        Ok(AgentResponse::ActionPredictionRecorded)
2615    }
2616
2617    fn executable_identity_key(
2618        &self,
2619        executable: PathBuf,
2620        environment: BTreeMap<String, Option<String>>,
2621    ) -> Result<ExecutableIdentityKey> {
2622        // Restricted to the variables that actually select what an identity
2623        // probe reports: the toolchain rustup resolves, and the SDK a linker
2624        // driver builds against. Anything else would let one key stand for two
2625        // different compilers.
2626        if !environment.keys().all(|name| {
2627            matches!(
2628                name.as_str(),
2629                "RUSTUP_HOME" | "RUSTUP_TOOLCHAIN" | "SDKROOT" | "MACOSX_DEPLOYMENT_TARGET"
2630            )
2631        }) {
2632            bail!("executable identity contains an unsupported environment variable");
2633        }
2634        Ok(ExecutableIdentityKey {
2635            executable,
2636            environment,
2637        })
2638    }
2639
2640    fn find_executable_identity(
2641        &self,
2642        executable: PathBuf,
2643        environment: BTreeMap<String, Option<String>>,
2644    ) -> Result<AgentResponse> {
2645        let key = self.executable_identity_key(executable, environment)?;
2646        let stdout = self
2647            .executable_identities
2648            .lock()
2649            .unwrap()
2650            .get(&key)
2651            .cloned();
2652        Ok(AgentResponse::ExecutableIdentity { stdout })
2653    }
2654
2655    fn store_executable_identity(
2656        &self,
2657        executable: PathBuf,
2658        environment: BTreeMap<String, Option<String>>,
2659        stdout: Vec<u8>,
2660    ) -> Result<AgentResponse> {
2661        if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
2662            bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
2663        }
2664        let key = self.executable_identity_key(executable, environment)?;
2665        let mut identities = self.executable_identities.lock().unwrap();
2666        let is_new = !identities.contains_key(&key);
2667        let previous_size = identities.get(&key).map_or(0, Vec::len);
2668        if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
2669            bail!("executable identity cache contains too many entries");
2670        }
2671        let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
2672        if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
2673            bail!("executable identity cache contains too many bytes");
2674        }
2675        identities.insert(key, stdout.clone());
2676        Ok(AgentResponse::ExecutableIdentity {
2677            stdout: Some(stdout),
2678        })
2679    }
2680
2681    /// Serve newline-delimited protocol requests on an authenticated session stream.
2682    pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
2683    where
2684        S: AsyncRead + AsyncWrite + Unpin,
2685    {
2686        let (reader, mut writer) = tokio::io::split(stream);
2687        let mut reader = BufReader::new(reader);
2688        let hello = read_request(&mut reader)
2689            .await?
2690            .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
2691        let request: AgentRequest = serde_json::from_str(&hello)?;
2692        match request {
2693            AgentRequest::Hello {
2694                protocol,
2695                client_version,
2696            } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
2697            AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
2698                send_response(
2699                    &mut writer,
2700                    &AgentResponse::Error {
2701                        message: format!(
2702                            "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
2703                        ),
2704                    },
2705                )
2706                .await?;
2707                return Ok(());
2708            }
2709            AgentRequest::Hello { client_version, .. } => {
2710                send_response(
2711                    &mut writer,
2712                    &AgentResponse::Error {
2713                        message: format!(
2714                            "cache client {client_version} does not match agent {}",
2715                            self.version
2716                        ),
2717                    },
2718                )
2719                .await?;
2720                return Ok(());
2721            }
2722            _ => bail!("the first agent request must be hello"),
2723        }
2724        send_response(
2725            &mut writer,
2726            &AgentResponse::Hello {
2727                protocol: AGENT_PROTOCOL_VERSION,
2728                agent_version: self.version.to_string(),
2729            },
2730        )
2731        .await?;
2732
2733        // Tickets accumulate for the life of the connection because a shim
2734        // publishes a compilation's blobs and the action result naming them over
2735        // one connection, in that order.
2736        let mut connection = ConnectionUploads::default();
2737        while let Some(line) = read_request(&mut reader).await? {
2738            let response = match serde_json::from_str(&line) {
2739                Ok(request) => self.respond_on(request, &mut connection).await,
2740                Err(error) => AgentResponse::Error {
2741                    message: format!("invalid agent request: {error}"),
2742                },
2743            };
2744            send_response(&mut writer, &response).await?;
2745        }
2746        Ok(())
2747    }
2748}
2749
2750/// Read one newline-delimited request, refusing one that grows past the cap.
2751///
2752/// Any process running as this user can open the session socket, so a request
2753/// that never terminates its line must not be able to grow the agent's memory
2754/// without bound.
2755async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
2756where
2757    R: AsyncBufRead + Unpin,
2758{
2759    let mut line = Vec::new();
2760    loop {
2761        let available = reader.fill_buf().await?;
2762        if available.is_empty() {
2763            break;
2764        }
2765        let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
2766            Some(index) => (index, true),
2767            None => (available.len(), false),
2768        };
2769        if line.len() + consumed > MAX_REQUEST_BYTES {
2770            bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
2771        }
2772        line.extend_from_slice(&available[..consumed]);
2773        // The newline itself is consumed but never kept.
2774        reader.consume(consumed + usize::from(complete));
2775        if complete {
2776            return Ok(Some(String::from_utf8(line)?));
2777        }
2778    }
2779    if line.is_empty() {
2780        Ok(None)
2781    } else {
2782        Ok(Some(String::from_utf8(line)?))
2783    }
2784}
2785
2786/// Whether `task` is a well-formed task action identity.
2787///
2788/// Identities name files and directories in the store, so anything that reads
2789/// the store back has to be able to tell an identity from whatever else a user
2790/// left lying there.
2791pub fn is_task_identity(task: &str) -> bool {
2792    task.len() == 64
2793        && task
2794            .bytes()
2795            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
2796}
2797
2798fn validate_task_identity(task: &str) -> Result<()> {
2799    if !is_task_identity(task) {
2800        bail!("invalid task action identity");
2801    }
2802    Ok(())
2803}
2804
2805/// Where a store keeps its task prediction manifests.
2806fn task_manifest_dir(store: &Path) -> PathBuf {
2807    store.join("task-manifests").join("v1")
2808}
2809
2810/// The action digests a task's prediction manifest recorded.
2811///
2812/// Read straight off disk rather than through an agent, because a collector
2813/// needs the action set of tasks no session is running. A manifest that is
2814/// missing or no longer parseable yields no actions rather than an error: this
2815/// is a prediction index, so the worst a thin answer costs is a cold prefetch,
2816/// or an object collected earlier than it deserved.
2817pub fn task_manifest_actions(store: &Path, task: &str) -> Result<Vec<CacheDigest>> {
2818    validate_task_identity(task)?;
2819    let path = task_manifest_dir(store).join(format!("{task}.json"));
2820    let bytes = match fs::read(&path) {
2821        Ok(bytes) => bytes,
2822        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
2823        Err(error) => {
2824            return Err(error).wrap_err_with(|| format!("failed to read {}", path.display()));
2825        }
2826    };
2827    let Ok(manifest) = serde_json::from_slice::<TaskActionManifest>(&bytes) else {
2828        return Ok(Vec::new());
2829    };
2830    if validate_task_manifest(&manifest, task).is_err() {
2831        return Ok(Vec::new());
2832    }
2833    Ok(manifest
2834        .predictions
2835        .into_iter()
2836        .map(|prediction| prediction.action)
2837        .collect())
2838}
2839
2840fn validate_action_prediction(prediction: &ActionPrediction) -> Result<()> {
2841    if prediction.validate() {
2842        Ok(())
2843    } else {
2844        bail!("invalid action prediction")
2845    }
2846}
2847
2848fn validate_task_manifest(manifest: &TaskActionManifest, task: &str) -> Result<()> {
2849    if manifest.task == task && manifest.validate() {
2850        Ok(())
2851    } else {
2852        bail!("invalid task action manifest")
2853    }
2854}
2855
2856fn merge_task_manifests(
2857    task: &str,
2858    base: Option<TaskActionManifest>,
2859    update: TaskActionManifest,
2860) -> Result<TaskActionManifest> {
2861    validate_task_manifest(&update, task)?;
2862    let mut predictions = BTreeMap::new();
2863    if let Some(base) = base {
2864        validate_task_manifest(&base, task)?;
2865        predictions.extend(
2866            base.predictions
2867                .into_iter()
2868                .map(|prediction| (prediction.invocation.clone(), prediction)),
2869        );
2870    }
2871    predictions.extend(
2872        update
2873            .predictions
2874            .into_iter()
2875            .map(|prediction| (prediction.invocation.clone(), prediction)),
2876    );
2877    let manifest = TaskActionManifest {
2878        version: TASK_ACTION_MANIFEST_VERSION,
2879        task: task.to_owned(),
2880        predictions: predictions.into_values().collect(),
2881    };
2882    validate_task_manifest(&manifest, task)?;
2883    Ok(manifest)
2884}
2885
2886fn merge_remote_task_manifest(
2887    task: &str,
2888    remote: TaskActionManifest,
2889    local: TaskActionManifest,
2890) -> (TaskActionManifest, bool) {
2891    match merge_task_manifests(task, Some(remote), local.clone()) {
2892        Ok(manifest) => (manifest, true),
2893        Err(error) => {
2894            warn!("remote task action manifest merge failed for {task}: {error}");
2895            (local, false)
2896        }
2897    }
2898}
2899
2900async fn send_response(
2901    writer: &mut (impl AsyncWrite + Unpin),
2902    response: &AgentResponse,
2903) -> Result<()> {
2904    let mut encoded = serde_json::to_vec(response)?;
2905    encoded.push(b'\n');
2906    writer.write_all(&encoded).await?;
2907    writer.flush().await?;
2908    Ok(())
2909}
2910
2911#[cfg(test)]
2912#[path = "agent_tests.rs"]
2913mod tests;