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: PathBuf,
30 modules: 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<OpenhclKernelPackageArch, (PathBuf, PathBuf)> =
74 BTreeMap::new();
75 let mut kernel_reqs: BTreeMap<
76 (OpenhclKernelPackageKind, OpenhclKernelPackageArch),
77 Vec<WriteVar<PathBuf>>,
78 > = BTreeMap::new();
79 let mut modules_reqs: BTreeMap<
80 (OpenhclKernelPackageKind, OpenhclKernelPackageArch),
81 Vec<WriteVar<PathBuf>>,
82 > = BTreeMap::new();
83 let mut pkg_reqs: BTreeMap<
84 (OpenhclKernelPackageKind, OpenhclKernelPackageArch),
85 Vec<WriteVar<PathBuf>>,
86 > = BTreeMap::new();
87 let mut metadata_reqs: BTreeMap<
88 (OpenhclKernelPackageKind, OpenhclKernelPackageArch),
89 Vec<WriteVar<PathBuf>>,
90 > = BTreeMap::new();
91
92 for req in requests {
93 match req {
94 Request::SetVersion(kind, v) => {
95 let mut old = versions.insert(kind, v.clone());
96 same_across_all_reqs("SetVersion", &mut old, v)?
97 }
98 Request::SetLocal {
99 arch,
100 kernel,
101 modules,
102 } => {
103 if let Some(existing) = local_paths.get(&arch) {
104 if existing != &(kernel.clone(), modules.clone()) {
105 anyhow::bail!("Conflicting local paths for {:?}", arch);
106 }
107 } else {
108 local_paths.insert(arch, (kernel, modules));
109 }
110 }
111 Request::GetKernel { kind, arch, kernel } => {
112 kernel_reqs.entry((kind, arch)).or_default().push(kernel);
113 }
114 Request::GetModules {
115 kind,
116 arch,
117 modules,
118 } => {
119 modules_reqs.entry((kind, arch)).or_default().push(modules);
120 }
121 Request::GetPackageRoot { kind, arch, pkg } => {
122 pkg_reqs.entry((kind, arch)).or_default().push(pkg);
123 }
124 Request::GetMetadata {
125 kind,
126 arch,
127 metadata,
128 } => {
129 metadata_reqs
130 .entry((kind, arch))
131 .or_default()
132 .push(metadata);
133 }
134 }
135 }
136
137 let all_reqs: std::collections::BTreeSet<(
139 OpenhclKernelPackageKind,
140 OpenhclKernelPackageArch,
141 )> = kernel_reqs
142 .keys()
143 .chain(modules_reqs.keys())
144 .chain(pkg_reqs.keys())
145 .chain(metadata_reqs.keys())
146 .cloned()
147 .collect();
148
149 for (kind, arch) in &all_reqs {
151 if !local_paths.contains_key(arch) && !versions.contains_key(kind) {
152 anyhow::bail!(
153 "Must provide either SetLocal for {:?} or SetVersion for {:?}",
154 arch,
155 kind
156 );
157 }
158 }
159
160 if all_reqs.is_empty() {
161 return Ok(());
162 }
163
164 let (local_reqs, download_reqs): (Vec<_>, Vec<_>) = all_reqs
166 .into_iter()
167 .partition(|(_, arch)| local_paths.contains_key(arch));
168
169 let (kernel_reqs_local, mut kernel_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) =
171 kernel_reqs
172 .into_iter()
173 .partition(|((_, arch), _)| local_paths.contains_key(arch));
174 let (modules_reqs_local, mut modules_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) =
175 modules_reqs
176 .into_iter()
177 .partition(|((_, arch), _)| local_paths.contains_key(arch));
178 let (pkg_reqs_local, mut pkg_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) = pkg_reqs
179 .into_iter()
180 .partition(|((_, arch), _)| local_paths.contains_key(arch));
181 let (metadata_reqs_local, mut metadata_reqs_download): (BTreeMap<_, _>, BTreeMap<_, _>) =
182 metadata_reqs
183 .into_iter()
184 .partition(|((_, arch), _)| local_paths.contains_key(arch));
185
186 if !local_reqs.is_empty() {
188 ctx.emit_rust_step("use local kernel package", |ctx| {
189 let mut kernel_reqs = kernel_reqs_local.claim(ctx);
190 let mut modules_reqs = modules_reqs_local.claim(ctx);
191 let mut pkg_reqs = pkg_reqs_local.claim(ctx);
192 let mut metadata_reqs = metadata_reqs_local.claim(ctx);
193 let local_paths = local_paths.clone();
194 let local_reqs = local_reqs.clone();
195
196 move |rt| {
197 for (_, arch) in local_reqs {
198 let (kernel_path, modules_path) = local_paths.get(&arch).unwrap();
199
200 log::info!(
201 "using local kernel at {:?} and modules at {:?}",
202 kernel_path,
203 modules_path
204 );
205
206 for kind in [
208 OpenhclKernelPackageKind::Main,
209 OpenhclKernelPackageKind::Dev,
210 OpenhclKernelPackageKind::Cvm,
211 OpenhclKernelPackageKind::CvmDev,
212 ] {
213 if let Some(vars) = kernel_reqs.remove(&(kind, arch)) {
214 rt.write_all(vars, kernel_path);
215 }
216 }
217
218 for kind in [
220 OpenhclKernelPackageKind::Main,
221 OpenhclKernelPackageKind::Dev,
222 OpenhclKernelPackageKind::Cvm,
223 OpenhclKernelPackageKind::CvmDev,
224 ] {
225 if let Some(vars) = modules_reqs.remove(&(kind, arch)) {
226 rt.write_all(vars, modules_path);
227 }
228 }
229
230 if let Some(parent) = kernel_path.parent() {
232 let parent_buf = parent.to_path_buf();
233 for kind in [
234 OpenhclKernelPackageKind::Main,
235 OpenhclKernelPackageKind::Dev,
236 OpenhclKernelPackageKind::Cvm,
237 OpenhclKernelPackageKind::CvmDev,
238 ] {
239 if let Some(vars) = pkg_reqs.remove(&(kind, arch)) {
240 rt.write_all(vars, &parent_buf);
241 }
242 }
243
244 let metadata_path = parent_buf.join("kernel_build_metadata.json");
246 for kind in [
247 OpenhclKernelPackageKind::Main,
248 OpenhclKernelPackageKind::Dev,
249 OpenhclKernelPackageKind::Cvm,
250 OpenhclKernelPackageKind::CvmDev,
251 ] {
252 if let Some(vars) = metadata_reqs.remove(&(kind, arch)) {
253 rt.write_all(vars, &metadata_path);
254 }
255 }
256 }
257 }
258 Ok(())
259 }
260 });
261 }
262
263 if download_reqs.is_empty() {
264 return Ok(());
265 }
266
267 let extract_zip_deps = flowey_lib_common::_util::extract::extract_zip_if_new_deps(ctx);
269
270 for (kind, arch) in download_reqs {
271 let version = versions.get(&kind).expect("checked above");
272 let tag = format!(
273 "rolling-lts/hcl-{}/{}",
274 match kind {
275 OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Cvm => "main",
276 OpenhclKernelPackageKind::Dev | OpenhclKernelPackageKind::CvmDev => "dev",
277 },
278 version
279 );
280
281 let file_name = format!(
282 "Microsoft.OHCL.Kernel{}.{}{}-{}.tar.gz",
283 match kind {
284 OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Cvm => "",
285 OpenhclKernelPackageKind::Dev | OpenhclKernelPackageKind::CvmDev => ".Dev",
286 },
287 version,
288 match kind {
289 OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Dev => "",
290 OpenhclKernelPackageKind::Cvm | OpenhclKernelPackageKind::CvmDev => "-cvm",
291 },
292 match arch {
293 OpenhclKernelPackageArch::X86_64 => "x64",
294 OpenhclKernelPackageArch::Aarch64 => "arm64",
295 },
296 );
297
298 let kernel_package_tar_gz =
299 ctx.reqv(|v| flowey_lib_common::download_gh_release::Request {
300 repo_owner: "microsoft".into(),
301 repo_name: "OHCL-Linux-Kernel".into(),
302 needs_auth: false,
303 tag,
304 file_name: file_name.clone(),
305 path: v,
306 });
307
308 let kernel_file_name = match arch {
309 OpenhclKernelPackageArch::X86_64 => "vmlinux",
310 OpenhclKernelPackageArch::Aarch64 => "Image",
311 };
312
313 let has_kernel_req = kernel_reqs_download.contains_key(&(kind, arch));
314 let has_modules_req = modules_reqs_download.contains_key(&(kind, arch));
315 let has_pkg_req = pkg_reqs_download.contains_key(&(kind, arch));
316 let has_metadata_req = metadata_reqs_download.contains_key(&(kind, arch));
317
318 ctx.emit_rust_step("extract and resolve kernel package", |ctx| {
319 let extract_zip_deps = extract_zip_deps.clone().claim(ctx);
320 let kernel_vars = if has_kernel_req {
321 Some(
322 kernel_reqs_download
323 .remove(&(kind, arch))
324 .unwrap()
325 .claim(ctx),
326 )
327 } else {
328 None
329 };
330 let modules_vars = if has_modules_req {
331 Some(
332 modules_reqs_download
333 .remove(&(kind, arch))
334 .unwrap()
335 .claim(ctx),
336 )
337 } else {
338 None
339 };
340 let pkg_vars = if has_pkg_req {
341 Some(pkg_reqs_download.remove(&(kind, arch)).unwrap().claim(ctx))
342 } else {
343 None
344 };
345 let metadata_vars = if has_metadata_req {
346 Some(
347 metadata_reqs_download
348 .remove(&(kind, arch))
349 .unwrap()
350 .claim(ctx),
351 )
352 } else {
353 None
354 };
355 let kernel_package_tar_gz = kernel_package_tar_gz.claim(ctx);
356 let file_name = file_name.clone();
357 let kernel_file_name = kernel_file_name.to_string();
358
359 move |rt| {
360 let kernel_package_tar_gz = rt.read(kernel_package_tar_gz);
361
362 let extract_dir = flowey_lib_common::_util::extract::extract_zip_if_new(
364 rt,
365 extract_zip_deps,
366 &kernel_package_tar_gz,
367 &file_name,
368 )?;
369
370 let kernel_path = extract_dir.join(&kernel_file_name);
372 let modules_path = extract_dir.join("modules");
373 let metadata_path = extract_dir.join("kernel_build_metadata.json");
374
375 if let Some(vars) = kernel_vars {
376 rt.write_all(vars, &kernel_path);
377 }
378 if let Some(vars) = modules_vars {
379 rt.write_all(vars, &modules_path);
380 }
381 if let Some(vars) = pkg_vars {
382 rt.write_all(vars, &extract_dir);
383 }
384 if let Some(vars) = metadata_vars {
385 rt.write_all(vars, &metadata_path);
386 }
387
388 Ok(())
389 }
390 });
391 }
392
393 Ok(())
394 }
395}