KV store shared by all the flows

This is first step toward a context used by flow scripts to access shared data,
such as entity metadata or measurement units. For this first step,
the store is not passed as an argument to the onMessage and onInterval method.
The scripts has to create a `new FlowStore()` providing a `get()` method
to get the value attached to some key - typically the metadata attached to an MQTT topic.

The scripts only have a read access to the shared store
which is populated by regular flows which output is configured to be the context.

Signed-off-by: Didier Wenzek <[email protected]>

Author: didier-wenzek

crates/extensions/tedge_flows/src/actor.rs

@@ -409,6 +409,12 @@ impl FlowsMapper {
                     error!(target: "flows", "{flow}: cannot flush {path}: {err}");
                 }
             }
+            FlowOutput::Context => {
+                // Any context output reaching this point is an error
+                for message in messages {
+                    error!(target: "flows", "{flow}: cannot store context value: {}", message.payload_str().unwrap_or_default());
+                }
+            }
         }
         Ok(())
     }

crates/extensions/tedge_flows/src/config.rs

@@ -88,6 +88,9 @@ pub enum OutputConfig {
 
     #[serde(rename = "file")]
     File { path: Utf8PathBuf },
+
+    #[serde(rename = "context")]
+    Context,
 }
 
 #[derive(thiserror::Error, Debug)]
@@ -269,6 +272,7 @@ impl TryFrom<OutputConfig> for FlowOutput {
                 topic: topic.map(into_topic).transpose()?,
             },
             OutputConfig::File { path } => FlowOutput::File { path },
+            OutputConfig::Context => FlowOutput::Context,
         })
     }
 }

crates/extensions/tedge_flows/src/flow.rs

@@ -87,6 +87,7 @@ pub enum FlowInput {
 pub enum FlowOutput {
     Mqtt { topic: Option<Topic> },
     File { path: Utf8PathBuf },
+    Context,
 }
 
 /// The final outcome of a sequence of transformations applied by a flow to a message

crates/extensions/tedge_flows/src/js_lib/kv_store.rs

@@ -0,0 +1,75 @@
+use crate::js_value::JsonValue;
+use rquickjs::class::Trace;
+use rquickjs::Class;
+use rquickjs::Ctx;
+use rquickjs::JsLifetime;
+use rquickjs::Result;
+use std::collections::BTreeMap;
+use std::ops::Deref;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+#[derive(Clone, Default, JsLifetime)]
+pub struct KVStore {
+    data: Arc<Mutex<BTreeMap<String, JsonValue>>>,
+}
+
+impl KVStore {
+    pub fn init(&self, ctx: &Ctx<'_>) {
+        let globals = ctx.globals();
+        let _ = Class::<FlowStore>::define(&globals);
+        self.store_as_userdata(ctx)
+    }
+
+    pub fn get(&self, key: &str) -> JsonValue {
+        let data = self.data.lock().unwrap();
+        data.get(key).cloned().unwrap_or(JsonValue::Null)
+    }
+
+    pub fn insert(&self, key: impl Into<String>, value: impl Into<JsonValue>) {
+        match value.into() {
+            JsonValue::Null => self.remove(&key.into()),
+            value => {
+                let mut data = self.data.lock().unwrap();
+                data.insert(key.into(), value);
+            }
+        }
+    }
+
+    pub fn remove(&self, key: &str) {
+        let mut data = self.data.lock().unwrap();
+        data.remove(key);
+    }
+
+    fn store_as_userdata(&self, ctx: &Ctx<'_>) {
+        let _ = ctx.store_userdata(self.clone());
+    }
+
+    fn get_from_userdata(ctx: &Ctx<'_>) -> Self {
+        match ctx.userdata::<Self>() {
+            None => {
+                let store = KVStore::default();
+                store.store_as_userdata(ctx);
+                store
+            }
+            Some(userdata) => userdata.deref().clone(),
+        }
+    }
+}
+
+#[derive(Clone, Trace, JsLifetime)]
+#[rquickjs::class(frozen)]
+pub struct FlowStore {}
+
+#[rquickjs::methods]
+impl<'js> FlowStore {
+    #[qjs(constructor)]
+    fn new(_ctx: Ctx<'js>) -> Result<FlowStore> {
+        Ok(FlowStore {})
+    }
+
+    fn get(&self, ctx: Ctx<'js>, key: String) -> Result<JsonValue> {
+        let data = KVStore::get_from_userdata(&ctx);
+        Ok(data.get(&key))
+    }
+}

crates/extensions/tedge_flows/src/js_lib/mod.rs

@@ -1,3 +1,4 @@
 pub mod console;
+pub mod kv_store;
 pub mod text_decoder;
 pub mod text_encoder;

crates/extensions/tedge_flows/src/js_runtime.rs

@@ -1,4 +1,5 @@
 use crate::js_lib;
+use crate::js_lib::kv_store::KVStore;
 use crate::js_script::JsScript;
 use crate::js_value::JsonValue;
 use crate::LoadError;
@@ -17,6 +18,7 @@ use tracing::debug;
 
 pub struct JsRuntime {
     runtime: rquickjs::AsyncRuntime,
+    pub(crate) store: KVStore,
     worker: mpsc::Sender<JsRequest>,
     execution_timeout: Duration,
 }
@@ -35,10 +37,12 @@ impl JsRuntime {
             })))
             .await;
         let context = rquickjs::AsyncContext::full(&runtime).await?;
