sk_core/
hooks.rs

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

sk_core/
hooks.rs
