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 for kind in [
202 OpenhclKernelPackageKind::Main,
203 OpenhclKernelPackageKind::Dev,
204 OpenhclKernelPackageKind::Cvm,
205 OpenhclKernelPackageKind::CvmDev,
206 ] {
207 if let Some(vars) = kernel_reqs.remove(&(kind, arch)) {
208 rt.write_all(vars, kernel_path);
209 }
210 }
211
212 for kind in [
214 OpenhclKernelPackageKind::Main,
215 OpenhclKernelPackageKind::Dev,
216 OpenhclKernelPackageKind::Cvm,
217 OpenhclKernelPackageKind::CvmDev,
218 ] {
219 if let Some(vars) = modules_reqs.remove(&(kind, arch)) {
220 rt.write_all(vars, modules_path);
221 }
222 }
223
224 if let Some(parent) = kernel_path.parent() {
226 let parent_buf = parent.to_path_buf();
227 for kind in [
228 OpenhclKernelPackageKind::Main,
229 OpenhclKernelPackageKind::Dev,
230 OpenhclKernelPackageKind::Cvm,
231 OpenhclKernelPackageKind::CvmDev,
232 ] {
233 if let Some(vars) = pkg_reqs.remove(&(kind, arch)) {
234 rt.write_all(vars, &parent_buf);
235 }
236 }
237
238 let metadata_path = parent_buf.join("kernel_build_metadata.json");
240 for kind in [
241 OpenhclKernelPackageKind::Main,
242 OpenhclKernelPackageKind::Dev,
243 OpenhclKernelPackageKind::Cvm,
244 OpenhclKernelPackageKind::CvmDev,
245 ] {
246 if let Some(vars) = metadata_reqs.remove(&(kind, arch)) {
247 rt.write_all(vars, &metadata_path);
248 }
249 }
250 }
251 }
252 Ok(())
253 }
254 });
255 }
256
257 if download_reqs.is_empty() {
258 return Ok(());
259 }
260
261 let extract_zip_deps = flowey_lib_common::_util::extract::extract_zip_if_new_deps(ctx);
263
264 for (kind, arch) in download_reqs {
265 let version = versions.get(&kind).expect("checked above");
266 let tag = format!(
267 "rolling-lts/hcl-{}/{}",
268 match kind {
269 OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Cvm => "main",
270 OpenhclKernelPackageKind::Dev | OpenhclKernelPackageKind::CvmDev => "dev",
271 },
272 version
273 );
274
275 let file_name = format!(
276 "Microsoft.OHCL.Kernel{}.{}{}-{}.tar.gz",
277 match kind {
278 OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Cvm => "",
279 OpenhclKernelPackageKind::Dev | OpenhclKernelPackageKind::CvmDev => ".Dev",
280 },
281 version,
282 match kind {
283 OpenhclKernelPackageKind::Main | OpenhclKernelPackageKind::Dev => "",
284 OpenhclKernelPackageKind::Cvm | OpenhclKernelPackageKind::CvmDev => "-cvm",
285 },
286 match arch {
287 OpenhclKernelPackageArch::X86_64 => "x64",
288 OpenhclKernelPackageArch::Aarch64 => "arm64",
289 },
290 );
291
292 let kernel_package_tar_gz =
293 ctx.reqv(|v| flowey_lib_common::download_gh_release::Request {
294 repo_owner: "microsoft".into(),
295 repo_name: "OHCL-Linux-Kernel".into(),
296 needs_auth: false,
297 tag,
298 file_name: file_name.clone(),
299 path: v,
300 });
301
302 let kernel_file_name = match arch {
303 OpenhclKernelPackageArch::X86_64 => "vmlinux",
304 OpenhclKernelPackageArch::Aarch64 => "Image",
305 };
306
307 let has_kernel_req = kernel_reqs_download.contains_key(&(kind, arch));
308 let has_modules_req = modules_reqs_download.contains_key(&(kind, arch));
309 let has_pkg_req = pkg_reqs_download.contains_key(&(kind, arch));
310 let has_metadata_req = metadata_reqs_download.contains_key(&(kind, arch));
311
312 ctx.emit_rust_step("extract and resolve kernel package", |ctx| {
313 let extract_zip_deps = extract_zip_deps.clone().claim(ctx);
314 let kernel_vars = if has_kernel_req {
315 Some(
316 kernel_reqs_download
317 .remove(&(kind, arch))
318 .unwrap()
319 .claim(ctx),
320 )
321 } else {
322 None
323 };
324 let modules_vars = if has_modules_req {
325 Some(
326 modules_reqs_download
327 .remove(&(kind, arch))
328 .unwrap()
329 .claim(ctx),
330 )
331 } else {
332 None
333 };
334 let pkg_vars = if has_pkg_req {
335 Some(pkg_reqs_download.remove(&(kind, arch)).unwrap().claim(ctx))
336 } else {
337 None
338 };
339 let metadata_vars = if has_metadata_req {
340 Some(
341 metadata_reqs_download
342 .remove(&(kind, arch))
343 .unwrap()
344 .claim(ctx),
345 )
346 } else {
347 None
348 };
349 let kernel_package_tar_gz = kernel_package_tar_gz.claim(ctx);
350 let file_name = file_name.clone();
351 let kernel_file_name = kernel_file_name.to_string();
352
353 move |rt| {
354 let kernel_package_tar_gz = rt.read(kernel_package_tar_gz);
355
356 let extract_dir = flowey_lib_common::_util::extract::extract_zip_if_new(
358 rt,
359 extract_zip_deps,
360 &kernel_package_tar_gz,
361 &file_name,
362 )?;
363
364 let kernel_path = extract_dir.join(&kernel_file_name);
366 let modules_path = extract_dir.join("modules");
367 let metadata_path = extract_dir.join("kernel_build_metadata.json");
368
369 if let Some(vars) = kernel_vars {
370 rt.write_all(vars, &kernel_path);
371 }
372 if let Some(vars) = modules_vars {
373 rt.write_all(vars, &modules_path);
374 }
375 if let Some(vars) = pkg_vars {
376 rt.write_all(vars, &extract_dir);
377 }
378 if let Some(vars) = metadata_vars {
379 rt.write_all(vars, &metadata_path);
380 }
381
382 Ok(())
383 }
384 });
385 }
386
387 Ok(())
388 }
389}