Merge pull request #214 from code0-tech/213-non-blocking-bug

Non-Blocking gRPC Server

Author: raphael-goetz

Cargo.lock

@@ -4,16 +4,16 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-tokio = { version = "1.44.1", features = ["rt-multi-thread"] }
+tokio = { version = "1.44.1", features = ["rt-multi-thread", "signal"] }
 futures = "0.3.31"
 log = "0.4.26"
 env_logger = "0.11.8"
 prost = "0.14.1"
 tonic = "0.14.1"
-tucana = { version = "0.0.39", features = ["all"] }
-code0-flow = { version = "0.0.18" }
+tucana = { version = "0.0.42", features = ["all"] }
+code0-flow = { version = "0.0.19", features = ["flow_health", "flow_validator"] }
 serde_json = "1.0.140"
-async-nats = "0.44.2"
+async-nats = "0.45.0"
 tonic-health = "0.14.1"
 tokio-stream = "0.1.17"
 uuid = { version = "1.18.0", features = ["v4"] }

Cargo.toml

@@ -1,4 +1,3 @@
-use crate::sagittarius::test_execution_client_impl::SagittariusTestExecutionServiceClient;
 use crate::{configuration::Config as AquilaConfig, flow::get_flow_identifier};
 use async_nats::jetstream::kv::Config;
 use code0_flow::flow_config::load_env_file;
