1use flowey::node::prelude::*;
8use std::collections::BTreeMap;
9
10#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Copy)]
11pub enum OpenhclKernelPackageKind {
12 Main,
13 Cvm,
14 Dev,
15 CvmDev,
16}
17
18#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Copy)]
19pub enum OpenhclKernelPackageArch {
20 X86_64,
21 Aarch64,
22}
23
24flowey_request! {
25 pub enum Request {
26 SetLocal {
28 arch: OpenhclKernelPackageArch,
29 kernel: ReadVar<PathBuf>,
30 modules: ReadVar<PathBuf>,
31 },
32 SetVersion(OpenhclKernelPackageKind, String),
34 GetKernel {
36 kind: OpenhclKernelPackageKind,
37 arch: OpenhclKernelPackageArch,
38 kernel: WriteVar<PathBuf>,
39 },
40 GetModules {
42 kind: OpenhclKernelPackageKind,
43 arch: OpenhclKernelPackageArch,
44 modules: WriteVar<PathBuf>,
45 },
46 GetPackageRoot {
48 kind: OpenhclKernelPackageKind,
49 arch: OpenhclKernelPackageArch,
50 pkg: WriteVar<PathBuf>,
51 },
52 GetMetadata {
54 kind: OpenhclKernelPackageKind,
55 arch: OpenhclKernelPackageArch,
56 metadata: WriteVar<PathBuf>,
57 },
58 }
59}
60
61new_flow_node!(struct Node);
62
63impl FlowNode for Node {
64 type Request = Request;
65
66 fn imports(ctx: &mut ImportCtx<'_>) {
67 ctx.import::<flowey_lib_common::install_dist_pkg::Node>();
68 ctx.import::<flowey_lib_common::download_gh_release::Node>();
69 }
70
71 fn emit(requests: Vec<Self::Request>, ctx: &mut NodeCtx<'_>) -> anyhow::Result<()> {
72 let mut versions: BTreeMap<OpenhclKernelPackageKind, String> = BTreeMap::new();
73 let mut local_paths: BTreeMap<
74 OpenhclKernelPackageArch,
75 (ReadVar<PathBuf>, ReadVar<PathBuf>),
76 > = BTreeMap::new();
77 let mut kernel_reqs: BTreeMap<
78 (OpenhclKernelPackageKind, OpenhclKernelPackageArch),
79 Vec<WriteVar<PathBuf>>,
80 > = BTreeMap::new();
81 let mut modules_reqs: BTreeMap<
82 (OpenhclKernelPackageKind, OpenhclKernelPackageArch),
83 Vec<WriteVar<PathBuf>>,
84 > = BTreeMap::new();
85 let mut pkg_reqs: BTreeMap<
86 (OpenhclKernelPackageKind, OpenhclKernelPackageArch),
87 Vec<WriteVar<PathBuf>>,
88 > = BTreeMap::new();
89 let mut metadata_reqs: BTreeMap<
90 (OpenhclKernelPackageKind, OpenhclKernelPackageArch),
91 Vec<WriteVar<PathBuf>>,
92 > = BTreeMap::new();
93
94 for req in requests {
95 match req {
96 Request::SetVersion(kind, v) => {
97 let mut old = versions.insert(kind, v.clone());
98 same_across_all_reqs("SetVersion", &mut old, v)?
99 }
100 Request::SetLocal {
101 arch,
102 kernel,
103 modules,
104 } => {
105 if local_paths.contains_key(&arch) {
106 anyhow::bail!("Duplicate local paths for {:?}", arch);
107 }
108 local_paths.insert(arch, (kernel, modules));
109 }
110 Request::GetKernel { kind, arch, kernel } => {
111 kernel_reqs.entry((kind, arch)).or_default().push(kernel);
112 }
113 Request::GetModules {
114 kind,
115 arch,
116 modules,
117 } => {
118 modules_reqs.entry((kind, arch)).or_default().push(modules);
119 }
120 Request::GetPackageRoot { kind, arch, pkg } => {
121 pkg_reqs.entry((kind, arch)).or_default().push(pkg);
122 }
123 Request::GetMetadata {
124 kind,
125 arch,
126 metadata,
127 } => {
128 metadata_reqs
129 .entry((kind, arch))
130 .or_default()
131 .push(metadata);
132 }
133 }
134 }
135
136 let all_reqs: std::collections::BTreeSet<(
138 OpenhclKernelPackageKind,
139 OpenhclKernelPackageArch,
140 )> = kernel_reqs
141 .keys()
142 .chain(modules_reqs.keys())
143 .chain(pkg_reqs.keys())
144 .chain(metadata_reqs.keys())
145 .cloned()
146 .collect();
147
148 for (kind, arch) in &all_reqs {
150 if !local_paths.contains_key(arch) && !versions.contains_key(kind) {
151 anyhow::bail!(
152 "Must provide either SetLocal for {:?} or SetVersion for {:?}",
153 arch,
154 kind
155 );
156 }
157 }
158
159 if all_reqs.is_empty() {
160 return Ok(());
161 }
162
163 let (local_reqs, download_reqs): (Vec<_>, Vec<_>) = all_reqs
165 .into_iter()
166 .partition(|(_, arch)| local_paths.contains_key(arch));
167
168 let (kernel_reqs_local, mut kernel_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) =
170 kernel_reqs
171 .into_iter()
172 .partition(|((_, arch), _)| local_paths.contains_key(arch));
173 let (modules_reqs_local, mut modules_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) =
174 modules_reqs
175 .into_iter()
176 .partition(|((_, arch), _)| local_paths.contains_key(arch));
177 let (pkg_reqs_local, mut pkg_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) = pkg_reqs
178 .into_iter()
179 .partition(|((_, arch), _)| local_paths.contains_key(arch));
180 let (metadata_reqs_local, mut metadata_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) =
181 metadata_reqs
182 .into_iter()
183 .partition(|((_, arch), _)| local_paths.contains_key(arch));
184
185 if !local_reqs.is_empty() {
187 ctx.emit_rust_step("use local kernel package", |ctx| {
188 let mut kernel_reqs = kernel_reqs_local.claim(ctx);
189 let mut modules_reqs = modules_reqs_local.claim(ctx);
190 let mut pkg_reqs = pkg_reqs_local.claim(ctx);
191 let mut metadata_reqs = metadata_reqs_local.claim(ctx);
192 let local_paths: BTreeMap<_, _> = local_paths
193 .into_iter()
194 .map(|(arch, (k, m))| (arch, (k.claim(ctx), m.claim(ctx))))
195 .collect();
196 let local_reqs = local_reqs.clone();
197
198 move |rt| {
199 for (_, arch) in local_reqs {
200 let (kernel_var, modules_var) = local_paths.get(&arch).unwrap();
201 let kernel_path = rt.read(kernel_var.clone());
202 let modules_path = rt.read(modules_var.clone());
203
204 log::info!(
205 "using local kernel at {:?} and modules at {:?}",
206 kernel_path,
207 modules_path
208 );
209
210 for kind in [
212 OpenhclKernelPackageKind::Main,
213 OpenhclKernelPackageKind::Dev,
214 OpenhclKernelPackageKind::Cvm,
215 OpenhclKernelPackageKind::CvmDev,
216 ] {
217 if let Some(vars) = kernel_reqs.remove(&(kind, arch)) {
218 rt.write_all(vars, &kernel_path);
219 }
220 }
221
222 for kind in [
224 OpenhclKernelPackageKind::Main,
225 OpenhclKernelPackageKind::Dev,
226 OpenhclKernelPackageKind::Cvm,
227 OpenhclKernelPackageKind::CvmDev,
228 ] {
229 if let Some(vars) = modules_reqs.remove(&(kind, arch)) {
230 rt.write_all(vars, &modules_path);
231 }
232 }
233
234 if let Some(parent) = kernel_path.parent() {
236 let parent_buf = parent.to_path_buf();
237 for kind in [
238 OpenhclKernelPackageKind::Main,
239 OpenhclKernelPackageKind::Dev,
240 OpenhclKernelPackageKind::Cvm,
241 OpenhclKernelPackageKind::CvmDev,
242 ] {
243 if let Some(vars) = pkg_reqs.remove(&(kind, arch)) {
244 rt.write_all(vars, &parent_buf);
245 }
246 }
247
248 let metadata_path = parent_buf.join("kernel_build_metadata.json");
250 for kind in [
251 OpenhclKernelPackageKind::Main,
252 OpenhclKernelPackageKind::Dev,
253 OpenhclKernelPackageKind::Cvm,
254 OpenhclKernelPackageKind::CvmDev,
255 ] {
256 if let Some(vars) = metadata_reqs.remove(&(kind, arch)) {
257 rt.write_all(vars, &metadata_path);
258 }
259 }
260 }
261 }
262 Ok(())
263 }
264 });
265 }
266
267 if download_reqs.is_empty() {
268 return Ok(());
269 }
270
271 let extract_zip_deps = flowey_lib_common::_util::extract::extract_zip_if_new_deps(ctx);
273
274 for (kind, arch) in download_reqs {
275 let version = versions.get(&kind).expect("checked above");
276 let tag = format!(
277 "rolling-lts/hcl-{}/{}",
278 match kind {
279 OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Cvm => "main",
280 OpenhclKernelPackageKind::Dev | OpenhclKernelPackageKind::CvmDev => "dev",
281 },
282 version
283 );
284
285 let file_name = format!(
286 "Microsoft.OHCL.Kernel{}.{}{}-{}.tar.gz",
287 match kind {
288 OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Cvm => "",
289 OpenhclKernelPackageKind::Dev | OpenhclKernelPackageKind::CvmDev => ".Dev",
290 },
291 version,
292 match kind {
293 OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Dev => "",
294 OpenhclKernelPackageKind::Cvm | OpenhclKernelPackageKind::CvmDev => "-cvm",
295 },
296 match arch {
297 OpenhclKernelPackageArch::X86_64 => "x64",
298 OpenhclKernelPackageArch::Aarch64 => "arm64",
299 },
300 );
301
302 let kernel_package_tar_gz =
303 ctx.reqv(|v| flowey_lib_common::download_gh_release::Request {
304 repo_owner: "microsoft".into(),
305 repo_name: "OHCL-Linux-Kernel".into(),
306 needs_auth: false,
307 tag,
308 file_name: file_name.clone(),
309 path: v,
310 });
311
312 let kernel_file_name = match arch {
313 OpenhclKernelPackageArch::X86_64 => "vmlinux",
314 OpenhclKernelPackageArch::Aarch64 => "Image",
315 };
316
317 let has_kernel_req = kernel_reqs_download.contains_key(&(kind, arch));
318 let has_modules_req = modules_reqs_download.contains_key(&(kind, arch));
319 let has_pkg_req = pkg_reqs_download.contains_key(&(kind, arch));
320 let has_metadata_req = metadata_reqs_download.contains_key(&(kind, arch));
321
322 ctx.emit_rust_step("extract and resolve kernel package", |ctx| {
323 let extract_zip_deps = extract_zip_deps.clone().claim(ctx);
324 let kernel_vars = if has_kernel_req {
325 Some(
326 kernel_reqs_download
327 .remove(&(kind, arch))
328 .unwrap()
329 .claim(ctx),
330 )
331 } else {
332 None
333 };
334 let modules_vars = if has_modules_req {
335 Some(
336 modules_reqs_download
337 .remove(&(kind, arch))
338 .unwrap()
339 .claim(ctx),
340 )
341 } else {
342 None
343 };
344 let pkg_vars = if has_pkg_req {
345 Some(pkg_reqs_download.remove(&(kind, arch)).unwrap().claim(ctx))
346 } else {
347 None
348 };
349 let metadata_vars = if has_metadata_req {
350 Some(
351 metadata_reqs_download
352 .remove(&(kind, arch))
353 .unwrap()
354 .claim(ctx),
355 )
356 } else {
357 None
358 };
359 let kernel_package_tar_gz = kernel_package_tar_gz.claim(ctx);
360 let file_name = file_name.clone();
361 let kernel_file_name = kernel_file_name.to_string();
362
363 move |rt| {
364 let kernel_package_tar_gz = rt.read(kernel_package_tar_gz);
365
366 let extract_dir = flowey_lib_common::_util::extract::extract_zip_if_new(
368 rt,
369 extract_zip_deps,
370 &kernel_package_tar_gz,
371 &file_name,
372 )?;
373
374 let kernel_path = extract_dir.join(&kernel_file_name);
376 let modules_path = extract_dir.join("modules");
377 let metadata_path = extract_dir.join("kernel_build_metadata.json");
378
379 if let Some(vars) = kernel_vars {
380 rt.write_all(vars, &kernel_path);
381 }
382 if let Some(vars) = modules_vars {
383 rt.write_all(vars, &modules_path);
384 }
385 if let Some(vars) = pkg_vars {
386 rt.write_all(vars, &extract_dir);
387 }
388 if let Some(vars) = metadata_vars {
389 rt.write_all(vars, &metadata_path);
390 }
391
392 Ok(())
393 }
394 });
395 }
396
397 Ok(())
398 }
399}