microsoft/openvmm

Public

mirrored from https://github.com/microsoft/openvmmAvailable

CodeCommitsIssuesPull requestsActionsInsightsSecurity
c0d6416eef547dcf25a7148b4d02f2c5fca3f8a3

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

flowey/flowey_lib_common/src/cache.rs

592lines · modecode

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Cache the contents of a particular directory between runs.
5//!
6//! The contents of the provided `dir` will be saved at the end of a run, using
7//! the user-defined `key` string to tag the contents of the cache.
8//!
9//! Subsequent runs will then use the `key` to restore the contents of the
10//! directory.
11//!
12//! # A note of file sizes
13//!
14//! This node is backed by the in-box Cache@2 Task on ADO, and the in-box
15//! actions/cache@v3 Action on Github Actions.
16//!
17//! These actions have limits on the size of data they can cache at any given
18//! time, and potentially have issues with particularly large artifacts (e.g:
19//! gigabytes in size).
20//!
21//! In cases where you're intending to cache large files, it is recommended to
22//! implement caching functionality directly using [`NodeCtx::persistent_dir`],
23//! which is guaranteed to be reliable (when running on a system where such
24//! persistent storage is available).
25//!
26//! # Clearing the cache
27//!
28//! Clearing the cache is done in different ways depending on the backend:
29//!
30//! - Local: just delete the cache folder on your machine
31//! - Github: use the cache tasks's web UI to manage cache entries
32//! - ADO: define a pipeline-level variable called `FloweyCacheGeneration`, and set
33//! it to an new arbitrary value.
34//! - This is because ADO doesn't have a native way to flush the cache
35//! outside of updating the cache key in the YAML file itself.
36
37use flowey::node::prelude::*;
38use std::collections::BTreeSet;
39use std::io::Seek;
40use std::io::Write;
41
42/// Status of a cache directory.
43#[derive(Debug, Serialize, Deserialize)]
44pub enum CacheHit {
45 /// Complete miss - cache is empty.
46 Miss,
47 /// Direct hit - cache is perfectly restored.
48 Hit,
49 /// Partial hit - cache was partially restored.
50 PartialHit,
51}
52
53/// How the result of the cache task should be reported.
54#[derive(Serialize, Deserialize)]
55pub enum CacheResult {
56 /// Don't care about the details, only care that the task was run.
57 SideEffect(WriteVar<SideEffect>),
58 /// Get details on the result of the cache restore.
59 HitVar(WriteVar<CacheHit>),
60}
61
62flowey_request! {
63 pub struct Request {
64 /// Friendly label for the directory being cached
65 pub label: String,
66 /// Absolute path to the directory that will be cached between runs
67 pub dir: ReadVar<PathBuf>,
68 /// The key created when saving a cache and the key used to search for a
69 /// cache.
70 pub key: ReadVar<String>,
71 /// An optional set of alternative restore keys.
72 ///
73 /// If no cache hit occurs for key, these restore keys are used
74 /// sequentially in the order provided to find and restore a cache.
75 pub restore_keys: Option<ReadVar<Vec<String>>>,
76 /// Variable to write the result of trying to restore the cache.
77 pub hitvar: CacheResult,
78 }
79}
80
81enum ClaimedCacheResult {
82 SideEffect,
83 HitVar(ClaimedWriteVar<CacheHit>),
84}
85
86new_flow_node!(struct Node);
87
88impl FlowNode for Node {
89 type Request = Request;
90
91 fn imports(_ctx: &mut ImportCtx<'_>) {}
92
93 fn emit(requests: Vec<Self::Request>, ctx: &mut NodeCtx<'_>) -> anyhow::Result<()> {
94 // -- end of req processing -- //
95
96 match ctx.backend() {
97 FlowBackend::Local => {
98 if !ctx.supports_persistent_dir() {
99 ctx.emit_minor_rust_step("Reporting cache misses", |ctx| {
100 let hitvars = requests
101 .into_iter()
102 .map(|v| match v.hitvar {
103 CacheResult::SideEffect(v) => {
104 v.claim(ctx);
105 ClaimedCacheResult::SideEffect
106 }
107 CacheResult::HitVar(v) => ClaimedCacheResult::HitVar(v.claim(ctx)),
108 })
109 .collect::<Vec<_>>();
110
111 |rt| {
112 for var in hitvars {
113 match var {
114 ClaimedCacheResult::SideEffect => {}
115 ClaimedCacheResult::HitVar(var) => {
116 rt.write(var, &CacheHit::Miss)
117 }
118 }
119 }
120 }
121 });
122
123 return Ok(());
124 };
125
126 for Request {
127 label,
128 dir,
129 key,
130 restore_keys,
131 hitvar,
132 } in requests
133 {
134 // work around a bug in how post-job nodes affect stage1 day
135 // culling...
136 let persistent_dir = ctx.persistent_dir().unwrap();
137
138 // regardless if we're reporting the hit back to the user, we'll
139 // want to record the hit status so we can efficiently skip
140 // saving to the cache in the post-job step
141 let (side_effect, (hitvar_reader, hitvar)) = match hitvar {
142 CacheResult::HitVar(hitvar) => (None, (hitvar.new_reader(), hitvar)),
143 CacheResult::SideEffect(var) => (Some(var), ctx.new_var()),
144 };
145
146 let (resolve_post_job, require_post_job) = ctx.new_post_job_side_effect();
147
148 ctx.emit_rust_step(format!("Restore cache: {label}"), |ctx| {
149 require_post_job.claim(ctx);
150 side_effect.claim(ctx);
151 let persistent_dir = persistent_dir.clone().claim(ctx);
152 let dir = dir.clone().claim(ctx);
153 let key = key.clone().claim(ctx);
154 let restore_keys = restore_keys.claim(ctx);
155 let hitvar = hitvar.claim(ctx);
156 |rt| {
157 let persistent_dir = rt.read(persistent_dir);
158 let dir = rt.read(dir);
159 let key = rt.read(key);
160 let restore_keys = restore_keys.map(|x| rt.read(x));
161
162 let set_hitvar = move |val| {
163 log::info!("cache status: {val:?}");
164 rt.write(hitvar, &val);
165 };
166
167 // figure out what cache entries are available to us
168 //
169 // (reading this entire file into memory seems fine at
170 // this juncture, given the sort of datasets we're
171 // working with)
172 let available_keys: BTreeSet<String> = if let Ok(s) =
173 fs_err::read_to_string(persistent_dir.join("cache_keys"))
174 {
175 s.split('\n').map(|s| s.trim().to_owned()).collect()
176 } else {
177 BTreeSet::new()
178 };
179
180 // using the keys the user provided us, check if there's
181 // a match
182 let mut existing_cache_dir = None;
183 for (idx, key) in Some(key)
184 .into_iter()
185 .chain(restore_keys.into_iter().flatten())
186 .enumerate()
187 {
188 if available_keys.contains(&key) {
189 existing_cache_dir = Some((idx == 0, hash_key_to_dir(&key)));
190 break;
191 }
192 }
193
194 let Some((direct_hit, existing_cache_dir)) = existing_cache_dir else {
195 set_hitvar(CacheHit::Miss);
196 return Ok(());
197 };
198
199 crate::_util::copy_dir_all(
200 persistent_dir.join(existing_cache_dir),
201 dir,
202 )
203 .context("while restoring cache")?;
204
205 set_hitvar(if direct_hit {
206 CacheHit::Hit
207 } else {
208 CacheHit::PartialHit
209 });
210
211 Ok(())
212 }
213 });
214
215 ctx.emit_rust_step(format!("Saving cache: {label}"), |ctx| {
216 resolve_post_job.claim(ctx);
217 let hitvar_reader = hitvar_reader.claim(ctx);
218 let persistent_dir = persistent_dir.clone().claim(ctx);
219 let dir = dir.claim(ctx);
220 let key = key.claim(ctx);
221 move |rt| {
222 let persistent_dir = rt.read(persistent_dir);
223 let dir = rt.read(dir);
224 let key = rt.read(key);
225 let hitvar_reader = rt.read(hitvar_reader);
226
227 let mut cache_keys_file = fs_err::OpenOptions::new()
228 .append(true)
229 .create(true)
230 .read(true)
231 .open(persistent_dir.join("cache_keys"))?;
232
233 if matches!(hitvar_reader, CacheHit::Hit) {
234 // no need to update the cache
235 log::info!("was direct hit - no updates needed");
236 return Ok(());
237 }
238
239 // otherwise, need to update the cache
240 crate::_util::copy_dir_all(
241 dir,
242 persistent_dir.join(hash_key_to_dir(&key)),
243 )?;
244
245 cache_keys_file.seek(std::io::SeekFrom::End(0))?;
246 writeln!(cache_keys_file, "{}", key)?;
247
248 log::info!("cache saved");
249
250 Ok(())
251 }
252 });
253 }
254 }
255 FlowBackend::Ado => {
256 for Request {
257 label,
258 dir,
259 key,
260 restore_keys,
261 hitvar,
262 } in requests
263 {
264 let (resolve_post_job, require_post_job) = ctx.new_post_job_side_effect();
265
266 let (dir_string, key, restore_keys) = {
267 let (processed_dir, write_processed_dir) = ctx.new_var();
268 let (processed_key, write_processed_key) = ctx.new_var();
269 let (processed_keys, write_processed_keys) = if restore_keys.is_some() {
270 let (a, b) = ctx.new_var();
271 (Some(a), Some(b))
272 } else {
273 (None, None)
274 };
275
276 ctx.emit_rust_step("Pre-processing cache vars", |ctx| {
277 require_post_job.claim(ctx);
278 let write_processed_dir = write_processed_dir.claim(ctx);
279 let write_processed_key = write_processed_key.claim(ctx);
280 let write_processed_keys = write_processed_keys.claim(ctx);
281
282 let dir = dir.clone().claim(ctx);
283 let key = key.claim(ctx);
284 let restore_keys = restore_keys.claim(ctx);
285
286 |rt| {
287 let dir = rt.read(dir);
288 // while we're here, we'll convert dir into a
289 // String, so it can be stuffed into an ADO var
290 rt.write(
291 write_processed_dir,
292 &dir.absolute()?.display().to_string(),
293 );
294
295 // Inject `$(FloweyCacheGeneration)` as part of the
296 // cache key to provide a non-intrusive mechanism to
297 // cycle the ADO cache when it gets into an
298 // inconsistent state.
299 //
300 // Deny cross-os caching by default (matching Github
301 // CI works by default).
302 //
303 // FUTURE: add toggle to request to allow cross-os
304 // caching?
305 let inject_extras = |s| {
306 format!(r#""$(FloweyCacheGeneration)" | "$(Agent.OS)" | "{s}""#)
307 };
308
309 let key = rt.read(key);
310 rt.write(write_processed_key, &inject_extras(key));
311
312 if let Some(write_processed_keys) = write_processed_keys {
313 let restore_keys = rt.read(restore_keys.unwrap());
314 rt.write(
315 write_processed_keys,
316 &restore_keys
317 .into_iter()
318 .map(inject_extras)
319 .collect::<Vec<_>>()
320 .join("\\n"),
321 );
322 }
323
324 Ok(())
325 }
326 });
327
328 (processed_dir, processed_key, processed_keys)
329 };
330
331 let (hitvar_str_reader, hitvar_str_writer) =
332 if matches!(hitvar, CacheResult::HitVar(..)) {
333 let (r, w) = ctx.new_var();
334 (Some(r), Some(w))
335 } else {
336 (None, None)
337 };
338
339 ctx.emit_ado_step(format!("Restore cache: {label}"), |ctx| {
340 let dir_string = dir_string.clone().claim(ctx);
341 let key = key.claim(ctx);
342 let restore_keys = restore_keys.claim(ctx);
343 let hitvar_str_writer = hitvar_str_writer.claim(ctx);
344 |rt| {
345 let dir_string = rt.get_var(dir_string).as_raw_var_name();
346 let key = rt.get_var(key).as_raw_var_name();
347 let restore_keys = if let Some(restore_keys) = restore_keys {
348 format!(
349 "restore_keys: $({})",
350 rt.get_var(restore_keys).as_raw_var_name()
351 )
352 } else {
353 String::new()
354 };
355
356 let hitvar_input = if let Some(hitvar_str_writer) = hitvar_str_writer {
357 let hitvar_ado = AdoRuntimeVar::dangerous_from_global(
358 "FLOWEY_CACHE_HITVAR",
359 false,
360 );
361 // note the _lack_ of $() around the var!
362 let input =
363 format!("cacheHitVar: {}", hitvar_ado.as_raw_var_name());
364 rt.set_var(hitvar_str_writer, hitvar_ado);
365 input
366 } else {
367 String::new()
368 };
369
370 format!(
371 r#"
372 - task: Cache@2
373 inputs:
374 key: '$({key})'
375 path: $({dir_string})
376 {restore_keys}
377 {hitvar_input}
378 "#
379 )
380 }
381 });
382
383 if let Some(hitvar_str_reader) = hitvar_str_reader {
384 ctx.emit_rust_step("map ADO hitvar to flowey", |ctx| {
385 let label = label.clone();
386 let dir = dir.clone().claim(ctx);
387
388 let CacheResult::HitVar(hitvar) = hitvar else {
389 unreachable!()
390 };
391
392 let hitvar = hitvar.claim(ctx);
393 let hitvar_str_reader = hitvar_str_reader.claim(ctx);
394 move |rt| {
395 let dir = rt.read(dir);
396 let hitvar_str = rt.read(hitvar_str_reader);
397 let mut var = match hitvar_str.as_str() {
398 "true" => CacheHit::Hit,
399 "false" => CacheHit::Miss,
400 "inexact" => CacheHit::PartialHit,
401 other => anyhow::bail!("unexpected cacheHitVar value: {other}"),
402 };
403
404 // WORKAROUND: ADO is really cool software, and
405 // sometimes, when it feels like it, i'll get into
406 // an inconsistent state where it reports a cache
407 // hit, but then the cache is actually empty.
408 if matches!(var, CacheHit::Hit | CacheHit::PartialHit) {
409 if dir.read_dir()?.next().is_none() {
410 log::error!("Detected inconsistent ADO cache entry: {label}");
411 log::error!("Please define/cycle the `FloweyCacheGeneration` pipeline variable");
412 var = CacheHit::Miss;
413 }
414 }
415
416 rt.write(hitvar, &var);
417 Ok(())
418 }
419 });
420 }
421
422 ctx.emit_rust_step(format!("validate cache entry: {label}"), |ctx| {
423 resolve_post_job.claim(ctx);
424 let dir = dir.clone().claim(ctx);
425 move |rt| {
426 let mut dir_contents = rt.read(dir).read_dir()?.peekable();
427
428 if dir_contents.peek().is_none() {
429 log::error!("Detected empty cache folder for entry: {label}");
430 log::error!("This is a bug - please update the pipeline code");
431 anyhow::bail!("cache error")
432 }
433
434 for entry in dir_contents {
435 let entry = entry?;
436 log::debug!("uploading: {}", entry.path().display());
437 }
438
439 Ok(())
440 }
441 });
442 }
443 }
444 FlowBackend::Github => {
445 for Request {
446 label,
447 dir,
448 key,
449 restore_keys,
450 hitvar,
451 } in requests
452 {
453 let (resolve_post_job, require_post_job) = ctx.new_post_job_side_effect();
454
455 let (dir_string, key, restore_keys) = {
456 let (processed_dir, write_processed_dir) = ctx.new_var();
457 let (processed_key, write_processed_key) = ctx.new_var();
458 let (processed_keys, write_processed_keys) = if restore_keys.is_some() {
459 let (a, b) = ctx.new_var();
460 (Some(a), Some(b))
461 } else {
462 (None, None)
463 };
464
465 ctx.emit_rust_step("Pre-processing cache vars", |ctx| {
466 require_post_job.claim(ctx);
467 let write_processed_dir = write_processed_dir.claim(ctx);
468 let write_processed_key = write_processed_key.claim(ctx);
469 let write_processed_keys = write_processed_keys.claim(ctx);
470
471 let dir = dir.clone().claim(ctx);
472 let key = key.claim(ctx);
473 let restore_keys = restore_keys.claim(ctx);
474
475 |rt| {
476 let dir = rt.read(dir);
477 rt.write(
478 write_processed_dir,
479 &dir.absolute()?.display().to_string(),
480 );
481
482 let key = rt.read(key);
483 let key = format!("{key}-{}-{}", rt.arch(), rt.platform());
484 rt.write(write_processed_key, &key);
485
486 if let Some(write_processed_keys) = write_processed_keys {
487 let restore_keys = rt.read(restore_keys.unwrap());
488 rt.write(
489 write_processed_keys,
490 &format!(
491 r#""[{}]""#,
492 restore_keys
493 .into_iter()
494 .map(|s| format!(
495 "'{s}-{}-{}'",
496 rt.arch(),
497 rt.platform()
498 ))
499 .collect::<Vec<_>>()
500 .join(", ")
501 ),
502 );
503 }
504
505 Ok(())
506 }
507 });
508
509 (processed_dir, processed_key, processed_keys)
510 };
511
512 let (hitvar_str_reader, hitvar_str_writer) =
513 if matches!(hitvar, CacheResult::HitVar(..)) {
514 let (r, w) = ctx.new_var();
515 (Some(r), Some(w))
516 } else {
517 (None, None)
518 };
519
520 let mut step = ctx
521 .emit_gh_step(format!("Restore cache: {label}"), "actions/cache@v4")
522 .with("key", key)
523 .with("path", dir_string);
524 if let Some(restore_keys) = restore_keys {
525 step = step.with("restore-keys", restore_keys);
526 }
527 if let Some(hitvar_str_writer) = hitvar_str_writer {
528 step = step.output("cache-hit", hitvar_str_writer);
529 }
530 step.finish(ctx);
531
532 if let Some(hitvar_str_reader) = hitvar_str_reader {
533 ctx.emit_minor_rust_step("map Github cache-hit to flowey", |ctx| {
534 let CacheResult::HitVar(hitvar) = hitvar else {
535 unreachable!()
536 };
537
538 let hitvar = hitvar.claim(ctx);
539 let hitvar_str_reader = hitvar_str_reader.claim(ctx);
540 // TODO: How do we distinguish between a partial hit and a miss?
541 move |rt| {
542 let hitvar_str = rt.read(hitvar_str_reader);
543 // Github's cache action brilliantly only reports "false" if missing a cache key that exists,
544 // and leaves it blank if its a miss in other cases.
545 let var = match hitvar_str.as_str() {
546 "true" => CacheHit::Hit,
547 _ => CacheHit::Miss,
548 };
549
550 rt.write(hitvar, &var);
551 }
552 });
553 }
554
555 ctx.emit_rust_step(format!("validate cache entry: {label}"), |ctx| {
556 resolve_post_job.claim(ctx);
557 let dir = dir.clone().claim(ctx);
558 move |rt| {
559 let mut dir_contents = rt.read(dir).read_dir()?.peekable();
560
561 if dir_contents.peek().is_none() {
562 log::error!("Detected empty cache folder for entry: {label}");
563 log::error!("This is a bug - please update the pipeline code");
564 anyhow::bail!("cache error")
565 }
566
567 for entry in dir_contents {
568 let entry = entry?;
569 log::debug!("uploading: {}", entry.path().display());
570 }
571
572 Ok(())
573 }
574 });
575 }
576 }
577 }
578
579 Ok(())
580 }
581}
582
583// _technically_, if we want to be _super_ sure we're not gonna have a hash
584// collision, we should also do a content-hash of the thing we're about to
585// cache... but this should be OK for now, given that we don't expect to have a
586// massive number of cache entries.
587fn hash_key_to_dir(key: &str) -> String {
588 let hasher = &mut rustc_hash::FxHasher::default();
589 std::hash::Hash::hash(&key, hasher);
590 let hash = std::hash::Hasher::finish(hasher);
591 format!("{:08x?}", hash)
592}
593