sk_core/
hooks.rs

1use std::fs;
2use std::process::Stdio;
3
4use anyhow::{
5    anyhow,
6    bail,
7};
8use derive_more::Display;
9use sk_api::v1::SimulationHooksConfig;
10use tokio::io::{
11    AsyncWriteExt,
12    BufWriter,
13};
14use tokio::process::Command;
15use tracing::*;
16
17use crate::prelude::*;
18
19#[derive(Debug, Display)]
20#[display("{_variant}")]
21pub enum Type {
22    PreStart,
23    PreRun,
24    PostRun,
25    PostStop,
26}
27
28pub fn merge_hooks(maybe_files: &Option<Vec<String>>) -> anyhow::Result<Option<SimulationHooksConfig>> {
29    let Some(files) = maybe_files else {
30        return Ok(None);
31    };
32    if files.is_empty() {
33        return Ok(None);
34    }
35
36    Some(files.iter().try_fold(SimulationHooksConfig::default(), |mut merged_hooks, f| {
37        let next = serde_yaml::from_slice::<SimulationHooksConfig>(
38            &fs::read(f).map_err(|e| anyhow!("error reading hook {f}: {e}"))?,
39        )
40        .map_err(|e| anyhow!("error parsing hook {f}: {e}"))?;
41        merge_vecs(&mut merged_hooks.pre_start_hooks, next.pre_start_hooks);
42        merge_vecs(&mut merged_hooks.pre_run_hooks, next.pre_run_hooks);
43        merge_vecs(&mut merged_hooks.post_run_hooks, next.post_run_hooks);
44        merge_vecs(&mut merged_hooks.post_stop_hooks, next.post_stop_hooks);
45        Ok(merged_hooks)
46    }))
47    .transpose()
48}
49
50pub async fn execute(sim: &Simulation, type_: Type, recorder: &SkEventRecorder) -> EmptyResult {
51    let maybe_hooks = match &sim.spec.hooks {
52        Some(hooks_config) => match type_ {
53            Type::PreStart => hooks_config.pre_start_hooks.as_ref(),
54            Type::PreRun => hooks_config.pre_run_hooks.as_ref(),
55            Type::PostRun => hooks_config.post_run_hooks.as_ref(),
56            Type::PostStop => hooks_config.post_stop_hooks.as_ref(),
57        },
58        _ => None,
59    };
60
61    if let Some(hooks) = maybe_hooks {
62        info!("Executing {:?} hooks", type_);
63
64        for hook in hooks {
65            info!("Running `{}` with args {:?}", hook.cmd, hook.args);
66            let mut child = Command::new(hook.cmd.clone())
67                .args(hook.args.clone().unwrap_or_default())
68                .stdin(Stdio::piped())
69                .stdout(Stdio::piped())
70                .stderr(Stdio::piped())
71                .spawn()?;
72            if let Some(true) = hook.send_sim {
73                let mut stdin = BufWriter::new(child.stdin.take().ok_or(anyhow!("could not take stdin"))?);
74                stdin.write_all(&serde_json::to_vec(sim)?).await?;
75                stdin.flush().await?;
76            }
77            let cmd_str = format!("{} {}", hook.cmd, hook.args.clone().unwrap_or_default().join(" "));
78            let output = child.wait_with_output().await?;
79            info!("Hook output: {:?}", output);
80
81            if !output.status.success()
82                && let Some(false) = hook.ignore_failure
83            {
84                recorder
85                    .send_hook_failed_event(&format!("{type_}"), &cmd_str, output.status.code())
86                    .await?;
87                bail!("hook failed");
88            }
89        }
90        recorder.send_hooks_succeeded_event(&format!("{type_}")).await?;
91        info!("Done executing {:?} hooks", type_);
92    };
93
94    Ok(())
95}
96
97fn merge_vecs<T>(maybe_v1: &mut Option<Vec<T>>, maybe_v2: Option<Vec<T>>) {
98    if let Some(v2) = maybe_v2 {
99        if let Some(v1) = maybe_v1 { v1.extend(v2) } else { *maybe_v1 = Some(v2) }
100    }
101}
102
103#[cfg(test)]
104#[cfg_attr(coverage, coverage(off))]
105mod test {
106    use assert_fs::prelude::*;
107    use sk_testutils::*;
108
109    use super::*;
110
111    const HOOK1: &str = r#"
112---
113preStartHooks:
114  - cmd: prestart1
115    args:
116      - prestart-arg1
117      - prestart-arg2
118preRunHooks:
119  - cmd: prerun1
120    args:
121      - prerun-arg1
122postRunHooks:
123  - cmd: postrun1
124    args:
125      - postrun-arg1
126"#;
127
128    const HOOK2: &str = r#"
129---
130preStartHooks:
131  - cmd: prestart2
132  - cmd: prestart3
133preRunHooks:
134  - cmd: prerun2
135    args:
136      - prerun-arg2
137postStopHooks:
138  - cmd: poststop1
139    args:
140      - poststop-arg1
141"#;
142
143    const HOOK3: &str = r#"
144---
145preRunHooks:
146  - cmd: prerun3
147    args:
148      - prerun-arg3
149postRunHooks:
150  - cmd: postrun2
151    args:
152      - prerun-arg2
153"#;
154
155    const EXPECTED_MERGED: &str = r#"
156---
157preStartHooks:
158  - cmd: prestart1
159    args:
160      - prestart-arg1
161      - prestart-arg2
162  - cmd: prestart2
163  - cmd: prestart3
164preRunHooks:
165  - cmd: prerun1
166    args:
167      - prerun-arg1
168  - cmd: prerun2
169    args:
170      - prerun-arg2
171  - cmd: prerun3
172    args:
173      - prerun-arg3
174postRunHooks:
175  - cmd: postrun1
176    args:
177      - postrun-arg1
178  - cmd: postrun2
179    args:
180      - prerun-arg2
181postStopHooks:
182  - cmd: poststop1
183    args:
184      - poststop-arg1
185"#;
186
187    #[rstest]
188    fn test_merge_hooks() {
189        let temp = assert_fs::TempDir::new().unwrap();
190        let hook1 = temp.child("hook1.yml");
191        hook1.write_str(HOOK1).unwrap();
192        let hook2 = temp.child("hook2.yml");
193        hook2.write_str(HOOK2).unwrap();
194        let hook3 = temp.child("hook3.yml");
195        hook3.write_str(HOOK3).unwrap();
196
197        let merged_config = merge_hooks(&Some(vec![
198            hook1.path().to_str().unwrap().into(),
199            hook2.path().to_str().unwrap().into(),
200            hook3.path().to_str().unwrap().into(),
201        ]))
202        .unwrap()
203        .unwrap();
204        assert_eq!(merged_config, serde_yaml::from_str(EXPECTED_MERGED).unwrap());
205    }
206
207    #[fixture]
208    fn event_recorder() -> SkEventRecorder {
209        SkEventRecorder::mock()
210    }
211
212    #[rstest(tokio::test)]
213    async fn test_execute_hooks(test_sim: Simulation, event_recorder: SkEventRecorder) {
214        // Should print "foo"
215        let res = execute(&test_sim, Type::PreStart, &event_recorder).await;
216        assert!(res.is_ok());
217
218        // No PreStop hook defined
219        let res = execute(&test_sim, Type::PostStop, &event_recorder).await;
220        assert!(res.is_ok());
221
222        // PreRun hook calls bad command
223        let res = execute(&test_sim, Type::PreRun, &event_recorder).await;
224        assert!(res.is_err());
225    }
226}
sk_core/hooks.rs

sk_core/
hooks.rs
