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