Move db draining logic out of actor in tedge-flows

Author: jarhodes314

crates/extensions/tedge_flows/src/actor.rs

@@ -44,9 +44,7 @@ impl Actor for FlowsMapper {
 
             tokio::select! {
                 _ = deadline_future => {
-                    let drained_messages = self.drain_db().await?;
-                    self.on_messages(MessageSource::MeaDB, drained_messages).await?;
-
+                    self.poll_ready_sources().await?;
                     self.on_interval().await?;
                 }
                 message = self.messages.recv() => {
@@ -132,17 +130,6 @@ impl FlowsMapper {
         Ok(())
     }
 
-    async fn on_messages(
-        &mut self,
-        source: MessageSource,
-        messages: Vec<(DateTime, Message)>,
-    ) -> Result<(), RuntimeError> {
-        for (timestamp, message) in messages {
-            self.on_message(source, timestamp, message).await?
-        }
-        Ok(())
-    }
-
     async fn on_interval(&mut self) -> Result<(), RuntimeError> {
         let now = Instant::now();
         let timestamp = DateTime::now();
@@ -196,6 +183,8 @@ impl FlowsMapper {
                     if let Err(err) = self
                         .processor
                         .database
+                        .lock()
+                        .await
                         .store_many(output_series, messages)
                         .await
                     {
@@ -207,22 +196,47 @@ impl FlowsMapper {
         Ok(())
     }
 
-    async fn drain_db(&mut self) -> Result<Vec<(DateTime, Message)>, RuntimeError> {
+    async fn poll_ready_sources(&mut self) -> Result<(), RuntimeError> {
         let timestamp = DateTime::now();
-        let mut messages = vec![];
-        for (flow_id, flow_messages) in self.processor.drain_db(timestamp).await {
-            match flow_messages {
-                Ok(flow_messages) => {
-                    for (t, m) in flow_messages.iter() {
-                        info!(target: "flows", "drained: @{}.{} [{}]", t.seconds, t.nanoseconds, m.topic);
+
+        // Collect flow IDs with ready sources
+        let ready_flows: Vec<String> = self
+            .processor
+            .flows
+            .iter()
+            .filter_map(|(flow_id, flow)| {
+                flow.input_source
+                    .as_ref()
+                    .filter(|source| source.is_ready(timestamp))
+                    .map(|_| flow_id.clone())
+            })
+            .collect();
+
+        // Poll each ready source and process messages
+        for flow_id in ready_flows {
+            if let Some(flow) = self.processor.flows.get_mut(&flow_id) {
+                if let Some(source) = &mut flow.input_source {
+                    match source.poll(timestamp).await {
+                        Ok(messages) => {
+                            for (t, m) in messages.iter() {
+                                info!(target: "flows", "drained: @{}.{} [{}]", t.seconds, t.nanoseconds, m.topic);
+                            }
+                            source.update_after_poll(timestamp);
+
+                            // Process the messages through the flow
+                            for (msg_timestamp, message) in messages {
+                                self.on_message(MessageSource::MeaDB, msg_timestamp, message)
+                                    .await?;
+                            }
+                        }
+                        Err(err) => {
+                            error!(target: "flows", "{flow_id}: Failed to poll source: {err}");
+                        }
                     }
-                    messages.extend(flow_messages);
-                }
-                Err(err) => {
-                    error!(target: "flows", "{flow_id}: {err}");
                 }
             }
         }
-        Ok(messages)
+
+        Ok(())
     }
 }

crates/extensions/tedge_flows/src/config.rs

@@ -1,7 +1,9 @@
+use crate::database::MeaDb;
 use crate::flow::Flow;
 use crate::flow::FlowInput;
 use crate::flow::FlowOutput;
 use crate::flow::FlowStep;
+use crate::input_source::DatabaseSource;
 use crate::js_runtime::JsRuntime;
 use crate::js_script::JsScript;
 use crate::LoadError;
@@ -10,8 +12,10 @@ use camino::Utf8PathBuf;
 use serde::Deserialize;
 use serde_json::Value;
 use std::fmt::Debug;
+use std::sync::Arc;
 use std::time::Duration;
 use tedge_mqtt_ext::TopicFilter;
+use tokio::sync::Mutex;
 
 #[derive(Deserialize)]
 pub struct FlowConfig {
@@ -110,8 +114,25 @@ impl FlowConfig {
         js_runtime: &mut JsRuntime,
         config_dir: &Utf8Path,
         source: Utf8PathBuf,
+        database: Arc<Mutex<Box<dyn MeaDb>>>,
     ) -> Result<Flow, ConfigError> {
         let input = self.input.try_into()?;
+
+        // Create input source for MeaDB inputs
+        let input_source = match &input {
+            FlowInput::MeaDB {
+                series,
+                frequency,
+                max_age,
+            } => Some(Box::new(DatabaseSource::new(
+                database.clone(),
+                series.clone(),
+                *frequency,
+                *max_age,
+            )) as Box<dyn crate::input_source::InputSource>),
+            FlowInput::Mqtt { .. } => None,
+        };
+
         let mut steps = vec![];
         for (i, step) in self.steps.into_iter().enumerate() {
             let mut step = step.compile(config_dir, i, &source).await?;
@@ -122,15 +143,13 @@ impl FlowConfig {
             steps.push(step);
         }
         let output = self.output.try_into()?;
-        let mut flow = Flow {
+        let flow = Flow {
             input,
+            input_source,
             steps,
             source,
             output,
-            next_drain: None,
-            last_drain: None,
         };
-        flow.init_next_drain();
         Ok(flow)
     }
 }

crates/extensions/tedge_flows/src/flow.rs

@@ -1,6 +1,7 @@
 use std::ops::Sub;
 use std::time::Duration;
 
+use crate::input_source::InputSource;
 use crate::js_runtime::JsRuntime;
 use crate::js_script::JsScript;
 use crate::stats::Counter;
@@ -16,24 +17,20 @@ use tokio::time::Instant;
 use tracing::warn;
 
 /// A chain of transformation of MQTT messages
-#[derive(Debug)]
 pub struct Flow {
     /// The source topics
     pub input: FlowInput,
 
+    /// Input source for polling (e.g., database drains)
+    pub input_source: Option<Box<dyn InputSource>>,
+
     /// Transformation steps to apply in order to the messages
     pub steps: Vec<FlowStep>,
 
     pub source: Utf8PathBuf,
 
     /// Target of the transformed messages
     pub output: FlowOutput,
-
-    /// Next time to drain database for MeaDB inputs (for deadline-based wakeup)
-    pub next_drain: Option<tokio::time::Instant>,
-
-    /// Last time database was drained (for frequency checking)
-    pub last_drain: Option<DateTime>,
 }
 
 /// A message transformation step
@@ -104,45 +101,6 @@ impl Flow {
         }
     }
 
-    pub fn init_next_drain(&mut self) {
-        if let FlowInput::MeaDB { frequency, .. } = &self.input {
-            if !frequency.is_zero() {
-                self.next_drain = Some(tokio::time::Instant::now() + *frequency);
-            }
-        }
-    }
-
-    pub fn should_drain_at(&mut self, timestamp: DateTime) -> bool {
-        if let FlowInput::MeaDB { frequency, .. } = &self.input {
-            if frequency.is_zero() {
-                return false;
-            }
-
-            // Check if enough time has passed since last drain
-            match self.last_drain {
-                Some(last_drain) => {
-                    let elapsed_secs = timestamp.seconds.saturating_sub(last_drain.seconds);
-                    let frequency_secs = frequency.as_secs();
-                    if elapsed_secs >= frequency_secs {
-                        self.last_drain = Some(timestamp);
-                        // Also update the deadline for the actor loop
-                        self.next_drain = Some(tokio::time::Instant::now() + *frequency);
-                        true
-                    } else {
-                        false
-                    }
-                }
-                None => {
-                    // First drain
-                    self.last_drain = Some(timestamp);
-                    self.next_drain = Some(tokio::time::Instant::now() + *frequency);
-                    true
-                }
-            }
-        } else {
-            false
-        }
-    }
 
     pub fn topics(&self) -> TopicFilter {
         let mut topics = self.input.topics();