Adds extra key-value string data storage to runtime objects

Allows external processes to store data in durable runtimes.

Signed-off-by: David Gilligan-Cook <dcook@imageworks.com>

crates/spfs-cli/main/Cargo.toml

@@ -27,6 +27,7 @@ nix = { workspace = true }
 number_prefix = "*"                                                           # we hope to match versions with indicatif
 relative-path = "1.3"
 serde_json = { workspace = true }
+serde_yaml = "0.8.17"
 spfs = { path = "../../spfs" }
 spfs-cli-common = { path = "../common" }
 strum = { workspace = true, features = ["derive"] }
@@ -47,3 +48,7 @@ features = [
     "Win32_System_SystemInformation",
     "Win32_System_Threading",
 ]
+
+[dev-dependencies]
+rstest = { workspace = true }
+tempfile = { workspace = true }

crates/spfs-cli/main/src/cmd_run.rs

@@ -2,16 +2,103 @@
 // SPDX-License-Identifier: Apache-2.0
 // https://github.com/imageworks/spk
 
+use std::collections::BTreeMap;
 use std::ffi::OsString;
+use std::io;
 use std::time::Instant;
 
-use anyhow::{Context, Result};
+use anyhow::{anyhow, Context, Result};
 use clap::{ArgGroup, Args};
 use spfs::runtime::BindMount;
 use spfs::storage::FromConfig;
 use spfs::tracking::{EnvSpec, EnvSpecItem};
 use spfs_cli_common as cli;
 
+#[cfg(test)]
+#[path = "./cmd_run_test.rs"]
+mod cmd_run_test;
+#[cfg(test)]
+#[path = "./fixtures.rs"]
+mod fixtures;
+
+#[derive(Args, Clone, Debug)]
+pub struct ExtraData {
+    /// Adds extra key-value string data to the new runtime.
+    ///
+    /// This allows external processes to store arbitrary data in the
+    /// runtimes they create. This is most useful with durable runtimes.
+    /// The data can be retrieved by parsing the output of `spfs runtime
+    /// info`
+    ///
+    /// Extra data is key/value string pairs separated by either an
+    /// equals sign or colon (--extra-data name=value --extra-data
+    /// other:value).  Additionally, many pair of extra data can be
+    /// specified at once in yaml or json format (--extra-data '{name:
+    /// value, other: value}').
+    ///
+    /// Extra data can also be given in a json or yaml file via the
+    /// --extra-data--file flag. If given, --extra-data will supersede
+    /// anything in the options file(s).
+    #[clap(long, value_name = "KEY:VALUE")]
+    pub extra_data: Vec<String>,
+
+    /// Specify extra key-value data from a json or yaml file (see --extra-data)
+    #[clap(long)]
+    pub extra_data_file: Vec<std::path::PathBuf>,
+}
+
+type KeyValuePair = (String, String);
+
+impl ExtraData {
+    // Returns a list of extra data key-value pairs gathered from all
+    // the extra data related command line arguments. The same keys,
+    // and values, can appear multiple times in the list if specified
+    // multiple times in various command line arguments.
+    pub fn get_data(&self) -> Result<Vec<KeyValuePair>> {
+        let mut data: Vec<KeyValuePair> = Vec::new();
+
+        for filename in self.extra_data_file.iter() {
+            let reader: Box<dyn io::Read> = if Ok("-".to_string())
+                == filename.clone().into_os_string().into_string()
+            {
+                // Treat '-' as "read from stdin"
+                Box::new(io::stdin())
+            } else {
+                Box::new(
+                    std::fs::File::open(filename)
+                        .with_context(|| format!("Failed to open extra data file: {filename:?}"))?,
+                )
+            };
+            let extra_data: BTreeMap<String, String> = serde_yaml::from_reader(reader)
+                .with_context(|| {
+                    format!("Failed to parse as extra data key-value pairs: {filename:?}")
+                })?;
+            data.extend(extra_data);
+        }
+
+        for pair in self.extra_data.iter() {
+            let pair = pair.trim();
+            if pair.starts_with('{') {
+                let given: BTreeMap<String, String> = serde_yaml::from_str(pair)
+                    .context("--extra-data value looked like yaml, but could not be parsed")?;
+                data.extend(given);
+                continue;
+            }
+
+            let (name, value) = pair
+                .split_once('=')
+                .or_else(|| pair.split_once(':'))
+                .ok_or_else(|| {
+                    anyhow!("Invalid option: -extra-data {pair} (should be in the form name=value)")
+                })?;
+
+            data.push((name.to_string(), value.to_string()));
+        }
+
+        Ok(data)
+    }
+}
+
 /// Run a program in a configured spfs environment
 #[derive(Debug, Args)]
 #[clap(group(
