Add RuntimeUsageService

crates/taurus/src/client/mod.rs

@@ -1 +1,2 @@
 pub mod runtime_status;
+pub mod runtime_usage;

crates/taurus/src/client/runtime_usage.rs

@@ -0,0 +1,39 @@
+use code0_flow::flow_service::retry::create_channel_with_retry;
+use tonic::transport::Channel;
+use tucana::{
+    aquila::{RuntimeUsageRequest, runtime_usage_service_client::RuntimeUsageServiceClient},
+    shared::RuntimeUsage,
+};
+
+pub struct TaurusRuntimeUsageService {
+    channel: Channel,
+}
+
+impl TaurusRuntimeUsageService {
+    pub async fn from_url(aquila_url: String) -> Self {
+        let channel = create_channel_with_retry("Aquila", aquila_url).await;
+        TaurusRuntimeUsageService { channel }
+    }
+
+    pub async fn update_runtime_usage(&self, runtime_usage: RuntimeUsage) {
+        log::info!("Updating the current Runtime Status!");
+        let mut client = RuntimeUsageServiceClient::new(self.channel.clone());
+
+        let request = RuntimeUsageRequest {
+            runtime_usage: vec![runtime_usage],
+        };
+
+        match client.update(request).await {
+            Ok(response) => {
+                log::info!(
+                    "Was the update of the RuntimeStatus accepted by Sagittarius? {}",
+                    response.into_inner().success
+                );
+            }
+            Err(err) => {
+                log::error!("Failed to update RuntimeStatus: {:?}", err);
+            }
+        }
+    }
+}
+

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 {