Skip to main content

kache_core/
lib.rs

1#[cfg(feature = "planning")]
2use std::collections::{HashMap, HashSet};
3
4#[cfg(feature = "planning")]
5use anyhow::Result;
6#[cfg(feature = "planning")]
7use async_trait::async_trait;
8use serde::{Deserialize, Serialize};
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
11pub struct BuildIntent {
12    #[serde(default)]
13    pub crate_names: Vec<String>,
14    #[serde(default)]
15    pub namespace: Option<String>,
16    #[serde(default)]
17    pub cargo_lock_deps: Vec<(String, String)>,
18}
19
20/// Which source produced a candidate, i.e. how much to trust it
21/// (kunobi-ninja/kache#617).
22///
23/// `Unknown` is the `#[serde(other)]` arm so a newer planner naming a source
24/// this build has never heard of degrades to "untrusted" instead of failing
25/// the whole plan.
26#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
27#[serde(rename_all = "snake_case")]
28pub enum CandidateSource {
29    /// Exact lockfile-shard match: this build's dependency set produced it.
30    Shard,
31    /// This machine built this crate before.
32    History,
33    /// A crate NAME matched something in the remote listing. A crate name is
34    /// not a build identity, so this is a guess.
35    KeyCache,
36    #[default]
37    #[serde(other)]
38    Unknown,
39}
40
41impl CandidateSource {
42    /// Confidence rank, lower is better. Only the ORDER matters; these are not
43    /// probabilities and must not be presented as any.
44    pub fn confidence_rank(self) -> u8 {
45        match self {
46            CandidateSource::Shard => 0,
47            CandidateSource::History => 1,
48            CandidateSource::KeyCache => 2,
49            CandidateSource::Unknown => 3,
50        }
51    }
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
55pub struct PrefetchCandidate {
56    pub cache_key: String,
57    pub crate_name: String,
58    /// What a miss would cost to rebuild. `None` = unknown, which is NOT the
59    /// same as zero: an un-backfilled store row reads 0, and treating that as
60    /// "free to fetch and worthless to have" would bury it (#617).
61    #[serde(default, skip_serializing_if = "Option::is_none")]
62    pub compile_time_ms: Option<u64>,
63    /// Stored artifact size. An admission and ranking ESTIMATE, never a
64    /// promise about compressed transfer bytes: anything enforced has to be
65    /// counted on the wire.
66    #[serde(default, skip_serializing_if = "Option::is_none")]
67    pub size_bytes: Option<u64>,
68    #[serde(default)]
69    pub source: CandidateSource,
70    /// Position in the build's dependency order, i.e. roughly when the build
71    /// will ask for it. `None` = not in the intent's crate list.
72    #[serde(default, skip_serializing_if = "Option::is_none")]
73    pub demand_index: Option<u32>,
74}
75
76impl PrefetchCandidate {
77    /// A candidate with no metadata, as the pre-#617 wire produced.
78    pub fn new(cache_key: String, crate_name: String) -> Self {
79        Self {
80            cache_key,
81            crate_name,
82            compile_time_ms: None,
83            size_bytes: None,
84            source: CandidateSource::Unknown,
85            demand_index: None,
86        }
87    }
88
89    pub fn with_source(mut self, source: CandidateSource) -> Self {
90        self.source = source;
91        self
92    }
93}
94
95#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
96#[serde(rename_all = "snake_case")]
97pub enum PrefetchDisposition {
98    Execute,
99    UseFallback,
100    DoNothing,
101}
102
103#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
104pub struct PrefetchPlan {
105    #[serde(default)]
106    pub plan_id: Option<String>,
107    #[serde(default)]
108    pub planner: Option<String>,
109    pub disposition: PrefetchDisposition,
110    #[serde(default)]
111    pub candidates: Vec<PrefetchCandidate>,
112}
113
114/// How many candidates each source may contribute to one plan
115/// (kunobi-ninja/kache#616).
116///
117/// These bound plan COMPOSITION, which is a different job from the daemon's
118/// key/byte/time budgets: those bound resource use and are the trust boundary,
119/// these stop one low-confidence source from crowding out better candidates
120/// before the budget is even reached. `0` disables a cap.
121#[derive(Debug, Clone, Copy, PartialEq, Eq)]
122pub struct PlanLimits {
123    /// Key-cache variants per crate. A crate name is not a build identity, so
124    /// of `n` variants at most one can be the right one; taking many is paying
125    /// `n` downloads for at most one hit.
126    pub key_cache_per_crate: usize,
127    /// Key-cache candidates across the whole plan, so a build with hundreds of
128    /// unresolved crates cannot fill the plan with guesses.
129    pub key_cache_total: usize,
130    /// History entries per crate. Unlike shards, where one crate legitimately
131    /// has several compile units, extra history rows for one crate are older
132    /// variants.
133    pub history_per_crate: usize,
134}
135
136impl Default for PlanLimits {
137    fn default() -> Self {
138        Self {
139            key_cache_per_crate: 2,
140            key_cache_total: 64,
141            history_per_crate: 2,
142        }
143    }
144}
145
146/// Urgency bucket width, in dependency-order positions (#617).
147///
148/// Demand order is bucketed rather than used exactly because it comes from a
149/// guppy graph traversal, which only approximates when cargo will actually ask
150/// (cargo reorders for parallelism, build scripts, proc macros, features).
151/// Treating position 40 and 45 as meaningfully different is false precision;
152/// 40 versus 400 is real. Roughly the prefetch concurrency, so one window is
153/// about one wave of downloads.
154#[cfg(feature = "planning")]
155const URGENCY_BUCKET: u32 = 16;
156
157/// Sort key for dispatch order, lowest first (#617).
158///
159/// Lexicographic, deliberately, rather than a weighted score: a weighted sum
160/// needs coefficients, and nothing can calibrate them until #618 makes
161/// "arrived before it was demanded" measurable. Every element here is an
162/// ordering, not a magnitude.
163///
164/// 1. Urgency bucket. Prefetch races the build, so a high-value artifact
165///    needed at minute eight loses to a medium-value one needed at second
166///    five. Candidates with no demand index sort last.
167/// 2. Confidence. Within one wave, prefer the source most likely to be right.
168/// 3. Value, descending. Expensive rebuilds first, so the limited slots buy
169///    the most avoided work. Unknown cost sorts after known cost rather than
170///    being scored as zero.
171///
172/// Callers must apply this as a STABLE sort: equal keys keep source order,
173/// which is the planner's confidence-merge order.
174#[cfg(feature = "planning")]
175pub fn dispatch_sort_key(candidate: &PrefetchCandidate) -> (u32, u8, std::cmp::Reverse<u64>) {
176    let bucket = candidate
177        .demand_index
178        .map(|index| index / URGENCY_BUCKET)
179        .unwrap_or(u32::MAX);
180    (
181        bucket,
182        candidate.source.confidence_rank(),
183        // `None` -> 0 -> sorts last under Reverse, without claiming the
184        // candidate is worthless.
185        std::cmp::Reverse(candidate.compile_time_ms.unwrap_or(0)),
186    )
187}
188
189/// Truncate `items` to `limit`, returning how many were dropped. `0` disables.
190///
191/// Shared by every composition cap so "0 means unlimited" is defined once
192/// rather than re-derived at each call site.
193#[cfg(feature = "planning")]
194fn cap_to(items: &mut Vec<PrefetchCandidate>, limit: usize) -> usize {
195    if limit == 0 || items.len() <= limit {
196        return 0;
197    }
198    let dropped = items.len() - limit;
199    items.truncate(limit);
200    dropped
201}
202
203/// What a plan left out, so a truncated plan is distinguishable from one that
204/// had nothing more to offer (#616).
205#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
206pub struct PlanComposition {
207    pub from_shards: usize,
208    pub from_history: usize,
209    pub from_key_cache: usize,
210    /// Dropped by [`PlanLimits`], by source.
211    pub dropped_history_per_crate: usize,
212    pub dropped_key_cache_per_crate: usize,
213    pub dropped_key_cache_total: usize,
214}
215
216impl PlanComposition {
217    pub fn dropped_total(&self) -> usize {
218        self.dropped_history_per_crate
219            + self.dropped_key_cache_per_crate
220            + self.dropped_key_cache_total
221    }
222}
223
224#[cfg(feature = "planning")]
225#[async_trait]
226pub trait PlannerDataSource {
227    async fn shard_candidates(
228        &self,
229        namespace: &str,
230        deps: &[(String, String)],
231    ) -> Result<Vec<PrefetchCandidate>>;
232
233    async fn history_candidates(&self, crate_names: &[String]) -> Result<Vec<PrefetchCandidate>>;
234
235    async fn key_cache_keys_for_crate(&self, crate_name: &str) -> Result<Vec<String>>;
236}
237
238#[cfg(feature = "planning")]
239pub async fn build_prefetch_plan<T>(
240    source: &T,
241    intent: &BuildIntent,
242    planner_name: &str,
243) -> Result<PrefetchPlan>
244where
245    T: PlannerDataSource + Sync + ?Sized,
246{
247    build_prefetch_plan_with_limits(source, intent, planner_name, PlanLimits::default())
248        .await
249        .map(|(plan, _composition)| plan)
250}
251
252/// [`build_prefetch_plan`] with explicit composition limits, also returning
253/// what the limits dropped so a caller can report it (#616).
254#[cfg(feature = "planning")]
255pub async fn build_prefetch_plan_with_limits<T>(
256    source: &T,
257    intent: &BuildIntent,
258    planner_name: &str,
259    limits: PlanLimits,
260) -> Result<(PrefetchPlan, PlanComposition)>
261where
262    T: PlannerDataSource + Sync + ?Sized,
263{
264    let crate_order = crate_query_order(intent);
265    let demand_index = demand_index_map(&crate_order);
266    let mut seen = HashSet::new();
267    let mut resolved_crates = HashSet::new();
268    let mut candidates = Vec::new();
269    let mut composition = PlanComposition::default();
270
271    // Sources are merged in descending order of confidence, each one filling
272    // only the crates the ones before it left unresolved. Shard lookups used
273    // to RETURN as soon as they produced anything (kunobi-ninja/kache#614),
274    // but a shard hit is exact per bucket: one dependency bump invalidates one
275    // of `NUM_SHARDS` buckets while the rest still match, so a single matching
276    // shard was enough to short-circuit history and key-cache recovery for
277    // every crate in every bucket that missed.
278    if let Some(namespace) = intent.namespace.as_deref()
279        && !intent.cargo_lock_deps.is_empty()
280    {
281        // A failed shard lookup is not fatal: the sources below still run.
282        if let Ok(shard_candidates) = source
283            .shard_candidates(namespace, &intent.cargo_lock_deps)
284            .await
285        {
286            for candidate in order_candidates_by_crate_order(shard_candidates, intent) {
287                resolved_crates.insert(candidate.crate_name.clone());
288                if seen.insert(candidate.cache_key.clone()) {
289                    candidates.push(candidate.with_source(CandidateSource::Shard));
290                }
291            }
292        }
293    }
294
295    let unresolved = |resolved: &HashSet<String>| -> Vec<String> {
296        crate_order
297            .iter()
298            .filter(|name| !resolved.contains(*name))
299            .cloned()
300            .collect()
301    };
302
303    composition.from_shards = candidates.len();
304
305    let history_query = unresolved(&resolved_crates);
306    if !history_query.is_empty() {
307        // Group by crate so the per-crate cap keeps the FIRST entries, which
308        // both data sources return most-recently-used first.
309        let mut by_crate: HashMap<String, Vec<PrefetchCandidate>> = HashMap::new();
310        for candidate in order_candidates_by_crate_order(
311            source.history_candidates(&history_query).await?,
312            intent,
313        ) {
314            by_crate
315                .entry(candidate.crate_name.clone())
316                .or_default()
317                .push(candidate.with_source(CandidateSource::History));
318        }
319
320        for crate_name in &history_query {
321            let Some(mut for_crate) = by_crate.remove(crate_name) else {
322                continue;
323            };
324            composition.dropped_history_per_crate +=
325                cap_to(&mut for_crate, limits.history_per_crate);
326            for candidate in for_crate {
327                resolved_crates.insert(candidate.crate_name.clone());
328                if seen.insert(candidate.cache_key.clone()) {
329                    composition.from_history += 1;
330                    candidates.push(candidate);
331                }
332            }
333        }
334    }
335
336    // The key cache is the weakest source: it maps a crate NAME to every cache
337    // key in the bucket, with no target, toolchain, profile, or feature
338    // dimension, so of `n` variants at most one can be right. Capping is not a
339    // fix for that (dimensioning the remote layout is, separately) but it stops
340    // the guesses crowding out better candidates (#616).
341    for crate_name in unresolved(&resolved_crates) {
342        if limits.key_cache_total > 0 && composition.from_key_cache >= limits.key_cache_total {
343            // Whatever this crate would have offered is dropped wholesale; count
344            // it so the truncation is visible rather than inferred.
345            composition.dropped_key_cache_total += source
346                .key_cache_keys_for_crate(&crate_name)
347                .await?
348                .into_iter()
349                .filter(|key| !seen.contains(key))
350                .count();
351            continue;
352        }
353
354        let mut for_crate: Vec<PrefetchCandidate> = source
355            .key_cache_keys_for_crate(&crate_name)
356            .await?
357            .into_iter()
358            .map(|cache_key| {
359                PrefetchCandidate::new(cache_key, crate_name.clone())
360                    .with_source(CandidateSource::KeyCache)
361            })
362            .collect();
363
364        composition.dropped_key_cache_per_crate +=
365            cap_to(&mut for_crate, limits.key_cache_per_crate);
366
367        for candidate in for_crate {
368            if limits.key_cache_total > 0 && composition.from_key_cache >= limits.key_cache_total {
369                composition.dropped_key_cache_total += 1;
370                continue;
371            }
372            if seen.insert(candidate.cache_key.clone()) {
373                composition.from_key_cache += 1;
374                candidates.push(candidate);
375            }
376        }
377    }
378
379    // Stamp demand position so the daemon can rank without the intent: it only
380    // receives the plan, and `PrefetchRequest::from_plan` drops everything else.
381    for candidate in &mut candidates {
382        candidate.demand_index = demand_index.get(&candidate.crate_name).copied();
383    }
384
385    // Dispatch order (#617). Stable, so equal keys keep the confidence-merge
386    // order the sources were appended in.
387    candidates.sort_by_key(dispatch_sort_key);
388
389    Ok((execute_plan(planner_name, candidates), composition))
390}
391
392#[cfg(feature = "planning")]
393fn execute_plan(planner_name: &str, candidates: Vec<PrefetchCandidate>) -> PrefetchPlan {
394    let planner = planner_name.trim();
395    PrefetchPlan {
396        plan_id: None,
397        planner: Some(if planner.is_empty() {
398            "planner".to_string()
399        } else {
400            planner.to_string()
401        }),
402        disposition: PrefetchDisposition::Execute,
403        candidates,
404    }
405}
406
407/// crate name -> position in dependency order, for [`dispatch_sort_key`].
408#[cfg(feature = "planning")]
409fn demand_index_map(crate_order: &[String]) -> HashMap<String, u32> {
410    crate_order
411        .iter()
412        .enumerate()
413        .map(|(index, name)| (name.clone(), index as u32))
414        .collect()
415}
416
417#[cfg(feature = "planning")]
418fn crate_query_order(intent: &BuildIntent) -> Vec<String> {
419    let mut seen = HashSet::new();
420    intent
421        .crate_names
422        .iter()
423        .filter(|crate_name| seen.insert((*crate_name).clone()))
424        .cloned()
425        .collect()
426}
427
428#[cfg(feature = "planning")]
429fn order_candidates_by_crate_order(
430    mut candidates: Vec<PrefetchCandidate>,
431    intent: &BuildIntent,
432) -> Vec<PrefetchCandidate> {
433    if intent.crate_names.is_empty() {
434        return candidates;
435    }
436
437    let mut priority = HashMap::new();
438    for (index, crate_name) in crate_query_order(intent).iter().enumerate() {
439        priority.entry(crate_name.clone()).or_insert(index);
440    }
441
442    let mut indexed_candidates = candidates.drain(..).enumerate().collect::<Vec<_>>();
443    indexed_candidates.sort_by_key(|(index, candidate)| {
444        (
445            priority
446                .get(&candidate.crate_name)
447                .copied()
448                .unwrap_or(usize::MAX),
449            *index,
450        )
451    });
452    indexed_candidates
453        .into_iter()
454        .map(|(_, candidate)| candidate)
455        .collect()
456}
457
458#[cfg(test)]
459mod tests {
460    use super::*;
461
462    #[cfg(feature = "planning")]
463    use std::collections::HashMap;
464
465    #[cfg(feature = "planning")]
466    use anyhow::anyhow;
467
468    #[cfg(feature = "planning")]
469    use proptest::prelude::*;
470
471    #[cfg(feature = "planning")]
472    #[derive(Default)]
473    struct FakePlannerDataSource {
474        shard_candidates: Vec<PrefetchCandidate>,
475        shard_error: bool,
476        history_candidates: Vec<PrefetchCandidate>,
477        history_by_crate: HashMap<String, String>,
478        key_cache: HashMap<String, Vec<String>>,
479    }
480
481    #[cfg(feature = "planning")]
482    #[async_trait]
483    impl PlannerDataSource for FakePlannerDataSource {
484        async fn shard_candidates(
485            &self,
486            _namespace: &str,
487            _deps: &[(String, String)],
488        ) -> Result<Vec<PrefetchCandidate>> {
489            if self.shard_error {
490                Err(anyhow!("shard lookup failed"))
491            } else {
492                Ok(self.shard_candidates.clone())
493            }
494        }
495
496        async fn history_candidates(
497            &self,
498            crate_names: &[String],
499        ) -> Result<Vec<PrefetchCandidate>> {
500            if !self.history_candidates.is_empty() {
501                return Ok(self.history_candidates.clone());
502            }
503
504            Ok(crate_names
505                .iter()
506                .filter_map(|crate_name| {
507                    self.history_by_crate.get(crate_name).map(|cache_key| {
508                        PrefetchCandidate::new(cache_key.clone(), crate_name.clone())
509                    })
510                })
511                .collect())
512        }
513
514        async fn key_cache_keys_for_crate(&self, crate_name: &str) -> Result<Vec<String>> {
515            Ok(self.key_cache.get(crate_name).cloned().unwrap_or_default())
516        }
517    }
518
519    #[test]
520    fn test_build_intent_serde_roundtrip() {
521        let intent = BuildIntent {
522            crate_names: vec!["serde".into(), "tokio".into()],
523            namespace: Some("x86_64/hash/release".into()),
524            cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
525        };
526
527        let json = serde_json::to_string(&intent).unwrap();
528        let parsed: BuildIntent = serde_json::from_str(&json).unwrap();
529        assert_eq!(parsed, intent);
530    }
531
532    #[test]
533    fn test_build_intent_defaults_missing_fields() {
534        let parsed: BuildIntent = serde_json::from_str(r#"{"crate_names":["serde"]}"#).unwrap();
535        assert_eq!(parsed.crate_names, vec!["serde"]);
536        assert!(parsed.namespace.is_none());
537        assert!(parsed.cargo_lock_deps.is_empty());
538    }
539
540    #[test]
541    fn test_prefetch_plan_serde_roundtrip() {
542        let plan = PrefetchPlan {
543            plan_id: Some("plan-1".into()),
544            planner: Some("local".into()),
545            disposition: PrefetchDisposition::Execute,
546            candidates: vec![PrefetchCandidate::new("abc".into(), "serde".into())],
547        };
548
549        let json = serde_json::to_string(&plan).unwrap();
550        let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
551        assert_eq!(parsed, plan);
552    }
553
554    #[test]
555    fn test_prefetch_plan_missing_disposition_is_rejected() {
556        let err = serde_json::from_str::<PrefetchPlan>(
557            r#"{"planner":"legacy","candidates":[{"cache_key":"abc","crate_name":"serde"}]}"#,
558        )
559        .unwrap_err();
560        assert!(err.to_string().contains("missing field"));
561    }
562
563    #[test]
564    fn test_prefetch_plan_do_nothing_roundtrip() {
565        let plan = PrefetchPlan {
566            plan_id: Some("plan-2".into()),
567            planner: Some("remote".into()),
568            disposition: PrefetchDisposition::DoNothing,
569            candidates: vec![],
570        };
571
572        let json = serde_json::to_string(&plan).unwrap();
573        let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
574        assert_eq!(parsed, plan);
575    }
576
577    #[cfg(feature = "planning")]
578    #[tokio::test]
579    async fn test_build_prefetch_plan_prefers_shard_candidates() {
580        let source = FakePlannerDataSource {
581            shard_candidates: vec![PrefetchCandidate::new("from-shard".into(), "serde".into())],
582            ..Default::default()
583        };
584        let intent = BuildIntent {
585            crate_names: vec!["serde".into()],
586            namespace: Some("linux/hash/release".into()),
587            cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
588        };
589
590        let plan = build_prefetch_plan(&source, &intent, "fallback")
591            .await
592            .unwrap();
593
594        assert_eq!(plan.disposition, PrefetchDisposition::Execute);
595        assert_eq!(plan.planner.as_deref(), Some("fallback"));
596        assert_eq!(plan.candidates.len(), 1);
597        assert_eq!(plan.candidates[0].cache_key, "from-shard");
598    }
599
600    #[cfg(feature = "planning")]
601    #[tokio::test]
602    async fn test_build_prefetch_plan_falls_back_to_history_and_key_cache() {
603        let mut source = FakePlannerDataSource {
604            shard_error: true,
605            history_candidates: vec![PrefetchCandidate::new("history-key".into(), "serde".into())],
606            ..Default::default()
607        };
608        source.key_cache.insert(
609            "tokio".into(),
610            vec!["tokio-key".into(), "history-key".into()],
611        );
612
613        let intent = BuildIntent {
614            crate_names: vec!["serde".into(), "tokio".into()],
615            namespace: Some("linux/hash/debug".into()),
616            cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
617        };
618
619        let plan = build_prefetch_plan(&source, &intent, "fallback")
620            .await
621            .unwrap();
622
623        assert_eq!(plan.disposition, PrefetchDisposition::Execute);
624        assert_eq!(plan.candidates.len(), 2);
625        assert_eq!(plan.candidates[0].cache_key, "history-key");
626        assert_eq!(plan.candidates[1].cache_key, "tokio-key");
627    }
628
629    #[cfg(feature = "planning")]
630    #[tokio::test]
631    async fn test_build_prefetch_plan_orders_shard_candidates_by_crate_order() {
632        let source = FakePlannerDataSource {
633            shard_candidates: vec![
634                PrefetchCandidate::new("app-key".into(), "app".into()),
635                PrefetchCandidate::new("dep-key".into(), "dep".into()),
636                PrefetchCandidate::new("middle-key".into(), "middle".into()),
637            ],
638            ..Default::default()
639        };
640        let intent = BuildIntent {
641            crate_names: vec!["dep".into(), "middle".into(), "app".into()],
642            namespace: Some("linux/hash/debug".into()),
643            cargo_lock_deps: vec![("dep".into(), "1.0.0".into())],
644        };
645
646        let plan = build_prefetch_plan(&source, &intent, "fallback")
647            .await
648            .unwrap();
649
650        let keys = plan
651            .candidates
652            .iter()
653            .map(|candidate| candidate.cache_key.as_str())
654            .collect::<Vec<_>>();
655        assert_eq!(keys, vec!["dep-key", "middle-key", "app-key"]);
656    }
657
658    /// A PARTIAL shard hit must not suppress the lower-confidence sources
659    /// (kunobi-ninja/kache#614).
660    ///
661    /// Shard matching is exact per bucket, so a dependency bump invalidates
662    /// one bucket while the rest still match. The planner used to return as
663    /// soon as shards produced anything, so the crates in the missed buckets
664    /// were dropped from the plan even though history and the key cache could
665    /// resolve them.
666    #[cfg(feature = "planning")]
667    #[tokio::test]
668    async fn test_build_prefetch_plan_fills_crates_a_partial_shard_hit_missed() {
669        let mut source = FakePlannerDataSource {
670            // Only `dep` is in a bucket that still matches.
671            shard_candidates: vec![PrefetchCandidate::new("dep-shard-key".into(), "dep".into())],
672            history_by_crate: HashMap::from([("middle".into(), "middle-history-key".into())]),
673            ..Default::default()
674        };
675        source
676            .key_cache
677            .insert("app".into(), vec!["app-key-cache-key".into()]);
678
679        let intent = BuildIntent {
680            crate_names: vec!["dep".into(), "middle".into(), "app".into()],
681            namespace: Some("linux/hash/debug".into()),
682            cargo_lock_deps: vec![("dep".into(), "1.0.0".into())],
683        };
684
685        let plan = build_prefetch_plan(&source, &intent, "fallback")
686            .await
687            .unwrap();
688
689        let keys = plan
690            .candidates
691            .iter()
692            .map(|candidate| candidate.cache_key.as_str())
693            .collect::<Vec<_>>();
694        assert_eq!(
695            keys,
696            vec!["dep-shard-key", "middle-history-key", "app-key-cache-key"],
697            "each source should fill the crates the higher-confidence ones left unresolved"
698        );
699    }
700
701    /// A crate the shards already resolved is not re-queried from the
702    /// lower-confidence sources (#614): shard keys are exact, history and the
703    /// key cache are not, so they only fill gaps.
704    #[cfg(feature = "planning")]
705    #[tokio::test]
706    async fn test_build_prefetch_plan_does_not_requery_shard_resolved_crates() {
707        let mut source = FakePlannerDataSource {
708            shard_candidates: vec![PrefetchCandidate::new(
709                "serde-shard-key".into(),
710                "serde".into(),
711            )],
712            history_by_crate: HashMap::from([("serde".into(), "serde-stale-history-key".into())]),
713            ..Default::default()
714        };
715        source
716            .key_cache
717            .insert("serde".into(), vec!["serde-stale-key-cache-key".into()]);
718
719        let intent = BuildIntent {
720            crate_names: vec!["serde".into()],
721            namespace: Some("linux/hash/debug".into()),
722            cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
723        };
724
725        let plan = build_prefetch_plan(&source, &intent, "fallback")
726            .await
727            .unwrap();
728
729        let keys = plan
730            .candidates
731            .iter()
732            .map(|candidate| candidate.cache_key.as_str())
733            .collect::<Vec<_>>();
734        assert_eq!(keys, vec!["serde-shard-key"]);
735    }
736
737    #[cfg(feature = "planning")]
738    #[tokio::test]
739    async fn test_build_prefetch_plan_queries_history_by_crate_order() {
740        let source = FakePlannerDataSource {
741            history_by_crate: HashMap::from([
742                ("app".into(), "app-key".into()),
743                ("dep".into(), "dep-key".into()),
744                ("middle".into(), "middle-key".into()),
745            ]),
746            ..Default::default()
747        };
748        let intent = BuildIntent {
749            crate_names: vec!["dep".into(), "middle".into(), "app".into()],
750            namespace: None,
751            cargo_lock_deps: vec![],
752        };
753
754        let plan = build_prefetch_plan(&source, &intent, "fallback")
755            .await
756            .unwrap();
757
758        let keys = plan
759            .candidates
760            .iter()
761            .map(|candidate| candidate.cache_key.as_str())
762            .collect::<Vec<_>>();
763        assert_eq!(keys, vec!["dep-key", "middle-key", "app-key"]);
764    }
765
766    #[cfg(feature = "planning")]
767    proptest! {
768        #[test]
769        fn property_crate_query_order_preserves_first_occurrence(
770            crate_ids in proptest::collection::vec(any::<u8>(), 1..64),
771        ) {
772            let mut crate_names = crate_ids
773                .into_iter()
774                .map(|id| format!("crate-{id}"))
775                .collect::<Vec<_>>();
776            crate_names.push(crate_names[0].clone());
777            let intent = BuildIntent {
778                crate_names: crate_names.clone(),
779                ..Default::default()
780            };
781
782            let mut expected = Vec::new();
783            for crate_name in crate_names {
784                if !expected.contains(&crate_name) {
785                    expected.push(crate_name);
786                }
787            }
788
789            prop_assert_eq!(crate_query_order(&intent), expected);
790        }
791
792        #[test]
793        fn property_candidate_ordering_is_a_stable_priority_permutation(
794            requested_ids in proptest::collection::vec(0u8..8, 1..24),
795            candidate_ids in proptest::collection::vec(0u8..16, 0..48),
796        ) {
797            let requested_names = requested_ids
798                .into_iter()
799                .map(|id| format!("crate-{id}"))
800                .collect::<Vec<_>>();
801            let intent = BuildIntent {
802                crate_names: requested_names.clone(),
803                ..Default::default()
804            };
805
806            // Interleave two unknown and two known candidates so even minimal
807            // shrunk inputs exercise both priority and stable tie ordering.
808            let mut candidate_names = Vec::with_capacity(candidate_ids.len() + 4);
809            candidate_names.push("not-requested".to_string());
810            candidate_names.push(requested_names[0].clone());
811            candidate_names.extend(
812                candidate_ids
813                    .into_iter()
814                    .map(|id| format!("crate-{id}")),
815            );
816            candidate_names.push("not-requested".to_string());
817            candidate_names.push(requested_names[0].clone());
818
819            let candidates = candidate_names
820                .into_iter()
821                .enumerate()
822                .map(|(index, crate_name)| {
823                    PrefetchCandidate::new(format!("key-{index}"), crate_name)
824                })
825                .collect::<Vec<_>>();
826
827            let mut expected = candidates.clone();
828            expected.sort_by_key(|candidate| {
829                requested_names
830                    .iter()
831                    .position(|name| name == &candidate.crate_name)
832                    .unwrap_or(usize::MAX)
833            });
834
835            prop_assert_eq!(
836                order_candidates_by_crate_order(candidates, &intent),
837                expected
838            );
839        }
840    }
841    // ── Composition caps and ranking (#616, #617) ────────────────────────
842
843    #[cfg(feature = "planning")]
844    fn candidate(
845        key: &str,
846        crate_name: &str,
847        source: CandidateSource,
848        compile_time_ms: Option<u64>,
849        demand_index: Option<u32>,
850    ) -> PrefetchCandidate {
851        PrefetchCandidate {
852            cache_key: key.into(),
853            crate_name: crate_name.into(),
854            compile_time_ms,
855            size_bytes: None,
856            source,
857            demand_index,
858        }
859    }
860
861    /// Confidence ordering is the whole point of the rank; assert the ORDER,
862    /// not the literal numbers.
863    #[test]
864    fn test_confidence_rank_orders_sources() {
865        assert!(
866            CandidateSource::Shard.confidence_rank() < CandidateSource::History.confidence_rank()
867        );
868        assert!(
869            CandidateSource::History.confidence_rank()
870                < CandidateSource::KeyCache.confidence_rank()
871        );
872        assert!(
873            CandidateSource::KeyCache.confidence_rank()
874                < CandidateSource::Unknown.confidence_rank()
875        );
876    }
877
878    /// An unrecognised source name degrades to `Unknown` instead of failing the
879    /// whole plan (forward compatibility with a newer planner).
880    #[test]
881    fn test_unknown_candidate_source_deserializes() {
882        let candidate: PrefetchCandidate =
883            serde_json::from_str(r#"{"cache_key":"k","crate_name":"c","source":"telepathy"}"#)
884                .unwrap();
885        assert_eq!(candidate.source, CandidateSource::Unknown);
886    }
887
888    /// A pre-#617 candidate has no metadata and must not be rejected, and the
889    /// missing fields must read as unknown rather than zero.
890    #[test]
891    fn test_legacy_candidate_wire_still_parses() {
892        let candidate: PrefetchCandidate =
893            serde_json::from_str(r#"{"cache_key":"k","crate_name":"c"}"#).unwrap();
894        assert_eq!(candidate.compile_time_ms, None);
895        assert_eq!(candidate.size_bytes, None);
896        assert_eq!(candidate.demand_index, None);
897        assert_eq!(candidate.source, CandidateSource::Unknown);
898    }
899
900    /// Urgency dominates value: prefetch races the build, so a costly artifact
901    /// needed late loses to a cheaper one needed now (#617).
902    #[cfg(feature = "planning")]
903    #[test]
904    fn test_dispatch_order_prefers_urgency_over_value() {
905        let urgent_cheap = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
906        let late_expensive = candidate("b", "b", CandidateSource::Shard, Some(100_000), Some(500));
907        assert!(dispatch_sort_key(&urgent_cheap) < dispatch_sort_key(&late_expensive));
908    }
909
910    /// Within one urgency window, value decides.
911    #[cfg(feature = "planning")]
912    #[test]
913    fn test_dispatch_order_prefers_value_within_a_bucket() {
914        let cheap = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
915        let costly = candidate("b", "b", CandidateSource::Shard, Some(5_000), Some(1));
916        assert!(
917            dispatch_sort_key(&costly) < dispatch_sort_key(&cheap),
918            "same bucket, same source: the expensive rebuild goes first"
919        );
920    }
921
922    /// Positions inside one bucket are treated as equally urgent: guppy order
923    /// only approximates demand time, so finer distinctions are false precision.
924    #[cfg(feature = "planning")]
925    #[test]
926    fn test_dispatch_order_buckets_nearby_positions_together() {
927        let first = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
928        let nearby = candidate("b", "b", CandidateSource::Shard, Some(10), Some(15));
929        let next_bucket = candidate("c", "c", CandidateSource::Shard, Some(10), Some(16));
930        assert_eq!(
931            dispatch_sort_key(&first),
932            dispatch_sort_key(&nearby),
933            "positions 0 and 15 share a bucket"
934        );
935        assert!(
936            dispatch_sort_key(&nearby) < dispatch_sort_key(&next_bucket),
937            "position 16 starts the next bucket"
938        );
939    }
940
941    /// Within a bucket, confidence beats value: a shard match is worth more
942    /// than a bigger number from a guess.
943    #[cfg(feature = "planning")]
944    #[test]
945    fn test_dispatch_order_prefers_confidence_within_a_bucket() {
946        let shard_cheap = candidate("a", "a", CandidateSource::Shard, Some(1), Some(0));
947        let guess_costly = candidate("b", "b", CandidateSource::KeyCache, Some(99_999), Some(0));
948        assert!(dispatch_sort_key(&shard_cheap) < dispatch_sort_key(&guess_costly));
949    }
950
951    /// Unknown cost sorts after known cost, but is not dropped and is not
952    /// treated as zero-value against a *different* bucket.
953    #[cfg(feature = "planning")]
954    #[test]
955    fn test_dispatch_order_places_unknown_cost_last_within_its_bucket() {
956        let known = candidate("a", "a", CandidateSource::Shard, Some(1), Some(0));
957        let unknown = candidate("b", "b", CandidateSource::Shard, None, Some(0));
958        assert!(dispatch_sort_key(&known) < dispatch_sort_key(&unknown));
959
960        // ...but an unknown-cost candidate needed NOW still beats a known-cost
961        // one needed much later.
962        let late_known = candidate("c", "c", CandidateSource::Shard, Some(50_000), Some(999));
963        assert!(dispatch_sort_key(&unknown) < dispatch_sort_key(&late_known));
964    }
965
966    /// A candidate outside the intent's crate list sorts last rather than first.
967    #[cfg(feature = "planning")]
968    #[test]
969    fn test_dispatch_order_places_unknown_demand_last() {
970        let known = candidate("a", "a", CandidateSource::Shard, Some(1), Some(100_000));
971        let no_demand = candidate("b", "b", CandidateSource::Shard, Some(50_000), None);
972        assert!(dispatch_sort_key(&known) < dispatch_sort_key(&no_demand));
973    }
974
975    /// The per-crate key-cache cap keeps the first N and reports the rest.
976    #[cfg(feature = "planning")]
977    #[tokio::test]
978    async fn test_key_cache_per_crate_cap() {
979        let mut source = FakePlannerDataSource::default();
980        source.key_cache.insert(
981            "serde".into(),
982            vec!["k1".into(), "k2".into(), "k3".into(), "k4".into()],
983        );
984        let intent = BuildIntent {
985            crate_names: vec!["serde".into()],
986            ..Default::default()
987        };
988
989        let (plan, composition) = build_prefetch_plan_with_limits(
990            &source,
991            &intent,
992            "fallback",
993            PlanLimits {
994                key_cache_per_crate: 2,
995                ..PlanLimits::default()
996            },
997        )
998        .await
999        .unwrap();
1000
1001        assert_eq!(plan.candidates.len(), 2, "capped to two variants");
1002        assert_eq!(composition.from_key_cache, 2);
1003        assert_eq!(composition.dropped_key_cache_per_crate, 2);
1004        assert!(composition.dropped_total() > 0, "truncation is visible");
1005    }
1006
1007    /// The plan-wide key-cache cap stops one weak source filling the plan, and
1008    /// counts what whole crates it skipped.
1009    #[cfg(feature = "planning")]
1010    #[tokio::test]
1011    async fn test_key_cache_total_cap_counts_skipped_crates() {
1012        let mut source = FakePlannerDataSource::default();
1013        for name in ["a", "b", "c"] {
1014            source
1015                .key_cache
1016                .insert(name.into(), vec![format!("{name}-k1")]);
1017        }
1018        let intent = BuildIntent {
1019            crate_names: vec!["a".into(), "b".into(), "c".into()],
1020            ..Default::default()
1021        };
1022
1023        let (plan, composition) = build_prefetch_plan_with_limits(
1024            &source,
1025            &intent,
1026            "fallback",
1027            PlanLimits {
1028                key_cache_total: 1,
1029                ..PlanLimits::default()
1030            },
1031        )
1032        .await
1033        .unwrap();
1034
1035        assert_eq!(plan.candidates.len(), 1);
1036        assert_eq!(composition.from_key_cache, 1);
1037        assert_eq!(
1038            composition.dropped_key_cache_total, 2,
1039            "the two skipped crates are counted, not silently missing"
1040        );
1041    }
1042
1043    /// `0` disables a cap rather than dropping everything.
1044    #[cfg(feature = "planning")]
1045    #[tokio::test]
1046    async fn test_zero_limit_disables_the_cap() {
1047        let mut source = FakePlannerDataSource::default();
1048        source
1049            .key_cache
1050            .insert("serde".into(), vec!["k1".into(), "k2".into(), "k3".into()]);
1051        let intent = BuildIntent {
1052            crate_names: vec!["serde".into()],
1053            ..Default::default()
1054        };
1055
1056        let (plan, composition) = build_prefetch_plan_with_limits(
1057            &source,
1058            &intent,
1059            "fallback",
1060            PlanLimits {
1061                key_cache_per_crate: 0,
1062                key_cache_total: 0,
1063                history_per_crate: 0,
1064            },
1065        )
1066        .await
1067        .unwrap();
1068
1069        assert_eq!(plan.candidates.len(), 3);
1070        assert_eq!(composition.dropped_total(), 0);
1071    }
1072
1073    /// Candidates are stamped with their source and demand position, so the
1074    /// daemon can rank without the intent (it only receives the plan).
1075    #[cfg(feature = "planning")]
1076    #[tokio::test]
1077    async fn test_plan_stamps_source_and_demand_index() {
1078        let source = FakePlannerDataSource {
1079            history_by_crate: HashMap::from([
1080                ("dep".into(), "dep-key".into()),
1081                ("app".into(), "app-key".into()),
1082            ]),
1083            ..Default::default()
1084        };
1085        let intent = BuildIntent {
1086            crate_names: vec!["dep".into(), "app".into()],
1087            ..Default::default()
1088        };
1089
1090        let (plan, _) =
1091            build_prefetch_plan_with_limits(&source, &intent, "fallback", PlanLimits::default())
1092                .await
1093                .unwrap();
1094
1095        let dep = plan
1096            .candidates
1097            .iter()
1098            .find(|c| c.crate_name == "dep")
1099            .expect("dep planned");
1100        let app = plan
1101            .candidates
1102            .iter()
1103            .find(|c| c.crate_name == "app")
1104            .expect("app planned");
1105        assert_eq!(dep.source, CandidateSource::History);
1106        assert_eq!(dep.demand_index, Some(0));
1107        assert_eq!(app.demand_index, Some(1), "position follows crate order");
1108    }
1109
1110    /// `cap_to` computes the DROPPED count, not a ratio. 5 items capped to 2
1111    /// drops 3; a divide would say 2. (Mutation-driven: `-` vs `/` agree on
1112    /// 4-and-2, so the earlier cap test could not tell them apart.)
1113    #[cfg(feature = "planning")]
1114    #[test]
1115    fn test_cap_to_reports_the_dropped_count() {
1116        let mut items: Vec<PrefetchCandidate> = (0..5)
1117            .map(|i| PrefetchCandidate::new(format!("k{i}"), "c".into()))
1118            .collect();
1119        assert_eq!(cap_to(&mut items, 2), 3);
1120        assert_eq!(items.len(), 2);
1121        assert_eq!(items[0].cache_key, "k0", "keeps the FIRST N");
1122
1123        let mut exact: Vec<PrefetchCandidate> =
1124            vec![PrefetchCandidate::new("a".into(), "c".into())];
1125        assert_eq!(cap_to(&mut exact, 1), 0, "exactly at the limit drops none");
1126        assert_eq!(cap_to(&mut exact, 0), 0, "0 disables the cap");
1127        assert_eq!(exact.len(), 1);
1128    }
1129
1130    /// `dropped_total` sums every drop class. Each field is distinct so a
1131    /// swapped operator or a missing term changes the result.
1132    #[test]
1133    fn test_dropped_total_sums_every_class() {
1134        let composition = PlanComposition {
1135            dropped_history_per_crate: 1,
1136            dropped_key_cache_per_crate: 2,
1137            dropped_key_cache_total: 4,
1138            ..PlanComposition::default()
1139        };
1140        assert_eq!(composition.dropped_total(), 7);
1141        assert_eq!(PlanComposition::default().dropped_total(), 0);
1142    }
1143
1144    /// The per-crate history cap keeps the most-recent entries and reports the
1145    /// rest. History is ordered most-recently-used first by both data sources,
1146    /// so the cap keeps the freshest variants.
1147    #[cfg(feature = "planning")]
1148    #[tokio::test]
1149    async fn test_history_per_crate_cap() {
1150        let source = FakePlannerDataSource {
1151            history_candidates: vec![
1152                PrefetchCandidate::new("serde-newest".into(), "serde".into()),
1153                PrefetchCandidate::new("serde-older".into(), "serde".into()),
1154                PrefetchCandidate::new("serde-oldest".into(), "serde".into()),
1155            ],
1156            ..Default::default()
1157        };
1158        let intent = BuildIntent {
1159            crate_names: vec!["serde".into()],
1160            ..Default::default()
1161        };
1162
1163        let (plan, composition) = build_prefetch_plan_with_limits(
1164            &source,
1165            &intent,
1166            "fallback",
1167            PlanLimits {
1168                history_per_crate: 1,
1169                ..PlanLimits::default()
1170            },
1171        )
1172        .await
1173        .unwrap();
1174
1175        assert_eq!(composition.from_history, 1);
1176        assert_eq!(composition.dropped_history_per_crate, 2);
1177        assert_eq!(
1178            plan.candidates
1179                .iter()
1180                .map(|c| c.cache_key.as_str())
1181                .collect::<Vec<_>>(),
1182            vec!["serde-newest"]
1183        );
1184    }
1185
1186    /// Once the plan-wide key-cache budget is spent, a later crate is skipped
1187    /// WHOLESALE, and everything it would have offered is counted, not just the
1188    /// per-crate-capped subset.
1189    #[cfg(feature = "planning")]
1190    #[tokio::test]
1191    async fn test_key_cache_total_skips_a_whole_crate_and_counts_all_of_it() {
1192        let mut source = FakePlannerDataSource::default();
1193        source.key_cache.insert("a".into(), vec!["a1".into()]);
1194        source.key_cache.insert(
1195            "b".into(),
1196            vec![
1197                "b1".into(),
1198                "b2".into(),
1199                "b3".into(),
1200                "b4".into(),
1201                "b5".into(),
1202            ],
1203        );
1204        let intent = BuildIntent {
1205            crate_names: vec!["a".into(), "b".into()],
1206            ..Default::default()
1207        };
1208
1209        let (_plan, composition) = build_prefetch_plan_with_limits(
1210            &source,
1211            &intent,
1212            "fallback",
1213            PlanLimits {
1214                key_cache_per_crate: 2,
1215                key_cache_total: 1,
1216                history_per_crate: 2,
1217            },
1218        )
1219        .await
1220        .unwrap();
1221
1222        assert_eq!(composition.from_key_cache, 1, "only crate `a` fits");
1223        assert_eq!(
1224            composition.dropped_key_cache_total, 5,
1225            "all five of `b`'s variants are counted, not the two a per-crate \
1226             cap would have admitted"
1227        );
1228        assert_eq!(
1229            composition.dropped_key_cache_per_crate, 0,
1230            "`b` was skipped before the per-crate cap applied"
1231        );
1232    }
1233
1234    /// The budget can also run out PART WAY through one crate's variants: the
1235    /// remainder is dropped and counted individually.
1236    #[cfg(feature = "planning")]
1237    #[tokio::test]
1238    async fn test_key_cache_total_can_run_out_mid_crate() {
1239        let mut source = FakePlannerDataSource::default();
1240        source
1241            .key_cache
1242            .insert("serde".into(), vec!["k1".into(), "k2".into(), "k3".into()]);
1243        let intent = BuildIntent {
1244            crate_names: vec!["serde".into()],
1245            ..Default::default()
1246        };
1247
1248        let (_plan, composition) = build_prefetch_plan_with_limits(
1249            &source,
1250            &intent,
1251            "fallback",
1252            PlanLimits {
1253                key_cache_per_crate: 2,
1254                key_cache_total: 1,
1255                history_per_crate: 2,
1256            },
1257        )
1258        .await
1259        .unwrap();
1260
1261        assert_eq!(composition.from_key_cache, 1);
1262        assert_eq!(
1263            composition.dropped_key_cache_per_crate, 1,
1264            "k3 dropped by the per-crate cap"
1265        );
1266        assert_eq!(
1267            composition.dropped_key_cache_total, 1,
1268            "k2 dropped by the plan-wide cap, mid-crate"
1269        );
1270    }
1271}