@@ -57,9 +144,13 @@ pub struct CmdRun {
     /// '/spfs/MOUNTPOINT' under /spfs. Mounts are specified as
     /// key/value pairs separated by either an equals sign or colon
     /// (--extra-mount path=mountpoint --extra-mount path:mountpoint).
+
     #[clap(long, value_name = "SRC_PATH:MOUNTPOINT_UNDER_SPFS")]
     pub extra_mount: Option<Vec<BindMount>>,
 
+    #[clap(flatten)]
+    pub external_data: ExtraData,
+
     /// The tag or id of the desired runtime
     ///
     /// Use '-' to request an empty environment
@@ -159,6 +250,14 @@ impl CmdRun {
                 tracing::debug!("with extra mounts: {extras:?}");
             }
 
+            let data = self.external_data.get_data()?;
+            if !data.is_empty() {
+                tracing::debug!("with extra data: {data:?}");
+                for (key, value) in data {
+                    runtime.insert_extra_data(key, value);
+                }
+            }
+
             let start_time = Instant::now();
             runtime.config.mount_backend = config.filesystem.backend;
             runtime.config.secondary_repositories = config.get_secondary_runtime_repositories();

crates/spfs-cli/main/src/cmd_run_test.rs

@@ -0,0 +1,196 @@
+// Copyright (c) Sony Pictures Imageworks, et al.
+// SPDX-License-Identifier: Apache-2.0
+// https://github.com/imageworks/spk
+
+use std::fs;
+
+use rstest::rstest;
+
+use super::fixtures::*;
+use super::ExtraData;
+
+// Test - --extra-data field1:value1 --extra-data field2:value2
+// Test - --extra-data field1=value1 --extra-data field2=value2
+// Test - --extra-data field1:value1 --extra-data field2=value2
+// Test - --extra-data field1=value1 --extra-data field2:value2
+// Test - --extra-data '{ field1: value1, field2: value2 }'
+#[rstest]
+#[case(vec!["field1:value1".to_string(), "field2:value2".to_string()])]
+#[case(vec!["field1=value1".to_string(), "field2=value2".to_string()])]
+#[case(vec!["field1:value1".to_string(), "field2=value2".to_string()])]
+#[case(vec!["field1=value1".to_string(), "field2:value2".to_string()])]
+#[case(vec!["{field1: value1, field2: value2}".to_string()])]
+fn test_cmd_run_create_extra_data(#[case] values: Vec<String>) {
+    // Setup some data for the key value pairs
+    let field1 = "field1".to_string();
+    let field2 = "field2".to_string();
+    let value1 = "value1".to_string();
+    let value2 = "value2".to_string();
+
+    let filenames = Vec::new();
+
+    // Test - --extra-data field1:value1 --extra-data field2:value2
+    let data = ExtraData {
+        extra_data: values,
+        extra_data_file: filenames,
+    };
+
+    let result = data
+        .get_data()
+        .expect("unable to parse cmd_run --extra-data values");
+    assert!(!result.is_empty());
+    assert!(result.len() == 2);
+    assert!(result[0] == (field1, value1));
+    assert!(result[1] == (field2, value2));
+}
+
+#[rstest]
+#[should_panic]
+#[case(vec!["field1,value1".to_string(), "field2/value2".to_string()])]
+#[should_panic]
+#[case(vec!["field1 value1".to_string(), "field2 value2".to_string()])]
+#[should_panic]
+#[case(vec!["{field1: value1, field2:    - value2    - value1}".to_string()])]
+fn test_cmd_run_create_extra_data_invalid(#[case] invalid_values: Vec<String>) {
+    let filenames = Vec::new();
+
+    // Test - --extra-data with_an_invalid_argument_value
+    let data = ExtraData {
+        extra_data: invalid_values,
+        extra_data_file: filenames,
+    };
+
+    let _result = data
+        .get_data()
+        .expect("unable to parse cmd_run --extra-data values");
+}
+
+#[rstest]
+fn test_cmd_run_create_extra_data_from_file(tmpdir: tempfile::TempDir) {
+    // Setup some data for the key value pairs
+    let field1 = "field1".to_string();
+    let field2 = "field2".to_string();
+    let value1 = "value1".to_string();
+    let value2 = "value2".to_string();
+    let extra_data = format!("{field1}: {value1}\n{field2}: {value2}\n");
+
+    let filename = tmpdir.path().join("filename.yaml");
+    fs::write(filename.clone(), extra_data)
+        .expect("Unable to write extra data to file during setup");
+    let filenames = vec![filename];
+
+    let values = Vec::new();
+
+    // Test - --extra-data-file filename.yaml
+    let data = ExtraData {
+        extra_data: values,
+        extra_data_file: filenames,
+    };
+
+    let result = data
+        .get_data()
+        .expect("unable to parse cmd_run --extra-data values");
+    assert!(!result.is_empty());
+    assert!(result.len() == 2);
+    assert!(result[0] == (field1, value1));
+    assert!(result[1] == (field2, value2));
+}
+
+#[rstest]
+#[should_panic]
+fn test_cmd_run_create_extra_data_from_file_not_exist(tmpdir: tempfile::TempDir) {
+    // Setup a file name that does not exist
+    let filename = tmpdir.path().join("nosuchfile.yaml");
+    let filenames = vec![filename];
+
+    let values = Vec::new();
+
+    // Test - --extra-data-file nosuchfile.yaml
+    let data = ExtraData {
+        extra_data: values,
+        extra_data_file: filenames,
+    };
+
+    let _result = data
+        .get_data()
+        .expect("unable to parse cmd_run --extra-data values");
+}
+
+#[rstest]
+#[should_panic]
+fn test_cmd_run_create_extra_data_from_file_invalid_keyvalues(tmpdir: tempfile::TempDir) {
+    // Setup some data for the key value pairs
+    let field1 = "field1".to_string();
+    let field2 = "field2".to_string();
+    let value1 = "value1".to_string();
+    let value2 = "value2".to_string();
+    let extra_data = format!("{field1}: {value1}\n{field2}:\n    - {value2}\n    - {value1}\n");
+
+    let filename = tmpdir.path().join("filename.yaml");
+    fs::write(filename.clone(), extra_data)
+        .expect("Unable to write extra data to file during setup");
+    let filenames = vec![filename];
+
+    let values = Vec::new();
+
+    // Test - --extra-data-file filename.yaml that contains more than key-value string pairs
+    let data = ExtraData {
+        extra_data: values,
+        extra_data_file: filenames,
+    };
+
+    let _result = data
+        .get_data()
+        .expect("unable to parse cmd_run --extra-data values");
+}
+
+#[rstest]
+fn test_cmd_run_create_extra_data_all(tmpdir: tempfile::TempDir) {
+    // Setup some data for the key value pairs
+    let field1 = "field1".to_string();
+    let field2 = "field2".to_string();
+    let field3 = "field3".to_string();
+    let field4 = "field4".to_string();
+
+    let value1 = "value1".to_string();
+    let value2 = "value2".to_string();
+    let value3 = "value3".to_string();
+    let value4 = "value4".to_string();
+
+    let values: Vec<String> = vec![
+        format!("{field1}:{value1}"),
+        format!("{field2}={value2}"),
+        "{field3: value3, field4: value4}".to_string(),
+    ];
+
+    let extra_data = format!("{field4}: {value4}\n{field3}: {value3}\n");
+
+    let filename = tmpdir.path().join("filename.yaml");
+    fs::write(filename.clone(), extra_data)
+        .expect("Unable to write extra data to file during setup");
+    let filenames = vec![filename];
+
+    // Test -  --extra-data field1:value1 --extra-data field2=value2
+    //         --extra-data '{ field1: value1, field2: value2 }'
+    //         --extra-data-file filename.yaml
+    let data = ExtraData {
+        extra_data: values,
+        extra_data_file: filenames,
+    };
+
+    let result = data
+        .get_data()
+        .expect("unable to parse cmd_run --extra-data values");
+
+    assert!(!result.is_empty());
+    assert!(result.len() == 6);
+    // from --extra-data-file filename.taml
+    assert!(result[0] == (field3.clone(), value3.clone()));
+    assert!(result[1] == (field4.clone(), value4.clone()));
+    // from --extra-data field1:value1 extra-data field2:value2
+    assert!(result[2] == (field1, value1));
+    assert!(result[3] == (field2, value2));
+    // from --extra-data '{field3: value3, field4: value4}'
+    assert!(result[4] == (field3, value3));
+    assert!(result[5] == (field4, value4));
+}

crates/spfs-cli/main/src/cmd_shell.rs

@@ -8,6 +8,7 @@ use spfs::runtime::BindMount;
 use spfs_cli_common as cli;
 
 use super::cmd_run;
+use super::cmd_run::ExtraData;
 
 /// Enter a subshell in a configured spfs environment
 #[derive(Debug, Args)]
@@ -54,6 +55,9 @@ pub struct CmdShell {
     #[clap(long, value_name = "SRCPATH:SPFSLOCATION")]
     pub extra_mount: Option<Vec<BindMount>>,
 
+    #[clap(flatten)]
+    pub external_data: ExtraData,
+
     /// The tag or id of the desired runtime
     ///
     /// Use '-' or nothing to request an empty environment
@@ -74,6 +78,7 @@ impl CmdShell {
             reference: self.reference.clone(),
             keep_runtime: self.keep_runtime,
             extra_mount: self.extra_mount.clone(),
+            external_data: self.external_data.clone(),
             command: Default::default(),
         };
         run_cmd.run(config).await

crates/spfs-cli/main/src/fixtures.rs

@@ -0,0 +1,10 @@
+use rstest::fixture;
+use tempfile::TempDir;
+
+#[fixture]
+pub fn tmpdir() -> TempDir {
+    tempfile::Builder::new()
+        .prefix("spfs-test-")
+        .tempdir()
+        .expect("failed to create dir for test")
+}