-        let worker = JsWorker::spawn(context).await;
+        let store = KVStore::default();
+        let worker = JsWorker::spawn(context, store.clone()).await;
         let execution_timeout = Duration::from_secs(5);
         Ok(JsRuntime {
             runtime,
+            store,
             worker,
             execution_timeout,
         })
@@ -164,20 +168,21 @@ struct JsWorker {
 }
 
 impl JsWorker {
-    pub async fn spawn(context: rquickjs::AsyncContext) -> mpsc::Sender<JsRequest> {
+    pub async fn spawn(context: rquickjs::AsyncContext, store: KVStore) -> mpsc::Sender<JsRequest> {
         let (sender, requests) = mpsc::channel(100);
         tokio::spawn(async move {
             let worker = JsWorker { context, requests };
-            worker.run().await
+            worker.run(store).await
         });
         sender
     }
 
-    async fn run(mut self) {
+    async fn run(mut self, store: KVStore) {
         rquickjs::async_with!(self.context => |ctx| {
             js_lib::console::init(&ctx);
             js_lib::text_decoder::init(&ctx);
             js_lib::text_encoder::init(&ctx);
+            store.init(&ctx);
             let mut modules = JsModules::new();
             while let Some(request) = self.requests.recv().await {
                 match request {

crates/extensions/tedge_flows/src/js_script.rs

@@ -555,6 +555,38 @@ export async function onMessage(message, config) {
         );
     }
 
+    #[tokio::test]
+    async fn using_the_context() {
+        let js = r#"
+export function onMessage(message) {
+    let data = new FlowStore()
+    return {
+        topic: message.topic,
+        payload: JSON.stringify(data.get(message.topic))
+    }
+}
+        "#;
+        let (runtime, script) = runtime_with(js).await;
+
+        runtime.store.insert(
+            "foo/bar",
+            serde_json::json!({
+                "hello": "world",
+                "guess": 42,
+            }),
+        );
+
+        let input = Message::new("foo/bar", "");
+        let output = Message::new("foo/bar", r#"{"guess":42,"hello":"world"}"#);
+        assert_eq!(
+            script
+                .on_message(&runtime, SystemTime::now(), &input)
+                .await
+                .unwrap(),
+            vec![output]
+        );
+    }
+
     async fn runtime_with(js: &str) -> (JsRuntime, JsScript) {
         let mut runtime = JsRuntime::try_new().await.unwrap();
         let mut script = JsScript::new("toml".into(), 1, "js".into());

crates/extensions/tedge_flows/src/runtime.rs

@@ -5,6 +5,8 @@ use crate::js_runtime::JsRuntime;
 use crate::registry::BaseFlowRegistry;
 use crate::registry::FlowRegistryExt;
 use crate::stats::Counter;
+use crate::FlowError;
+use crate::FlowOutput;
 use crate::LoadError;
 use camino::Utf8Path;
 use camino::Utf8PathBuf;
@@ -121,6 +123,9 @@ impl<Registry: FlowRegistryExt + Send> MessageProcessor<Registry> {
 
         self.stats.runtime_on_message_done(started_at);
         out_messages
+            .into_iter()
+            .filter_map(|flow_output| self.store_context_values(flow_output))
+            .collect()
     }
 
     pub async fn on_interval(&mut self, timestamp: SystemTime, now: Instant) -> Vec<FlowResult> {
@@ -135,6 +140,45 @@ impl<Registry: FlowRegistryExt + Send> MessageProcessor<Registry> {
         out_messages
     }
 
+    fn store_context_values(&mut self, messages: FlowResult) -> Option<FlowResult> {
+        match messages {
+            FlowResult::Ok {
+                messages,
+                output: FlowOutput::Context,
+                flow,
+            } => {
+                for message in messages {
+                    if let Err(error) = self.store_context_value(&message) {
+                        return Some(FlowResult::Err {
+                            flow,
+                            error,
+                            output: FlowOutput::Context,
+                        });
+                    }
+                }
+                None
+            }
+            messages => Some(messages),
+        }
+    }
+
+    pub fn store_context_value(&mut self, message: &Message) -> Result<(), FlowError> {
+        if message.payload.is_empty() {
+            self.js_runtime.store.remove(&message.topic)
+        } else {
+            let payload = message.payload_str().ok_or(FlowError::UnsupportedMessage(
+                "Non UFT8 payload".to_string(),
+            ))?;
+            let value: serde_json::Value = serde_json::from_str(payload)
+                .map_err(|err| FlowError::UnsupportedMessage(format!("Non JSON payload: {err}")))?;
+            self.js_runtime
+                .store
+                .insert(message.topic.to_owned(), value);
+        }
+
+        Ok(())
+    }
+
     pub async fn dump_processing_stats(&self) {
         self.stats.dump_processing_stats();
     }

tests/RobotFramework/tests/tedge_flows/flows/load_metadata.toml

@@ -0,0 +1,7 @@
+#
+# Update flows context with measurement metadata
+#
+[input.mqtt]
+topics = ["te/+/+/+/+/m/+/meta"]
+
+[output.context]

tests/RobotFramework/tests/tedge_flows/flows/measurements.toml

@@ -2,5 +2,5 @@ input.mqtt.topics = ["te/+/+/+/+/m/+"]
 
 steps = [
     { script = "add_timestamp.js" },
-    { script = "te_to_c8y.js", meta_topics = ["te/+/+/+/+/m/+/meta"] }
+    { script = "te_to_c8y.js" }
 ]