grass_driver/
job.rs

1use std::{
2    collections::HashMap,
3    fs::File,
4    io::Result,
5    io::{Cursor, Error, Write},
6    path::{Path, PathBuf},
7    process::{Child, Command},
8};
9
10use grass_ir::GrassIR;
11use serde::Deserialize;
12use sha::{sha256::Sha256, utils::DigestExt};
13use tempfile::TempDir;
14
15use crate::{
16    cache::CacheState,
17    dependency::{Dependency, DependencySource},
18};
19
20use crate::return_true;
21
22#[derive(Deserialize)]
23pub enum BuildFlavor {
24    Debug,
25    Release,
26    ReleaseWithDebugInfo,
27}
28
29impl Default for BuildFlavor {
30    fn default() -> Self {
31        Self::Release
32    }
33}
34
35#[derive(Deserialize)]
36#[allow(unused)]
37pub struct JobDefinition {
38    ir: Vec<GrassIR>,
39
40    // ################### Runtime confiruation ################
41    working_dir: PathBuf,
42    #[serde(default)]
43    cmdline_args: Vec<String>,
44    #[serde(default)]
45    env_vars: HashMap<String, String>,
46    #[serde(default)]
47    const_bag_types: Vec<String>,
48
49    // ############# Runtime Configuration ######################
50    #[serde(default = "default_runtime")]
51    runtime_package: String,
52    #[serde(default = "default_macro")]
53    macro_package: String,
54
55    runtime_source: DependencySource,
56    #[serde(default = "latest_grass_runtime")]
57    runtime_version: String,
58
59    macro_source: DependencySource,
60    #[serde(default = "latest_grass_runtime")]
61    macro_version: String,
62
63    // ################### Package Information ####################
64    #[serde(default)]
65    deps: Vec<Dependency>,
66    #[serde(default = "default_version")]
67    version: String,
68    #[serde(default)]
69    build_flavor: BuildFlavor,
70    #[serde(default)]
71    temp_dir: Option<PathBuf>,
72
73    // ################### Cache Control #########################
74    #[serde(default = "return_true")]
75    use_cache: bool,
76    #[serde(default = "return_true")]
77    update_cache: bool,
78    #[serde(default = "default_cache_root")]
79    cache_root: PathBuf,
80
81    // ################ Runtime values ####################
82    #[serde(default)]
83    tool_chain_path: Option<PathBuf>,
84    #[serde(skip)]
85    compilation_dir: Option<TempDir>,
86    #[serde(skip)]
87    ir_hash: String,
88    #[serde(skip)]
89    artifact_path: Option<PathBuf>,
90}
91
92fn default_version() -> String {
93    "0.0.1".to_string()
94}
95
96fn default_runtime() -> String {
97    "grass-runtime".to_string()
98}
99
100fn default_macro() -> String {
101    "grass-macro".to_string()
102}
103
104fn latest_grass_runtime() -> String {
105    "*".to_string()
106}
107
108fn default_cache_root() -> PathBuf {
109    let mut ret: PathBuf = std::env::var("HOME")
110        .unwrap_or_else(|_| "/".to_string())
111        .into();
112    ret.push(".grass-ql");
113    ret.push("cache");
114    ret
115}
116
117impl JobDefinition {
118    pub fn get_stderr_log(&mut self) -> Result<File> {
119        File::open(self.get_compilation_dir()?.join("stderr.log"))
120    }
121    pub fn cargo(&mut self, args: &[&str], capture_stdout: bool) -> Result<Child> {
122        let mut cmd = Command::new("cargo");
123
124        if let Some(tool_chain_path) = self.tool_chain_path.as_ref() {
125            cmd.env("PATH", tool_chain_path);
126        }
127
128        cmd.args(args);
129
130        let comp_dir = self.get_compilation_dir()?;
131        cmd.current_dir(comp_dir);
132
133        let err_log = File::create(comp_dir.join("stderr.log"))?;
134        cmd.stderr(err_log);
135
136        if capture_stdout {
137            let out_log = File::create(comp_dir.join("stdout.txt"))?;
138            cmd.stdout(out_log);
139        }
140
141        cmd.spawn()
142    }
143    fn populate_ir_hash(&mut self) -> Result<()> {
144        let mut buffer = Vec::new();
145        {
146            let mut buffer_writer = Cursor::new(&mut buffer);
147            serde_json::to_writer(&mut buffer_writer, &self.ir)?;
148
149            let mut deps = self.deps.clone();
150            deps.sort_by_key(|x| x.name.clone());
151
152            serde_json::to_writer(&mut buffer_writer, &deps)?;
153
154            serde_json::to_writer(&mut buffer_writer, &self.macro_package)?;
155            serde_json::to_writer(&mut buffer_writer, &self.macro_version)?;
156            serde_json::to_writer(&mut buffer_writer, &self.macro_source)?;
157
158            serde_json::to_writer(&mut buffer_writer, &self.runtime_package)?;
159            serde_json::to_writer(&mut buffer_writer, &self.runtime_version)?;
160            serde_json::to_writer(&mut buffer_writer, &self.runtime_source)?;
161        }
162        let input_str = String::from_utf8(buffer).unwrap();
163        log::debug!("Job hasher input: {}", input_str);
164        let mut hasher = Sha256::default();
165
166        hasher.write_all(input_str.as_bytes())?;
167
168        hasher.flush()?;
169
170        let ir_hash = hasher.to_hex();
171
172        self.ir_hash = ir_hash;
173        Ok(())
174    }
175    fn get_artifact_hash(&mut self) -> Result<&str> {
176        if self.ir_hash.len() == 0 {
177            self.populate_ir_hash()?;
178        }
179        Ok(self.ir_hash.as_str())
180    }
181    fn get_artifact_name(&mut self) -> Result<String> {
182        Ok(format!("grass-artifact-{}", self.get_artifact_hash()?))
183    }
184    fn get_artifact_path(&mut self) -> Result<PathBuf> {
185        let mut ret = self.get_compilation_dir()?.to_path_buf();
186        ret.push("target");
187        match self.build_flavor {
188            BuildFlavor::Debug => ret.push("debug"),
189            BuildFlavor::Release | BuildFlavor::ReleaseWithDebugInfo => ret.push("release"),
190        }
191        ret.push(self.get_artifact_name()?);
192        Ok(ret)
193    }
194    fn write_manifest_file(&mut self, root: &Path) -> Result<()> {
195        let manifest_path = root.join("Cargo.toml");
196        let mut manifest_file = File::create(&manifest_path)?;
197        writeln!(&mut manifest_file, "[package]")?;
198        writeln!(
199            &mut manifest_file,
200            "name = \"{}\"",
201            self.get_artifact_name()?
202        )?;
203        writeln!(&mut manifest_file, "version = \"{}\"", self.version)?;
204        writeln!(&mut manifest_file, "edition = \"2021\"")?;
205        writeln!(&mut manifest_file, "[dependencies]")?;
206        for dep in self.deps.iter().chain([
207            &Dependency::create_grass_dep(
208                &self.runtime_package,
209                &self.runtime_source,
210                &self.runtime_version,
211            ),
212            &Dependency::create_grass_dep(
213                &self.macro_package,
214                &self.macro_source,
215                &self.macro_version,
216            ),
217        ]) {
218            dep.write_dependency_line(&mut manifest_file)?;
219        }
220        match self.build_flavor {
221            BuildFlavor::Release => {
222                writeln!(&mut manifest_file, "[profile.release]")?;
223                writeln!(&mut manifest_file, "strip = true")?;
224            }
225            BuildFlavor::ReleaseWithDebugInfo => {
226                writeln!(&mut manifest_file, "[profile.release]")?;
227                writeln!(&mut manifest_file, "debug = true")?;
228            }
229            _ => (),
230        }
231        Ok(())
232    }
233    fn write_source_code(&self, root: &Path) -> Result<()> {
234        let source_dir = root.join("src");
235        std::fs::create_dir(&source_dir)?;
236
237        let source_path = source_dir.join("main.rs");
238        let mut source_file = File::create(source_path)?;
239
240        writeln!(&mut source_file, "#[allow(unused_imports)]")?;
241        writeln!(
242            &mut source_file,
243            "use grass_runtime::const_bag::{{ConstBagRef, ConstBagType}};"
244        )?;
245
246        for (id, t) in self.const_bag_types.iter().enumerate() {
247            let ty = match t.as_str() {
248                "str" => "String",
249                "i64" => "i64",
250                "f64" => "f64",
251                _ => {
252                    panic!("Unsupported const bag type: {}", t);
253                }
254            };
255            writeln!(
256                &mut source_file,
257                "const __CONST_BAG_VALUE_{id} : ConstBagRef<{ty}> = ConstBagRef::<{ty}>::new({id});",
258                id = id,
259                ty = ty,
260            )?;
261        }
262
263        for (id, ir) in self.ir.iter().enumerate() {
264            let ir_path = source_dir.as_path().join(format!("grass_ir_{}.json", id));
265            let ir_file = File::create(&ir_path)?;
266            serde_json::to_writer(ir_file, ir)?;
267
268            writeln!(
269                &mut source_file,
270                "fn grass_query_{id}(cmd_args: &[&str]) -> Result<(), Box<dyn std::error::Error>> {{",
271                id = id
272            )?;
273            writeln!(
274                &mut source_file,
275                "    grass_macro::import_grass_ir_from_file!(\"{ir_file}\");",
276                ir_file = ir_path.as_os_str().to_string_lossy()
277            )?;
278            writeln!(&mut source_file, "    Ok(())")?;
279            writeln!(&mut source_file, "}}")?;
280        }
281
282        writeln!(
283            &mut source_file,
284            "fn main() -> Result<(), Box<dyn std::error::Error>> {{"
285        )?;
286
287        writeln!(
288            &mut source_file,
289            "    let owned_cmd_args: Vec<_> = std::env::args().collect();"
290        )?;
291
292        writeln!(
293            &mut source_file,
294            "    let cmd_args: Vec<_> = owned_cmd_args.iter().map(|a| a.as_str()).collect();"
295        )?;
296
297        for id in 0..self.ir.len() {
298            writeln!(
299                &mut source_file,
300                "    grass_query_{id}(&cmd_args)?;",
301                id = id
302            )?;
303        }
304        writeln!(&mut source_file, "    Ok(())")?;
305        writeln!(&mut source_file, "}}")?;
306
307        Ok(())
308    }
309    pub fn get_compilation_dir(&mut self) -> Result<&Path> {
310        if self.compilation_dir.is_some() {
311            Ok(self.compilation_dir.as_ref().unwrap().as_ref())
312        } else {
313            let mut root_dir = tempfile::Builder::new();
314
315            root_dir.prefix("grass-workspace-");
316            root_dir.rand_bytes(5);
317            let compilation_dir = if let Some(temp_dir) = &self.temp_dir {
318                root_dir.tempdir_in(temp_dir.as_path())?
319            } else {
320                root_dir.tempdir()?
321            };
322
323            self.write_manifest_file(compilation_dir.path())?;
324            self.write_source_code(compilation_dir.path())?;
325
326            self.compilation_dir = Some(compilation_dir);
327            Ok(self.compilation_dir.as_ref().unwrap().path())
328        }
329    }
330    pub fn build_artifact(&mut self) -> Result<&Path> {
331        log::info!("Building artifact {}", self.get_artifact_name()?);
332        let mut child = match self.build_flavor {
333            BuildFlavor::Debug => self.cargo(&["build"], true),
334            BuildFlavor::Release | BuildFlavor::ReleaseWithDebugInfo => {
335                self.cargo(&["build", "--release"], true)
336            }
337        }?;
338        let status = child.wait()?;
339        if status.success() {
340            self.artifact_path = Some(self.get_artifact_path()?);
341            return self.get_artifact();
342        }
343
344        Err(Error::new(
345            std::io::ErrorKind::Other,
346            format!(
347                "Cargo exited with error code {}",
348                status.code().unwrap_or(0)
349            ),
350        ))
351    }
352    pub fn get_artifact(&mut self) -> Result<&Path> {
353        if self.artifact_path.is_some() {
354            return Ok(self.artifact_path.as_ref().unwrap());
355        }
356
357        let mut cache = if self.use_cache || self.update_cache {
358            Some(CacheState::load_cache(&self.cache_root)?)
359        } else {
360            None
361        };
362
363        if self.use_cache {
364            log::info!(
365                "Checking binary cache for binary {}",
366                self.get_artifact_hash()?
367            );
368            let mut cached_artifact = PathBuf::new();
369            if cache
370                .as_mut()
371                .unwrap()
372                .query_cache_entry(self.get_artifact_hash()?, &mut cached_artifact)?
373            {
374                log::info!(
375                    "Found cached binary at {}",
376                    cached_artifact.to_str().unwrap()
377                );
378                self.artifact_path = Some(cached_artifact);
379                return self.get_artifact();
380            }
381        }
382
383        if self.update_cache {
384            let hash = self.get_artifact_hash()?.to_string();
385            let mut cached_path = PathBuf::new();
386            cache.as_mut().unwrap().update_cache(
387                &hash,
388                |buf| {
389                    let path = self.build_artifact()?;
390                    *buf = path.to_path_buf();
391                    Ok(())
392                },
393                &mut cached_path,
394            )?;
395            self.artifact_path = Some(cached_path);
396            return self.get_artifact();
397        }
398        self.build_artifact()
399    }
400
401    pub fn execute_artifact(&mut self) -> Result<Child> {
402        let working_dir = self.working_dir.clone();
403        let environment = self.env_vars.clone();
404        let artifact_path = self.get_artifact()?;
405        log::info!(
406            "Launching artifact {}",
407            artifact_path.as_os_str().to_string_lossy()
408        );
409        log::info!("Working directory = {:?}", working_dir);
410        log::info!("Environment vars: {:?}", environment);
411        Ok(Command::new(artifact_path)
412            .current_dir(&self.working_dir)
413            .envs(&self.env_vars)
414            .args(&self.cmdline_args)
415            .spawn()?)
416    }
417    pub fn print_expanded_code(&mut self) -> Result<()> {
418        self.cargo(&["expand"], false)?.wait()?;
419        Ok(())
420    }
421}
grass_driver/job.rs

grass_driver/
job.rs
