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, Serialize, Deserialize, PartialEq, Eq)]
21pub struct PrefetchCandidate {
22 pub cache_key: String,
23 pub crate_name: String,
24}
25
26#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
27#[serde(rename_all = "snake_case")]
28pub enum PrefetchDisposition {
29 Execute,
30 UseFallback,
31 DoNothing,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
35pub struct PrefetchPlan {
36 #[serde(default)]
37 pub plan_id: Option<String>,
38 #[serde(default)]
39 pub planner: Option<String>,
40 pub disposition: PrefetchDisposition,
41 #[serde(default)]
42 pub candidates: Vec<PrefetchCandidate>,
43}
44
45#[cfg(feature = "planning")]
46#[async_trait]
47pub trait PlannerDataSource {
48 async fn shard_candidates(
49 &self,
50 namespace: &str,
51 deps: &[(String, String)],
52 ) -> Result<Vec<PrefetchCandidate>>;
53
54 async fn history_candidates(&self, crate_names: &[String]) -> Result<Vec<PrefetchCandidate>>;
55
56 async fn key_cache_keys_for_crate(&self, crate_name: &str) -> Result<Vec<String>>;
57}
58
59#[cfg(feature = "planning")]
60pub async fn build_prefetch_plan<T>(
61 source: &T,
62 intent: &BuildIntent,
63 planner_name: &str,
64) -> Result<PrefetchPlan>
65where
66 T: PlannerDataSource + Sync + ?Sized,
67{
68 if let Some(namespace) = intent.namespace.as_deref()
69 && !intent.cargo_lock_deps.is_empty()
70 {
71 match source
72 .shard_candidates(namespace, &intent.cargo_lock_deps)
73 .await
74 {
75 Ok(candidates) if !candidates.is_empty() => {
76 return Ok(execute_plan(
77 planner_name,
78 order_candidates_by_crate_order(candidates, intent),
79 ));
80 }
81 Ok(_) | Err(_) => {}
82 }
83 }
84
85 let crate_order = crate_query_order(intent);
86 let mut seen = HashSet::new();
87 let mut resolved_crates = HashSet::new();
88 let mut candidates = Vec::new();
89
90 for candidate in
91 order_candidates_by_crate_order(source.history_candidates(&crate_order).await?, intent)
92 {
93 resolved_crates.insert(candidate.crate_name.clone());
94 if seen.insert(candidate.cache_key.clone()) {
95 candidates.push(candidate);
96 }
97 }
98
99 for crate_name in crate_order
100 .iter()
101 .filter(|name| !resolved_crates.contains(*name))
102 {
103 for cache_key in source.key_cache_keys_for_crate(crate_name).await? {
104 if seen.insert(cache_key.clone()) {
105 candidates.push(PrefetchCandidate {
106 cache_key,
107 crate_name: crate_name.clone(),
108 });
109 }
110 }
111 }
112
113 Ok(execute_plan(planner_name, candidates))
114}
115
116#[cfg(feature = "planning")]
117fn execute_plan(planner_name: &str, candidates: Vec<PrefetchCandidate>) -> PrefetchPlan {
118 let planner = planner_name.trim();
119 PrefetchPlan {
120 plan_id: None,
121 planner: Some(if planner.is_empty() {
122 "planner".to_string()
123 } else {
124 planner.to_string()
125 }),
126 disposition: PrefetchDisposition::Execute,
127 candidates,
128 }
129}
130
131#[cfg(feature = "planning")]
132fn crate_query_order(intent: &BuildIntent) -> Vec<String> {
133 let mut seen = HashSet::new();
134 intent
135 .crate_names
136 .iter()
137 .filter(|crate_name| seen.insert((*crate_name).clone()))
138 .cloned()
139 .collect()
140}
141
142#[cfg(feature = "planning")]
143fn order_candidates_by_crate_order(
144 mut candidates: Vec<PrefetchCandidate>,
145 intent: &BuildIntent,
146) -> Vec<PrefetchCandidate> {
147 if intent.crate_names.is_empty() {
148 return candidates;
149 }
150
151 let mut priority = HashMap::new();
152 for (index, crate_name) in crate_query_order(intent).iter().enumerate() {
153 priority.entry(crate_name.clone()).or_insert(index);
154 }
155
156 let mut indexed_candidates = candidates.drain(..).enumerate().collect::<Vec<_>>();
157 indexed_candidates.sort_by_key(|(index, candidate)| {
158 (
159 priority
160 .get(&candidate.crate_name)
161 .copied()
162 .unwrap_or(usize::MAX),
163 *index,
164 )
165 });
166 indexed_candidates
167 .into_iter()
168 .map(|(_, candidate)| candidate)
169 .collect()
170}
171
172#[cfg(test)]
173mod tests {
174 use super::*;
175
176 #[cfg(feature = "planning")]
177 use std::collections::HashMap;
178
179 #[cfg(feature = "planning")]
180 use anyhow::anyhow;
181
182 #[cfg(feature = "planning")]
183 #[derive(Default)]
184 struct FakePlannerDataSource {
185 shard_candidates: Vec<PrefetchCandidate>,
186 shard_error: bool,
187 history_candidates: Vec<PrefetchCandidate>,
188 history_by_crate: HashMap<String, String>,
189 key_cache: HashMap<String, Vec<String>>,
190 }
191
192 #[cfg(feature = "planning")]
193 #[async_trait]
194 impl PlannerDataSource for FakePlannerDataSource {
195 async fn shard_candidates(
196 &self,
197 _namespace: &str,
198 _deps: &[(String, String)],
199 ) -> Result<Vec<PrefetchCandidate>> {
200 if self.shard_error {
201 Err(anyhow!("shard lookup failed"))
202 } else {
203 Ok(self.shard_candidates.clone())
204 }
205 }
206
207 async fn history_candidates(
208 &self,
209 crate_names: &[String],
210 ) -> Result<Vec<PrefetchCandidate>> {
211 if !self.history_candidates.is_empty() {
212 return Ok(self.history_candidates.clone());
213 }
214
215 Ok(crate_names
216 .iter()
217 .filter_map(|crate_name| {
218 self.history_by_crate
219 .get(crate_name)
220 .map(|cache_key| PrefetchCandidate {
221 cache_key: cache_key.clone(),
222 crate_name: crate_name.clone(),
223 })
224 })
225 .collect())
226 }
227
228 async fn key_cache_keys_for_crate(&self, crate_name: &str) -> Result<Vec<String>> {
229 Ok(self.key_cache.get(crate_name).cloned().unwrap_or_default())
230 }
231 }
232
233 #[test]
234 fn test_build_intent_serde_roundtrip() {
235 let intent = BuildIntent {
236 crate_names: vec!["serde".into(), "tokio".into()],
237 namespace: Some("x86_64/hash/release".into()),
238 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
239 };
240
241 let json = serde_json::to_string(&intent).unwrap();
242 let parsed: BuildIntent = serde_json::from_str(&json).unwrap();
243 assert_eq!(parsed, intent);
244 }
245
246 #[test]
247 fn test_build_intent_defaults_missing_fields() {
248 let parsed: BuildIntent = serde_json::from_str(r#"{"crate_names":["serde"]}"#).unwrap();
249 assert_eq!(parsed.crate_names, vec!["serde"]);
250 assert!(parsed.namespace.is_none());
251 assert!(parsed.cargo_lock_deps.is_empty());
252 }
253
254 #[test]
255 fn test_prefetch_plan_serde_roundtrip() {
256 let plan = PrefetchPlan {
257 plan_id: Some("plan-1".into()),
258 planner: Some("local".into()),
259 disposition: PrefetchDisposition::Execute,
260 candidates: vec![PrefetchCandidate {
261 cache_key: "abc".into(),
262 crate_name: "serde".into(),
263 }],
264 };
265
266 let json = serde_json::to_string(&plan).unwrap();
267 let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
268 assert_eq!(parsed, plan);
269 }
270
271 #[test]
272 fn test_prefetch_plan_missing_disposition_is_rejected() {
273 let err = serde_json::from_str::<PrefetchPlan>(
274 r#"{"planner":"legacy","candidates":[{"cache_key":"abc","crate_name":"serde"}]}"#,
275 )
276 .unwrap_err();
277 assert!(err.to_string().contains("missing field"));
278 }
279
280 #[test]
281 fn test_prefetch_plan_do_nothing_roundtrip() {
282 let plan = PrefetchPlan {
283 plan_id: Some("plan-2".into()),
284 planner: Some("remote".into()),
285 disposition: PrefetchDisposition::DoNothing,
286 candidates: vec![],
287 };
288
289 let json = serde_json::to_string(&plan).unwrap();
290 let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
291 assert_eq!(parsed, plan);
292 }
293
294 #[cfg(feature = "planning")]
295 #[tokio::test]
296 async fn test_build_prefetch_plan_prefers_shard_candidates() {
297 let source = FakePlannerDataSource {
298 shard_candidates: vec![PrefetchCandidate {
299 cache_key: "from-shard".into(),
300 crate_name: "serde".into(),
301 }],
302 ..Default::default()
303 };
304 let intent = BuildIntent {
305 crate_names: vec!["serde".into()],
306 namespace: Some("linux/hash/release".into()),
307 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
308 };
309
310 let plan = build_prefetch_plan(&source, &intent, "fallback")
311 .await
312 .unwrap();
313
314 assert_eq!(plan.disposition, PrefetchDisposition::Execute);
315 assert_eq!(plan.planner.as_deref(), Some("fallback"));
316 assert_eq!(plan.candidates.len(), 1);
317 assert_eq!(plan.candidates[0].cache_key, "from-shard");
318 }
319
320 #[cfg(feature = "planning")]
321 #[tokio::test]
322 async fn test_build_prefetch_plan_falls_back_to_history_and_key_cache() {
323 let mut source = FakePlannerDataSource {
324 shard_error: true,
325 history_candidates: vec![PrefetchCandidate {
326 cache_key: "history-key".into(),
327 crate_name: "serde".into(),
328 }],
329 ..Default::default()
330 };
331 source.key_cache.insert(
332 "tokio".into(),
333 vec!["tokio-key".into(), "history-key".into()],
334 );
335
336 let intent = BuildIntent {
337 crate_names: vec!["serde".into(), "tokio".into()],
338 namespace: Some("linux/hash/debug".into()),
339 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
340 };
341
342 let plan = build_prefetch_plan(&source, &intent, "fallback")
343 .await
344 .unwrap();
345
346 assert_eq!(plan.disposition, PrefetchDisposition::Execute);
347 assert_eq!(plan.candidates.len(), 2);
348 assert_eq!(plan.candidates[0].cache_key, "history-key");
349 assert_eq!(plan.candidates[1].cache_key, "tokio-key");
350 }
351
352 #[cfg(feature = "planning")]
353 #[tokio::test]
354 async fn test_build_prefetch_plan_orders_shard_candidates_by_crate_order() {
355 let source = FakePlannerDataSource {
356 shard_candidates: vec![
357 PrefetchCandidate {
358 cache_key: "app-key".into(),
359 crate_name: "app".into(),
360 },
361 PrefetchCandidate {
362 cache_key: "dep-key".into(),
363 crate_name: "dep".into(),
364 },
365 PrefetchCandidate {
366 cache_key: "middle-key".into(),
367 crate_name: "middle".into(),
368 },
369 ],
370 ..Default::default()
371 };
372 let intent = BuildIntent {
373 crate_names: vec!["dep".into(), "middle".into(), "app".into()],
374 namespace: Some("linux/hash/debug".into()),
375 cargo_lock_deps: vec![("dep".into(), "1.0.0".into())],
376 };
377
378 let plan = build_prefetch_plan(&source, &intent, "fallback")
379 .await
380 .unwrap();
381
382 let keys = plan
383 .candidates
384 .iter()
385 .map(|candidate| candidate.cache_key.as_str())
386 .collect::<Vec<_>>();
387 assert_eq!(keys, vec!["dep-key", "middle-key", "app-key"]);
388 }
389
390 #[cfg(feature = "planning")]
391 #[tokio::test]
392 async fn test_build_prefetch_plan_queries_history_by_crate_order() {
393 let source = FakePlannerDataSource {
394 history_by_crate: HashMap::from([
395 ("app".into(), "app-key".into()),
396 ("dep".into(), "dep-key".into()),
397 ("middle".into(), "middle-key".into()),
398 ]),
399 ..Default::default()
400 };
401 let intent = BuildIntent {
402 crate_names: vec!["dep".into(), "middle".into(), "app".into()],
403 namespace: None,
404 cargo_lock_deps: vec![],
405 };
406
407 let plan = build_prefetch_plan(&source, &intent, "fallback")
408 .await
409 .unwrap();
410
411 let keys = plan
412 .candidates
413 .iter()
414 .map(|candidate| candidate.cache_key.as_str())
415 .collect::<Vec<_>>();
416 assert_eq!(keys, vec!["dep-key", "middle-key", "app-key"]);
417 }
418}