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
10pub mod timeline;
11
12#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
13pub struct BuildIntent {
14 #[serde(default)]
15 pub crate_names: Vec<String>,
16 #[serde(default)]
17 pub namespace: Option<String>,
18 #[serde(default)]
19 pub cargo_lock_deps: Vec<(String, String)>,
20 #[serde(default)]
23 pub identity_key: Option<String>,
24 #[serde(default)]
27 pub lock_path: Option<String>,
28}
29
30#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
37#[cfg_attr(kani, derive(kani::Arbitrary))]
38#[serde(rename_all = "snake_case")]
39pub enum CandidateSource {
40 Manifest,
42 Shard,
44 History,
46 KeyCache,
49 #[default]
50 #[serde(other)]
51 Unknown,
52}
53
54impl CandidateSource {
55 pub fn confidence_rank(self) -> u8 {
58 match self {
59 CandidateSource::Manifest => 0,
60 CandidateSource::Shard => 1,
61 CandidateSource::History => 2,
62 CandidateSource::KeyCache => 3,
63 CandidateSource::Unknown => 4,
64 }
65 }
66}
67
68#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
69pub struct PrefetchCandidate {
70 pub cache_key: String,
71 pub crate_name: String,
72 #[serde(default, skip_serializing_if = "Option::is_none")]
76 pub compile_time_ms: Option<u64>,
77 #[serde(default, skip_serializing_if = "Option::is_none")]
81 pub size_bytes: Option<u64>,
82 #[serde(default)]
83 pub source: CandidateSource,
84 #[serde(default, skip_serializing_if = "Option::is_none")]
87 pub demand_index: Option<u32>,
88}
89
90impl PrefetchCandidate {
91 pub fn new(cache_key: String, crate_name: String) -> Self {
93 Self {
94 cache_key,
95 crate_name,
96 compile_time_ms: None,
97 size_bytes: None,
98 source: CandidateSource::Unknown,
99 demand_index: None,
100 }
101 }
102
103 pub fn with_source(mut self, source: CandidateSource) -> Self {
104 self.source = source;
105 self
106 }
107}
108
109#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
110#[serde(rename_all = "snake_case")]
111pub enum PrefetchDisposition {
112 Execute,
113 UseFallback,
114 DoNothing,
115}
116
117#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
118pub struct PrefetchPlan {
119 #[serde(default)]
120 pub plan_id: Option<String>,
121 #[serde(default)]
122 pub planner: Option<String>,
123 pub disposition: PrefetchDisposition,
124 #[serde(default)]
125 pub candidates: Vec<PrefetchCandidate>,
126}
127
128#[derive(Debug, Clone, Copy, PartialEq, Eq)]
136pub struct PlanLimits {
137 pub key_cache_per_crate: usize,
141 pub key_cache_total: usize,
144 pub history_per_crate: usize,
148}
149
150impl Default for PlanLimits {
151 fn default() -> Self {
152 Self {
153 key_cache_per_crate: 2,
154 key_cache_total: 64,
155 history_per_crate: 2,
156 }
157 }
158}
159
160#[cfg(feature = "planning")]
170const URGENCY_BUCKET: u32 = 16;
171
172#[cfg(feature = "planning")]
190pub fn dispatch_sort_key(
191 candidate: &PrefetchCandidate,
192) -> (u32, u8, std::cmp::Reverse<Option<u64>>) {
193 let bucket = candidate
194 .demand_index
195 .map(|index| index / URGENCY_BUCKET)
196 .unwrap_or(u32::MAX);
197 (
198 bucket,
199 candidate.source.confidence_rank(),
200 std::cmp::Reverse(candidate.compile_time_ms),
204 )
205}
206
207#[cfg(feature = "planning")]
208fn sort_candidates_for_dispatch(candidates: &mut [PrefetchCandidate]) {
209 candidates.sort_by_key(dispatch_sort_key);
210}
211
212#[cfg(feature = "planning")]
217fn cap_to(items: &mut Vec<PrefetchCandidate>, limit: usize) -> usize {
218 if limit == 0 || items.len() <= limit {
219 return 0;
220 }
221 let dropped = items.len() - limit;
222 items.truncate(limit);
223 dropped
224}
225
226#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
229pub struct PlanComposition {
230 pub from_identity: usize,
231 pub from_shards: usize,
232 pub from_history: usize,
233 pub from_key_cache: usize,
234 pub dropped_history_per_crate: usize,
236 pub dropped_key_cache_per_crate: usize,
237 pub dropped_key_cache_total: usize,
238}
239
240impl PlanComposition {
241 pub fn dropped_total(&self) -> usize {
242 self.dropped_history_per_crate
243 + self.dropped_key_cache_per_crate
244 + self.dropped_key_cache_total
245 }
246}
247
248#[cfg(feature = "planning")]
249#[async_trait]
250pub trait PlannerDataSource {
251 async fn shard_candidates(
252 &self,
253 namespace: &str,
254 deps: &[(String, String)],
255 ) -> Result<Vec<PrefetchCandidate>>;
256
257 async fn history_candidates(&self, crate_names: &[String]) -> Result<Vec<PrefetchCandidate>>;
258
259 async fn key_cache_keys_for_crate(&self, crate_name: &str) -> Result<Vec<String>>;
260
261 async fn identity_candidates(&self, identity_key: &str) -> Result<Vec<PrefetchCandidate>>;
264}
265
266#[cfg(feature = "planning")]
267pub async fn build_prefetch_plan<T>(
268 source: &T,
269 intent: &BuildIntent,
270 planner_name: &str,
271) -> Result<PrefetchPlan>
272where
273 T: PlannerDataSource + Sync + ?Sized,
274{
275 build_prefetch_plan_with_limits(source, intent, planner_name, PlanLimits::default())
276 .await
277 .map(|(plan, _composition)| plan)
278}
279
280#[cfg(feature = "planning")]
283pub async fn build_prefetch_plan_with_limits<T>(
284 source: &T,
285 intent: &BuildIntent,
286 planner_name: &str,
287 limits: PlanLimits,
288) -> Result<(PrefetchPlan, PlanComposition)>
289where
290 T: PlannerDataSource + Sync + ?Sized,
291{
292 let crate_order = crate_query_order(intent);
293 let demand_index = demand_index_map(&crate_order);
294 let mut seen = HashSet::new();
295 let mut resolved_crates = HashSet::new();
296 let mut candidates = Vec::new();
297 let mut composition = PlanComposition::default();
298
299 if let Some(identity_key) = intent
302 .identity_key
303 .as_deref()
304 .map(str::trim)
305 .filter(|key| !key.is_empty())
306 {
307 match source.identity_candidates(identity_key).await {
308 Ok(identity_candidates) => {
309 for (index, candidate) in identity_candidates.into_iter().enumerate() {
310 resolved_crates.insert(candidate.crate_name.clone());
311 if seen.insert(candidate.cache_key.clone()) {
312 let mut candidate = candidate.with_source(CandidateSource::Manifest);
313 if candidate.demand_index.is_none() {
314 candidate.demand_index = u32::try_from(index).ok();
315 }
316 candidates.push(candidate);
317 }
318 }
319 composition.from_identity = candidates.len();
320 }
321 Err(_) => {
322 }
324 }
325 }
326 let identity_resolved = resolved_crates.clone();
327
328 if let Some(namespace) = intent.namespace.as_deref()
335 && !intent.cargo_lock_deps.is_empty()
336 {
337 if let Ok(shard_candidates) = source
339 .shard_candidates(namespace, &intent.cargo_lock_deps)
340 .await
341 {
342 for candidate in order_candidates_by_crate_order(shard_candidates, intent) {
343 if identity_resolved.contains(&candidate.crate_name) {
344 continue;
345 }
346 resolved_crates.insert(candidate.crate_name.clone());
347 if seen.insert(candidate.cache_key.clone()) {
348 composition.from_shards += 1;
349 candidates.push(candidate.with_source(CandidateSource::Shard));
350 }
351 }
352 }
353 }
354
355 let unresolved = |resolved: &HashSet<String>| -> Vec<String> {
356 crate_order
357 .iter()
358 .filter(|name| !resolved.contains(*name))
359 .cloned()
360 .collect()
361 };
362
363 let history_query = unresolved(&resolved_crates);
364 if !history_query.is_empty() {
365 let mut by_crate: HashMap<String, Vec<PrefetchCandidate>> = HashMap::new();
368 for candidate in order_candidates_by_crate_order(
369 source.history_candidates(&history_query).await?,
370 intent,
371 ) {
372 by_crate
373 .entry(candidate.crate_name.clone())
374 .or_default()
375 .push(candidate.with_source(CandidateSource::History));
376 }
377
378 for crate_name in &history_query {
379 let Some(mut for_crate) = by_crate.remove(crate_name) else {
380 continue;
381 };
382 composition.dropped_history_per_crate +=
383 cap_to(&mut for_crate, limits.history_per_crate);
384 for candidate in for_crate {
385 resolved_crates.insert(candidate.crate_name.clone());
386 if seen.insert(candidate.cache_key.clone()) {
387 composition.from_history += 1;
388 candidates.push(candidate);
389 }
390 }
391 }
392 }
393
394 for crate_name in unresolved(&resolved_crates) {
400 if limits.key_cache_total > 0 && composition.from_key_cache >= limits.key_cache_total {
401 composition.dropped_key_cache_total += source
404 .key_cache_keys_for_crate(&crate_name)
405 .await?
406 .into_iter()
407 .filter(|key| !seen.contains(key))
408 .count();
409 continue;
410 }
411
412 let mut for_crate: Vec<PrefetchCandidate> = source
413 .key_cache_keys_for_crate(&crate_name)
414 .await?
415 .into_iter()
416 .map(|cache_key| {
417 PrefetchCandidate::new(cache_key, crate_name.clone())
418 .with_source(CandidateSource::KeyCache)
419 })
420 .collect();
421
422 composition.dropped_key_cache_per_crate +=
423 cap_to(&mut for_crate, limits.key_cache_per_crate);
424
425 for candidate in for_crate {
426 if limits.key_cache_total > 0 && composition.from_key_cache >= limits.key_cache_total {
427 composition.dropped_key_cache_total += 1;
428 continue;
429 }
430 if seen.insert(candidate.cache_key.clone()) {
431 composition.from_key_cache += 1;
432 candidates.push(candidate);
433 }
434 }
435 }
436
437 for candidate in &mut candidates {
441 if candidate.demand_index.is_none() {
442 candidate.demand_index = demand_index.get(&candidate.crate_name).copied();
443 }
444 }
445
446 sort_candidates_for_dispatch(&mut candidates);
449
450 Ok((execute_plan(planner_name, candidates), composition))
451}
452
453#[cfg(feature = "planning")]
454fn execute_plan(planner_name: &str, candidates: Vec<PrefetchCandidate>) -> PrefetchPlan {
455 let planner = planner_name.trim();
456 PrefetchPlan {
457 plan_id: None,
458 planner: Some(if planner.is_empty() {
459 "planner".to_string()
460 } else {
461 planner.to_string()
462 }),
463 disposition: PrefetchDisposition::Execute,
464 candidates,
465 }
466}
467
468#[cfg(feature = "planning")]
470fn demand_index_map(crate_order: &[String]) -> HashMap<String, u32> {
471 crate_order
472 .iter()
473 .enumerate()
474 .map(|(index, name)| (name.clone(), index as u32))
475 .collect()
476}
477
478#[cfg(feature = "planning")]
479fn crate_query_order(intent: &BuildIntent) -> Vec<String> {
480 let mut seen = HashSet::new();
481 intent
482 .crate_names
483 .iter()
484 .filter(|crate_name| seen.insert((*crate_name).clone()))
485 .cloned()
486 .collect()
487}
488
489#[cfg(feature = "planning")]
490fn order_candidates_by_crate_order(
491 mut candidates: Vec<PrefetchCandidate>,
492 intent: &BuildIntent,
493) -> Vec<PrefetchCandidate> {
494 if intent.crate_names.is_empty() {
495 return candidates;
496 }
497
498 let mut priority = HashMap::new();
499 for (index, crate_name) in crate_query_order(intent).iter().enumerate() {
500 priority.entry(crate_name.clone()).or_insert(index);
501 }
502
503 let mut indexed_candidates = candidates.drain(..).enumerate().collect::<Vec<_>>();
504 indexed_candidates.sort_by_key(|(index, candidate)| {
505 (
506 priority
507 .get(&candidate.crate_name)
508 .copied()
509 .unwrap_or(usize::MAX),
510 *index,
511 )
512 });
513 indexed_candidates
514 .into_iter()
515 .map(|(_, candidate)| candidate)
516 .collect()
517}
518
519#[cfg(all(kani, feature = "planning"))]
526mod kani_proofs {
527 use super::*;
528
529 fn candidate(
530 source: CandidateSource,
531 compile_time_ms: Option<u64>,
532 demand_index: Option<u32>,
533 ) -> PrefetchCandidate {
534 PrefetchCandidate {
535 cache_key: String::new(),
536 crate_name: String::new(),
537 compile_time_ms,
538 size_bytes: None,
539 source,
540 demand_index,
541 }
542 }
543
544 #[kani::proof]
545 fn known_demand_always_precedes_unknown_demand() {
546 let demand_index: u32 = kani::any();
547 let known = candidate(kani::any(), kani::any(), Some(demand_index));
548 let unknown = candidate(kani::any(), kani::any(), None);
549 kani::cover!();
550
551 assert!(dispatch_sort_key(&known) < dispatch_sort_key(&unknown));
552 }
553
554 #[kani::proof]
555 fn earlier_urgency_bucket_always_wins() {
556 let earlier: u32 = kani::any();
557 let later: u32 = kani::any();
558 kani::assume(earlier / URGENCY_BUCKET < later / URGENCY_BUCKET);
559 let earlier_candidate = candidate(kani::any(), kani::any(), Some(earlier));
560 let later_candidate = candidate(kani::any(), kani::any(), Some(later));
561 kani::cover!();
562
563 assert!(dispatch_sort_key(&earlier_candidate) < dispatch_sort_key(&later_candidate));
564 }
565
566 #[kani::proof]
567 fn confidence_always_wins_within_a_bucket() {
568 let better_source: CandidateSource = kani::any();
569 let worse_source: CandidateSource = kani::any();
570 kani::assume(better_source.confidence_rank() < worse_source.confidence_rank());
571 let better_demand: u32 = kani::any();
572 let worse_demand: u32 = kani::any();
573 kani::assume(better_demand / URGENCY_BUCKET == worse_demand / URGENCY_BUCKET);
574 let better = candidate(better_source, kani::any(), Some(better_demand));
575 let worse = candidate(worse_source, kani::any(), Some(worse_demand));
576 kani::cover!();
577
578 assert!(dispatch_sort_key(&better) < dispatch_sort_key(&worse));
579 }
580
581 #[kani::proof]
582 fn higher_compile_cost_wins_after_urgency_and_confidence() {
583 let cheaper: u64 = kani::any();
584 let costlier: u64 = kani::any();
585 kani::assume(cheaper < costlier);
586 let cheaper_demand: u32 = kani::any();
587 let costlier_demand: u32 = kani::any();
588 kani::assume(cheaper_demand / URGENCY_BUCKET == costlier_demand / URGENCY_BUCKET);
589 let source: CandidateSource = kani::any();
590 let cheaper_candidate = candidate(source, Some(cheaper), Some(cheaper_demand));
591 let costlier_candidate = candidate(source, Some(costlier), Some(costlier_demand));
592 kani::cover!();
593
594 assert!(dispatch_sort_key(&costlier_candidate) < dispatch_sort_key(&cheaper_candidate));
595 }
596
597 #[kani::proof]
598 fn known_compile_cost_including_zero_precedes_unknown_cost() {
599 let known_cost: u64 = kani::any();
600 let demand_index: u32 = kani::any();
601 let source: CandidateSource = kani::any();
602 let known = candidate(source, Some(known_cost), Some(demand_index));
603 let unknown = candidate(source, None, Some(demand_index));
604 kani::cover!(known_cost == 0);
605
606 assert!(dispatch_sort_key(&known) < dispatch_sort_key(&unknown));
607 }
608
609 #[kani::proof]
610 fn dispatch_sort_preserves_equal_key_order() {
611 let source: CandidateSource = kani::any();
612 let compile_time_ms: Option<u64> = kani::any();
613 let demand_index: Option<u32> = kani::any();
614 let mut first = candidate(source, compile_time_ms, demand_index);
615 first.size_bytes = Some(1);
616 let mut second = candidate(source, compile_time_ms, demand_index);
617 second.size_bytes = Some(2);
618 let mut candidates = [first, second];
619
620 sort_candidates_for_dispatch(&mut candidates);
621 kani::cover!();
622
623 assert_eq!(candidates[0].size_bytes, Some(1));
624 assert_eq!(candidates[1].size_bytes, Some(2));
625 }
626}
627
628#[cfg(test)]
629mod tests {
630 use super::*;
631
632 #[cfg(feature = "planning")]
633 use std::collections::HashMap;
634
635 #[cfg(feature = "planning")]
636 use anyhow::anyhow;
637
638 #[cfg(feature = "planning")]
639 use proptest::prelude::*;
640
641 #[cfg(feature = "planning")]
642 #[derive(Default)]
643 struct FakePlannerDataSource {
644 identity_candidates: Vec<PrefetchCandidate>,
645 shard_candidates: Vec<PrefetchCandidate>,
646 shard_error: bool,
647 history_candidates: Vec<PrefetchCandidate>,
648 history_by_crate: HashMap<String, String>,
649 key_cache: HashMap<String, Vec<String>>,
650 }
651
652 #[cfg(feature = "planning")]
653 #[async_trait]
654 impl PlannerDataSource for FakePlannerDataSource {
655 async fn shard_candidates(
656 &self,
657 _namespace: &str,
658 _deps: &[(String, String)],
659 ) -> Result<Vec<PrefetchCandidate>> {
660 if self.shard_error {
661 Err(anyhow!("shard lookup failed"))
662 } else {
663 Ok(self.shard_candidates.clone())
664 }
665 }
666
667 async fn history_candidates(
668 &self,
669 crate_names: &[String],
670 ) -> Result<Vec<PrefetchCandidate>> {
671 if !self.history_candidates.is_empty() {
672 return Ok(self.history_candidates.clone());
673 }
674
675 Ok(crate_names
676 .iter()
677 .filter_map(|crate_name| {
678 self.history_by_crate.get(crate_name).map(|cache_key| {
679 PrefetchCandidate::new(cache_key.clone(), crate_name.clone())
680 })
681 })
682 .collect())
683 }
684
685 async fn key_cache_keys_for_crate(&self, crate_name: &str) -> Result<Vec<String>> {
686 Ok(self.key_cache.get(crate_name).cloned().unwrap_or_default())
687 }
688
689 async fn identity_candidates(&self, _identity_key: &str) -> Result<Vec<PrefetchCandidate>> {
690 Ok(self.identity_candidates.clone())
691 }
692 }
693
694 #[test]
695 fn test_build_intent_serde_roundtrip() {
696 let intent = BuildIntent {
697 crate_names: vec!["serde".into(), "tokio".into()],
698 namespace: Some("x86_64/hash/release".into()),
699 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
700 identity_key: None,
701 lock_path: None,
702 };
703
704 let json = serde_json::to_string(&intent).unwrap();
705 let parsed: BuildIntent = serde_json::from_str(&json).unwrap();
706 assert_eq!(parsed, intent);
707 }
708
709 #[test]
710 fn test_build_intent_defaults_missing_fields() {
711 let parsed: BuildIntent = serde_json::from_str(r#"{"crate_names":["serde"]}"#).unwrap();
712 assert_eq!(parsed.crate_names, vec!["serde"]);
713 assert!(parsed.namespace.is_none());
714 assert!(parsed.cargo_lock_deps.is_empty());
715 assert!(parsed.identity_key.is_none());
716 }
717
718 #[test]
719 fn test_prefetch_plan_serde_roundtrip() {
720 let plan = PrefetchPlan {
721 plan_id: Some("plan-1".into()),
722 planner: Some("local".into()),
723 disposition: PrefetchDisposition::Execute,
724 candidates: vec![PrefetchCandidate::new("abc".into(), "serde".into())],
725 };
726
727 let json = serde_json::to_string(&plan).unwrap();
728 let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
729 assert_eq!(parsed, plan);
730 }
731
732 #[test]
733 fn test_prefetch_plan_missing_disposition_is_rejected() {
734 let err = serde_json::from_str::<PrefetchPlan>(
735 r#"{"planner":"legacy","candidates":[{"cache_key":"abc","crate_name":"serde"}]}"#,
736 )
737 .unwrap_err();
738 assert!(err.to_string().contains("missing field"));
739 }
740
741 #[test]
742 fn test_prefetch_plan_do_nothing_roundtrip() {
743 let plan = PrefetchPlan {
744 plan_id: Some("plan-2".into()),
745 planner: Some("remote".into()),
746 disposition: PrefetchDisposition::DoNothing,
747 candidates: vec![],
748 };
749
750 let json = serde_json::to_string(&plan).unwrap();
751 let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
752 assert_eq!(parsed, plan);
753 }
754
755 #[cfg(feature = "planning")]
756 #[tokio::test]
757 async fn test_build_prefetch_plan_prefers_shard_candidates() {
758 let source = FakePlannerDataSource {
759 shard_candidates: vec![PrefetchCandidate::new("from-shard".into(), "serde".into())],
760 ..Default::default()
761 };
762 let intent = BuildIntent {
763 crate_names: vec!["serde".into()],
764 namespace: Some("linux/hash/release".into()),
765 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
766 identity_key: None,
767 lock_path: None,
768 };
769
770 let (plan, composition) =
771 build_prefetch_plan_with_limits(&source, &intent, "fallback", PlanLimits::default())
772 .await
773 .unwrap();
774
775 assert_eq!(plan.disposition, PrefetchDisposition::Execute);
776 assert_eq!(plan.planner.as_deref(), Some("fallback"));
777 assert_eq!(plan.candidates.len(), 1);
778 assert_eq!(plan.candidates[0].cache_key, "from-shard");
779 assert_eq!(composition.from_shards, 1);
780 }
781
782 #[cfg(feature = "planning")]
783 #[tokio::test]
784 async fn test_build_prefetch_plan_prefers_identity_manifest_over_shards() {
785 let source = FakePlannerDataSource {
786 identity_candidates: vec![PrefetchCandidate::new(
787 "from-identity".into(),
788 "serde".into(),
789 )],
790 shard_candidates: vec![PrefetchCandidate::new("from-shard".into(), "serde".into())],
791 ..Default::default()
792 };
793 let intent = BuildIntent {
794 crate_names: vec!["serde".into()],
795 namespace: Some("linux/hash/release".into()),
796 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
797 identity_key: Some("id/abcd/x86_64-unknown-linux-gnu/release".into()),
798 lock_path: None,
799 };
800
801 let plan = build_prefetch_plan(&source, &intent, "fallback")
802 .await
803 .unwrap();
804
805 assert_eq!(plan.candidates.len(), 1);
806 assert_eq!(plan.candidates[0].cache_key, "from-identity");
807 assert_eq!(plan.candidates[0].source, CandidateSource::Manifest);
808 assert_eq!(plan.candidates[0].demand_index, Some(0));
809 }
810
811 #[cfg(feature = "planning")]
812 #[tokio::test]
813 async fn test_build_prefetch_plan_falls_back_to_history_and_key_cache() {
814 let mut source = FakePlannerDataSource {
815 shard_error: true,
816 history_candidates: vec![PrefetchCandidate::new("history-key".into(), "serde".into())],
817 ..Default::default()
818 };
819 source.key_cache.insert(
820 "tokio".into(),
821 vec!["tokio-key".into(), "history-key".into()],
822 );
823
824 let intent = BuildIntent {
825 crate_names: vec!["serde".into(), "tokio".into()],
826 namespace: Some("linux/hash/debug".into()),
827 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
828 identity_key: None,
829 lock_path: None,
830 };
831
832 let plan = build_prefetch_plan(&source, &intent, "fallback")
833 .await
834 .unwrap();
835
836 assert_eq!(plan.disposition, PrefetchDisposition::Execute);
837 assert_eq!(plan.candidates.len(), 2);
838 assert_eq!(plan.candidates[0].cache_key, "history-key");
839 assert_eq!(plan.candidates[1].cache_key, "tokio-key");
840 }
841
842 #[cfg(feature = "planning")]
843 #[tokio::test]
844 async fn test_build_prefetch_plan_orders_shard_candidates_by_crate_order() {
845 let source = FakePlannerDataSource {
846 shard_candidates: vec![
847 PrefetchCandidate::new("app-key".into(), "app".into()),
848 PrefetchCandidate::new("dep-key".into(), "dep".into()),
849 PrefetchCandidate::new("middle-key".into(), "middle".into()),
850 ],
851 ..Default::default()
852 };
853 let intent = BuildIntent {
854 crate_names: vec!["dep".into(), "middle".into(), "app".into()],
855 namespace: Some("linux/hash/debug".into()),
856 cargo_lock_deps: vec![("dep".into(), "1.0.0".into())],
857 identity_key: None,
858 lock_path: None,
859 };
860
861 let plan = build_prefetch_plan(&source, &intent, "fallback")
862 .await
863 .unwrap();
864
865 let keys = plan
866 .candidates
867 .iter()
868 .map(|candidate| candidate.cache_key.as_str())
869 .collect::<Vec<_>>();
870 assert_eq!(keys, vec!["dep-key", "middle-key", "app-key"]);
871 }
872
873 #[cfg(feature = "planning")]
882 #[tokio::test]
883 async fn test_build_prefetch_plan_fills_crates_a_partial_shard_hit_missed() {
884 let mut source = FakePlannerDataSource {
885 shard_candidates: vec![PrefetchCandidate::new("dep-shard-key".into(), "dep".into())],
887 history_by_crate: HashMap::from([("middle".into(), "middle-history-key".into())]),
888 ..Default::default()
889 };
890 source
891 .key_cache
892 .insert("app".into(), vec!["app-key-cache-key".into()]);
893
894 let intent = BuildIntent {
895 crate_names: vec!["dep".into(), "middle".into(), "app".into()],
896 namespace: Some("linux/hash/debug".into()),
897 cargo_lock_deps: vec![("dep".into(), "1.0.0".into())],
898 identity_key: None,
899 lock_path: None,
900 };
901
902 let plan = build_prefetch_plan(&source, &intent, "fallback")
903 .await
904 .unwrap();
905
906 let keys = plan
907 .candidates
908 .iter()
909 .map(|candidate| candidate.cache_key.as_str())
910 .collect::<Vec<_>>();
911 assert_eq!(
912 keys,
913 vec!["dep-shard-key", "middle-history-key", "app-key-cache-key"],
914 "each source should fill the crates the higher-confidence ones left unresolved"
915 );
916 }
917
918 #[cfg(feature = "planning")]
922 #[tokio::test]
923 async fn test_build_prefetch_plan_does_not_requery_shard_resolved_crates() {
924 let mut source = FakePlannerDataSource {
925 shard_candidates: vec![PrefetchCandidate::new(
926 "serde-shard-key".into(),
927 "serde".into(),
928 )],
929 history_by_crate: HashMap::from([("serde".into(), "serde-stale-history-key".into())]),
930 ..Default::default()
931 };
932 source
933 .key_cache
934 .insert("serde".into(), vec!["serde-stale-key-cache-key".into()]);
935
936 let intent = BuildIntent {
937 crate_names: vec!["serde".into()],
938 namespace: Some("linux/hash/debug".into()),
939 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
940 identity_key: None,
941 lock_path: None,
942 };
943
944 let plan = build_prefetch_plan(&source, &intent, "fallback")
945 .await
946 .unwrap();
947
948 let keys = plan
949 .candidates
950 .iter()
951 .map(|candidate| candidate.cache_key.as_str())
952 .collect::<Vec<_>>();
953 assert_eq!(keys, vec!["serde-shard-key"]);
954 }
955
956 #[cfg(feature = "planning")]
957 #[tokio::test]
958 async fn test_build_prefetch_plan_queries_history_by_crate_order() {
959 let source = FakePlannerDataSource {
960 history_by_crate: HashMap::from([
961 ("app".into(), "app-key".into()),
962 ("dep".into(), "dep-key".into()),
963 ("middle".into(), "middle-key".into()),
964 ]),
965 ..Default::default()
966 };
967 let intent = BuildIntent {
968 crate_names: vec!["dep".into(), "middle".into(), "app".into()],
969 namespace: None,
970 cargo_lock_deps: vec![],
971 identity_key: None,
972 lock_path: None,
973 };
974
975 let plan = build_prefetch_plan(&source, &intent, "fallback")
976 .await
977 .unwrap();
978
979 let keys = plan
980 .candidates
981 .iter()
982 .map(|candidate| candidate.cache_key.as_str())
983 .collect::<Vec<_>>();
984 assert_eq!(keys, vec!["dep-key", "middle-key", "app-key"]);
985 }
986
987 #[cfg(feature = "planning")]
988 proptest! {
989 #[test]
990 fn property_crate_query_order_preserves_first_occurrence(
991 crate_ids in proptest::collection::vec(any::<u8>(), 1..64),
992 ) {
993 let mut crate_names = crate_ids
994 .into_iter()
995 .map(|id| format!("crate-{id}"))
996 .collect::<Vec<_>>();
997 crate_names.push(crate_names[0].clone());
998 let intent = BuildIntent {
999 crate_names: crate_names.clone(),
1000 ..Default::default()
1001 };
1002
1003 let mut expected = Vec::new();
1004 for crate_name in crate_names {
1005 if !expected.contains(&crate_name) {
1006 expected.push(crate_name);
1007 }
1008 }
1009
1010 prop_assert_eq!(crate_query_order(&intent), expected);
1011 }
1012
1013 #[test]
1014 fn property_candidate_ordering_is_a_stable_priority_permutation(
1015 requested_ids in proptest::collection::vec(0u8..8, 1..24),
1016 candidate_ids in proptest::collection::vec(0u8..16, 0..48),
1017 ) {
1018 let requested_names = requested_ids
1019 .into_iter()
1020 .map(|id| format!("crate-{id}"))
1021 .collect::<Vec<_>>();
1022 let intent = BuildIntent {
1023 crate_names: requested_names.clone(),
1024 ..Default::default()
1025 };
1026
1027 let mut candidate_names = Vec::with_capacity(candidate_ids.len() + 4);
1030 candidate_names.push("not-requested".to_string());
1031 candidate_names.push(requested_names[0].clone());
1032 candidate_names.extend(
1033 candidate_ids
1034 .into_iter()
1035 .map(|id| format!("crate-{id}")),
1036 );
1037 candidate_names.push("not-requested".to_string());
1038 candidate_names.push(requested_names[0].clone());
1039
1040 let candidates = candidate_names
1041 .into_iter()
1042 .enumerate()
1043 .map(|(index, crate_name)| {
1044 PrefetchCandidate::new(format!("key-{index}"), crate_name)
1045 })
1046 .collect::<Vec<_>>();
1047
1048 let mut expected = candidates.clone();
1049 expected.sort_by_key(|candidate| {
1050 requested_names
1051 .iter()
1052 .position(|name| name == &candidate.crate_name)
1053 .unwrap_or(usize::MAX)
1054 });
1055
1056 prop_assert_eq!(
1057 order_candidates_by_crate_order(candidates, &intent),
1058 expected
1059 );
1060 }
1061 }
1062 #[cfg(feature = "planning")]
1065 fn candidate(
1066 key: &str,
1067 crate_name: &str,
1068 source: CandidateSource,
1069 compile_time_ms: Option<u64>,
1070 demand_index: Option<u32>,
1071 ) -> PrefetchCandidate {
1072 PrefetchCandidate {
1073 cache_key: key.into(),
1074 crate_name: crate_name.into(),
1075 compile_time_ms,
1076 size_bytes: None,
1077 source,
1078 demand_index,
1079 }
1080 }
1081
1082 #[test]
1085 fn test_confidence_rank_orders_sources() {
1086 assert!(
1087 CandidateSource::Manifest.confidence_rank() < CandidateSource::Shard.confidence_rank()
1088 );
1089 assert!(
1090 CandidateSource::Shard.confidence_rank() < CandidateSource::History.confidence_rank()
1091 );
1092 assert!(
1093 CandidateSource::History.confidence_rank()
1094 < CandidateSource::KeyCache.confidence_rank()
1095 );
1096 assert!(
1097 CandidateSource::KeyCache.confidence_rank()
1098 < CandidateSource::Unknown.confidence_rank()
1099 );
1100 }
1101
1102 #[test]
1105 fn test_unknown_candidate_source_deserializes() {
1106 let candidate: PrefetchCandidate =
1107 serde_json::from_str(r#"{"cache_key":"k","crate_name":"c","source":"telepathy"}"#)
1108 .unwrap();
1109 assert_eq!(candidate.source, CandidateSource::Unknown);
1110 }
1111
1112 #[test]
1115 fn test_legacy_candidate_wire_still_parses() {
1116 let candidate: PrefetchCandidate =
1117 serde_json::from_str(r#"{"cache_key":"k","crate_name":"c"}"#).unwrap();
1118 assert_eq!(candidate.compile_time_ms, None);
1119 assert_eq!(candidate.size_bytes, None);
1120 assert_eq!(candidate.demand_index, None);
1121 assert_eq!(candidate.source, CandidateSource::Unknown);
1122 }
1123
1124 #[cfg(feature = "planning")]
1127 #[test]
1128 fn test_dispatch_order_prefers_urgency_over_value() {
1129 let urgent_cheap = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
1130 let late_expensive = candidate("b", "b", CandidateSource::Shard, Some(100_000), Some(500));
1131 assert!(dispatch_sort_key(&urgent_cheap) < dispatch_sort_key(&late_expensive));
1132 }
1133
1134 #[cfg(feature = "planning")]
1136 #[test]
1137 fn test_dispatch_order_prefers_value_within_a_bucket() {
1138 let cheap = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
1139 let costly = candidate("b", "b", CandidateSource::Shard, Some(5_000), Some(1));
1140 assert!(
1141 dispatch_sort_key(&costly) < dispatch_sort_key(&cheap),
1142 "same bucket, same source: the expensive rebuild goes first"
1143 );
1144 }
1145
1146 #[cfg(feature = "planning")]
1150 #[test]
1151 fn test_dispatch_order_buckets_nearby_positions_together() {
1152 let first = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
1153 let nearby = candidate("b", "b", CandidateSource::Shard, Some(10), Some(15));
1154 let next_bucket = candidate("c", "c", CandidateSource::Shard, Some(10), Some(16));
1155 assert_eq!(
1156 dispatch_sort_key(&first),
1157 dispatch_sort_key(&nearby),
1158 "positions 0 and 15 share a bucket"
1159 );
1160 assert!(
1161 dispatch_sort_key(&nearby) < dispatch_sort_key(&next_bucket),
1162 "position 16 starts the next bucket"
1163 );
1164 }
1165
1166 #[cfg(feature = "planning")]
1169 #[test]
1170 fn test_dispatch_order_prefers_confidence_within_a_bucket() {
1171 let shard_cheap = candidate("a", "a", CandidateSource::Shard, Some(1), Some(0));
1172 let guess_costly = candidate("b", "b", CandidateSource::KeyCache, Some(99_999), Some(0));
1173 assert!(dispatch_sort_key(&shard_cheap) < dispatch_sort_key(&guess_costly));
1174 }
1175
1176 #[cfg(feature = "planning")]
1179 #[test]
1180 fn test_dispatch_order_places_unknown_cost_last_within_its_bucket() {
1181 let known = candidate("a", "a", CandidateSource::Shard, Some(1), Some(0));
1182 let known_zero = candidate("zero", "zero", CandidateSource::Shard, Some(0), Some(0));
1183 let unknown = candidate("b", "b", CandidateSource::Shard, None, Some(0));
1184 assert!(dispatch_sort_key(&known) < dispatch_sort_key(&unknown));
1185 assert!(dispatch_sort_key(&known_zero) < dispatch_sort_key(&unknown));
1186
1187 let late_known = candidate("c", "c", CandidateSource::Shard, Some(50_000), Some(999));
1190 assert!(dispatch_sort_key(&unknown) < dispatch_sort_key(&late_known));
1191 }
1192
1193 #[cfg(feature = "planning")]
1195 #[test]
1196 fn test_dispatch_order_places_unknown_demand_last() {
1197 let known = candidate("a", "a", CandidateSource::Shard, Some(1), Some(100_000));
1198 let no_demand = candidate("b", "b", CandidateSource::Shard, Some(50_000), None);
1199 assert!(dispatch_sort_key(&known) < dispatch_sort_key(&no_demand));
1200 }
1201
1202 #[cfg(feature = "planning")]
1204 #[test]
1205 fn test_dispatch_sort_orders_candidates_and_preserves_ties() {
1206 let late = candidate("late", "late", CandidateSource::Shard, Some(1), Some(32));
1207 let first_tie = candidate("first", "first", CandidateSource::Shard, Some(10), Some(0));
1208 let second_tie = candidate(
1209 "second",
1210 "second",
1211 CandidateSource::Shard,
1212 Some(10),
1213 Some(0),
1214 );
1215 let mut candidates = vec![late, first_tie, second_tie];
1216
1217 sort_candidates_for_dispatch(&mut candidates);
1218
1219 let keys = candidates
1220 .iter()
1221 .map(|candidate| candidate.cache_key.as_str())
1222 .collect::<Vec<_>>();
1223 assert_eq!(keys, vec!["first", "second", "late"]);
1224 }
1225
1226 #[cfg(feature = "planning")]
1228 #[tokio::test]
1229 async fn test_key_cache_per_crate_cap() {
1230 let mut source = FakePlannerDataSource::default();
1231 source.key_cache.insert(
1232 "serde".into(),
1233 vec!["k1".into(), "k2".into(), "k3".into(), "k4".into()],
1234 );
1235 let intent = BuildIntent {
1236 crate_names: vec!["serde".into()],
1237 ..Default::default()
1238 };
1239
1240 let (plan, composition) = build_prefetch_plan_with_limits(
1241 &source,
1242 &intent,
1243 "fallback",
1244 PlanLimits {
1245 key_cache_per_crate: 2,
1246 ..PlanLimits::default()
1247 },
1248 )
1249 .await
1250 .unwrap();
1251
1252 assert_eq!(plan.candidates.len(), 2, "capped to two variants");
1253 assert_eq!(composition.from_key_cache, 2);
1254 assert_eq!(composition.dropped_key_cache_per_crate, 2);
1255 assert!(composition.dropped_total() > 0, "truncation is visible");
1256 }
1257
1258 #[cfg(feature = "planning")]
1261 #[tokio::test]
1262 async fn test_key_cache_total_cap_counts_skipped_crates() {
1263 let mut source = FakePlannerDataSource::default();
1264 for name in ["a", "b", "c"] {
1265 source
1266 .key_cache
1267 .insert(name.into(), vec![format!("{name}-k1")]);
1268 }
1269 let intent = BuildIntent {
1270 crate_names: vec!["a".into(), "b".into(), "c".into()],
1271 ..Default::default()
1272 };
1273
1274 let (plan, composition) = build_prefetch_plan_with_limits(
1275 &source,
1276 &intent,
1277 "fallback",
1278 PlanLimits {
1279 key_cache_total: 1,
1280 ..PlanLimits::default()
1281 },
1282 )
1283 .await
1284 .unwrap();
1285
1286 assert_eq!(plan.candidates.len(), 1);
1287 assert_eq!(composition.from_key_cache, 1);
1288 assert_eq!(
1289 composition.dropped_key_cache_total, 2,
1290 "the two skipped crates are counted, not silently missing"
1291 );
1292 }
1293
1294 #[cfg(feature = "planning")]
1296 #[tokio::test]
1297 async fn test_zero_limit_disables_the_cap() {
1298 let mut source = FakePlannerDataSource::default();
1299 source
1300 .key_cache
1301 .insert("serde".into(), vec!["k1".into(), "k2".into(), "k3".into()]);
1302 let intent = BuildIntent {
1303 crate_names: vec!["serde".into()],
1304 ..Default::default()
1305 };
1306
1307 let (plan, composition) = build_prefetch_plan_with_limits(
1308 &source,
1309 &intent,
1310 "fallback",
1311 PlanLimits {
1312 key_cache_per_crate: 0,
1313 key_cache_total: 0,
1314 history_per_crate: 0,
1315 },
1316 )
1317 .await
1318 .unwrap();
1319
1320 assert_eq!(plan.candidates.len(), 3);
1321 assert_eq!(composition.dropped_total(), 0);
1322 }
1323
1324 #[cfg(feature = "planning")]
1327 #[tokio::test]
1328 async fn test_plan_stamps_source_and_demand_index() {
1329 let source = FakePlannerDataSource {
1330 history_by_crate: HashMap::from([
1331 ("dep".into(), "dep-key".into()),
1332 ("app".into(), "app-key".into()),
1333 ]),
1334 ..Default::default()
1335 };
1336 let intent = BuildIntent {
1337 crate_names: vec!["dep".into(), "app".into()],
1338 ..Default::default()
1339 };
1340
1341 let (plan, _) =
1342 build_prefetch_plan_with_limits(&source, &intent, "fallback", PlanLimits::default())
1343 .await
1344 .unwrap();
1345
1346 let dep = plan
1347 .candidates
1348 .iter()
1349 .find(|c| c.crate_name == "dep")
1350 .expect("dep planned");
1351 let app = plan
1352 .candidates
1353 .iter()
1354 .find(|c| c.crate_name == "app")
1355 .expect("app planned");
1356 assert_eq!(dep.source, CandidateSource::History);
1357 assert_eq!(dep.demand_index, Some(0));
1358 assert_eq!(app.demand_index, Some(1), "position follows crate order");
1359 }
1360
1361 #[cfg(feature = "planning")]
1365 #[test]
1366 fn test_cap_to_reports_the_dropped_count() {
1367 let mut items: Vec<PrefetchCandidate> = (0..5)
1368 .map(|i| PrefetchCandidate::new(format!("k{i}"), "c".into()))
1369 .collect();
1370 assert_eq!(cap_to(&mut items, 2), 3);
1371 assert_eq!(items.len(), 2);
1372 assert_eq!(items[0].cache_key, "k0", "keeps the FIRST N");
1373
1374 let mut exact: Vec<PrefetchCandidate> =
1375 vec![PrefetchCandidate::new("a".into(), "c".into())];
1376 assert_eq!(cap_to(&mut exact, 1), 0, "exactly at the limit drops none");
1377 assert_eq!(cap_to(&mut exact, 0), 0, "0 disables the cap");
1378 assert_eq!(exact.len(), 1);
1379 }
1380
1381 #[test]
1384 fn test_dropped_total_sums_every_class() {
1385 let composition = PlanComposition {
1386 dropped_history_per_crate: 1,
1387 dropped_key_cache_per_crate: 2,
1388 dropped_key_cache_total: 4,
1389 ..PlanComposition::default()
1390 };
1391 assert_eq!(composition.dropped_total(), 7);
1392 assert_eq!(PlanComposition::default().dropped_total(), 0);
1393 }
1394
1395 #[cfg(feature = "planning")]
1399 #[tokio::test]
1400 async fn test_history_per_crate_cap() {
1401 let source = FakePlannerDataSource {
1402 history_candidates: vec![
1403 PrefetchCandidate::new("serde-newest".into(), "serde".into()),
1404 PrefetchCandidate::new("serde-older".into(), "serde".into()),
1405 PrefetchCandidate::new("serde-oldest".into(), "serde".into()),
1406 ],
1407 ..Default::default()
1408 };
1409 let intent = BuildIntent {
1410 crate_names: vec!["serde".into()],
1411 ..Default::default()
1412 };
1413
1414 let (plan, composition) = build_prefetch_plan_with_limits(
1415 &source,
1416 &intent,
1417 "fallback",
1418 PlanLimits {
1419 history_per_crate: 1,
1420 ..PlanLimits::default()
1421 },
1422 )
1423 .await
1424 .unwrap();
1425
1426 assert_eq!(composition.from_history, 1);
1427 assert_eq!(composition.dropped_history_per_crate, 2);
1428 assert_eq!(
1429 plan.candidates
1430 .iter()
1431 .map(|c| c.cache_key.as_str())
1432 .collect::<Vec<_>>(),
1433 vec!["serde-newest"]
1434 );
1435 }
1436
1437 #[cfg(feature = "planning")]
1441 #[tokio::test]
1442 async fn test_key_cache_total_skips_a_whole_crate_and_counts_all_of_it() {
1443 let mut source = FakePlannerDataSource::default();
1444 source.key_cache.insert("a".into(), vec!["a1".into()]);
1445 source.key_cache.insert(
1446 "b".into(),
1447 vec![
1448 "b1".into(),
1449 "b2".into(),
1450 "b3".into(),
1451 "b4".into(),
1452 "b5".into(),
1453 ],
1454 );
1455 let intent = BuildIntent {
1456 crate_names: vec!["a".into(), "b".into()],
1457 ..Default::default()
1458 };
1459
1460 let (_plan, composition) = build_prefetch_plan_with_limits(
1461 &source,
1462 &intent,
1463 "fallback",
1464 PlanLimits {
1465 key_cache_per_crate: 2,
1466 key_cache_total: 1,
1467 history_per_crate: 2,
1468 },
1469 )
1470 .await
1471 .unwrap();
1472
1473 assert_eq!(composition.from_key_cache, 1, "only crate `a` fits");
1474 assert_eq!(
1475 composition.dropped_key_cache_total, 5,
1476 "all five of `b`'s variants are counted, not the two a per-crate \
1477 cap would have admitted"
1478 );
1479 assert_eq!(
1480 composition.dropped_key_cache_per_crate, 0,
1481 "`b` was skipped before the per-crate cap applied"
1482 );
1483 }
1484
1485 #[cfg(feature = "planning")]
1488 #[tokio::test]
1489 async fn test_key_cache_total_can_run_out_mid_crate() {
1490 let mut source = FakePlannerDataSource::default();
1491 source
1492 .key_cache
1493 .insert("serde".into(), vec!["k1".into(), "k2".into(), "k3".into()]);
1494 let intent = BuildIntent {
1495 crate_names: vec!["serde".into()],
1496 ..Default::default()
1497 };
1498
1499 let (_plan, composition) = build_prefetch_plan_with_limits(
1500 &source,
1501 &intent,
1502 "fallback",
1503 PlanLimits {
1504 key_cache_per_crate: 2,
1505 key_cache_total: 1,
1506 history_per_crate: 2,
1507 },
1508 )
1509 .await
1510 .unwrap();
1511
1512 assert_eq!(composition.from_key_cache, 1);
1513 assert_eq!(
1514 composition.dropped_key_cache_per_crate, 1,
1515 "k3 dropped by the per-crate cap"
1516 );
1517 assert_eq!(
1518 composition.dropped_key_cache_total, 1,
1519 "k2 dropped by the plan-wide cap, mid-crate"
1520 );
1521 }
1522}