ref: split taurus main file into app module

raphael-goetz · raphael-goetz · commit b5ba1738d6b4 · 2026-04-23T19:08:15.000+02:00
diff --git a/crates/taurus/src/app/mod.rs b/crates/taurus/src/app/mod.rs
@@ -0,0 +1,218 @@
+mod worker;
+
+use std::time::Duration;
+
+use code0_flow::flow_config::load_env_file;
+use code0_flow::flow_config::mode::Mode::DYNAMIC;
+use code0_flow::flow_service::FlowUpdateService;
+use taurus_core::runtime::engine::ExecutionEngine;
+use taurus_provider::providers::remote::nats_remote_runtime::NATSRemoteRuntime;
+use tokio::signal;
+use tokio::task::JoinHandle;
+use tokio::time::sleep;
+use tonic_health::pb::health_server::HealthServer;
+use tucana::shared::{RuntimeFeature, Translation};
+
+use crate::client::runtime_status::TaurusRuntimeStatusService;
+use crate::client::runtime_usage::TaurusRuntimeUsageService;
+use crate::config::Config;
+
+pub async fn run() {
+    init_logging();
+    load_env_file();
+
+    let config = Config::new();
+    let engine = ExecutionEngine::new();
+    let client = connect_nats(&config).await;
+
+    let mut health_task = spawn_health_task(&config);
+    let (runtime_status_service, runtime_usage_service) =
+        setup_dynamic_services_if_needed(&config).await;
+
+    let nats_remote = NATSRemoteRuntime::new(client.clone());
+    let mut worker_task = worker::spawn_worker(client, engine, nats_remote, runtime_usage_service);
+
+    wait_for_shutdown(&mut worker_task, &mut health_task).await;
+    update_stopped_status(runtime_status_service.as_ref()).await;
+
+    log::info!("Taurus shutdown complete");
+}
+
+fn init_logging() {
+    env_logger::Builder::from_default_env()
+        .filter_level(log::LevelFilter::Debug)
+        .init();
+}
+
+async fn connect_nats(config: &Config) -> async_nats::Client {
+    match async_nats::connect(config.nats_url.clone()).await {
+        Ok(client) => {
+            log::info!("Connected to NATS server");
+            client
+        }
+        Err(err) => {
+            panic!("Failed to connect to NATS server: {}", err);
+        }
+    }
+}
+
+fn spawn_health_task(config: &Config) -> Option<JoinHandle<()>> {
+    if !config.with_health_service {
+        return None;
+    }
+
+    let health_service = code0_flow::flow_health::HealthService::new(config.nats_url.clone());
+    let address = match format!("{}:{}", config.grpc_host, config.grpc_port).parse() {
+        Ok(address) => address,
+        Err(err) => {
+            log::error!("Failed to parse gRPC address: {:?}", err);
+            return None;
+        }
+    };
+
+    log::info!("Health server starting at {}", address);
+    Some(tokio::spawn(async move {
+        if let Err(err) = tonic::transport::Server::builder()
+            .add_service(HealthServer::new(health_service))
+            .serve(address)
+            .await
+        {
+            log::error!("Health server error: {:?}", err);
+        } else {
+            log::info!("Health server stopped gracefully");
+        }
+    }))
+}
+
+async fn setup_dynamic_services_if_needed(
+    config: &Config,
+) -> (
+    Option<TaurusRuntimeStatusService>,
+    Option<TaurusRuntimeUsageService>,
+) {
+    if config.mode != DYNAMIC {
+        return (None, None);
+    }
+
+    push_definitions_until_success(config).await;
+
+    let runtime_usage_service = Some(
+        TaurusRuntimeUsageService::from_url(config.aquila_url.clone(), config.aquila_token.clone())
+            .await,
+    );
+
+    let runtime_status_service = Some(
+        TaurusRuntimeStatusService::from_url(
+            config.aquila_url.clone(),
+            config.aquila_token.clone(),
+            "taurus".into(),
+            runtime_features(),
+        )
+        .await,
+    );
+
+    if let Some(status_service) = runtime_status_service.as_ref() {
+        status_service
+            .update_runtime_status(tucana::shared::execution_runtime_status::Status::Running)
+            .await;
+    }
+
+    (runtime_status_service, runtime_usage_service)
+}
+
+async fn push_definitions_until_success(config: &Config) {
+    let definition_service = FlowUpdateService::from_url(
+        config.aquila_url.clone(),
+        config.definitions.as_str(),
+        config.aquila_token.clone(),
+    )
+    .await;
+
+    let mut retry_count = 1;
+    loop {
+        if definition_service.send_with_status().await {
+            break;
+        }
+
+        log::warn!(
+            "Updating definitions failed, trying again in 3 seconds (retry #{})",
+            retry_count
+        );
+        retry_count += 1;
+        sleep(Duration::from_secs(3)).await;
+    }
+}
+
+fn runtime_features() -> Vec<RuntimeFeature> {
+    vec![RuntimeFeature {
+        name: vec![Translation {
+            code: "en-US".to_string(),
+            content: "Runtime".to_string(),
+        }],
+        description: vec![Translation {
+            code: "en-US".to_string(),
+            content: "Will execute incoming flows.".to_string(),
+        }],
+    }]
+}
+
+async fn update_stopped_status(runtime_status_service: Option<&TaurusRuntimeStatusService>) {
+    if let Some(status_service) = runtime_status_service {
+        status_service
+            .update_runtime_status(tucana::shared::execution_runtime_status::Status::Stopped)
+            .await;
+    }
+}
+
+async fn wait_for_shutdown(
+    worker_task: &mut JoinHandle<()>,
+    health_task: &mut Option<JoinHandle<()>>,
+) {
+    #[cfg(unix)]
+    let sigterm = async {
+        use tokio::signal::unix::{SignalKind, signal};
+
+        let mut term = signal(SignalKind::terminate()).expect("failed to install SIGTERM handler");
+        term.recv().await;
+    };
+
+    #[cfg(not(unix))]
+    let sigterm = std::future::pending::<()>();
+
+    if let Some(health_task) = health_task.as_mut() {
+        tokio::select! {
+            _ = &mut *worker_task => {
+                log::warn!("NATS worker task finished, shutting down");
+                health_task.abort();
+            }
+            _ = &mut *health_task => {
+                log::warn!("Health server task finished, shutting down");
+                worker_task.abort();
+            }
+            _ = signal::ctrl_c() => {
+                log::info!("Ctrl+C/Exit signal received, shutting down");
+                worker_task.abort();
+                health_task.abort();
+            }
+            _ = sigterm => {
+                log::info!("SIGTERM received, shutting down");
+                worker_task.abort();
+                health_task.abort();
+            }
+        }
+    } else {
+        tokio::select! {
+            _ = &mut *worker_task => {
+                log::warn!("NATS worker task finished, shutting down");
+            }
+            _ = signal::ctrl_c() => {
+                log::info!("Ctrl+C/Exit signal received, shutting down");
+                worker_task.abort();
+            }
+            _ = sigterm => {
+                log::info!("SIGTERM received, shutting down");
+                worker_task.abort();
+            }
+        }
+    }
+}
diff --git a/crates/taurus/src/app/worker.rs b/crates/taurus/src/app/worker.rs
@@ -0,0 +1,161 @@
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::time::Instant;
+
+use futures_lite::StreamExt;
+use prost::Message;
+use taurus_core::runtime::engine::{EmitType, ExecutionEngine, RespondEmitter};
+use taurus_core::types::signal::Signal;
+use taurus_provider::providers::emitter::nats_emitter::NATSRespondEmitter;
+use taurus_provider::providers::remote::nats_remote_runtime::NATSRemoteRuntime;
+use tokio::task::JoinHandle;
+use tucana::shared::value::Kind;
+use tucana::shared::{ExecutionFlow, RuntimeUsage, Value};
+
+use crate::client::runtime_usage::TaurusRuntimeUsageService;
+
+pub fn spawn_worker(
+    client: async_nats::Client,
+    engine: ExecutionEngine,
+    nats_remote: NATSRemoteRuntime,
+    runtime_usage_service: Option<TaurusRuntimeUsageService>,
+) -> JoinHandle<()> {
+    tokio::spawn(async move {
+        let runtime_emitter = NATSRespondEmitter::new(client.clone());
+
+        let mut subscription = match client
+            .queue_subscribe(String::from("execution.*"), "taurus".into())
+            .await
+        {
+            Ok(subscription) => {
+                log::info!("Subscribed to 'execution.*'");
+                subscription
+            }
+            Err(err) => {
+                log::error!("Failed to subscribe to 'execution.*': {:?}", err);
+                return;
+            }
+        };
+
+        while let Some(message) = subscription.next().await {
+            process_message(
+                message,
+                &client,
+                &engine,
+                &nats_remote,
+                &runtime_emitter,
+                runtime_usage_service.as_ref(),
+            )
+            .await;
+        }
+
+        log::info!("NATS worker loop ended");
+    })
+}
+
+async fn process_message(
+    message: async_nats::Message,
+    client: &async_nats::Client,
+    engine: &ExecutionEngine,
+    nats_remote: &NATSRemoteRuntime,
+    runtime_emitter: &NATSRespondEmitter,
+    runtime_usage_service: Option<&TaurusRuntimeUsageService>,
+) {
+    let flow: ExecutionFlow = match ExecutionFlow::decode(&*message.payload) {
+        Ok(flow) => flow,
+        Err(err) => {
+            log::error!(
+                "Failed to deserialize flow: {:?}, payload: {:?}",
+                err,
+                &message.payload
+            );
+            return;
+        }
+    };
+
+    let flow_id = flow.flow_id;
+    let reply_subject = message.reply.clone();
+    let respond_count = Arc::new(AtomicUsize::new(0));
+    let respond_count_for_emitter = respond_count.clone();
+    let respond_emitter = |execution_id, emit_type: EmitType, value: Value| {
+        match emit_type {
+            EmitType::OngoingExec => {
+                respond_count_for_emitter.fetch_add(1, Ordering::Relaxed);
+            }
+            EmitType::StartingExec => log::debug!("Flow execution started"),
+            EmitType::FinishedExec => log::debug!("Flow execution finished"),
+            EmitType::FailedExec => log::debug!("Flow execution failed"),
+        }
+        runtime_emitter.emit(execution_id, emit_type, value);
+    };
+    let (signal, runtime_usage) = execute_flow(flow, engine, nats_remote, Some(&respond_emitter));
+
+    let has_responded = respond_count.load(Ordering::Relaxed) > 0;
+    let final_value = signal_to_terminal_value(signal);
+
+    // Stream contract: if we already emitted intermediate values, do not send an extra terminal
+    // payload on the same reply subject.
+    if let Some(value) = final_value
+        && !has_responded
+        && let Some(reply_subject) = reply_subject
+    {
+        log::info!("Returning value for flow_id {}: {:?}", flow_id, value);
+        if let Err(err) = client
+            .publish(reply_subject, value.encode_to_vec().into())
+            .await
+        {
+            log::error!("Failed to send response: {:?}", err);
+        }
+    }
+
+    if let Some(usage_service) = runtime_usage_service {
+        usage_service.update_runtime_usage(runtime_usage).await;
+    }
+}
+
+fn execute_flow(
+    flow: ExecutionFlow,
+    engine: &ExecutionEngine,
+    nats_remote: &NATSRemoteRuntime,
+    respond_emitter: Option<&dyn RespondEmitter>,
+) -> (Signal, RuntimeUsage) {
+    let start = Instant::now();
+    let flow_id = flow.flow_id;
+    let (signal, _reason) = engine.execute_flow(flow, Some(nats_remote), respond_emitter, true);
+    let duration_millis = start.elapsed().as_millis() as i64;
+
+    (
+        signal,
+        RuntimeUsage {
+            flow_id,
+            duration: duration_millis,
+        },
+    )
+}
+
+fn signal_to_terminal_value(signal: Signal) -> Option<Value> {
+    match signal {
+        Signal::Failure(error) => {
+            log::error!("Runtime error occurred: {:?}", error);
+            Some(error.as_value())
+        }
+        Signal::Success(value) => {
+            log::debug!("Execution ended with success signal");
+            Some(value)
+        }
+        Signal::Return(value) => {
+            log::debug!("Execution ended with return signal");
+            Some(value)
+        }
+        Signal::Respond(_) => {
+            log::debug!("Execution ended with respond signal");
+            None
+        }
+        Signal::Stop => {
+            log::debug!("Received stop signal as last signal");
+            Some(Value {
+                kind: Some(Kind::NullValue(0)),
+            })
+        }
+    }
+}
diff --git a/crates/taurus/src/main.rs b/crates/taurus/src/main.rs