@@ -31,8 +30,14 @@ async fn main() {
 
     //Create connection to JetStream
     let client = match async_nats::connect(config.nats_url.clone()).await {
-        Ok(client) => client,
-        Err(err) => panic!("Failed to connect to NATS server: {}", err),
+        Ok(client) => {
+            log::info!("Connected to NATS");
+            client
+        }
+        Err(err) => {
+            log::error!("Failed to connect to NATS: {:?}", err);
+            panic!("Failed to connect to NATS server: {:?}", err)
+        }
     };
 
     let jet_stream = async_nats::jetstream::new(client.clone());
@@ -45,47 +50,60 @@ async fn main() {
         .await;
 
     let kv_store = match jet_stream.get_key_value(config.nats_bucket.clone()).await {
-        Ok(kv) => Arc::new(kv),
-        Err(err) => panic!("Failed to get key-value store: {}", err),
+        Ok(kv) => {
+            log::info!("Connected to JetStream");
+            Arc::new(kv)
+        }
+        Err(err) => {
+            log::error!("Failed to get key-value store: {:?}", err);
+            panic!("Failed to get key-value store: {:?}", err)
+        }
     };
 
-    //Create connection to Sagittarius if the type is hybrid
-    if !config.is_static() {
-        let server = AquilaGRPCServer::new(&config);
-
-        match server.start().await {
-            Ok(_) => {
-                log::info!("Server started successfully");
-            }
-            Err(err) => {
-                log::error!("Failed to start server: {:?}", err);
-                panic!("Failed to start server");
-            }
-        };
+    if config.is_static() {
+        log::info!("Starting with static configuration");
+        init_flows_from_json(config.flow_fallback_path, kv_store).await;
+        return;
+    }
 
-        // Connect to Sagittarius Flow Endpoint
-        SagittariusFlowClient::new(
-            config.backend_url.clone(),
-            kv_store.clone(),
-            config.runtime_token.clone(),
-        )
-        .await
-        .init_flow_stream()
-        .await;
+    let server = AquilaGRPCServer::new(&config);
+    let backend_url_flow = config.backend_url.clone();
+    let runtime_token_flow = config.runtime_token.clone();
+    let kv_for_flow = kv_store.clone();
 
-        // Connect to Sagittarius Execution Endpoint
-        SagittariusTestExecutionServiceClient::new(
-            client,
-            kv_store,
-            config.backend_url,
-            config.runtime_token,
-        )
-        .await
-        .logon()
-        .await;
-    } else {
-        init_flows_from_json(config.flow_fallback_path, kv_store).await
+    let mut server_task = tokio::spawn(async move {
+        if let Err(err) = server.start().await {
+            log::error!("gRPC server error: {:?}", err);
+        } else {
+            log::info!("gRPC server stopped gracefully");
+        }
+    });
+
+    let mut flow_task = tokio::spawn(async move {
+        let mut flow_client =
+            SagittariusFlowClient::new(backend_url_flow, kv_for_flow, runtime_token_flow).await;
+
+        flow_client.init_flow_stream().await;
+        log::warn!("Flow stream task exited");
+    });
+
+    tokio::select! {
+        _ = &mut server_task => {
+            log::warn!("gRPC server task finished, shutting down");
+            flow_task.abort();
+        }
+        _ = &mut flow_task => {
+            log::warn!("Flow stream task finished, shutting down");
+            server_task.abort();
+        }
+        _ = tokio::signal::ctrl_c() => {
+            log::info!("Ctrl+C/Exit signal received, shutting down");
+            server_task.abort();
+            flow_task.abort();
+        }
     }
+
+    log::info!("Aquila shutdown complete");
 }
 
 async fn init_flows_from_json(

src/main.rs

@@ -42,7 +42,7 @@ impl SagittariusFlowClient {
     async fn handle_response(&mut self, response: FlowResponse) {
         let data = match response.data {
             Some(data) => {
-                log::info!("Recieved a FlowResponse");
+                log::info!("Received a FlowResponse");
                 data
             }
             None => {
@@ -115,7 +115,7 @@ impl SagittariusFlowClient {
 
         let response = match self.client.update(request).await {
             Ok(res) => {
-                log::info!("Succesfully established a Stream (for Flows)");
+                log::info!("Successfully established a Stream (for Flows)");
                 res
             }
             Err(status) => {

src/sagittarius/flow_service_client_impl.rs

@@ -18,12 +18,11 @@ impl DataTypeService for AquilaDataTypeServiceServer {
     async fn update(
         &self,
         request: tonic::Request<tucana::aquila::DataTypeUpdateRequest>,
-    ) -> std::result::Result<tonic::Response<tucana::aquila::DataTypeUpdateResponse>, tonic::Status>
-    {
+    ) -> Result<tonic::Response<tucana::aquila::DataTypeUpdateResponse>, tonic::Status> {
         let data_type_update_request = request.into_inner();
 
         log::info!(
-            "Recieved DataTypes: {:?}",
+            "Received DataTypes: {:?}",
             data_type_update_request.data_types
         );
 

src/server/data_type_service_server_impl.rs

@@ -19,12 +19,11 @@ impl FlowTypeService for AquilaFlowTypeServiceServer {
     async fn update(
         &self,
         request: tonic::Request<tucana::aquila::FlowTypeUpdateRequest>,
-    ) -> std::result::Result<tonic::Response<tucana::aquila::FlowTypeUpdateResponse>, tonic::Status>
-    {
+    ) -> Result<tonic::Response<tucana::aquila::FlowTypeUpdateResponse>, tonic::Status> {
         let flow_type_update_request = request.into_inner();
 
         log::info!(
-            "Recieved FlowTypes: {:?}",
+            "Received FlowTypes: {:?}",
             flow_type_update_request.flow_types
         );
 

src/server/flow_type_service_server_impl.rs

@@ -20,14 +20,14 @@ impl RuntimeFunctionDefinitionService for AquilaRuntimeFunctionServiceServer {
     async fn update(
         &self,
         request: tonic::Request<tucana::aquila::RuntimeFunctionDefinitionUpdateRequest>,
-    ) -> std::result::Result<
+    ) -> Result<
         tonic::Response<tucana::aquila::RuntimeFunctionDefinitionUpdateResponse>,
         tonic::Status,
     > {
         let runtime_function_definition_update_request = request.into_inner();
 
         log::info!(
-            "Recieved RuntimeFunctions: {:?}",
+            "Received RuntimeFunctions: {:?}",
             runtime_function_definition_update_request.runtime_functions
         );