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