1#[cfg(feature = "planning")]
2use std::collections::{HashMap, HashSet};
3
4#[cfg(feature = "planning")]
5use anyhow::Result;
6#[cfg(feature = "planning")]
7use async_trait::async_trait;
8use serde::{Deserialize, Serialize};
9
10#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
11pub struct BuildIntent {
12 #[serde(default)]
13 pub crate_names: Vec<String>,
14 #[serde(default)]
15 pub namespace: Option<String>,
16 #[serde(default)]
17 pub cargo_lock_deps: Vec<(String, String)>,
18}
19
20#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
27#[serde(rename_all = "snake_case")]
28pub enum CandidateSource {
29 Shard,
31 History,
33 KeyCache,
36 #[default]
37 #[serde(other)]
38 Unknown,
39}
40
41impl CandidateSource {
42 pub fn confidence_rank(self) -> u8 {
45 match self {
46 CandidateSource::Shard => 0,
47 CandidateSource::History => 1,
48 CandidateSource::KeyCache => 2,
49 CandidateSource::Unknown => 3,
50 }
51 }
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
55pub struct PrefetchCandidate {
56 pub cache_key: String,
57 pub crate_name: String,
58 #[serde(default, skip_serializing_if = "Option::is_none")]
62 pub compile_time_ms: Option<u64>,
63 #[serde(default, skip_serializing_if = "Option::is_none")]
67 pub size_bytes: Option<u64>,
68 #[serde(default)]
69 pub source: CandidateSource,
70 #[serde(default, skip_serializing_if = "Option::is_none")]
73 pub demand_index: Option<u32>,
74}
75
76impl PrefetchCandidate {
77 pub fn new(cache_key: String, crate_name: String) -> Self {
79 Self {
80 cache_key,
81 crate_name,
82 compile_time_ms: None,
83 size_bytes: None,
84 source: CandidateSource::Unknown,
85 demand_index: None,
86 }
87 }
88
89 pub fn with_source(mut self, source: CandidateSource) -> Self {
90 self.source = source;
91 self
92 }
93}
94
95#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
96#[serde(rename_all = "snake_case")]
97pub enum PrefetchDisposition {
98 Execute,
99 UseFallback,
100 DoNothing,
101}
102
103#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
104pub struct PrefetchPlan {
105 #[serde(default)]
106 pub plan_id: Option<String>,
107 #[serde(default)]
108 pub planner: Option<String>,
109 pub disposition: PrefetchDisposition,
110 #[serde(default)]
111 pub candidates: Vec<PrefetchCandidate>,
112}
113
114#[derive(Debug, Clone, Copy, PartialEq, Eq)]
122pub struct PlanLimits {
123 pub key_cache_per_crate: usize,
127 pub key_cache_total: usize,
130 pub history_per_crate: usize,
134}
135
136impl Default for PlanLimits {
137 fn default() -> Self {
138 Self {
139 key_cache_per_crate: 2,
140 key_cache_total: 64,
141 history_per_crate: 2,
142 }
143 }
144}
145
146#[cfg(feature = "planning")]
155const URGENCY_BUCKET: u32 = 16;
156
157#[cfg(feature = "planning")]
175pub fn dispatch_sort_key(candidate: &PrefetchCandidate) -> (u32, u8, std::cmp::Reverse<u64>) {
176 let bucket = candidate
177 .demand_index
178 .map(|index| index / URGENCY_BUCKET)
179 .unwrap_or(u32::MAX);
180 (
181 bucket,
182 candidate.source.confidence_rank(),
183 std::cmp::Reverse(candidate.compile_time_ms.unwrap_or(0)),
186 )
187}
188
189#[cfg(feature = "planning")]
194fn cap_to(items: &mut Vec<PrefetchCandidate>, limit: usize) -> usize {
195 if limit == 0 || items.len() <= limit {
196 return 0;
197 }
198 let dropped = items.len() - limit;
199 items.truncate(limit);
200 dropped
201}
202
203#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
206pub struct PlanComposition {
207 pub from_shards: usize,
208 pub from_history: usize,
209 pub from_key_cache: usize,
210 pub dropped_history_per_crate: usize,
212 pub dropped_key_cache_per_crate: usize,
213 pub dropped_key_cache_total: usize,
214}
215
216impl PlanComposition {
217 pub fn dropped_total(&self) -> usize {
218 self.dropped_history_per_crate
219 + self.dropped_key_cache_per_crate
220 + self.dropped_key_cache_total
221 }
222}
223
224#[cfg(feature = "planning")]
225#[async_trait]
226pub trait PlannerDataSource {
227 async fn shard_candidates(
228 &self,
229 namespace: &str,
230 deps: &[(String, String)],
231 ) -> Result<Vec<PrefetchCandidate>>;
232
233 async fn history_candidates(&self, crate_names: &[String]) -> Result<Vec<PrefetchCandidate>>;
234
235 async fn key_cache_keys_for_crate(&self, crate_name: &str) -> Result<Vec<String>>;
236}
237
238#[cfg(feature = "planning")]
239pub async fn build_prefetch_plan<T>(
240 source: &T,
241 intent: &BuildIntent,
242 planner_name: &str,
243) -> Result<PrefetchPlan>
244where
245 T: PlannerDataSource + Sync + ?Sized,
246{
247 build_prefetch_plan_with_limits(source, intent, planner_name, PlanLimits::default())
248 .await
249 .map(|(plan, _composition)| plan)
250}
251
252#[cfg(feature = "planning")]
255pub async fn build_prefetch_plan_with_limits<T>(
256 source: &T,
257 intent: &BuildIntent,
258 planner_name: &str,
259 limits: PlanLimits,
260) -> Result<(PrefetchPlan, PlanComposition)>
261where
262 T: PlannerDataSource + Sync + ?Sized,
263{
264 let crate_order = crate_query_order(intent);
265 let demand_index = demand_index_map(&crate_order);
266 let mut seen = HashSet::new();
267 let mut resolved_crates = HashSet::new();
268 let mut candidates = Vec::new();
269 let mut composition = PlanComposition::default();
270
271 if let Some(namespace) = intent.namespace.as_deref()
279 && !intent.cargo_lock_deps.is_empty()
280 {
281 if let Ok(shard_candidates) = source
283 .shard_candidates(namespace, &intent.cargo_lock_deps)
284 .await
285 {
286 for candidate in order_candidates_by_crate_order(shard_candidates, intent) {
287 resolved_crates.insert(candidate.crate_name.clone());
288 if seen.insert(candidate.cache_key.clone()) {
289 candidates.push(candidate.with_source(CandidateSource::Shard));
290 }
291 }
292 }
293 }
294
295 let unresolved = |resolved: &HashSet<String>| -> Vec<String> {
296 crate_order
297 .iter()
298 .filter(|name| !resolved.contains(*name))
299 .cloned()
300 .collect()
301 };
302
303 composition.from_shards = candidates.len();
304
305 let history_query = unresolved(&resolved_crates);
306 if !history_query.is_empty() {
307 let mut by_crate: HashMap<String, Vec<PrefetchCandidate>> = HashMap::new();
310 for candidate in order_candidates_by_crate_order(
311 source.history_candidates(&history_query).await?,
312 intent,
313 ) {
314 by_crate
315 .entry(candidate.crate_name.clone())
316 .or_default()
317 .push(candidate.with_source(CandidateSource::History));
318 }
319
320 for crate_name in &history_query {
321 let Some(mut for_crate) = by_crate.remove(crate_name) else {
322 continue;
323 };
324 composition.dropped_history_per_crate +=
325 cap_to(&mut for_crate, limits.history_per_crate);
326 for candidate in for_crate {
327 resolved_crates.insert(candidate.crate_name.clone());
328 if seen.insert(candidate.cache_key.clone()) {
329 composition.from_history += 1;
330 candidates.push(candidate);
331 }
332 }
333 }
334 }
335
336 for crate_name in unresolved(&resolved_crates) {
342 if limits.key_cache_total > 0 && composition.from_key_cache >= limits.key_cache_total {
343 composition.dropped_key_cache_total += source
346 .key_cache_keys_for_crate(&crate_name)
347 .await?
348 .into_iter()
349 .filter(|key| !seen.contains(key))
350 .count();
351 continue;
352 }
353
354 let mut for_crate: Vec<PrefetchCandidate> = source
355 .key_cache_keys_for_crate(&crate_name)
356 .await?
357 .into_iter()
358 .map(|cache_key| {
359 PrefetchCandidate::new(cache_key, crate_name.clone())
360 .with_source(CandidateSource::KeyCache)
361 })
362 .collect();
363
364 composition.dropped_key_cache_per_crate +=
365 cap_to(&mut for_crate, limits.key_cache_per_crate);
366
367 for candidate in for_crate {
368 if limits.key_cache_total > 0 && composition.from_key_cache >= limits.key_cache_total {
369 composition.dropped_key_cache_total += 1;
370 continue;
371 }
372 if seen.insert(candidate.cache_key.clone()) {
373 composition.from_key_cache += 1;
374 candidates.push(candidate);
375 }
376 }
377 }
378
379 for candidate in &mut candidates {
382 candidate.demand_index = demand_index.get(&candidate.crate_name).copied();
383 }
384
385 candidates.sort_by_key(dispatch_sort_key);
388
389 Ok((execute_plan(planner_name, candidates), composition))
390}
391
392#[cfg(feature = "planning")]
393fn execute_plan(planner_name: &str, candidates: Vec<PrefetchCandidate>) -> PrefetchPlan {
394 let planner = planner_name.trim();
395 PrefetchPlan {
396 plan_id: None,
397 planner: Some(if planner.is_empty() {
398 "planner".to_string()
399 } else {
400 planner.to_string()
401 }),
402 disposition: PrefetchDisposition::Execute,
403 candidates,
404 }
405}
406
407#[cfg(feature = "planning")]
409fn demand_index_map(crate_order: &[String]) -> HashMap<String, u32> {
410 crate_order
411 .iter()
412 .enumerate()
413 .map(|(index, name)| (name.clone(), index as u32))
414 .collect()
415}
416
417#[cfg(feature = "planning")]
418fn crate_query_order(intent: &BuildIntent) -> Vec<String> {
419 let mut seen = HashSet::new();
420 intent
421 .crate_names
422 .iter()
423 .filter(|crate_name| seen.insert((*crate_name).clone()))
424 .cloned()
425 .collect()
426}
427
428#[cfg(feature = "planning")]
429fn order_candidates_by_crate_order(
430 mut candidates: Vec<PrefetchCandidate>,
431 intent: &BuildIntent,
432) -> Vec<PrefetchCandidate> {
433 if intent.crate_names.is_empty() {
434 return candidates;
435 }
436
437 let mut priority = HashMap::new();
438 for (index, crate_name) in crate_query_order(intent).iter().enumerate() {
439 priority.entry(crate_name.clone()).or_insert(index);
440 }
441
442 let mut indexed_candidates = candidates.drain(..).enumerate().collect::<Vec<_>>();
443 indexed_candidates.sort_by_key(|(index, candidate)| {
444 (
445 priority
446 .get(&candidate.crate_name)
447 .copied()
448 .unwrap_or(usize::MAX),
449 *index,
450 )
451 });
452 indexed_candidates
453 .into_iter()
454 .map(|(_, candidate)| candidate)
455 .collect()
456}
457
458#[cfg(test)]
459mod tests {
460 use super::*;
461
462 #[cfg(feature = "planning")]
463 use std::collections::HashMap;
464
465 #[cfg(feature = "planning")]
466 use anyhow::anyhow;
467
468 #[cfg(feature = "planning")]
469 use proptest::prelude::*;
470
471 #[cfg(feature = "planning")]
472 #[derive(Default)]
473 struct FakePlannerDataSource {
474 shard_candidates: Vec<PrefetchCandidate>,
475 shard_error: bool,
476 history_candidates: Vec<PrefetchCandidate>,
477 history_by_crate: HashMap<String, String>,
478 key_cache: HashMap<String, Vec<String>>,
479 }
480
481 #[cfg(feature = "planning")]
482 #[async_trait]
483 impl PlannerDataSource for FakePlannerDataSource {
484 async fn shard_candidates(
485 &self,
486 _namespace: &str,
487 _deps: &[(String, String)],
488 ) -> Result<Vec<PrefetchCandidate>> {
489 if self.shard_error {
490 Err(anyhow!("shard lookup failed"))
491 } else {
492 Ok(self.shard_candidates.clone())
493 }
494 }
495
496 async fn history_candidates(
497 &self,
498 crate_names: &[String],
499 ) -> Result<Vec<PrefetchCandidate>> {
500 if !self.history_candidates.is_empty() {
501 return Ok(self.history_candidates.clone());
502 }
503
504 Ok(crate_names
505 .iter()
506 .filter_map(|crate_name| {
507 self.history_by_crate.get(crate_name).map(|cache_key| {
508 PrefetchCandidate::new(cache_key.clone(), crate_name.clone())
509 })
510 })
511 .collect())
512 }
513
514 async fn key_cache_keys_for_crate(&self, crate_name: &str) -> Result<Vec<String>> {
515 Ok(self.key_cache.get(crate_name).cloned().unwrap_or_default())
516 }
517 }
518
519 #[test]
520 fn test_build_intent_serde_roundtrip() {
521 let intent = BuildIntent {
522 crate_names: vec!["serde".into(), "tokio".into()],
523 namespace: Some("x86_64/hash/release".into()),
524 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
525 };
526
527 let json = serde_json::to_string(&intent).unwrap();
528 let parsed: BuildIntent = serde_json::from_str(&json).unwrap();
529 assert_eq!(parsed, intent);
530 }
531
532 #[test]
533 fn test_build_intent_defaults_missing_fields() {
534 let parsed: BuildIntent = serde_json::from_str(r#"{"crate_names":["serde"]}"#).unwrap();
535 assert_eq!(parsed.crate_names, vec!["serde"]);
536 assert!(parsed.namespace.is_none());
537 assert!(parsed.cargo_lock_deps.is_empty());
538 }
539
540 #[test]
541 fn test_prefetch_plan_serde_roundtrip() {
542 let plan = PrefetchPlan {
543 plan_id: Some("plan-1".into()),
544 planner: Some("local".into()),
545 disposition: PrefetchDisposition::Execute,
546 candidates: vec![PrefetchCandidate::new("abc".into(), "serde".into())],
547 };
548
549 let json = serde_json::to_string(&plan).unwrap();
550 let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
551 assert_eq!(parsed, plan);
552 }
553
554 #[test]
555 fn test_prefetch_plan_missing_disposition_is_rejected() {
556 let err = serde_json::from_str::<PrefetchPlan>(
557 r#"{"planner":"legacy","candidates":[{"cache_key":"abc","crate_name":"serde"}]}"#,
558 )
559 .unwrap_err();
560 assert!(err.to_string().contains("missing field"));
561 }
562
563 #[test]
564 fn test_prefetch_plan_do_nothing_roundtrip() {
565 let plan = PrefetchPlan {
566 plan_id: Some("plan-2".into()),
567 planner: Some("remote".into()),
568 disposition: PrefetchDisposition::DoNothing,
569 candidates: vec![],
570 };
571
572 let json = serde_json::to_string(&plan).unwrap();
573 let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
574 assert_eq!(parsed, plan);
575 }
576
577 #[cfg(feature = "planning")]
578 #[tokio::test]
579 async fn test_build_prefetch_plan_prefers_shard_candidates() {
580 let source = FakePlannerDataSource {
581 shard_candidates: vec![PrefetchCandidate::new("from-shard".into(), "serde".into())],
582 ..Default::default()
583 };
584 let intent = BuildIntent {
585 crate_names: vec!["serde".into()],
586 namespace: Some("linux/hash/release".into()),
587 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
588 };
589
590 let plan = build_prefetch_plan(&source, &intent, "fallback")
591 .await
592 .unwrap();
593
594 assert_eq!(plan.disposition, PrefetchDisposition::Execute);
595 assert_eq!(plan.planner.as_deref(), Some("fallback"));
596 assert_eq!(plan.candidates.len(), 1);
597 assert_eq!(plan.candidates[0].cache_key, "from-shard");
598 }
599
600 #[cfg(feature = "planning")]
601 #[tokio::test]
602 async fn test_build_prefetch_plan_falls_back_to_history_and_key_cache() {
603 let mut source = FakePlannerDataSource {
604 shard_error: true,
605 history_candidates: vec![PrefetchCandidate::new("history-key".into(), "serde".into())],
606 ..Default::default()
607 };
608 source.key_cache.insert(
609 "tokio".into(),
610 vec!["tokio-key".into(), "history-key".into()],
611 );
612
613 let intent = BuildIntent {
614 crate_names: vec!["serde".into(), "tokio".into()],
615 namespace: Some("linux/hash/debug".into()),
616 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
617 };
618
619 let plan = build_prefetch_plan(&source, &intent, "fallback")
620 .await
621 .unwrap();
622
623 assert_eq!(plan.disposition, PrefetchDisposition::Execute);
624 assert_eq!(plan.candidates.len(), 2);
625 assert_eq!(plan.candidates[0].cache_key, "history-key");
626 assert_eq!(plan.candidates[1].cache_key, "tokio-key");
627 }
628
629 #[cfg(feature = "planning")]
630 #[tokio::test]
631 async fn test_build_prefetch_plan_orders_shard_candidates_by_crate_order() {
632 let source = FakePlannerDataSource {
633 shard_candidates: vec![
634 PrefetchCandidate::new("app-key".into(), "app".into()),
635 PrefetchCandidate::new("dep-key".into(), "dep".into()),
636 PrefetchCandidate::new("middle-key".into(), "middle".into()),
637 ],
638 ..Default::default()
639 };
640 let intent = BuildIntent {
641 crate_names: vec!["dep".into(), "middle".into(), "app".into()],
642 namespace: Some("linux/hash/debug".into()),
643 cargo_lock_deps: vec![("dep".into(), "1.0.0".into())],
644 };
645
646 let plan = build_prefetch_plan(&source, &intent, "fallback")
647 .await
648 .unwrap();
649
650 let keys = plan
651 .candidates
652 .iter()
653 .map(|candidate| candidate.cache_key.as_str())
654 .collect::<Vec<_>>();
655 assert_eq!(keys, vec!["dep-key", "middle-key", "app-key"]);
656 }
657
658 #[cfg(feature = "planning")]
667 #[tokio::test]
668 async fn test_build_prefetch_plan_fills_crates_a_partial_shard_hit_missed() {
669 let mut source = FakePlannerDataSource {
670 shard_candidates: vec![PrefetchCandidate::new("dep-shard-key".into(), "dep".into())],
672 history_by_crate: HashMap::from([("middle".into(), "middle-history-key".into())]),
673 ..Default::default()
674 };
675 source
676 .key_cache
677 .insert("app".into(), vec!["app-key-cache-key".into()]);
678
679 let intent = BuildIntent {
680 crate_names: vec!["dep".into(), "middle".into(), "app".into()],
681 namespace: Some("linux/hash/debug".into()),
682 cargo_lock_deps: vec![("dep".into(), "1.0.0".into())],
683 };
684
685 let plan = build_prefetch_plan(&source, &intent, "fallback")
686 .await
687 .unwrap();
688
689 let keys = plan
690 .candidates
691 .iter()
692 .map(|candidate| candidate.cache_key.as_str())
693 .collect::<Vec<_>>();
694 assert_eq!(
695 keys,
696 vec!["dep-shard-key", "middle-history-key", "app-key-cache-key"],
697 "each source should fill the crates the higher-confidence ones left unresolved"
698 );
699 }
700
701 #[cfg(feature = "planning")]
705 #[tokio::test]
706 async fn test_build_prefetch_plan_does_not_requery_shard_resolved_crates() {
707 let mut source = FakePlannerDataSource {
708 shard_candidates: vec![PrefetchCandidate::new(
709 "serde-shard-key".into(),
710 "serde".into(),
711 )],
712 history_by_crate: HashMap::from([("serde".into(), "serde-stale-history-key".into())]),
713 ..Default::default()
714 };
715 source
716 .key_cache
717 .insert("serde".into(), vec!["serde-stale-key-cache-key".into()]);
718
719 let intent = BuildIntent {
720 crate_names: vec!["serde".into()],
721 namespace: Some("linux/hash/debug".into()),
722 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
723 };
724
725 let plan = build_prefetch_plan(&source, &intent, "fallback")
726 .await
727 .unwrap();
728
729 let keys = plan
730 .candidates
731 .iter()
732 .map(|candidate| candidate.cache_key.as_str())
733 .collect::<Vec<_>>();
734 assert_eq!(keys, vec!["serde-shard-key"]);
735 }
736
737 #[cfg(feature = "planning")]
738 #[tokio::test]
739 async fn test_build_prefetch_plan_queries_history_by_crate_order() {
740 let source = FakePlannerDataSource {
741 history_by_crate: HashMap::from([
742 ("app".into(), "app-key".into()),
743 ("dep".into(), "dep-key".into()),
744 ("middle".into(), "middle-key".into()),
745 ]),
746 ..Default::default()
747 };
748 let intent = BuildIntent {
749 crate_names: vec!["dep".into(), "middle".into(), "app".into()],
750 namespace: None,
751 cargo_lock_deps: vec![],
752 };
753
754 let plan = build_prefetch_plan(&source, &intent, "fallback")
755 .await
756 .unwrap();
757
758 let keys = plan
759 .candidates
760 .iter()
761 .map(|candidate| candidate.cache_key.as_str())
762 .collect::<Vec<_>>();
763 assert_eq!(keys, vec!["dep-key", "middle-key", "app-key"]);
764 }
765
766 #[cfg(feature = "planning")]
767 proptest! {
768 #[test]
769 fn property_crate_query_order_preserves_first_occurrence(
770 crate_ids in proptest::collection::vec(any::<u8>(), 1..64),
771 ) {
772 let mut crate_names = crate_ids
773 .into_iter()
774 .map(|id| format!("crate-{id}"))
775 .collect::<Vec<_>>();
776 crate_names.push(crate_names[0].clone());
777 let intent = BuildIntent {
778 crate_names: crate_names.clone(),
779 ..Default::default()
780 };
781
782 let mut expected = Vec::new();
783 for crate_name in crate_names {
784 if !expected.contains(&crate_name) {
785 expected.push(crate_name);
786 }
787 }
788
789 prop_assert_eq!(crate_query_order(&intent), expected);
790 }
791
792 #[test]
793 fn property_candidate_ordering_is_a_stable_priority_permutation(
794 requested_ids in proptest::collection::vec(0u8..8, 1..24),
795 candidate_ids in proptest::collection::vec(0u8..16, 0..48),
796 ) {
797 let requested_names = requested_ids
798 .into_iter()
799 .map(|id| format!("crate-{id}"))
800 .collect::<Vec<_>>();
801 let intent = BuildIntent {
802 crate_names: requested_names.clone(),
803 ..Default::default()
804 };
805
806 let mut candidate_names = Vec::with_capacity(candidate_ids.len() + 4);
809 candidate_names.push("not-requested".to_string());
810 candidate_names.push(requested_names[0].clone());
811 candidate_names.extend(
812 candidate_ids
813 .into_iter()
814 .map(|id| format!("crate-{id}")),
815 );
816 candidate_names.push("not-requested".to_string());
817 candidate_names.push(requested_names[0].clone());
818
819 let candidates = candidate_names
820 .into_iter()
821 .enumerate()
822 .map(|(index, crate_name)| {
823 PrefetchCandidate::new(format!("key-{index}"), crate_name)
824 })
825 .collect::<Vec<_>>();
826
827 let mut expected = candidates.clone();
828 expected.sort_by_key(|candidate| {
829 requested_names
830 .iter()
831 .position(|name| name == &candidate.crate_name)
832 .unwrap_or(usize::MAX)
833 });
834
835 prop_assert_eq!(
836 order_candidates_by_crate_order(candidates, &intent),
837 expected
838 );
839 }
840 }
841 #[cfg(feature = "planning")]
844 fn candidate(
845 key: &str,
846 crate_name: &str,
847 source: CandidateSource,
848 compile_time_ms: Option<u64>,
849 demand_index: Option<u32>,
850 ) -> PrefetchCandidate {
851 PrefetchCandidate {
852 cache_key: key.into(),
853 crate_name: crate_name.into(),
854 compile_time_ms,
855 size_bytes: None,
856 source,
857 demand_index,
858 }
859 }
860
861 #[test]
864 fn test_confidence_rank_orders_sources() {
865 assert!(
866 CandidateSource::Shard.confidence_rank() < CandidateSource::History.confidence_rank()
867 );
868 assert!(
869 CandidateSource::History.confidence_rank()
870 < CandidateSource::KeyCache.confidence_rank()
871 );
872 assert!(
873 CandidateSource::KeyCache.confidence_rank()
874 < CandidateSource::Unknown.confidence_rank()
875 );
876 }
877
878 #[test]
881 fn test_unknown_candidate_source_deserializes() {
882 let candidate: PrefetchCandidate =
883 serde_json::from_str(r#"{"cache_key":"k","crate_name":"c","source":"telepathy"}"#)
884 .unwrap();
885 assert_eq!(candidate.source, CandidateSource::Unknown);
886 }
887
888 #[test]
891 fn test_legacy_candidate_wire_still_parses() {
892 let candidate: PrefetchCandidate =
893 serde_json::from_str(r#"{"cache_key":"k","crate_name":"c"}"#).unwrap();
894 assert_eq!(candidate.compile_time_ms, None);
895 assert_eq!(candidate.size_bytes, None);
896 assert_eq!(candidate.demand_index, None);
897 assert_eq!(candidate.source, CandidateSource::Unknown);
898 }
899
900 #[cfg(feature = "planning")]
903 #[test]
904 fn test_dispatch_order_prefers_urgency_over_value() {
905 let urgent_cheap = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
906 let late_expensive = candidate("b", "b", CandidateSource::Shard, Some(100_000), Some(500));
907 assert!(dispatch_sort_key(&urgent_cheap) < dispatch_sort_key(&late_expensive));
908 }
909
910 #[cfg(feature = "planning")]
912 #[test]
913 fn test_dispatch_order_prefers_value_within_a_bucket() {
914 let cheap = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
915 let costly = candidate("b", "b", CandidateSource::Shard, Some(5_000), Some(1));
916 assert!(
917 dispatch_sort_key(&costly) < dispatch_sort_key(&cheap),
918 "same bucket, same source: the expensive rebuild goes first"
919 );
920 }
921
922 #[cfg(feature = "planning")]
925 #[test]
926 fn test_dispatch_order_buckets_nearby_positions_together() {
927 let first = candidate("a", "a", CandidateSource::Shard, Some(10), Some(0));
928 let nearby = candidate("b", "b", CandidateSource::Shard, Some(10), Some(15));
929 let next_bucket = candidate("c", "c", CandidateSource::Shard, Some(10), Some(16));
930 assert_eq!(
931 dispatch_sort_key(&first),
932 dispatch_sort_key(&nearby),
933 "positions 0 and 15 share a bucket"
934 );
935 assert!(
936 dispatch_sort_key(&nearby) < dispatch_sort_key(&next_bucket),
937 "position 16 starts the next bucket"
938 );
939 }
940
941 #[cfg(feature = "planning")]
944 #[test]
945 fn test_dispatch_order_prefers_confidence_within_a_bucket() {
946 let shard_cheap = candidate("a", "a", CandidateSource::Shard, Some(1), Some(0));
947 let guess_costly = candidate("b", "b", CandidateSource::KeyCache, Some(99_999), Some(0));
948 assert!(dispatch_sort_key(&shard_cheap) < dispatch_sort_key(&guess_costly));
949 }
950
951 #[cfg(feature = "planning")]
954 #[test]
955 fn test_dispatch_order_places_unknown_cost_last_within_its_bucket() {
956 let known = candidate("a", "a", CandidateSource::Shard, Some(1), Some(0));
957 let unknown = candidate("b", "b", CandidateSource::Shard, None, Some(0));
958 assert!(dispatch_sort_key(&known) < dispatch_sort_key(&unknown));
959
960 let late_known = candidate("c", "c", CandidateSource::Shard, Some(50_000), Some(999));
963 assert!(dispatch_sort_key(&unknown) < dispatch_sort_key(&late_known));
964 }
965
966 #[cfg(feature = "planning")]
968 #[test]
969 fn test_dispatch_order_places_unknown_demand_last() {
970 let known = candidate("a", "a", CandidateSource::Shard, Some(1), Some(100_000));
971 let no_demand = candidate("b", "b", CandidateSource::Shard, Some(50_000), None);
972 assert!(dispatch_sort_key(&known) < dispatch_sort_key(&no_demand));
973 }
974
975 #[cfg(feature = "planning")]
977 #[tokio::test]
978 async fn test_key_cache_per_crate_cap() {
979 let mut source = FakePlannerDataSource::default();
980 source.key_cache.insert(
981 "serde".into(),
982 vec!["k1".into(), "k2".into(), "k3".into(), "k4".into()],
983 );
984 let intent = BuildIntent {
985 crate_names: vec!["serde".into()],
986 ..Default::default()
987 };
988
989 let (plan, composition) = build_prefetch_plan_with_limits(
990 &source,
991 &intent,
992 "fallback",
993 PlanLimits {
994 key_cache_per_crate: 2,
995 ..PlanLimits::default()
996 },
997 )
998 .await
999 .unwrap();
1000
1001 assert_eq!(plan.candidates.len(), 2, "capped to two variants");
1002 assert_eq!(composition.from_key_cache, 2);
1003 assert_eq!(composition.dropped_key_cache_per_crate, 2);
1004 assert!(composition.dropped_total() > 0, "truncation is visible");
1005 }
1006
1007 #[cfg(feature = "planning")]
1010 #[tokio::test]
1011 async fn test_key_cache_total_cap_counts_skipped_crates() {
1012 let mut source = FakePlannerDataSource::default();
1013 for name in ["a", "b", "c"] {
1014 source
1015 .key_cache
1016 .insert(name.into(), vec![format!("{name}-k1")]);
1017 }
1018 let intent = BuildIntent {
1019 crate_names: vec!["a".into(), "b".into(), "c".into()],
1020 ..Default::default()
1021 };
1022
1023 let (plan, composition) = build_prefetch_plan_with_limits(
1024 &source,
1025 &intent,
1026 "fallback",
1027 PlanLimits {
1028 key_cache_total: 1,
1029 ..PlanLimits::default()
1030 },
1031 )
1032 .await
1033 .unwrap();
1034
1035 assert_eq!(plan.candidates.len(), 1);
1036 assert_eq!(composition.from_key_cache, 1);
1037 assert_eq!(
1038 composition.dropped_key_cache_total, 2,
1039 "the two skipped crates are counted, not silently missing"
1040 );
1041 }
1042
1043 #[cfg(feature = "planning")]
1045 #[tokio::test]
1046 async fn test_zero_limit_disables_the_cap() {
1047 let mut source = FakePlannerDataSource::default();
1048 source
1049 .key_cache
1050 .insert("serde".into(), vec!["k1".into(), "k2".into(), "k3".into()]);
1051 let intent = BuildIntent {
1052 crate_names: vec!["serde".into()],
1053 ..Default::default()
1054 };
1055
1056 let (plan, composition) = build_prefetch_plan_with_limits(
1057 &source,
1058 &intent,
1059 "fallback",
1060 PlanLimits {
1061 key_cache_per_crate: 0,
1062 key_cache_total: 0,
1063 history_per_crate: 0,
1064 },
1065 )
1066 .await
1067 .unwrap();
1068
1069 assert_eq!(plan.candidates.len(), 3);
1070 assert_eq!(composition.dropped_total(), 0);
1071 }
1072
1073 #[cfg(feature = "planning")]
1076 #[tokio::test]
1077 async fn test_plan_stamps_source_and_demand_index() {
1078 let source = FakePlannerDataSource {
1079 history_by_crate: HashMap::from([
1080 ("dep".into(), "dep-key".into()),
1081 ("app".into(), "app-key".into()),
1082 ]),
1083 ..Default::default()
1084 };
1085 let intent = BuildIntent {
1086 crate_names: vec!["dep".into(), "app".into()],
1087 ..Default::default()
1088 };
1089
1090 let (plan, _) =
1091 build_prefetch_plan_with_limits(&source, &intent, "fallback", PlanLimits::default())
1092 .await
1093 .unwrap();
1094
1095 let dep = plan
1096 .candidates
1097 .iter()
1098 .find(|c| c.crate_name == "dep")
1099 .expect("dep planned");
1100 let app = plan
1101 .candidates
1102 .iter()
1103 .find(|c| c.crate_name == "app")
1104 .expect("app planned");
1105 assert_eq!(dep.source, CandidateSource::History);
1106 assert_eq!(dep.demand_index, Some(0));
1107 assert_eq!(app.demand_index, Some(1), "position follows crate order");
1108 }
1109
1110 #[cfg(feature = "planning")]
1114 #[test]
1115 fn test_cap_to_reports_the_dropped_count() {
1116 let mut items: Vec<PrefetchCandidate> = (0..5)
1117 .map(|i| PrefetchCandidate::new(format!("k{i}"), "c".into()))
1118 .collect();
1119 assert_eq!(cap_to(&mut items, 2), 3);
1120 assert_eq!(items.len(), 2);
1121 assert_eq!(items[0].cache_key, "k0", "keeps the FIRST N");
1122
1123 let mut exact: Vec<PrefetchCandidate> =
1124 vec![PrefetchCandidate::new("a".into(), "c".into())];
1125 assert_eq!(cap_to(&mut exact, 1), 0, "exactly at the limit drops none");
1126 assert_eq!(cap_to(&mut exact, 0), 0, "0 disables the cap");
1127 assert_eq!(exact.len(), 1);
1128 }
1129
1130 #[test]
1133 fn test_dropped_total_sums_every_class() {
1134 let composition = PlanComposition {
1135 dropped_history_per_crate: 1,
1136 dropped_key_cache_per_crate: 2,
1137 dropped_key_cache_total: 4,
1138 ..PlanComposition::default()
1139 };
1140 assert_eq!(composition.dropped_total(), 7);
1141 assert_eq!(PlanComposition::default().dropped_total(), 0);
1142 }
1143
1144 #[cfg(feature = "planning")]
1148 #[tokio::test]
1149 async fn test_history_per_crate_cap() {
1150 let source = FakePlannerDataSource {
1151 history_candidates: vec![
1152 PrefetchCandidate::new("serde-newest".into(), "serde".into()),
1153 PrefetchCandidate::new("serde-older".into(), "serde".into()),
1154 PrefetchCandidate::new("serde-oldest".into(), "serde".into()),
1155 ],
1156 ..Default::default()
1157 };
1158 let intent = BuildIntent {
1159 crate_names: vec!["serde".into()],
1160 ..Default::default()
1161 };
1162
1163 let (plan, composition) = build_prefetch_plan_with_limits(
1164 &source,
1165 &intent,
1166 "fallback",
1167 PlanLimits {
1168 history_per_crate: 1,
1169 ..PlanLimits::default()
1170 },
1171 )
1172 .await
1173 .unwrap();
1174
1175 assert_eq!(composition.from_history, 1);
1176 assert_eq!(composition.dropped_history_per_crate, 2);
1177 assert_eq!(
1178 plan.candidates
1179 .iter()
1180 .map(|c| c.cache_key.as_str())
1181 .collect::<Vec<_>>(),
1182 vec!["serde-newest"]
1183 );
1184 }
1185
1186 #[cfg(feature = "planning")]
1190 #[tokio::test]
1191 async fn test_key_cache_total_skips_a_whole_crate_and_counts_all_of_it() {
1192 let mut source = FakePlannerDataSource::default();
1193 source.key_cache.insert("a".into(), vec!["a1".into()]);
1194 source.key_cache.insert(
1195 "b".into(),
1196 vec![
1197 "b1".into(),
1198 "b2".into(),
1199 "b3".into(),
1200 "b4".into(),
1201 "b5".into(),
1202 ],
1203 );
1204 let intent = BuildIntent {
1205 crate_names: vec!["a".into(), "b".into()],
1206 ..Default::default()
1207 };
1208
1209 let (_plan, composition) = build_prefetch_plan_with_limits(
1210 &source,
1211 &intent,
1212 "fallback",
1213 PlanLimits {
1214 key_cache_per_crate: 2,
1215 key_cache_total: 1,
1216 history_per_crate: 2,
1217 },
1218 )
1219 .await
1220 .unwrap();
1221
1222 assert_eq!(composition.from_key_cache, 1, "only crate `a` fits");
1223 assert_eq!(
1224 composition.dropped_key_cache_total, 5,
1225 "all five of `b`'s variants are counted, not the two a per-crate \
1226 cap would have admitted"
1227 );
1228 assert_eq!(
1229 composition.dropped_key_cache_per_crate, 0,
1230 "`b` was skipped before the per-crate cap applied"
1231 );
1232 }
1233
1234 #[cfg(feature = "planning")]
1237 #[tokio::test]
1238 async fn test_key_cache_total_can_run_out_mid_crate() {
1239 let mut source = FakePlannerDataSource::default();
1240 source
1241 .key_cache
1242 .insert("serde".into(), vec!["k1".into(), "k2".into(), "k3".into()]);
1243 let intent = BuildIntent {
1244 crate_names: vec!["serde".into()],
1245 ..Default::default()
1246 };
1247
1248 let (_plan, composition) = build_prefetch_plan_with_limits(
1249 &source,
1250 &intent,
1251 "fallback",
1252 PlanLimits {
1253 key_cache_per_crate: 2,
1254 key_cache_total: 1,
1255 history_per_crate: 2,
1256 },
1257 )
1258 .await
1259 .unwrap();
1260
1261 assert_eq!(composition.from_key_cache, 1);
1262 assert_eq!(
1263 composition.dropped_key_cache_per_crate, 1,
1264 "k3 dropped by the per-crate cap"
1265 );
1266 assert_eq!(
1267 composition.dropped_key_cache_total, 1,
1268 "k2 dropped by the plan-wide cap, mid-crate"
1269 );
1270 }
1271}