feat: execution will now calculate total time and send it to Aquila

Author: raphael-goetz

crates/taurus/src/main.rs

@@ -2,6 +2,7 @@ mod client;
 mod config;
 
 use crate::client::runtime_status::TaurusRuntimeStatusService;
+use crate::client::runtime_usage::TaurusRuntimeUsageService;
 use crate::config::Config;
 use code0_flow::flow_service::FlowUpdateService;
 
@@ -11,16 +12,20 @@ use futures_lite::StreamExt;
 use log::error;
 use prost::Message;
 use std::collections::HashMap;
+use std::time::{SystemTime, UNIX_EPOCH};
 use taurus_core::context::context::Context;
 use taurus_core::context::executor::Executor;
 use taurus_core::context::registry::FunctionStore;
 use taurus_core::context::signal::Signal;
 use tokio::signal;
 use tonic_health::pb::health_server::HealthServer;
 use tucana::shared::value::Kind;
-use tucana::shared::{ExecutionFlow, NodeFunction, RuntimeFeature, Translation, Value};
+use tucana::shared::{
+    ExecutionFlow, NodeFunction, RuntimeFeature, RuntimeUsage, Translation, Value,
+};
 
-fn handle_message(flow: ExecutionFlow, store: &FunctionStore) -> Signal {
+fn handle_message(flow: ExecutionFlow, store: &FunctionStore) -> (Signal, RuntimeUsage) {
+    let now = SystemTime::now();
     let mut context = Context::default();
 
     let node_functions: HashMap<i64, NodeFunction> = flow
@@ -29,7 +34,23 @@ fn handle_message(flow: ExecutionFlow, store: &FunctionStore) -> Signal {
         .map(|node| (node.database_id, node))
         .collect();
 
-    Executor::new(store, node_functions).execute(flow.starting_node_id, &mut context, true)
+    let signal =
+        Executor::new(store, node_functions).execute(flow.starting_node_id, &mut context, true);
+    let timestamp = match now.duration_since(UNIX_EPOCH) {
+        Ok(time) => time.as_millis(),
+        Err(err) => {
+            log::error!("cannot get current system time: {:?}", err);
+            0
+        }
+    };
+
+    (
+        signal,
+        RuntimeUsage {
+            flow_id: flow.flow_id,
+            duration: timestamp as i64,
+        },
+    )
 }
 
 #[tokio::main]
@@ -43,6 +64,7 @@ async fn main() {
     let config = Config::new();
     let store = FunctionStore::default();
     let mut runtime_status_service: Option<TaurusRuntimeStatusService> = None;
+    let mut runtime_usage_service: Option<TaurusRuntimeUsageService> = None;
 
     let client = match async_nats::connect(config.nats_url.clone()).await {
         Ok(client) => {
@@ -92,6 +114,9 @@ async fn main() {
         .send()
         .await;
 
+        let usage_service = TaurusRuntimeUsageService::from_url(config.aquila_url.clone()).await;
+        runtime_usage_service = Some(usage_service);
+
         let status_service = TaurusRuntimeStatusService::from_url(
             config.aquila_url.clone(),
             "taurus".into(),
@@ -143,7 +168,8 @@ async fn main() {
             };
 
             let flow_id = flow.flow_id;
-            let value = match handle_message(flow, &store) {
+            let result = handle_message(flow, &store);
+            let value = match result.0 {
                 Signal::Failure(error) => {
                     log::error!(
                         "RuntimeError occurred, execution failed because: {:?}",
@@ -173,6 +199,10 @@ async fn main() {
 
             log::info!("For the flow_id {} returing the value {:?}", flow_id, value);
 
+            if let Some(usage_service) = &runtime_usage_service {
+                usage_service.update_runtime_usage(result.1).await;
+            }
+
             // Send a response to the reply subject
             if let Some(reply) = msg.reply {
                 match client.publish(reply, value.encode_to_vec().into()).await {