1#[cfg(feature = "planning")]
2use std::collections::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(planner_name, candidates));
77 }
78 Ok(_) | Err(_) => {}
79 }
80 }
81
82 let mut seen = HashSet::new();
83 let mut resolved_crates = HashSet::new();
84 let mut candidates = Vec::new();
85
86 for candidate in source.history_candidates(&intent.crate_names).await? {
87 resolved_crates.insert(candidate.crate_name.clone());
88 if seen.insert(candidate.cache_key.clone()) {
89 candidates.push(candidate);
90 }
91 }
92
93 for crate_name in intent
94 .crate_names
95 .iter()
96 .filter(|name| !resolved_crates.contains(*name))
97 {
98 for cache_key in source.key_cache_keys_for_crate(crate_name).await? {
99 if seen.insert(cache_key.clone()) {
100 candidates.push(PrefetchCandidate {
101 cache_key,
102 crate_name: crate_name.clone(),
103 });
104 }
105 }
106 }
107
108 Ok(execute_plan(planner_name, candidates))
109}
110
111#[cfg(feature = "planning")]
112fn execute_plan(planner_name: &str, candidates: Vec<PrefetchCandidate>) -> PrefetchPlan {
113 let planner = planner_name.trim();
114 PrefetchPlan {
115 plan_id: None,
116 planner: Some(if planner.is_empty() {
117 "planner".to_string()
118 } else {
119 planner.to_string()
120 }),
121 disposition: PrefetchDisposition::Execute,
122 candidates,
123 }
124}
125
126#[cfg(test)]
127mod tests {
128 use super::*;
129
130 #[cfg(feature = "planning")]
131 use std::collections::HashMap;
132
133 #[cfg(feature = "planning")]
134 use anyhow::anyhow;
135
136 #[cfg(feature = "planning")]
137 #[derive(Default)]
138 struct FakePlannerDataSource {
139 shard_candidates: Vec<PrefetchCandidate>,
140 shard_error: bool,
141 history_candidates: Vec<PrefetchCandidate>,
142 key_cache: HashMap<String, Vec<String>>,
143 }
144
145 #[cfg(feature = "planning")]
146 #[async_trait]
147 impl PlannerDataSource for FakePlannerDataSource {
148 async fn shard_candidates(
149 &self,
150 _namespace: &str,
151 _deps: &[(String, String)],
152 ) -> Result<Vec<PrefetchCandidate>> {
153 if self.shard_error {
154 Err(anyhow!("shard lookup failed"))
155 } else {
156 Ok(self.shard_candidates.clone())
157 }
158 }
159
160 async fn history_candidates(
161 &self,
162 _crate_names: &[String],
163 ) -> Result<Vec<PrefetchCandidate>> {
164 Ok(self.history_candidates.clone())
165 }
166
167 async fn key_cache_keys_for_crate(&self, crate_name: &str) -> Result<Vec<String>> {
168 Ok(self.key_cache.get(crate_name).cloned().unwrap_or_default())
169 }
170 }
171
172 #[test]
173 fn test_build_intent_serde_roundtrip() {
174 let intent = BuildIntent {
175 crate_names: vec!["serde".into(), "tokio".into()],
176 namespace: Some("x86_64/hash/release".into()),
177 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
178 };
179
180 let json = serde_json::to_string(&intent).unwrap();
181 let parsed: BuildIntent = serde_json::from_str(&json).unwrap();
182 assert_eq!(parsed, intent);
183 }
184
185 #[test]
186 fn test_prefetch_plan_serde_roundtrip() {
187 let plan = PrefetchPlan {
188 plan_id: Some("plan-1".into()),
189 planner: Some("local".into()),
190 disposition: PrefetchDisposition::Execute,
191 candidates: vec![PrefetchCandidate {
192 cache_key: "abc".into(),
193 crate_name: "serde".into(),
194 }],
195 };
196
197 let json = serde_json::to_string(&plan).unwrap();
198 let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
199 assert_eq!(parsed, plan);
200 }
201
202 #[test]
203 fn test_prefetch_plan_missing_disposition_is_rejected() {
204 let err = serde_json::from_str::<PrefetchPlan>(
205 r#"{"planner":"legacy","candidates":[{"cache_key":"abc","crate_name":"serde"}]}"#,
206 )
207 .unwrap_err();
208 assert!(err.to_string().contains("missing field"));
209 }
210
211 #[test]
212 fn test_prefetch_plan_do_nothing_roundtrip() {
213 let plan = PrefetchPlan {
214 plan_id: Some("plan-2".into()),
215 planner: Some("remote".into()),
216 disposition: PrefetchDisposition::DoNothing,
217 candidates: vec![],
218 };
219
220 let json = serde_json::to_string(&plan).unwrap();
221 let parsed: PrefetchPlan = serde_json::from_str(&json).unwrap();
222 assert_eq!(parsed, plan);
223 }
224
225 #[cfg(feature = "planning")]
226 #[tokio::test]
227 async fn test_build_prefetch_plan_prefers_shard_candidates() {
228 let source = FakePlannerDataSource {
229 shard_candidates: vec![PrefetchCandidate {
230 cache_key: "from-shard".into(),
231 crate_name: "serde".into(),
232 }],
233 ..Default::default()
234 };
235 let intent = BuildIntent {
236 crate_names: vec!["serde".into()],
237 namespace: Some("linux/hash/release".into()),
238 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
239 };
240
241 let plan = build_prefetch_plan(&source, &intent, "fallback")
242 .await
243 .unwrap();
244
245 assert_eq!(plan.disposition, PrefetchDisposition::Execute);
246 assert_eq!(plan.planner.as_deref(), Some("fallback"));
247 assert_eq!(plan.candidates.len(), 1);
248 assert_eq!(plan.candidates[0].cache_key, "from-shard");
249 }
250
251 #[cfg(feature = "planning")]
252 #[tokio::test]
253 async fn test_build_prefetch_plan_falls_back_to_history_and_key_cache() {
254 let mut source = FakePlannerDataSource {
255 shard_error: true,
256 history_candidates: vec![PrefetchCandidate {
257 cache_key: "history-key".into(),
258 crate_name: "serde".into(),
259 }],
260 ..Default::default()
261 };
262 source.key_cache.insert(
263 "tokio".into(),
264 vec!["tokio-key".into(), "history-key".into()],
265 );
266
267 let intent = BuildIntent {
268 crate_names: vec!["serde".into(), "tokio".into()],
269 namespace: Some("linux/hash/debug".into()),
270 cargo_lock_deps: vec![("serde".into(), "1.0.0".into())],
271 };
272
273 let plan = build_prefetch_plan(&source, &intent, "fallback")
274 .await
275 .unwrap();
276
277 assert_eq!(plan.disposition, PrefetchDisposition::Execute);
278 assert_eq!(plan.candidates.len(), 2);
279 assert_eq!(plan.candidates[0].cache_key, "history-key");
280 assert_eq!(plan.candidates[1].cache_key, "tokio-key");
281 }
282}