diff --git a/src/authorization/mod.rs b/src/authorization/mod.rs index 22656cb..a112574 100644 --- a/src/authorization/mod.rs +++ b/src/authorization/mod.rs @@ -1,3 +1,8 @@ +//! Bearer-token helpers shared by every gRPC client and server in Aquila: +//! [`authorization::get_authorization_metadata`] to attach a token to an +//! outgoing request, [`authorization::extract_token`] to read one back off +//! an incoming request. + pub mod authorization { use std::str::FromStr; use tonic::{ @@ -32,6 +37,9 @@ pub mod authorization { map } + /// Reads and validates the bearer token off an incoming request's + /// `authorization` header. Generic over the request body type so it + /// works for both unary and streaming gRPC requests. pub fn extract_token(request: &Request) -> Result<&str, Status> { let header = request.metadata().get("authorization").ok_or_else(|| { log::warn!("Missing authorization header"); diff --git a/src/configuration/config/display.rs b/src/configuration/config/display.rs new file mode 100644 index 0000000..45245cc --- /dev/null +++ b/src/configuration/config/display.rs @@ -0,0 +1,111 @@ +//! Human-readable rendering of [`Config`] for the startup log line, kept +//! separate from the struct definitions so the config *shape* isn't buried +//! under formatting code. Must stay in sync with the redactions in +//! [`super::DynamicConfig`]'s `Debug` impl — anything secret here needs the +//! same `[FILTERED]` treatment. + +use std::fmt; + +use super::Config; + +impl fmt::Display for Config { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(formatter, "Aquila configuration")?; + writeln!(formatter, " Environment: {}", self.environment)?; + writeln!(formatter, " Mode: {}", self.mode)?; + writeln!(formatter, " Log level: {}", self.log_level)?; + writeln!(formatter, " OpenTelemetry")?; + writeln!(formatter, " Enabled: {}", self.opentelemetry.enabled)?; + writeln!( + formatter, + " Service: {}", + self.opentelemetry.service_name + )?; + writeln!( + formatter, + " Logs: {}", + display_optional_url(&self.opentelemetry.logs_endpoint) + )?; + writeln!( + formatter, + " Metrics: {}", + display_optional_url(&self.opentelemetry.metrics_endpoint) + )?; + writeln!( + formatter, + " Traces: {}", + display_optional_url(&self.opentelemetry.traces_endpoint) + )?; + writeln!(formatter, " NATS")?; + writeln!(formatter, " URL: {}", self.nats.url)?; + writeln!(formatter, " Bucket: {}", self.nats.bucket)?; + writeln!(formatter, " gRPC")?; + writeln!( + formatter, + " Address: {}:{}", + self.grpc.host, self.grpc.port + )?; + writeln!( + formatter, + " Health service: {}", + self.grpc.health_service + )?; + writeln!(formatter, " Static mode")?; + writeln!(formatter, " Flow path: {}", self.static_config.flow_path)?; + writeln!(formatter, " Dynamic mode")?; + writeln!( + formatter, + " Backend URL: {}", + self.dynamic_config.backend_url + )?; + writeln!(formatter, " Backend token: [FILTERED]")?; + writeln!( + formatter, + " Request timeout: {}s", + self.dynamic_config.backend_unary_timeout_secs + )?; + writeln!(formatter, " Runtime status")?; + writeln!( + formatter, + " Not responding after: {}s", + self.runtime_status.not_responding_after_secs + )?; + writeln!( + formatter, + " Stopped after: {}s", + self.runtime_status.stopped_after_not_responding_secs + )?; + write!( + formatter, + " Monitor interval: {}s", + self.runtime_status.monitor_interval_secs + ) + } +} + +fn display_optional_url(url: &Option) -> &str { + url.as_deref() + .filter(|value| !value.trim().is_empty()) + .unwrap_or("") +} + +#[cfg(test)] +mod tests { + use super::Config; + + #[test] + fn display_output_is_readable_and_filters_backend_token() { + let mut config = Config::default(); + config.dynamic_config.backend_token = "super-secret".into(); + + let output = config.to_string(); + + assert!(output.starts_with("Aquila configuration\n")); + assert!(output.contains(" Environment: development")); + assert!(output.contains(" Address: 127.0.0.1:8081")); + assert!(output.contains(" Request timeout: 5s")); + assert!(output.contains(" Backend token: [FILTERED]")); + assert!(!output.contains("super-secret")); + assert!(!output.contains("Config {")); + } +} diff --git a/src/configuration/config.rs b/src/configuration/config/mod.rs similarity index 68% rename from src/configuration/config.rs rename to src/configuration/config/mod.rs index 55b0a14..e9faaa9 100644 --- a/src/configuration/config.rs +++ b/src/configuration/config/mod.rs @@ -1,4 +1,12 @@ -use std::{fmt, path::Path}; +//! Aquila's top-level configuration: loaded from `aquila.{yaml,toml,json,...}` +//! (via the `config` crate), overridable by environment variables, and +//! falling back to the [`Default`] impls below when no file is present. +//! +//! See [`display`] for how this renders in the startup log line. + +mod display; + +use std::path::Path; use code0_flow::flow_telemetry::OpenTelemetry; use config::{Config as ConfigLoader, ConfigError, File}; @@ -181,87 +189,6 @@ impl Config { } } -impl fmt::Display for Config { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - writeln!(formatter, "Aquila configuration")?; - writeln!(formatter, " Environment: {}", self.environment)?; - writeln!(formatter, " Mode: {}", self.mode)?; - writeln!(formatter, " Log level: {}", self.log_level)?; - writeln!(formatter, " OpenTelemetry")?; - writeln!(formatter, " Enabled: {}", self.opentelemetry.enabled)?; - writeln!( - formatter, - " Service: {}", - self.opentelemetry.service_name - )?; - writeln!( - formatter, - " Logs: {}", - display_optional_url(&self.opentelemetry.logs_endpoint) - )?; - writeln!( - formatter, - " Metrics: {}", - display_optional_url(&self.opentelemetry.metrics_endpoint) - )?; - writeln!( - formatter, - " Traces: {}", - display_optional_url(&self.opentelemetry.traces_endpoint) - )?; - writeln!(formatter, " NATS")?; - writeln!(formatter, " URL: {}", self.nats.url)?; - writeln!(formatter, " Bucket: {}", self.nats.bucket)?; - writeln!(formatter, " gRPC")?; - writeln!( - formatter, - " Address: {}:{}", - self.grpc.host, self.grpc.port - )?; - writeln!( - formatter, - " Health service: {}", - self.grpc.health_service - )?; - writeln!(formatter, " Static mode")?; - writeln!(formatter, " Flow path: {}", self.static_config.flow_path)?; - writeln!(formatter, " Dynamic mode")?; - writeln!( - formatter, - " Backend URL: {}", - self.dynamic_config.backend_url - )?; - writeln!(formatter, " Backend token: [FILTERED]")?; - writeln!( - formatter, - " Request timeout: {}s", - self.dynamic_config.backend_unary_timeout_secs - )?; - writeln!(formatter, " Runtime status")?; - writeln!( - formatter, - " Not responding after: {}s", - self.runtime_status.not_responding_after_secs - )?; - writeln!( - formatter, - " Stopped after: {}s", - self.runtime_status.stopped_after_not_responding_secs - )?; - write!( - formatter, - " Monitor interval: {}s", - self.runtime_status.monitor_interval_secs - ) - } -} - -fn display_optional_url(url: &Option) -> &str { - url.as_deref() - .filter(|value| !value.trim().is_empty()) - .unwrap_or("") -} - #[cfg(test)] mod tests { use std::sync::Mutex; @@ -304,22 +231,6 @@ mod tests { assert!(!output.contains("super-secret")); } - #[test] - fn display_output_is_readable_and_filters_backend_token() { - let mut config = Config::default(); - config.dynamic_config.backend_token = "super-secret".into(); - - let output = config.to_string(); - - assert!(output.starts_with("Aquila configuration\n")); - assert!(output.contains(" Environment: development")); - assert!(output.contains(" Address: 127.0.0.1:8081")); - assert!(output.contains(" Request timeout: 5s")); - assert!(output.contains(" Backend token: [FILTERED]")); - assert!(!output.contains("super-secret")); - assert!(!output.contains("Config {")); - } - #[test] fn opentelemetry_endpoints_are_enabled_by_presence() { let config: OpenTelemetry = ConfigLoader::builder() diff --git a/src/configuration/env.rs b/src/configuration/env.rs index a3f6aa9..772e248 100644 --- a/src/configuration/env.rs +++ b/src/configuration/env.rs @@ -1,7 +1,12 @@ +//! The deployment environment Aquila believes it's running in, used to gate +//! environment-specific behavior (like the dev-only flow JSON export in +//! [`crate::sagittarius::flow_service_client_impl`]). + use std::fmt; use serde::{Deserialize, Serialize}; +/// Which deployment environment Aquila is running in. #[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum Environment { diff --git a/src/configuration/mod.rs b/src/configuration/mod.rs index ffed714..4594158 100644 --- a/src/configuration/mod.rs +++ b/src/configuration/mod.rs @@ -1,3 +1,8 @@ +//! Everything Aquila loads before it can start serving: the process-wide +//! [`config::Config`] (file + env), the static [`service::ServiceConfiguration`] +//! allowlist, and small shared value types ([`env::Environment`], [`mode::Mode`], +//! [`state::AppReadiness`]) used across both. + pub mod config; pub mod env; pub mod mode; diff --git a/src/configuration/mode.rs b/src/configuration/mode.rs index bab984e..71a8073 100644 --- a/src/configuration/mode.rs +++ b/src/configuration/mode.rs @@ -1,3 +1,8 @@ +//! Aquila's two run modes: [`Mode::Static`] serves a fixed local flow export +//! with no Sagittarius dependency, [`Mode::Dynamic`] syncs flows and module +//! state live over gRPC. See [`crate::startup::static_mode`] and +//! [`crate::startup::dynamic_mode`] for what each mode actually wires up. + use std::fmt; use serde::{Deserialize, Serialize}; diff --git a/src/configuration/service/dto.rs b/src/configuration/service/dto.rs new file mode 100644 index 0000000..da23a72 --- /dev/null +++ b/src/configuration/service/dto.rs @@ -0,0 +1,95 @@ +//! Wire format for the service configuration file (`AQUILA_SERVICE_CONFIG_PATH`) +//! and its conversion into the domain types in the parent module. +//! +//! The file format is intentionally a bit friendlier than the domain model: +//! actions are declared with a plain `identifier` and per-project `configs`, +//! then expanded here into the [`tucana::shared::ModuleConfigurations`] +//! shape the rest of Aquila works with. Runtimes need no such expansion, so +//! [`RuntimeServiceConfiguration`] doubles as both the wire format and the +//! domain type. + +use serde::{Deserialize, Serialize}; +use tucana::shared::{ModuleConfigurations, helper::value::from_json_value}; + +use super::{ActionServiceConfiguration, ServiceConfiguration}; + +#[derive(Serialize, Deserialize, Clone)] +pub(super) struct SerializableModuleConfiguration { + pub(super) identifier: String, + pub(super) value: serde_json::Value, +} + +#[derive(Serialize, Deserialize, Clone)] +pub(super) struct SerializableModuleProjectConfiguration { + pub(super) project_id: i64, + #[serde(default)] + pub(super) configs: Vec, +} + +#[derive(Serialize, Deserialize, Clone)] +pub(super) struct SerializableActionServiceConfiguration { + pub(super) token: String, + pub(super) identifier: String, + #[serde(default)] + pub(super) configs: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Default)] +pub(super) struct SerializableServiceConfiguration { + #[serde(default)] + pub(super) actions: Vec, + #[serde(default)] + pub(super) runtimes: Vec, +} + +/// A pre-provisioned runtime's token and identity. Used directly as both the +/// file format and the domain type since, unlike actions, no expansion is needed. +#[derive(Serialize, Deserialize, Clone)] +pub struct RuntimeServiceConfiguration { + pub(super) token: String, + pub(super) identifier: String, + #[serde(default)] + pub(super) resolved_modules: Vec, +} + +impl From for tucana::shared::ModuleConfiguration { + fn from(value: SerializableModuleConfiguration) -> Self { + Self { + identifier: value.identifier, + value: Some(from_json_value(value.value)), + } + } +} + +impl From for tucana::shared::ModuleProjectConfigurations { + fn from(value: SerializableModuleProjectConfiguration) -> Self { + Self { + project_id: value.project_id, + module_configurations: value.configs.into_iter().map(Into::into).collect(), + } + } +} + +impl From for ActionServiceConfiguration { + fn from(value: SerializableActionServiceConfiguration) -> Self { + let module_identifier = value.identifier.clone(); + + Self { + token: value.token, + service_name: value.identifier, + config: vec![ModuleConfigurations { + module_identifier, + module_configurations: value.configs.into_iter().map(Into::into).collect(), + }], + } + } +} + +impl From for ServiceConfiguration { + fn from(value: SerializableServiceConfiguration) -> Self { + Self { + actions: value.actions.into_iter().map(Into::into).collect(), + runtimes: value.runtimes.into_iter().collect(), + } + } +} diff --git a/src/configuration/service.rs b/src/configuration/service/mod.rs similarity index 77% rename from src/configuration/service.rs rename to src/configuration/service/mod.rs index f9fff1d..16410d8 100644 --- a/src/configuration/service.rs +++ b/src/configuration/service/mod.rs @@ -1,101 +1,40 @@ -use serde::{Deserialize, Serialize}; -use serde_json::from_str; -use std::{fs::File, io::Read, path::Path}; -use tucana::shared::{ModuleConfigurations, helper::value::from_json_value}; +//! Domain model for Aquila's local service configuration file +//! (`AQUILA_SERVICE_CONFIG_PATH`): which action and runtime tokens are +//! pre-provisioned, and what configuration each action should receive. +//! +//! This is a static, file-backed allowlist — separate from the dynamic +//! module configuration Sagittarius pushes at runtime. See [`dto`] for the +//! on-disk format and how it's expanded into these types. -#[derive(Serialize, Deserialize, Clone)] -struct SerializableModuleConfiguration { - identifier: String, - value: serde_json::Value, -} +mod dto; -#[derive(Serialize, Deserialize, Clone)] -struct SerializableModuleProjectConfiguration { - project_id: i64, - #[serde(default)] - configs: Vec, -} +pub use dto::RuntimeServiceConfiguration; -#[derive(Serialize, Deserialize, Clone)] -pub struct SerializableActionServiceConfiguration { - token: String, - identifier: String, - #[serde(default)] - configs: Vec, -} +use std::{fs::File, io::Read, path::Path}; -#[derive(Serialize, Deserialize, Clone, Default)] -pub struct SerializableServiceConfiguration { - #[serde(default)] - actions: Vec, - #[serde(default)] - runtimes: Vec, -} +use serde_json::from_str; +use tucana::shared::ModuleConfigurations; -#[derive(Serialize, Deserialize, Clone)] +use dto::SerializableServiceConfiguration; + +#[derive(Clone)] pub struct ActionServiceConfiguration { token: String, service_name: String, config: Vec, } -#[derive(Serialize, Deserialize, Clone)] -pub struct RuntimeServiceConfiguration { - token: String, - identifier: String, - #[serde(default)] - resolved_modules: Vec, -} - -#[derive(Serialize, Deserialize, Clone, Default)] +#[derive(Clone, Default)] pub struct ServiceConfiguration { actions: Vec, runtimes: Vec, } -impl From for tucana::shared::ModuleConfiguration { - fn from(value: SerializableModuleConfiguration) -> Self { - Self { - identifier: value.identifier, - value: Some(from_json_value(value.value)), - } - } -} - -impl From for tucana::shared::ModuleProjectConfigurations { - fn from(value: SerializableModuleProjectConfiguration) -> Self { - Self { - project_id: value.project_id, - module_configurations: value.configs.into_iter().map(Into::into).collect(), - } - } -} - -impl From for ActionServiceConfiguration { - fn from(value: SerializableActionServiceConfiguration) -> Self { - let module_identifier = value.identifier.clone(); - - Self { - token: value.token, - service_name: value.identifier, - config: vec![ModuleConfigurations { - module_identifier, - module_configurations: value.configs.into_iter().map(Into::into).collect(), - }], - } - } -} - -impl From for ServiceConfiguration { - fn from(value: SerializableServiceConfiguration) -> Self { - Self { - actions: value.actions.into_iter().map(Into::into).collect(), - runtimes: value.runtimes.into_iter().collect(), - } - } -} impl ServiceConfiguration { - // This function is used to extract the real service name via modules + /// Maps a runtime's advertised identifier to the family it belongs to, + /// since individual `taurus-*` runtime instances all share one + /// provisioned token under the `taurus` identifier while `draco-*` + /// runtimes are provisioned individually. pub fn extract_service_name(name: &String) -> Option { if name.starts_with("draco") { return Some(name.clone()); @@ -146,6 +85,9 @@ impl ServiceConfiguration { } } + /// Every module identifier Aquila should advertise as available, + /// combining each action's own identifier with each runtime's resolved + /// module list (or its own identifier, if it hasn't resolved any yet). pub fn collect_modules(&self) -> Vec { let actions: Vec = self .actions @@ -164,6 +106,9 @@ impl ServiceConfiguration { vec![actions, runtime].concat() } + /// Loads the service configuration file at `path`, falling back to an + /// empty configuration (rather than failing startup) if the file is + /// missing, unreadable, or malformed — this file is optional. pub fn from_path(path: impl AsRef) -> Self { let mut data = String::new(); @@ -207,9 +152,11 @@ impl ServiceConfiguration { #[cfg(test)] mod tests { use super::{ - RuntimeServiceConfiguration, SerializableActionServiceConfiguration, - SerializableModuleConfiguration, SerializableModuleProjectConfiguration, - SerializableServiceConfiguration, ServiceConfiguration, + RuntimeServiceConfiguration, ServiceConfiguration, + dto::{ + SerializableActionServiceConfiguration, SerializableModuleConfiguration, + SerializableModuleProjectConfiguration, SerializableServiceConfiguration, + }, }; fn fixture() -> ServiceConfiguration { diff --git a/src/configuration/state.rs b/src/configuration/state.rs index ed5e9b2..1a66e8e 100644 --- a/src/configuration/state.rs +++ b/src/configuration/state.rs @@ -1,10 +1,14 @@ +//! Shared readiness flags consulted by [`crate::server::create_readiness_interceptor`] +//! to reject gRPC requests before Aquila's upstream dependencies are usable, +//! rather than accepting them and failing partway through a handler. + use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; /// Tracks readiness of each external service. #[derive(Clone)] pub struct AppReadiness { - // Readiness state of Sagittarius service + /// Whether the Sagittarius gRPC channel has completed its initial connect. pub sagittarius_ready: Arc, } diff --git a/src/flow/mod.rs b/src/flow/mod.rs index 91d716f..97db29a 100644 --- a/src/flow/mod.rs +++ b/src/flow/mod.rs @@ -1,7 +1,12 @@ +//! The key scheme Aquila's flow KV store uses, and the two operations built +//! on it: computing a flow's key and checking whether a key belongs to a +//! given flow id. There is no secondary index, so every "find by flow id" +//! lookup elsewhere in the codebase is a scan using [`key_has_flow_id`]. + use tucana::shared::ValidationFlow; /// Every flow identifier has this key -/// ... +/// `...` pub fn get_flow_identifier(flow: &ValidationFlow) -> String { format!( "{}.{}.{}.{}", diff --git a/src/main.rs b/src/main.rs index 76e5be3..8b89e20 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,12 @@ +//! Aquila is the runtime gateway that sits between Sagittarius and the +//! actions/runtimes that execute flows: it authenticates them, relays +//! execution requests and results over NATS, and keeps a local flow store in +//! sync (or, in static mode, serves a fixed export with no Sagittarius at all). +//! +//! See [`startup`] for the static vs. dynamic mode split, [`server`] for the +//! gRPC surface actions/runtimes talk to, and [`sagittarius`] for the +//! clients that talk back to Sagittarius. + use crate::configuration::{ config::Config as AquilaConfig, service::ServiceConfiguration, state::AppReadiness, }; @@ -62,6 +71,9 @@ async fn main() { telemetry.shutdown(); } +/// Routes panic messages through the telemetry error pipeline in addition to +/// the default stderr output, so a panic in a spawned task still shows up +/// wherever Aquila's other errors are reported. fn install_panic_logging() { std::panic::set_hook(Box::new(move |panic_info| { let message = if let Some(message) = panic_info.payload().downcast_ref::<&str>() { diff --git a/src/sagittarius/flow_service_client_impl/dev_export.rs b/src/sagittarius/flow_service_client_impl/dev_export.rs new file mode 100644 index 0000000..5ed66b4 --- /dev/null +++ b/src/sagittarius/flow_service_client_impl/dev_export.rs @@ -0,0 +1,87 @@ +//! Development-only convenience export of the current flow set to a local +//! JSON file, so a developer can inspect what Sagittarius last synced +//! without needing NATS/KV tooling. Only called when running with +//! `environment == DEVELOPMENT`. + +use std::path::Path; + +use tokio::fs; +use tucana::shared::Flows; + +const EXPORT_PATH: &str = "flowExport.json"; +const EXPORT_TMP_PATH: &str = "flowExport.json.tmp"; + +/// Writes `flows` to `flowExport.json`, replacing any previous export. +/// +/// Goes through a temp file plus rename so a reader never observes a +/// partially written export. If the rename fails because the destination +/// already exists, the destination is removed and the rename retried once +/// before giving up. +pub(super) async fn overwrite(flows: Flows) { + let json = match serde_json::to_vec_pretty(&flows) { + Ok(bytes) => bytes, + Err(error) => { + log::error!( + "Failed to serialize development flow export flow_count={} error={:?}", + flows.flows.len(), + error + ); + return; + } + }; + + let final_path = Path::new(EXPORT_PATH); + let tmp_path = Path::new(EXPORT_TMP_PATH); + + if let Err(error) = fs::write(tmp_path, &json).await { + log::error!( + "Failed to write development flow export path={} error={}", + tmp_path.display(), + error + ); + return; + } + + if let Err(error) = fs::rename(tmp_path, final_path).await { + log::warn!( + "Could not atomically replace development flow export path={} error={}; retrying after removing destination", + final_path.display(), + error + ); + match fs::remove_file(final_path).await { + Ok(()) => log::debug!( + "Removed previous development flow export path={}", + final_path.display() + ), + Err(remove_error) if remove_error.kind() == std::io::ErrorKind::NotFound => {} + Err(remove_error) => log::warn!( + "Failed to remove previous development flow export path={} error={}", + final_path.display(), + remove_error + ), + } + if let Err(retry_error) = fs::rename(tmp_path, final_path).await { + log::error!( + "Failed to replace development flow export source_path={} destination_path={} initial_rename_error={} retry_error={}", + tmp_path.display(), + final_path.display(), + error, + retry_error + ); + if let Err(cleanup_error) = fs::remove_file(tmp_path).await { + log::warn!( + "Failed to clean up temporary development flow export path={} error={}", + tmp_path.display(), + cleanup_error + ); + } + return; + } + } + + log::info!( + "Exported {} flows to {}", + flows.flows.len(), + final_path.display() + ); +} diff --git a/src/sagittarius/flow_service_client_impl/flow_store.rs b/src/sagittarius/flow_service_client_impl/flow_store.rs new file mode 100644 index 0000000..bf0384f --- /dev/null +++ b/src/sagittarius/flow_service_client_impl/flow_store.rs @@ -0,0 +1,100 @@ +//! KV-store operations backing the flows Aquila keeps synchronized with +//! Sagittarius: delete-by-id, wholesale replace, and single-flow upsert. +//! Every flow is keyed by [`get_flow_identifier`], so lookups by id require +//! a full key scan rather than a direct get. + +use futures::{StreamExt, TryStreamExt}; +use prost::Message; +use tucana::shared::{Flows, ValidationFlow}; + +use crate::flow::{get_flow_identifier, key_has_flow_id}; + +/// Deletes every stored key matching `flow_id`, returning how many were removed. +pub(super) async fn delete_flow(store: &async_nats::jetstream::kv::Store, flow_id: i64) -> usize { + let mut keys = match store.keys().await { + Ok(keys) => keys.boxed(), + Err(err) => { + log::error!( + "Failed to list stored flows for deletion flow_id={} error={:?}", + flow_id, + err + ); + return 0; + } + }; + + let mut deleted_count = 0; + while let Ok(Some(key)) = keys.try_next().await { + if !key_has_flow_id(&key, flow_id) { + continue; + } + + match store.delete(&key).await { + Ok(_) => deleted_count += 1, + Err(err) => log::error!( + "Failed to delete stored flow flow_id={} key={} error={:?}", + flow_id, + key, + err + ), + } + } + + deleted_count +} + +/// Stores a single flow update under its computed key, returning that key +/// for the caller's own logging/metrics. +pub(super) async fn store_flow( + store: &async_nats::jetstream::kv::Store, + flow: &ValidationFlow, +) -> (String, Result<(), async_nats::jetstream::kv::PutError>) { + let key = get_flow_identifier(flow); + let bytes = flow.encode_to_vec(); + let result = store.put(key.clone(), bytes.into()).await.map(|_| ()); + (key, result) +} + +/// Purges every currently stored flow and replaces it with `flows`. +/// Returns `(purged_count, stored_count)` so the caller can report metrics. +pub(super) async fn replace_all( + store: &async_nats::jetstream::kv::Store, + flows: Flows, +) -> (usize, usize) { + let mut keys = match store.keys().await { + Ok(keys) => keys.boxed(), + Err(err) => { + log::error!( + "Failed to list stored flows before replacement error={:?}", + err + ); + return (0, 0); + } + }; + + let mut purged_count = 0; + while let Ok(Some(key)) = keys.try_next().await { + match store.purge(&key).await { + Ok(_) => purged_count += 1, + Err(err) => log::error!("Failed to purge stored flow key={} error={}", key, err), + } + } + + let mut stored_count = 0; + for flow in flows.flows { + let (key, result) = store_flow(store, &flow).await; + match result { + Ok(()) => { + stored_count += 1; + log::debug!("Stored replacement flow key={}", key); + } + Err(err) => log::error!( + "Failed to store replacement flow key={} error={:?}", + key, + err + ), + } + } + + (purged_count, stored_count) +} diff --git a/src/sagittarius/flow_service_client_impl.rs b/src/sagittarius/flow_service_client_impl/mod.rs similarity index 52% rename from src/sagittarius/flow_service_client_impl.rs rename to src/sagittarius/flow_service_client_impl/mod.rs index e277d90..7ba3848 100644 --- a/src/sagittarius/flow_service_client_impl.rs +++ b/src/sagittarius/flow_service_client_impl/mod.rs @@ -1,22 +1,25 @@ -use crate::{ - authorization::authorization::get_authorization_metadata, - flow::{get_flow_identifier, key_has_flow_id}, - telemetry::metrics, -}; -use futures::{StreamExt, TryStreamExt}; -use prost::Message; -use std::{path::Path, sync::Arc}; -use tokio::fs; +//! Client for Sagittarius' flow synchronization stream: keeps Aquila's local +//! flow KV store in sync with whatever Sagittarius considers the current +//! set, and rebroadcasts action module configuration updates it receives +//! along the way. +//! +//! - [`flow_store`] applies the sync operations (delete/replace/update) to the KV store. +//! - [`dev_export`] mirrors the synced flows to a local JSON file, development only. + +mod dev_export; +mod flow_store; + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use futures::StreamExt; use tokio::sync::broadcast; use tonic::{Extensions, Request, transport::Channel}; -use tucana::{ - sagittarius::{ - FlowLogonRequest, FlowResponse, flow_response::Data, flow_service_client::FlowServiceClient, - }, - shared::Flows, +use tucana::sagittarius::{ + FlowLogonRequest, FlowResponse, flow_response::Data, flow_service_client::FlowServiceClient, }; -use std::sync::atomic::{AtomicBool, Ordering}; +use crate::{authorization::authorization::get_authorization_metadata, telemetry::metrics}; fn module_config_stats(configs: &tucana::shared::ModuleConfigurations) -> (usize, usize) { let project_count = configs.module_configurations.len(); @@ -35,6 +38,9 @@ pub struct SagittariusFlowClient { client: FlowServiceClient, env: String, token: String, + /// Flipped to `true` once the sync stream is established, so other + /// components can hold off on work that depends on Sagittarius state + /// actually being loaded. sagittarius_ready: Arc, action_config_tx: broadcast::Sender, } @@ -64,79 +70,7 @@ impl SagittariusFlowClient { self.env == "DEVELOPMENT" } - async fn export_flows_json_overwrite(&self, flows: Flows) { - if !self.is_development() { - return; - } - - let json = match serde_json::to_vec_pretty(&flows) { - Ok(b) => b, - Err(e) => { - log::error!( - "Failed to serialize development flow export flow_count={} error={:?}", - flows.flows.len(), - e - ); - return; - } - }; - - let final_path = Path::new("flowExport.json"); - let tmp_path = Path::new("flowExport.json.tmp"); - - if let Err(e) = fs::write(tmp_path, &json).await { - log::error!( - "Failed to write development flow export path={} error={}", - tmp_path.display(), - e - ); - return; - } - - if let Err(e) = fs::rename(tmp_path, final_path).await { - log::warn!( - "Could not atomically replace development flow export path={} error={}; retrying after removing destination", - final_path.display(), - e - ); - match fs::remove_file(final_path).await { - Ok(()) => log::debug!( - "Removed previous development flow export path={}", - final_path.display() - ), - Err(remove_error) if remove_error.kind() == std::io::ErrorKind::NotFound => {} - Err(remove_error) => log::warn!( - "Failed to remove previous development flow export path={} error={}", - final_path.display(), - remove_error - ), - } - if let Err(e2) = fs::rename(tmp_path, final_path).await { - log::error!( - "Failed to replace development flow export source_path={} destination_path={} initial_rename_error={} retry_error={}", - tmp_path.display(), - final_path.display(), - e, - e2 - ); - if let Err(cleanup_error) = fs::remove_file(tmp_path).await { - log::warn!( - "Failed to clean up temporary development flow export path={} error={}", - tmp_path.display(), - cleanup_error - ); - } - return; - } - } - - log::info!( - "Exported {} flows to {}", - flows.flows.len(), - final_path.display() - ); - } - + /// Applies one message from the flow sync stream to local state. async fn handle_response(&mut self, response: FlowResponse) { let data = match response.data { Some(data) => data, @@ -149,34 +83,7 @@ impl SagittariusFlowClient { match data { Data::DeletedFlowId(id) => { log::debug!("Applying flow deletion flow_id={}", id); - let mut keys = match self.store.keys().await { - Ok(keys) => keys.boxed(), - Err(err) => { - log::error!( - "Failed to list stored flows for deletion flow_id={} error={:?}", - id, - err - ); - return; - } - }; - - let mut deleted_count = 0; - while let Ok(Some(key)) = keys.try_next().await { - if !key_has_flow_id(&key, id) { - continue; - } - - match self.store.delete(&key).await { - Ok(_) => deleted_count += 1, - Err(err) => log::error!( - "Failed to delete stored flow flow_id={} key={} error={:?}", - id, - key, - err - ), - } - } + let deleted_count = flow_store::delete_flow(&self.store, id).await; if deleted_count == 0 { metrics::flow_operation("delete", "not_found", 1); @@ -191,11 +98,10 @@ impl SagittariusFlowClient { } } Data::UpdatedFlow(flow) => { - let key = get_flow_identifier(&flow); let flow_id = flow.flow_id; - let bytes = flow.encode_to_vec(); - match self.store.put(key.clone(), bytes.into()).await { - Ok(_) => { + let (key, result) = flow_store::store_flow(&self.store, &flow).await; + match result { + Ok(()) => { metrics::flow_operation("update", "success", 1); log::info!("Stored flow update flow_id={} key={}", flow_id, key) } @@ -217,45 +123,11 @@ impl SagittariusFlowClient { received_count ); - self.export_flows_json_overwrite(flows.clone()).await; - - let mut keys = match self.store.keys().await { - Ok(keys) => keys.boxed(), - Err(err) => { - log::error!( - "Failed to list stored flows before replacement error={:?}", - err - ); - return; - } - }; - - let mut purged_count = 0; - while let Ok(Some(key)) = keys.try_next().await { - match self.store.purge(&key).await { - Ok(_) => purged_count += 1, - Err(e) => { - log::error!("Failed to purge stored flow key={} error={}", key, e) - } - } + if self.is_development() { + dev_export::overwrite(flows.clone()).await; } - let mut stored_count = 0; - for flow in flows.flows { - let key = get_flow_identifier(&flow); - let bytes = flow.encode_to_vec(); - match self.store.put(key.clone(), bytes.into()).await { - Ok(_) => { - stored_count += 1; - log::debug!("Stored replacement flow key={}", key); - } - Err(err) => log::error!( - "Failed to store replacement flow key={} error={:?}", - key, - err - ), - }; - } + let (purged_count, stored_count) = flow_store::replace_all(&self.store, flows).await; log::info!( "Finished replacing stored flows received_count={} purged_count={} stored_count={}", received_count, @@ -291,6 +163,8 @@ impl SagittariusFlowClient { } } + /// Opens the flow sync stream and services it until it ends or errors, + /// at which point the caller is expected to reconnect. pub async fn init_flow_stream(&mut self) -> Result<(), tonic::Status> { self.sagittarius_ready.store(false, Ordering::SeqCst); diff --git a/src/sagittarius/mod.rs b/src/sagittarius/mod.rs index 0e5e5a3..ae86b39 100644 --- a/src/sagittarius/mod.rs +++ b/src/sagittarius/mod.rs @@ -1,3 +1,8 @@ +//! Clients Aquila uses to talk *to* Sagittarius: flow synchronization, +//! module registration, runtime status forwarding, and the test/live +//! execution stream. See [`retry`] for the shared reconnect-with-backoff +//! logic every long-lived stream client here is built on. + pub mod flow_service_client_impl; pub mod module_service_client_impl; pub mod retry; diff --git a/src/sagittarius/module_service_client_impl.rs b/src/sagittarius/module_service_client_impl.rs index f730f07..c2b930f 100644 --- a/src/sagittarius/module_service_client_impl.rs +++ b/src/sagittarius/module_service_client_impl.rs @@ -1,3 +1,7 @@ +//! Client for forwarding module registrations (from connected actions and +//! runtimes) to Sagittarius, so Sagittarius knows what's available to +//! reference from a flow. + use crate::configuration::service::ServiceConfiguration; use crate::{authorization::authorization::get_authorization_metadata, telemetry::errors}; use std::time::Duration; @@ -33,6 +37,9 @@ impl SagittariusModuleServiceClient { skip_all, fields(rpc.system = "grpc", rpc.service = "ModuleService", rpc.method = "Update") )] + /// Forwards a module update to Sagittarius, alongside the full list of + /// definition sources this Aquila instance currently has available — + /// Sagittarius needs that list to validate references in flows. pub async fn update_modules( &mut self, modules_update_request: tucana::aquila::ModuleUpdateRequest, diff --git a/src/sagittarius/retry.rs b/src/sagittarius/retry.rs index bc3eb93..6343c28 100644 --- a/src/sagittarius/retry.rs +++ b/src/sagittarius/retry.rs @@ -1,3 +1,8 @@ +//! Shared connection logic for every long-lived Sagittarius client: dial +//! with exponential backoff, and flip a shared readiness flag as the +//! connection comes up or drops so [`crate::server::create_readiness_interceptor`] +//! can reject gRPC traffic while Sagittarius is unreachable. + use std::sync::{ Arc, atomic::{AtomicBool, Ordering}, @@ -9,7 +14,14 @@ use tonic::transport::{Channel, Endpoint}; const MAX_BACKOFF: u64 = 2000 * 60; const MAX_RETRIES: i8 = 10; -// Will create a channel and retry if its not possible +/// Connects to `url`, retrying with exponential backoff (capped at +/// [`MAX_BACKOFF`] ms) up to [`MAX_RETRIES`] times before giving up. +/// `ready` is cleared before each attempt and set once connected, so callers +/// can gate readiness on it without polling the channel state directly. +/// +/// Panics rather than returning an error on exhaustion or a malformed URL, +/// since none of Aquila's callers have a meaningful way to run without a +/// Sagittarius connection. pub async fn create_channel_with_retry( channel_name: &str, url: String, diff --git a/src/sagittarius/runtime_status_service_client_impl.rs b/src/sagittarius/runtime_status_service_client_impl.rs index 3f2d5c3..03e4766 100644 --- a/src/sagittarius/runtime_status_service_client_impl.rs +++ b/src/sagittarius/runtime_status_service_client_impl.rs @@ -1,3 +1,8 @@ +//! Client for forwarding a runtime's status heartbeats to Sagittarius. +//! Currently unused end-to-end — see the disabled forwarding call in +//! `server::runtime_status_service_server_impl`, tracked in #360 — but kept +//! ready for when that's re-enabled. + use crate::{authorization::authorization::get_authorization_metadata, telemetry::errors}; use std::time::Duration; use tonic::{Extensions, Request, transport::Channel}; diff --git a/src/sagittarius/test_execution_client_impl/flow_id_cache.rs b/src/sagittarius/test_execution_client_impl/flow_id_cache.rs new file mode 100644 index 0000000..0fc28a1 --- /dev/null +++ b/src/sagittarius/test_execution_client_impl/flow_id_cache.rs @@ -0,0 +1,187 @@ +//! A bounded, TTL-expiring cache from execution id to flow id. +//! +//! Sagittarius execution results don't always carry a `flow_id` back (some +//! runtimes only echo the execution id), so Aquila remembers the mapping it +//! made when it first dispatched the execution and re-attaches the flow id +//! when the result comes back. The cache is bounded and TTL'd purely as a +//! safety net against executions that never send a result — nothing here +//! bounds it in the happy path, since [`ExecutionFlowIdCache::take`] removes +//! the entry as soon as it's used. + +use std::{ + collections::HashMap, + sync::Arc, + time::{Duration, Instant}, +}; + +use tokio::sync::Mutex; + +/// How long a mapping is kept if its execution never sends a result. +const EXECUTION_FLOW_ID_TTL: Duration = Duration::from_secs(30 * 60); +/// Upper bound on cache size, enforced by evicting the entry closest to +/// expiry when a new one would push the cache over the limit. +const MAX_EXECUTION_FLOW_IDS: usize = 10_000; + +struct ExecutionFlowIdMapping { + flow_id: i64, + expires_at: Instant, +} + +#[derive(Clone, Default)] +pub(super) struct ExecutionFlowIdCache { + entries: Arc>>, +} + +impl ExecutionFlowIdCache { + /// Remembers `flow_id` under `execution_id` for [`EXECUTION_FLOW_ID_TTL`]. + /// A no-op for an empty `execution_id`, since that can never be looked + /// back up by [`take`](Self::take). + pub(super) async fn remember(&self, execution_id: &str, flow_id: i64) { + if execution_id.is_empty() { + log::warn!("Cannot remember execution flow_id because execution_id is empty"); + return; + } + + let mut entries = self.entries.lock().await; + let now = Instant::now(); + let expired_count = prune_expired(&mut entries, now); + let replacing_existing = entries.contains_key(execution_id); + let evicted_execution_id = if !replacing_existing && entries.len() >= MAX_EXECUTION_FLOW_IDS + { + remove_soonest_to_expire(&mut entries) + } else { + None + }; + + entries.insert( + execution_id.to_string(), + ExecutionFlowIdMapping { + flow_id, + expires_at: now + EXECUTION_FLOW_ID_TTL, + }, + ); + + if let Some(evicted_execution_id) = evicted_execution_id { + log::warn!( + "Evicted execution flow mapping because cache is full evicted_execution_id={} max_entries={}", + evicted_execution_id, + MAX_EXECUTION_FLOW_IDS + ); + } + log::debug!( + "Remembered execution flow mapping execution_id={} flow_id={} cached_entries={} expired_entries={}", + execution_id, + flow_id, + entries.len(), + expired_count + ); + } + + /// Drops the mapping for `execution_id` without returning it, used once + /// a result has arrived carrying its own `flow_id` and the cached one is + /// no longer needed. + pub(super) async fn forget(&self, execution_id: &str) { + if execution_id.is_empty() { + return; + } + + let mut entries = self.entries.lock().await; + let removed = entries.remove(execution_id).is_some(); + log::debug!( + "Forgot execution flow mapping execution_id={} removed={}", + execution_id, + removed + ); + } + + /// Removes and returns the flow id for `execution_id`, unless it has + /// already expired. + pub(super) async fn take(&self, execution_id: &str) -> Option { + if execution_id.is_empty() { + return None; + } + + let mut entries = self.entries.lock().await; + let mapping = entries.remove(execution_id)?; + if mapping.expires_at > Instant::now() { + Some(mapping.flow_id) + } else { + log::debug!( + "Dropped expired execution flow mapping execution_id={}", + execution_id + ); + None + } + } +} + +fn prune_expired(entries: &mut HashMap, now: Instant) -> usize { + let initial_len = entries.len(); + entries.retain(|_, mapping| mapping.expires_at > now); + initial_len - entries.len() +} + +/// Evicts whichever entry is closest to expiring, since that's the best +/// approximation of "oldest" without tracking insertion order separately. +fn remove_soonest_to_expire(entries: &mut HashMap) -> Option { + let soonest_execution_id = entries + .iter() + .min_by_key(|(_, mapping)| mapping.expires_at) + .map(|(execution_id, _)| execution_id.clone())?; + + entries.remove(&soonest_execution_id); + Some(soonest_execution_id) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn take_returns_and_removes_the_mapping() { + futures::executor::block_on(async { + let cache = ExecutionFlowIdCache::default(); + cache.remember("execution-1", 42).await; + + assert_eq!(cache.take("execution-1").await, Some(42)); + assert_eq!(cache.take("execution-1").await, None); + }); + } + + #[test] + fn forget_removes_without_returning() { + futures::executor::block_on(async { + let cache = ExecutionFlowIdCache::default(); + cache.remember("execution-1", 42).await; + cache.forget("execution-1").await; + + assert_eq!(cache.take("execution-1").await, None); + }); + } + + #[test] + fn empty_execution_id_is_ignored() { + futures::executor::block_on(async { + let cache = ExecutionFlowIdCache::default(); + cache.remember("", 42).await; + + assert_eq!(cache.take("").await, None); + }); + } + + #[test] + fn take_drops_expired_mappings() { + futures::executor::block_on(async { + let cache = ExecutionFlowIdCache::default(); + cache.remember("execution-1", 42).await; + + { + let mut entries = cache.entries.lock().await; + entries.get_mut("execution-1").unwrap().expires_at = + Instant::now() - Duration::from_secs(1); + } + + assert_eq!(cache.take("execution-1").await, None); + }); + } +} diff --git a/src/sagittarius/test_execution_client_impl.rs b/src/sagittarius/test_execution_client_impl/mod.rs similarity index 53% rename from src/sagittarius/test_execution_client_impl.rs rename to src/sagittarius/test_execution_client_impl/mod.rs index c204e68..aae3de9 100644 --- a/src/sagittarius/test_execution_client_impl.rs +++ b/src/sagittarius/test_execution_client_impl/mod.rs @@ -1,254 +1,35 @@ -/* - Why is Aquila a client when Sagittarius wants a result of Aquila? +//! Client for Sagittarius' `test` execution stream. +//! +//! Why is Aquila a client here when Sagittarius wants a result *from* +//! Aquila? In some deployments Sagittarius can't open a connection to +//! Aquila, so Aquila dials out and sends a `Logon` request first — +//! establishing the connection from Aquila's side even though Sagittarius +//! is the one issuing execution requests on it. +//! +//! - [`response_sender`] owns the outgoing half of the stream and is shared +//! with every task that needs to report an execution result. +//! - [`flow_id_cache`] backs the piece of that sender that recovers a +//! result's `flow_id` when a runtime doesn't echo it. + +mod flow_id_cache; +mod response_sender; + +pub use response_sender::SagittariusExecutionResponseSender; + +use std::sync::Arc; - In some conditions Sagittarius can't connect to Aquila - Thus Aquila sends a `Logon` request to connect to Sagittarius establishing the connection -*/ use futures::{StreamExt, TryStreamExt}; use prost::Message; -use std::{ - collections::HashMap, - sync::Arc, - time::{Duration, Instant}, -}; -use tokio::sync::Mutex; use tokio_stream::wrappers::ReceiverStream; use tonic::transport::Channel; -use tonic::{Extensions, Request, Status}; +use tonic::{Extensions, Request}; use tucana::sagittarius::execution_logon_request::Data; use tucana::sagittarius::execution_service_client::ExecutionServiceClient; use tucana::sagittarius::{ExecutionLogonRequest, Logon}; -use tucana::shared::{ExecutionFlow, ExecutionResult, ValidationFlow}; +use tucana::shared::{ExecutionFlow, ValidationFlow}; use crate::{authorization::authorization::get_authorization_metadata, flow::key_has_flow_id}; -const EXECUTION_FLOW_ID_TTL: Duration = Duration::from_secs(30 * 60); -const MAX_EXECUTION_FLOW_IDS: usize = 10_000; - -struct ExecutionFlowIdMapping { - flow_id: i64, - expires_at: Instant, -} - -#[derive(Clone, Default)] -pub struct SagittariusExecutionResponseSender { - sender: Arc>>>, - execution_flow_ids: Arc>>, -} - -impl SagittariusExecutionResponseSender { - pub fn new() -> Self { - Self::default() - } - - async fn attach(&self, sender: tokio::sync::mpsc::Sender) { - let mut current = self.sender.lock().await; - let replacing_existing = current.is_some(); - *current = Some(sender); - log::debug!( - "Attached Sagittarius execution response sender replacing_existing={}", - replacing_existing - ); - } - - async fn clear(&self) { - let mut current = self.sender.lock().await; - let had_sender = current.is_some(); - *current = None; - log::debug!( - "Cleared Sagittarius execution response sender had_sender={}", - had_sender - ); - } - - async fn remember_execution_flow(&self, execution_id: &str, flow_id: i64) { - if execution_id.is_empty() { - log::warn!("Cannot remember execution flow_id because execution_id is empty"); - return; - } - - let mut execution_flow_ids = self.execution_flow_ids.lock().await; - let now = Instant::now(); - let expired_count = prune_expired_execution_flow_ids(&mut execution_flow_ids, now); - let replacing_existing = execution_flow_ids.contains_key(execution_id); - let evicted_execution_id = - if !replacing_existing && execution_flow_ids.len() >= MAX_EXECUTION_FLOW_IDS { - remove_oldest_execution_flow_id(&mut execution_flow_ids) - } else { - None - }; - - execution_flow_ids.insert( - execution_id.to_string(), - ExecutionFlowIdMapping { - flow_id, - expires_at: now + EXECUTION_FLOW_ID_TTL, - }, - ); - - if let Some(evicted_execution_id) = evicted_execution_id { - log::warn!( - "Evicted execution flow mapping because cache is full evicted_execution_id={} max_entries={}", - evicted_execution_id, - MAX_EXECUTION_FLOW_IDS - ); - } - log::debug!( - "Remembered execution flow mapping execution_id={} flow_id={} cached_entries={} expired_entries={}", - execution_id, - flow_id, - execution_flow_ids.len(), - expired_count - ); - } - - async fn forget_execution_flow(&self, execution_id: &str) { - if execution_id.is_empty() { - return; - } - - let mut execution_flow_ids = self.execution_flow_ids.lock().await; - let removed = execution_flow_ids.remove(execution_id).is_some(); - log::debug!( - "Forgot execution flow mapping execution_id={} removed={}", - execution_id, - removed - ); - } - - async fn take_execution_flow_id(&self, execution_id: &str) -> Option { - if execution_id.is_empty() { - return None; - } - - let mut execution_flow_ids = self.execution_flow_ids.lock().await; - let mapping = execution_flow_ids.remove(execution_id)?; - if mapping.expires_at > Instant::now() { - Some(mapping.flow_id) - } else { - log::debug!( - "Dropped expired execution flow mapping execution_id={}", - execution_id - ); - None - } - } - - pub async fn send_execution_result( - &self, - mut execution_result: ExecutionResult, - ) -> Result { - let execution_id = execution_result.execution_identifier.clone(); - - if execution_result.flow_id == 0 { - match self.take_execution_flow_id(&execution_id).await { - Some(flow_id) if flow_id != 0 => { - log::warn!( - "Filled missing execution result flow_id from Aquila mapping execution_id={} flow_id={}", - execution_id, - flow_id - ); - execution_result.flow_id = flow_id; - } - _ => { - log::warn!( - "Execution result has flow_id=0 and no Aquila mapping execution_id={}", - execution_id - ); - } - } - } else { - self.forget_execution_flow(&execution_id).await; - } - - let flow_id = execution_result.flow_id; - let node_result_count = execution_result.node_execution_results.len(); - let result_status = execution_result_status(&execution_result); - - log::debug!( - "Queueing execution result for Sagittarius stream execution_id={} flow_id={} result_status={} node_results={}", - execution_id, - flow_id, - result_status, - node_result_count - ); - - let sender = { - let current = self.sender.lock().await; - current.clone() - }; - - let Some(sender) = sender else { - log::error!( - "Cannot queue execution result for Sagittarius stream reason=stream_not_connected execution_id={} flow_id={} result_status={}", - execution_id, - flow_id, - result_status - ); - return Err(Status::unavailable( - "sagittarius execution stream is not connected", - )); - }; - - let remaining_capacity = sender.capacity(); - match sender - .send(ExecutionLogonRequest { - data: Some(Data::Response(execution_result)), - }) - .await - { - Ok(()) => { - log::debug!( - "Queued execution result for Sagittarius stream execution_id={} flow_id={} remaining_capacity_before_send={}", - execution_id, - flow_id, - remaining_capacity - ); - Ok(flow_id) - } - Err(_) => { - log::error!( - "Cannot queue execution result for Sagittarius stream reason=stream_closed execution_id={} flow_id={}", - execution_id, - flow_id - ); - Err(Status::unavailable( - "sagittarius execution stream is closed", - )) - } - } - } -} - -fn prune_expired_execution_flow_ids( - execution_flow_ids: &mut HashMap, - now: Instant, -) -> usize { - let initial_len = execution_flow_ids.len(); - execution_flow_ids.retain(|_, mapping| mapping.expires_at > now); - initial_len - execution_flow_ids.len() -} - -fn remove_oldest_execution_flow_id( - execution_flow_ids: &mut HashMap, -) -> Option { - let oldest_execution_id = execution_flow_ids - .iter() - .min_by_key(|(_, mapping)| mapping.expires_at) - .map(|(execution_id, _)| execution_id.clone())?; - - execution_flow_ids.remove(&oldest_execution_id); - Some(oldest_execution_id) -} - -fn execution_result_status(execution_result: &ExecutionResult) -> &'static str { - match execution_result.result.as_ref() { - Some(tucana::shared::execution_result::Result::Success(_)) => "success", - Some(tucana::shared::execution_result::Result::Error(_)) => "error", - None => "missing", - } -} - pub struct SagittariusTestExecutionServiceClient { nats_client: async_nats::Client, store: Arc, @@ -275,6 +56,9 @@ impl SagittariusTestExecutionServiceClient { } } + /// Scans the flow KV bucket for the entry whose key encodes `flow_id`. + /// There is no secondary index from flow id to key, so this is a linear + /// scan of every stored flow. async fn load_validation_flow(&self, flow_id: i64) -> Option { let mut keys = match self.store.keys().await { Ok(keys) => keys, @@ -354,6 +138,12 @@ impl SagittariusTestExecutionServiceClient { } } + /// Opens the Sagittarius execution stream and services it until it ends. + /// + /// The stream's outgoing half is driven by an mpsc channel rather than a + /// plain async generator so that [`SagittariusExecutionResponseSender`] + /// can push execution results onto it from other tasks while this loop + /// is busy reading incoming requests. pub async fn logon(&mut self) { let (tx, rx) = tokio::sync::mpsc::channel::(10000); let logon = ExecutionLogonRequest { diff --git a/src/sagittarius/test_execution_client_impl/response_sender.rs b/src/sagittarius/test_execution_client_impl/response_sender.rs new file mode 100644 index 0000000..b5e4eb6 --- /dev/null +++ b/src/sagittarius/test_execution_client_impl/response_sender.rs @@ -0,0 +1,155 @@ +//! Queues execution results onto the outgoing half of the Sagittarius +//! execution logon stream, and fills in a result's `flow_id` from the +//! [`ExecutionFlowIdCache`] when a runtime didn't echo it back. + +use std::sync::Arc; + +use tokio::sync::Mutex; +use tonic::Status; +use tucana::sagittarius::execution_logon_request::Data; +use tucana::sagittarius::ExecutionLogonRequest; +use tucana::shared::ExecutionResult; + +use super::flow_id_cache::ExecutionFlowIdCache; + +/// Handle shared between the task driving the Sagittarius execution stream +/// and every task that needs to report an execution result back to it. +/// +/// The stream sender is `None` whenever no stream is currently connected +/// (before the first logon, or after a disconnect), in which case results +/// are rejected rather than buffered. +#[derive(Clone, Default)] +pub struct SagittariusExecutionResponseSender { + sender: Arc>>>, + execution_flow_ids: ExecutionFlowIdCache, +} + +impl SagittariusExecutionResponseSender { + pub fn new() -> Self { + Self::default() + } + + pub(super) async fn attach(&self, sender: tokio::sync::mpsc::Sender) { + let mut current = self.sender.lock().await; + let replacing_existing = current.is_some(); + *current = Some(sender); + log::debug!( + "Attached Sagittarius execution response sender replacing_existing={}", + replacing_existing + ); + } + + pub(super) async fn clear(&self) { + let mut current = self.sender.lock().await; + let had_sender = current.is_some(); + *current = None; + log::debug!( + "Cleared Sagittarius execution response sender had_sender={}", + had_sender + ); + } + + pub(super) async fn remember_execution_flow(&self, execution_id: &str, flow_id: i64) { + self.execution_flow_ids.remember(execution_id, flow_id).await; + } + + pub(super) async fn forget_execution_flow(&self, execution_id: &str) { + self.execution_flow_ids.forget(execution_id).await; + } + + /// Sends `execution_result` to Sagittarius over the attached stream, + /// recovering its `flow_id` from the cache first if the result didn't + /// carry one. + pub async fn send_execution_result( + &self, + mut execution_result: ExecutionResult, + ) -> Result { + let execution_id = execution_result.execution_identifier.clone(); + + if execution_result.flow_id == 0 { + match self.execution_flow_ids.take(&execution_id).await { + Some(flow_id) if flow_id != 0 => { + log::warn!( + "Filled missing execution result flow_id from Aquila mapping execution_id={} flow_id={}", + execution_id, + flow_id + ); + execution_result.flow_id = flow_id; + } + _ => { + log::warn!( + "Execution result has flow_id=0 and no Aquila mapping execution_id={}", + execution_id + ); + } + } + } else { + self.forget_execution_flow(&execution_id).await; + } + + let flow_id = execution_result.flow_id; + let node_result_count = execution_result.node_execution_results.len(); + let result_status = execution_result_status(&execution_result); + + log::debug!( + "Queueing execution result for Sagittarius stream execution_id={} flow_id={} result_status={} node_results={}", + execution_id, + flow_id, + result_status, + node_result_count + ); + + let sender = { + let current = self.sender.lock().await; + current.clone() + }; + + let Some(sender) = sender else { + log::error!( + "Cannot queue execution result for Sagittarius stream reason=stream_not_connected execution_id={} flow_id={} result_status={}", + execution_id, + flow_id, + result_status + ); + return Err(Status::unavailable( + "sagittarius execution stream is not connected", + )); + }; + + let remaining_capacity = sender.capacity(); + match sender + .send(ExecutionLogonRequest { + data: Some(Data::Response(execution_result)), + }) + .await + { + Ok(()) => { + log::debug!( + "Queued execution result for Sagittarius stream execution_id={} flow_id={} remaining_capacity_before_send={}", + execution_id, + flow_id, + remaining_capacity + ); + Ok(flow_id) + } + Err(_) => { + log::error!( + "Cannot queue execution result for Sagittarius stream reason=stream_closed execution_id={} flow_id={}", + execution_id, + flow_id + ); + Err(Status::unavailable( + "sagittarius execution stream is closed", + )) + } + } + } +} + +fn execution_result_status(execution_result: &ExecutionResult) -> &'static str { + match execution_result.result.as_ref() { + Some(tucana::shared::execution_result::Result::Success(_)) => "success", + Some(tucana::shared::execution_result::Result::Error(_)) => "error", + None => "missing", + } +} diff --git a/src/server/action_transfer/logon.rs b/src/server/action_transfer/logon.rs new file mode 100644 index 0000000..2335969 --- /dev/null +++ b/src/server/action_transfer/logon.rs @@ -0,0 +1,301 @@ +//! Handles the first message of an action's transfer stream: authenticating +//! the token, registering the action's module with Sagittarius, and wiring +//! up the NATS subscriptions that feed the rest of the stream. + +use std::sync::Arc; + +use tokio::sync::Mutex; +use tonic::Status; +use tucana::aquila::{ActionLogon, ActionTransferResponse}; + +use crate::{ + configuration::service::ServiceConfiguration, + sagittarius::module_service_client_impl::SagittariusModuleServiceClient, + telemetry::{errors, metrics}, +}; + +use super::{nats_bridge::forward_nats_to_action, pending_replies::PendingReplyStore}; + +/// Extracts the bearer token from gRPC metadata. +pub(super) fn extract_token( + request: &tonic::Request>, +) -> Result { + log::debug!("Extracting authorization token from metadata"); + match request.metadata().get("authorization") { + Some(ascii) => match ascii.to_str() { + Ok(tk) => { + if tk.is_empty() { + log::error!("Authorization token is empty"); + return Err(Status::unauthenticated("authorization token is empty")); + } + + Ok(tk.to_string()) + } + Err(err) => { + log::error!("Cannot read authorization header because: {:?}", err); + Err(Status::unauthenticated("invalid authorization header")) + } + }, + None => { + log::error!("Missing authorization token"); + Err(Status::unauthenticated("missing authorization token")) + } + } +} + +/// Whether a broadcasted config update is meant for `action_identifier`, since +/// [`spawn_cfg_forwarder`] subscribes to a single broadcast channel shared by +/// every connected action. +fn applies_to_action(configs: &tucana::shared::ModuleConfigurations, action_identifier: &str) -> bool { + configs.module_identifier == action_identifier +} + +/// Rewrites every definition's `definition_source` on the module an action +/// logs on with, so downstream consumers can tell it came from this action +/// rather than from whatever source the action's module definition was +/// authored against. +fn overwrite_module_definition_sources(module: &mut tucana::shared::Module, action_identifier: &str) { + let source = format!("action.{}", action_identifier); + + for flow_type in &mut module.flow_types { + flow_type.definition_source = Some(source.clone()); + } + for runtime_flow_type in &mut module.runtime_flow_types { + runtime_flow_type.definition_source = Some(source.clone()); + } + for function_definition in &mut module.function_definitions { + function_definition.definition_source = source.clone(); + } + for runtime_function_definition in &mut module.runtime_function_definitions { + runtime_function_definition.definition_source = source.clone(); + } + for definition_data_type in &mut module.definition_data_types { + definition_data_type.definition_source = source.clone(); + } +} + +/// Validates the logon request, starts NATS + config forwarders, and returns the accepted logon. +#[allow(clippy::too_many_arguments)] +pub(super) async fn handle_logon( + token: &str, + mut action_logon: ActionLogon, + actions: Arc>, + module_service: Option>>, + client: async_nats::Client, + cfg_tx: tokio::sync::broadcast::Sender, + tx: tokio::sync::mpsc::Sender>, + pending_replies: PendingReplyStore, + cfg_forwarder_started: &mut bool, +) -> Result { + let module = match action_logon.module.as_mut() { + Some(m) => m, + None => { + log::warn!("Rejected action logon reason=missing_module"); + return Err(Status::aborted("Please provide a module configuration.")); + } + }; + let identifier = module.identifier.clone(); + log::info!("Action logon attempt identifier={}", identifier); + + { + let lock = actions.lock().await; + if !lock.has_action(&token.to_string(), &identifier) { + metrics::action_connection(&identifier, "rejected"); + metrics::action_failure(&identifier, "authentication"); + log::warn!( + "Rejected action logon identifier={} reason=token_not_registered", + identifier + ); + return Err(Status::unauthenticated( + "token not matching to action identifier", + )); + } + } + + overwrite_module_definition_sources(module, &identifier); + + if let Some(module_service) = module_service { + let mut client = module_service.lock().await; + let response = client + .update_modules(tucana::aquila::ModuleUpdateRequest { + modules: vec![module.clone()], + }) + .await; + + if !response.success { + metrics::action_connection(&identifier, "rejected"); + metrics::action_failure(&identifier, "module_update"); + errors::record_message( + "dependency", + "action.logon", + "Sagittarius rejected the action module update", + format!("action.identifier={identifier}"), + ); + return Err(Status::internal( + "could not update action module via Sagittarius", + )); + } + } + + log::debug!("Action connected identifier={}", identifier); + + let sub = match client.subscribe(format!("action.{}.*", identifier)).await { + Ok(s) => s, + Err(err) => { + metrics::action_connection(&identifier, "rejected"); + metrics::action_failure(&identifier, "subscription"); + errors::record( + "messaging", + "action.subscribe", + &err, + format!("action.identifier={identifier} subject=action.{identifier}.*"), + ); + return Err(Status::internal( + "could not register action into execution loop", + )); + } + }; + + if let Err(err) = client.flush().await { + metrics::action_connection(&identifier, "rejected"); + metrics::action_failure(&identifier, "subscription_flush"); + errors::record( + "messaging", + "action.subscribe.flush", + &err, + format!("action.identifier={identifier}"), + ); + return Err(Status::internal( + "could not register action subscription with NATS", + )); + } + + log::debug!("Subscribed to action subject action.{}.*", identifier); + + let tx_clone = tx.clone(); + let pending_replies_clone = pending_replies.clone(); + let forwarder_identifier = identifier.clone(); + tokio::spawn(async move { + forward_nats_to_action(forwarder_identifier, sub, tx_clone, pending_replies_clone).await; + }); + + // A logon is only the first message on the stream, but `handle_logon` can't + // assume it's only ever called once per stream, so the caller-owned flag + // guards against starting a second, redundant forwarder task. + if !*cfg_forwarder_started { + *cfg_forwarder_started = true; + log::debug!("Starting config forwarder action={}", identifier); + spawn_cfg_forwarder(identifier.clone(), cfg_tx, tx.clone()); + } + + Ok(action_logon) +} + +/// Forwards config updates for the given action identifier to the gRPC stream. +pub(super) fn spawn_cfg_forwarder( + action_identifier: String, + cfg_tx: tokio::sync::broadcast::Sender, + tx: tokio::sync::mpsc::Sender>, +) { + let mut cfg_rx = cfg_tx.subscribe(); + tokio::spawn(async move { + while let Ok(cfgs) = cfg_rx.recv().await { + if !applies_to_action(&cfgs, &action_identifier) { + log::debug!( + "Config update does not apply to action {}", + action_identifier + ); + continue; + } + + log::debug!("Forwarding config update to action {}", action_identifier); + let resp = ActionTransferResponse { + data: Some( + tucana::aquila::action_transfer_response::Data::ModuleConfigurations(cfgs), + ), + }; + + if tx.send(Ok(resp)).await.is_err() { + metrics::action_config_update(&action_identifier, "failed"); + metrics::action_failure(&action_identifier, "configuration_forward"); + log::debug!("Config forwarder channel closed for {}", action_identifier); + break; + } + metrics::action_config_update(&action_identifier, "success"); + } + + log::debug!("Config forwarder stopped for {}", action_identifier); + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn module_configurations_apply_by_module_identifier() { + let configs = tucana::shared::ModuleConfigurations { + module_identifier: "gls-action".to_string(), + module_configurations: vec![tucana::shared::ModuleProjectConfigurations { + project_id: 1, + module_configurations: vec![tucana::shared::ModuleConfiguration { + identifier: "username".to_string(), + value: None, + }], + }], + }; + + assert!(applies_to_action(&configs, "gls-action")); + assert!(!applies_to_action(&configs, "another-action")); + } + + #[test] + fn overwrite_module_definition_sources_uses_action_source() { + let mut module = tucana::shared::Module { + flow_types: vec![tucana::shared::FlowType { + definition_source: Some("module.old".to_string()), + ..Default::default() + }], + runtime_flow_types: vec![tucana::shared::RuntimeFlowType { + definition_source: Some("module.old".to_string()), + ..Default::default() + }], + function_definitions: vec![tucana::shared::FunctionDefinition { + definition_source: "module.old".to_string(), + ..Default::default() + }], + runtime_function_definitions: vec![tucana::shared::RuntimeFunctionDefinition { + definition_source: "module.old".to_string(), + ..Default::default() + }], + definition_data_types: vec![tucana::shared::DefinitionDataType { + definition_source: "module.old".to_string(), + ..Default::default() + }], + ..Default::default() + }; + + overwrite_module_definition_sources(&mut module, "send-email"); + + assert_eq!( + module.flow_types[0].definition_source.as_deref(), + Some("action.send-email") + ); + assert_eq!( + module.runtime_flow_types[0].definition_source.as_deref(), + Some("action.send-email") + ); + assert_eq!( + module.function_definitions[0].definition_source, + "action.send-email" + ); + assert_eq!( + module.runtime_function_definitions[0].definition_source, + "action.send-email" + ); + assert_eq!( + module.definition_data_types[0].definition_source, + "action.send-email" + ); + } +} diff --git a/src/server/action_transfer/mod.rs b/src/server/action_transfer/mod.rs new file mode 100644 index 0000000..d09af65 --- /dev/null +++ b/src/server/action_transfer/mod.rs @@ -0,0 +1,291 @@ +//! gRPC server for the bidirectional `ActionTransfer` stream: authenticates +//! an action's logon, then bridges execution requests/results and config +//! updates between NATS and the connected action for the lifetime of the +//! stream. +//! +//! The module is split by responsibility: +//! - [`logon`] validates the initial logon message and registers the action. +//! - [`nats_bridge`] moves execution requests/results between NATS and gRPC. +//! - [`pending_replies`] tracks which NATS reply subject an execution result belongs to. + +mod logon; +mod nats_bridge; +mod pending_replies; + +use std::{pin::Pin, sync::Arc}; + +use futures::StreamExt; +use futures_core::Stream; +use tokio::sync::Mutex; +use tokio_stream::wrappers::ReceiverStream; +use tonic::Status; +use tracing::Instrument; +use tucana::aquila::{ + ActionTransferRequest, ActionTransferResponse, action_transfer_service_server::ActionTransferService, +}; + +use crate::{ + configuration::service::ServiceConfiguration, + sagittarius::module_service_client_impl::SagittariusModuleServiceClient, + telemetry::metrics, +}; + +use logon::{extract_token, handle_logon}; +use nats_bridge::{handle_event, handle_result, send_stream_error}; +use pending_replies::PendingReplyStore; + +/// Implements the `ActionTransfer` gRPC service that a connected action +/// speaks to for the lifetime of its bidirectional stream. +/// +/// One instance is shared across all connected actions; per-connection state +/// (pending replies, logon status, ...) lives inside the `transfer` task +/// spawned for each stream instead of on `self`. +pub struct AquilaActionTransferServiceServer { + client: async_nats::Client, + kv: async_nats::jetstream::kv::Store, + /// Static, pre-provisioned action tokens/configuration loaded at startup. + actions: ServiceConfiguration, + /// Present only in dynamic mode, where module updates must be relayed to Sagittarius. + module_service: Option>>, + /// Broadcasts module configuration updates to every connected action's config forwarder. + action_config_tx: tokio::sync::broadcast::Sender, + /// Whether Aquila is running in static mode, which changes how config updates are sourced. + is_static: bool, +} + +impl AquilaActionTransferServiceServer { + pub fn new( + client: async_nats::Client, + kv: async_nats::jetstream::kv::Store, + actions: ServiceConfiguration, + module_service: Option>>, + action_config_tx: tokio::sync::broadcast::Sender, + is_static: bool, + ) -> Self { + Self { + client, + kv, + actions, + module_service, + action_config_tx, + is_static, + } + } +} + +#[tonic::async_trait] +impl ActionTransferService for AquilaActionTransferServiceServer { + type TransferStream = + Pin> + Send + 'static>>; + + /// Opens the bidirectional stream an action uses to log on and then + /// exchange events, execution requests, and results with Aquila. + /// + /// The gRPC protocol for this stream is a strict handshake: the first + /// message must be a [`Logon`](tucana::aquila::action_transfer_request::Data::Logon), + /// every message after that must not be, and everything is driven from a + /// single spawned task so the stream can be read and written to concurrently. + #[tracing::instrument( + name = "aquila.action.transfer", + skip_all, + fields(rpc.system = "grpc", rpc.service = "ActionTransferService", rpc.method = "Transfer") + )] + async fn transfer( + &self, + request: tonic::Request>, + ) -> std::result::Result, tonic::Status> { + let token = extract_token(&request)?; + log::debug!("Action transfer stream opened"); + + let mut first_request = true; + let mut action_props: Option = None; + let mut stream = request.into_inner(); + + let actions = Arc::new(Mutex::new(self.actions.clone())); + let kv = self.kv.clone(); + let client = self.client.clone(); + let module_service = self.module_service.clone(); + let cfg_tx = self.action_config_tx.clone(); + let is_static = self.is_static; + let pending_replies = PendingReplyStore::new(); + + let (tx, rx) = + tokio::sync::mpsc::channel::>(32); + + let stream_span = tracing::info_span!( + "aquila.action.stream", + action.identifier = tracing::field::Empty + ); + tokio::spawn(async move { + let mut cfg_forwarder_started = false; + let mut connected_at = None; + let mut connected_identifier = None; + log::debug!("Action transfer stream started"); + + while let Some(next) = stream.next().await { + let transfer_request = match next { + Ok(tr) => tr, + Err(status) => { + log::warn!("Action transfer input stream failed status={:?}", status); + break; + } + }; + + let data = match transfer_request.data { + Some(d) => d, + None => { + log::warn!("Received empty action transfer request"); + continue; + } + }; + + if first_request { + first_request = false; + + match data { + tucana::aquila::action_transfer_request::Data::Logon(action_logon) => { + let identifier = match action_logon.module { + Some(ref m) => m.identifier.clone(), + None => { + log::warn!("Rejected action logon reason=missing_module"); + send_stream_error( + &tx, + Status::aborted("Please provide a module configuration."), + ) + .await; + break; + } + }; + + log::debug!("Received logon for action {}", identifier); + + let accepted = match handle_logon( + &token, + action_logon, + actions.clone(), + module_service.clone(), + client.clone(), + cfg_tx.clone(), + tx.clone(), + pending_replies.clone(), + &mut cfg_forwarder_started, + ) + .await + { + Ok(v) => v, + Err(status) => { + log::warn!( + "Action logon failed identifier={} code={:?} message={}", + identifier, + status.code(), + status.message() + ); + send_stream_error(&tx, status).await; + break; + } + }; + + action_props = Some(accepted); + tracing::Span::current() + .record("action.identifier", identifier.as_str()); + metrics::action_connection(&identifier, "accepted"); + metrics::action_active(&identifier, 1); + connected_at = Some(std::time::Instant::now()); + connected_identifier = Some(identifier); + } + _ => { + log::error!("Action stream protocol violation expected=logon"); + send_stream_error( + &tx, + Status::failed_precondition("first action stream message must be logon"), + ) + .await; + break; + } + } + + continue; + } + + let props = match action_props.clone() { + Some(p) => p, + None => { + log::error!("Missing action properties after logon"); + break; + } + }; + + let identifier = match props.module { + Some(ref m) => m.identifier.clone(), + None => { + log::error!("Logon state missing module"); + break; + } + }; + + // Static mode has no Sagittarius to push config updates from, so + // re-broadcast the action's configured values on every message it + // sends instead of relying on a one-shot delivery at logon time. + if is_static { + let lock = actions.lock().await; + let configs = lock.get_action_configuration(&token, &identifier); + for conf in configs { + if let Err(err) = cfg_tx.send(conf) { + log::warn!("No action configuration receivers available: {:?}", err); + } + } + }; + + match data { + tucana::aquila::action_transfer_request::Data::Logon(_) => { + log::warn!( + "Action stream protocol violation identifier={} reason=duplicate_logon", + identifier + ); + send_stream_error( + &tx, + Status::failed_precondition("action stream logon was already accepted"), + ) + .await; + break; + } + tucana::aquila::action_transfer_request::Data::Event(event) => { + log::debug!("Received event action={}", identifier); + metrics::action_event(&identifier); + handle_event(&identifier, event, kv.clone(), client.clone()).await; + } + tucana::aquila::action_transfer_request::Data::Result(execution_result) => { + log::debug!( + "Received execution result execution_id={} action={}", + execution_result.execution_identifier, + identifier + ); + + handle_result( + &identifier, + execution_result, + client.clone(), + pending_replies.clone(), + ) + .await; + } + } + } + + if let Some(identifier) = connected_identifier { + metrics::action_active(&identifier, -1); + metrics::action_connection(&identifier, "closed"); + if let Some(connected_at) = connected_at { + metrics::action_connection_duration( + &identifier, + connected_at.elapsed().as_secs_f64(), + ); + } + } + log::debug!("Action transfer stream ended"); + } + .instrument(stream_span)); + + Ok(tonic::Response::new(Box::pin(ReceiverStream::new(rx)))) + } +} diff --git a/src/server/action_transfer/nats_bridge.rs b/src/server/action_transfer/nats_bridge.rs new file mode 100644 index 0000000..2e26361 --- /dev/null +++ b/src/server/action_transfer/nats_bridge.rs @@ -0,0 +1,451 @@ +//! Bridges NATS execution requests/results to and from a connected action's +//! gRPC stream: looks up matching flows for incoming events, forwards +//! execution requests to the action, and publishes results back to NATS. + +use async_nats::{Subject, Subscriber}; +use futures::StreamExt; +use prost::Message; +use tucana::{ + aquila::{ActionEvent, ActionExecutionRequest, ActionExecutionResponse, ActionTransferResponse}, + shared::{ExecutionFlow, Flows, ValidationFlow, Value}, +}; + +use crate::telemetry::{errors, metrics}; + +use super::pending_replies::{PendingReplyStore, pending_reply_keys}; + +/// Wraps the underlying NATS/KV error from a failed flow lookup so callers +/// get a stable, human-readable message while [`std::error::Error::source`] +/// still exposes the original cause for logging. +#[derive(Debug)] +struct FlowIdentificationError { + source: Box, +} + +impl std::fmt::Display for FlowIdentificationError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("failed to identify flows") + } +} + +impl std::error::Error for FlowIdentificationError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + Some(self.source.as_ref()) + } +} + +/// Scans the flow KV bucket for every entry whose key matches `pattern`, +/// decoding each match. There is no secondary index for flows, so a full key +/// scan is the only lookup path available. +async fn get_flows( + pattern: String, + kv: async_nats::jetstream::kv::Store, +) -> Result { + log::debug!("Scanning flows with pattern: {}", pattern); + let mut collector = Vec::new(); + let mut keys = match kv.keys().await { + Ok(keys) => keys.boxed(), + Err(err) => { + return Err(FlowIdentificationError { + source: Box::new(err), + }); + } + }; + + while let Ok(Some(key)) = tokio_stream::StreamExt::try_next(&mut keys).await { + if !is_matching_key(&pattern, &key) { + continue; + } + + match kv.get(key.clone()).await { + Ok(Some(bytes)) => { + let decoded_flow = ValidationFlow::decode(bytes); + match decoded_flow { + Ok(flow) => collector.push(flow), + Err(err) => { + errors::record( + "flow_storage", + "flow.decode", + &err, + format!("flow.key={key}"), + ); + } + } + } + Ok(None) => { + log::debug!("Flow key disappeared while reading: {}", key); + } + Err(err) => { + errors::record( + "flow_storage", + "flow.fetch", + &err, + format!("flow.key={key}"), + ); + } + } + } + + log::debug!("Matched {} flows for pattern {}", collector.len(), pattern); + Ok(Flows { flows: collector }) +} + +/// Matches a dot-separated KV key against a dot-separated pattern where `*` +/// matches any single segment. A pattern shorter than the key still counts +/// as a match on its own segments (the key's trailing segments are ignored). +fn is_matching_key(pattern: &String, key: &String) -> bool { + let split_pattern = pattern.split("."); + let split_key = key.split(".").collect::>(); + let zip = split_pattern.into_iter().zip(split_key); + + for (pattern_part, key_part) in zip { + if pattern_part == "*" { + continue; + } + + if pattern_part != key_part { + log::debug!("Key {} does not match pattern {}", key, pattern); + return false; + } + } + + true +} + +/// Turns a stored, pre-validated flow into the executable form sent to a +/// runtime, binding the triggering event's payload as its input value. +fn convert_validation_flow(flow: ValidationFlow, input_value: Option) -> ExecutionFlow { + ExecutionFlow { + flow_id: flow.flow_id, + starting_node_id: flow.starting_node_id, + input_value, + node_functions: flow.node_functions, + project_id: flow.project_id, + } +} + +/// Recovers the execution id from a NATS subject of the form +/// `action..`, used as a fallback when the +/// execution request's own payload doesn't carry one. +fn subject_execution_identifier(subject: &Subject) -> Option { + subject + .as_str() + .rsplit('.') + .next() + .filter(|execution_id| !execution_id.is_empty()) + .map(ToString::to_string) +} + +/// Classifies an execution result for metrics purposes without cloning or +/// otherwise touching the underlying node result payload. +fn action_result_outcome(result: &ActionExecutionResponse) -> &'static str { + match result + .node_result + .as_ref() + .and_then(|node_result| node_result.result.as_ref()) + { + Some(tucana::shared::node_execution_result::Result::Success(_)) => "success", + Some(tucana::shared::node_execution_result::Result::Error(_)) => "error", + None => "missing", + } +} + +/// Sends a terminal error down the gRPC response stream. Best-effort: if the +/// receiver already dropped the stream there's nothing left to notify. +pub(super) async fn send_stream_error( + tx: &tokio::sync::mpsc::Sender>, + status: tonic::Status, +) { + if tx.send(Err(status)).await.is_err() { + log::debug!("Action transfer response stream closed before error could be sent"); + } +} + +/// Looks up matching flows for an event and requests their execution. +/// +/// Each match is dispatched as an independent NATS request-reply exchange on +/// its own `execution.` subject, so one flow failing to find a runtime +/// doesn't block the others, and a runtime's response arrives correlated to +/// exactly the flow that triggered it. +pub(super) async fn handle_event( + action_identifier: &str, + event: ActionEvent, + kv: async_nats::jetstream::kv::Store, + client: async_nats::Client, +) { + let pattern = format!("{}.*.{}.*", event.event_type, event.project_id); + log::debug!( + "Handling action event event_type={} project_id={}", + event.event_type, + event.project_id + ); + + let flows = match get_flows(pattern.clone(), kv).await { + Ok(f) => f, + Err(err) => { + errors::record( + "flow_storage", + "action.event.find_flows", + &err, + format!( + "action.identifier={} event_type={} project_id={} pattern={}", + action_identifier, event.event_type, event.project_id, pattern + ), + ); + return; + } + }; + + let matched_count = flows.flows.len(); + log::info!( + "Matched flows for action event event_type={} project_id={} flow_count={}", + event.event_type, + event.project_id, + matched_count + ); + for flow in flows.flows { + let uuid = uuid::Uuid::new_v4().to_string(); + let flow_id = flow.flow_id; + let execution_flow: ExecutionFlow = convert_validation_flow(flow, event.payload.clone()); + let bytes = execution_flow.encode_to_vec(); + let topic = format!("execution.{}", uuid); + + log::info!( + "Requesting execution flow_id={} execution_id={} event_type={} project_id={}", + flow_id, + uuid, + event.event_type, + event.project_id + ); + + if let Err(err) = client.request(topic.clone(), bytes.into()).await { + errors::record( + "messaging", + "action.event.request_execution", + &err, + format!( + "action.identifier={} flow_id={} execution_id={} topic={}", + action_identifier, flow_id, uuid, topic + ), + ); + } + } +} + +/// Publishes execution results back to the original NATS reply subject. +/// +/// The reply subject was stashed by [`forward_nats_to_action`] when the +/// execution request first came in; the action only knows its own execution +/// id, not the NATS subject that's waiting for a reply, so this is the only +/// place that reunites the two. +pub(super) async fn handle_result( + action_identifier: &str, + execution_result: ActionExecutionResponse, + client: async_nats::Client, + pending_replies: PendingReplyStore, +) { + let execution_id = execution_result.execution_identifier.clone(); + metrics::action_result(action_identifier, action_result_outcome(&execution_result)); + + let Some(pending_reply) = pending_replies.remove(&execution_id).await else { + metrics::action_failure(action_identifier, "result_unmatched"); + errors::record_message( + "protocol", + "action.result.match", + "No pending NATS reply subject found", + format!("action.identifier={action_identifier} execution_id={execution_id}"), + ); + return; + }; + metrics::action_execution_duration( + action_identifier, + pending_reply.started_at.elapsed().as_secs_f64(), + ); + + log::debug!( + "Publishing execution result execution_id={} reply_subject={}", + execution_id, + pending_reply.reply_subject + ); + + let payload = execution_result.encode_to_vec(); + if let Err(err) = client + .publish(pending_reply.reply_subject.clone(), payload.into()) + .await + { + metrics::action_failure(action_identifier, "result_publish"); + errors::record( + "messaging", + "action.result.publish", + &err, + format!( + "action.identifier={} execution_id={} reply_subject={}", + action_identifier, execution_id, pending_reply.reply_subject + ), + ); + return; + } + + if let Err(err) = client.flush().await { + metrics::action_failure(action_identifier, "result_flush"); + errors::record( + "messaging", + "action.result.flush", + &err, + format!("action.identifier={action_identifier} execution_id={execution_id}"), + ); + } +} + +/// Forwards NATS execution requests to the connected action via gRPC and stores reply subjects. +/// +/// Runs for as long as the action's `action..*` subscription is +/// alive, which is the same lifetime as the action's logon; a fresh task is +/// spawned per logon rather than reused across reconnects. +pub(super) async fn forward_nats_to_action( + action_identifier: String, + mut sub: Subscriber, + tx: tokio::sync::mpsc::Sender>, + pending_replies: PendingReplyStore, +) { + log::debug!("Waiting for incoming action execution request"); + + while let Some(msg) = sub.next().await { + let mut execution = match ActionExecutionRequest::decode(msg.payload.as_ref()) { + Ok(req) => req, + Err(err) => { + metrics::action_execution(&action_identifier, "invalid"); + metrics::action_failure(&action_identifier, "execution_decode"); + errors::record( + "protocol", + "action.execution.decode", + &err, + format!( + "action.identifier={} subject={} payload_bytes={}", + action_identifier, + msg.subject, + msg.payload.len() + ), + ); + continue; + } + }; + + let subject_execution_id = subject_execution_identifier(&msg.subject); + if execution.execution_identifier.is_empty() + && let Some(subject_execution_id) = subject_execution_id.as_ref() + { + log::warn!( + "Filled missing action execution identifier from NATS subject subject={} execution_id={}", + msg.subject, + subject_execution_id + ); + execution.execution_identifier = subject_execution_id.clone(); + } + + let execution_id = execution.execution_identifier.clone(); + + let Some(reply_subject) = msg.reply.clone() else { + metrics::action_execution(&action_identifier, "invalid"); + metrics::action_failure(&action_identifier, "missing_reply_subject"); + log::error!( + "Received request without NATS reply subject execution_id={}", + execution_id + ); + continue; + }; + + let keys = pending_reply_keys(&execution_id, subject_execution_id.as_deref()); + if keys.is_empty() { + metrics::action_execution(&action_identifier, "invalid"); + metrics::action_failure(&action_identifier, "missing_execution_identifier"); + log::error!( + "Cannot store NATS reply subject without execution identifier subject={} reply_subject={}", + msg.subject, + reply_subject + ); + continue; + } + + pending_replies + .insert(reply_subject.clone(), keys.clone()) + .await; + + log::debug!( + "Stored reply subject reply_subject={} execution_id={} keys={:?}", + reply_subject, + execution_id, + keys + ); + + log::debug!( + "Forwarding execution request to action execution_id={} subject={}", + execution_id, + msg.subject + ); + + let resp = ActionTransferResponse { + data: Some(tucana::aquila::action_transfer_response::Data::Execution( + execution, + )), + }; + + if tx.send(Ok(resp)).await.is_err() { + metrics::action_execution(&action_identifier, "forward_failed"); + metrics::action_failure(&action_identifier, "execution_forward"); + log::debug!("Execution forwarder channel closed"); + + // cleanup, since the request can no longer be delivered to the action + pending_replies.remove(&execution_id).await; + + break; + } + metrics::action_execution(&action_identifier, "forwarded"); + } + + log::debug!("Execution forwarder stopped"); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn subject_execution_identifier_uses_last_subject_token() { + assert_eq!( + subject_execution_identifier(&Subject::from("action.example.execution-id")), + Some("execution-id".to_string()) + ); + } + + #[test] + fn action_result_outcome_distinguishes_success_error_and_missing() { + use tucana::shared::{Error, NodeExecutionResult, Value, node_execution_result}; + + let response = |result| ActionExecutionResponse { + execution_identifier: "execution-id".into(), + node_result: Some(NodeExecutionResult { + result, + ..Default::default() + }), + }; + + assert_eq!( + action_result_outcome(&response(Some(node_execution_result::Result::Success( + Value::default() + )))), + "success" + ); + assert_eq!( + action_result_outcome(&response(Some(node_execution_result::Result::Error( + Error::default() + )))), + "error" + ); + assert_eq!(action_result_outcome(&response(None)), "missing"); + assert_eq!( + action_result_outcome(&ActionExecutionResponse::default()), + "missing" + ); + } +} diff --git a/src/server/action_transfer/pending_replies.rs b/src/server/action_transfer/pending_replies.rs new file mode 100644 index 0000000..74ef36c --- /dev/null +++ b/src/server/action_transfer/pending_replies.rs @@ -0,0 +1,129 @@ +//! Tracks in-flight action executions so their NATS reply subject can be +//! located again once the connected action sends back a result. + +use async_nats::Subject; +use std::{collections::HashMap, sync::Arc, time::Instant}; +use tokio::sync::Mutex; + +/// A single execution awaiting a result, plus everything needed to route +/// that result and measure how long it took. +#[derive(Clone)] +pub(super) struct PendingReply { + /// Where to publish the `ActionExecutionResponse` once it arrives. + pub(super) reply_subject: Subject, + /// Every key this entry is filed under in the store, so removal by any + /// one alias can also clean up the others. + keys: Vec, + /// When the request was forwarded to the action, for execution-duration metrics. + pub(super) started_at: Instant, +} + +/// Shared, lock-protected registry mapping execution identifiers to the NATS +/// reply subject a result must eventually be published to. +#[derive(Clone, Default)] +pub(super) struct PendingReplyStore { + inner: Arc>>, +} + +impl PendingReplyStore { + pub(super) fn new() -> Self { + Self::default() + } + + /// Registers `reply_subject` under every alias in `keys` (typically the + /// payload execution id and the one derived from the NATS subject). + pub(super) async fn insert(&self, reply_subject: Subject, keys: Vec) { + let pending_reply = PendingReply { + reply_subject, + keys: keys.clone(), + started_at: Instant::now(), + }; + + let mut pending = self.inner.lock().await; + for key in keys { + pending.insert(key, pending_reply.clone()); + } + } + + /// Removes and returns the reply registered under `execution_id`, along + /// with any of its aliases. + pub(super) async fn remove(&self, execution_id: &str) -> Option { + let mut pending = self.inner.lock().await; + let pending_reply = pending.remove(execution_id)?; + + for key in &pending_reply.keys { + if key != execution_id { + pending.remove(key); + } + } + + Some(pending_reply) + } +} + +/// Determines which keys a pending reply should be filed under: the +/// execution id from the request payload and, if different, the one derived +/// from the NATS subject it arrived on. +pub(super) fn pending_reply_keys( + request_execution_id: &str, + subject_execution_id: Option<&str>, +) -> Vec { + let mut keys = Vec::new(); + + if !request_execution_id.is_empty() { + keys.push(request_execution_id.to_string()); + } + + if let Some(subject_execution_id) = subject_execution_id + && !subject_execution_id.is_empty() + && !keys.iter().any(|key| key == subject_execution_id) + { + keys.push(subject_execution_id.to_string()); + } + + keys +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn pending_reply_keys_include_payload_and_subject_ids_once() { + assert_eq!( + pending_reply_keys("payload-id", Some("subject-id")), + vec!["payload-id".to_string(), "subject-id".to_string()] + ); + assert_eq!( + pending_reply_keys("same-id", Some("same-id")), + vec!["same-id".to_string()] + ); + assert_eq!( + pending_reply_keys("", Some("subject-id")), + vec!["subject-id".to_string()] + ); + } + + #[test] + fn remove_removes_all_aliases() { + futures::executor::block_on(async { + let reply_subject = Subject::from("_INBOX.reply"); + let store = PendingReplyStore::new(); + + store + .insert( + reply_subject.clone(), + vec!["payload-id".to_string(), "subject-id".to_string()], + ) + .await; + + let removed = store + .remove("subject-id") + .await + .expect("pending reply should be found by alias"); + + assert_eq!(removed.reply_subject, reply_subject); + assert!(store.remove("payload-id").await.is_none()); + }); + } +} diff --git a/src/server/action_transfer_service_server_impl.rs b/src/server/action_transfer_service_server_impl.rs deleted file mode 100644 index 0f423b2..0000000 --- a/src/server/action_transfer_service_server_impl.rs +++ /dev/null @@ -1,1037 +0,0 @@ -use crate::{ - configuration::service::ServiceConfiguration, - sagittarius::module_service_client_impl::SagittariusModuleServiceClient, - telemetry::{errors, metrics}, -}; -use async_nats::{Subject, Subscriber}; -use futures::StreamExt; -use futures_core::Stream; -use prost::Message; -use std::{collections::HashMap, pin::Pin, sync::Arc, time::Instant}; -use tokio::sync::Mutex; -use tokio_stream::wrappers::ReceiverStream; -use tonic::Status; -use tracing::Instrument; -use tucana::{ - aquila::{ - ActionEvent, ActionExecutionRequest, ActionExecutionResponse, ActionLogon, - ActionTransferRequest, ActionTransferResponse, - action_transfer_service_server::ActionTransferService, - }, - shared::{ExecutionFlow, Flows, ValidationFlow, Value}, -}; - -type PendingReplies = Arc>>; - -#[derive(Clone)] -struct PendingReply { - reply_subject: Subject, - keys: Vec, - started_at: Instant, -} - -pub struct AquilaActionTransferServiceServer { - client: async_nats::Client, - kv: async_nats::jetstream::kv::Store, - actions: ServiceConfiguration, - module_service: Option>>, - action_config_tx: tokio::sync::broadcast::Sender, - is_static: bool, -} - -impl AquilaActionTransferServiceServer { - pub fn new( - client: async_nats::Client, - kv: async_nats::jetstream::kv::Store, - actions: ServiceConfiguration, - module_service: Option>>, - action_config_tx: tokio::sync::broadcast::Sender, - is_static: bool, - ) -> Self { - Self { - client, - kv, - actions, - module_service, - action_config_tx, - is_static, - } - } -} - -#[derive(Debug)] -struct FlowIdentificationError { - source: Box, -} - -impl std::fmt::Display for FlowIdentificationError { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("failed to identify flows") - } -} - -impl std::error::Error for FlowIdentificationError { - fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { - Some(self.source.as_ref()) - } -} - -async fn get_flows( - pattern: String, - kv: async_nats::jetstream::kv::Store, -) -> Result { - log::debug!("Scanning flows with pattern: {}", pattern); - let mut collector = Vec::new(); - let mut keys = match kv.keys().await { - Ok(keys) => keys.boxed(), - Err(err) => { - return Err(FlowIdentificationError { - source: Box::new(err), - }); - } - }; - - while let Ok(Some(key)) = tokio_stream::StreamExt::try_next(&mut keys).await { - if !is_matching_key(&pattern, &key) { - continue; - } - - match kv.get(key.clone()).await { - Ok(Some(bytes)) => { - let decoded_flow = ValidationFlow::decode(bytes); - match decoded_flow { - Ok(flow) => collector.push(flow), - Err(err) => { - errors::record( - "flow_storage", - "flow.decode", - &err, - format!("flow.key={key}"), - ); - } - } - } - Ok(None) => { - log::debug!("Flow key disappeared while reading: {}", key); - } - Err(err) => { - errors::record( - "flow_storage", - "flow.fetch", - &err, - format!("flow.key={key}"), - ); - } - } - } - - log::debug!("Matched {} flows for pattern {}", collector.len(), pattern); - Ok(Flows { flows: collector }) -} - -fn is_matching_key(pattern: &String, key: &String) -> bool { - let split_pattern = pattern.split("."); - let split_key = key.split(".").collect::>(); - let zip = split_pattern.into_iter().zip(split_key); - - for (pattern_part, key_part) in zip { - if pattern_part == "*" { - continue; - } - - if pattern_part != key_part { - log::debug!("Key {} does not match pattern {}", key, pattern); - return false; - } - } - - true -} - -fn convert_validation_flow(flow: ValidationFlow, input_value: Option) -> ExecutionFlow { - ExecutionFlow { - flow_id: flow.flow_id, - starting_node_id: flow.starting_node_id, - input_value, - node_functions: flow.node_functions, - project_id: flow.project_id, - } -} - -fn applies_to_action( - configs: &tucana::shared::ModuleConfigurations, - action_identifier: &str, -) -> bool { - configs.module_identifier == action_identifier -} - -fn overwrite_module_definition_sources( - module: &mut tucana::shared::Module, - action_identifier: &str, -) { - let source = format!("action.{}", action_identifier); - - for flow_type in &mut module.flow_types { - flow_type.definition_source = Some(source.clone()); - } - for runtime_flow_type in &mut module.runtime_flow_types { - runtime_flow_type.definition_source = Some(source.clone()); - } - for function_definition in &mut module.function_definitions { - function_definition.definition_source = source.clone(); - } - for runtime_function_definition in &mut module.runtime_function_definitions { - runtime_function_definition.definition_source = source.clone(); - } - for definition_data_type in &mut module.definition_data_types { - definition_data_type.definition_source = source.clone(); - } -} - -fn subject_execution_identifier(subject: &Subject) -> Option { - subject - .as_str() - .rsplit('.') - .next() - .filter(|execution_id| !execution_id.is_empty()) - .map(ToString::to_string) -} - -fn action_result_outcome(result: &ActionExecutionResponse) -> &'static str { - match result - .node_result - .as_ref() - .and_then(|node_result| node_result.result.as_ref()) - { - Some(tucana::shared::node_execution_result::Result::Success(_)) => "success", - Some(tucana::shared::node_execution_result::Result::Error(_)) => "error", - None => "missing", - } -} - -fn pending_reply_keys( - request_execution_id: &str, - subject_execution_id: Option<&str>, -) -> Vec { - let mut keys = Vec::new(); - - if !request_execution_id.is_empty() { - keys.push(request_execution_id.to_string()); - } - - if let Some(subject_execution_id) = subject_execution_id - && !subject_execution_id.is_empty() - && !keys.iter().any(|key| key == subject_execution_id) - { - keys.push(subject_execution_id.to_string()); - } - - keys -} - -async fn send_stream_error( - tx: &tokio::sync::mpsc::Sender>, - status: tonic::Status, -) { - if tx.send(Err(status)).await.is_err() { - log::debug!("Action transfer response stream closed before error could be sent"); - } -} - -fn insert_pending_reply( - pending: &mut HashMap, - reply_subject: Subject, - keys: Vec, -) { - let pending_reply = PendingReply { - reply_subject, - keys: keys.clone(), - started_at: Instant::now(), - }; - - for key in keys { - pending.insert(key, pending_reply.clone()); - } -} - -fn remove_pending_reply( - pending: &mut HashMap, - execution_id: &str, -) -> Option { - let pending_reply = pending.remove(execution_id)?; - - for key in &pending_reply.keys { - if key != execution_id { - pending.remove(key); - } - } - - Some(pending_reply) -} - -/// Extracts the bearer token from gRPC metadata. -fn extract_token( - request: &tonic::Request>, -) -> Result { - log::debug!("Extracting authorization token from metadata"); - match request.metadata().get("authorization") { - Some(ascii) => match ascii.to_str() { - Ok(tk) => { - if tk.is_empty() { - log::error!("Authorization token is empty"); - return Err(Status::unauthenticated("authorization token is empty")); - } - - Ok(tk.to_string()) - } - Err(err) => { - log::error!("Cannot read authorization header because: {:?}", err); - Err(Status::unauthenticated("invalid authorization header")) - } - }, - None => { - log::error!("Missing authorization token"); - Err(Status::unauthenticated("missing authorization token")) - } - } -} - -/// Validates the logon request, starts NATS + config forwarders, and returns the accepted logon. -async fn handle_logon( - token: &str, - mut action_logon: ActionLogon, - actions: Arc>, - module_service: Option>>, - client: async_nats::Client, - cfg_tx: tokio::sync::broadcast::Sender, - tx: tokio::sync::mpsc::Sender>, - pending_replies: PendingReplies, - cfg_forwarder_started: &mut bool, -) -> Result { - let module = match action_logon.module.as_mut() { - Some(m) => m, - None => { - log::warn!("Rejected action logon reason=missing_module"); - return Err(Status::aborted("Please provide a module configuration.")); - } - }; - let identifier = module.identifier.clone(); - log::info!("Action logon attempt identifier={}", identifier); - - { - let lock = actions.lock().await; - if !lock.has_action(&token.to_string(), &identifier) { - metrics::action_connection(&identifier, "rejected"); - metrics::action_failure(&identifier, "authentication"); - log::warn!( - "Rejected action logon identifier={} reason=token_not_registered", - identifier - ); - return Err(Status::unauthenticated( - "token not matching to action identifier", - )); - } - } - - overwrite_module_definition_sources(module, &identifier); - - if let Some(module_service) = module_service { - let mut client = module_service.lock().await; - let response = client - .update_modules(tucana::aquila::ModuleUpdateRequest { - modules: vec![module.clone()], - }) - .await; - - if !response.success { - metrics::action_connection(&identifier, "rejected"); - metrics::action_failure(&identifier, "module_update"); - errors::record_message( - "dependency", - "action.logon", - "Sagittarius rejected the action module update", - format!("action.identifier={identifier}"), - ); - return Err(Status::internal( - "could not update action module via Sagittarius", - )); - } - } - - log::debug!("Action connected identifier={}", identifier); - - let sub = match client.subscribe(format!("action.{}.*", identifier)).await { - Ok(s) => s, - Err(err) => { - metrics::action_connection(&identifier, "rejected"); - metrics::action_failure(&identifier, "subscription"); - errors::record( - "messaging", - "action.subscribe", - &err, - format!("action.identifier={identifier} subject=action.{identifier}.*"), - ); - return Err(Status::internal( - "could not register action into execution loop", - )); - } - }; - - if let Err(err) = client.flush().await { - metrics::action_connection(&identifier, "rejected"); - metrics::action_failure(&identifier, "subscription_flush"); - errors::record( - "messaging", - "action.subscribe.flush", - &err, - format!("action.identifier={identifier}"), - ); - return Err(Status::internal( - "could not register action subscription with NATS", - )); - } - - log::debug!("Subscribed to action subject action.{}.*", identifier); - - let tx_clone = tx.clone(); - let pending_replies_clone = pending_replies.clone(); - let forwarder_identifier = identifier.clone(); - tokio::spawn(async move { - forward_nats_to_action(forwarder_identifier, sub, tx_clone, pending_replies_clone).await; - }); - - if !*cfg_forwarder_started { - *cfg_forwarder_started = true; - log::debug!("Starting config forwarder action={}", identifier); - spawn_cfg_forwarder(identifier.clone(), cfg_tx, tx.clone()); - } - - Ok(action_logon) -} - -/// Forwards config updates for the given action identifier to the gRPC stream. -fn spawn_cfg_forwarder( - action_identifier: String, - cfg_tx: tokio::sync::broadcast::Sender, - tx: tokio::sync::mpsc::Sender>, -) { - let mut cfg_rx = cfg_tx.subscribe(); - tokio::spawn(async move { - while let Ok(cfgs) = cfg_rx.recv().await { - if !applies_to_action(&cfgs, &action_identifier) { - log::debug!( - "Config update does not apply to action {}", - action_identifier - ); - continue; - } - - log::debug!("Forwarding config update to action {}", action_identifier); - let resp = ActionTransferResponse { - data: Some( - tucana::aquila::action_transfer_response::Data::ModuleConfigurations(cfgs), - ), - }; - - if tx.send(Ok(resp)).await.is_err() { - metrics::action_config_update(&action_identifier, "failed"); - metrics::action_failure(&action_identifier, "configuration_forward"); - log::debug!("Config forwarder channel closed for {}", action_identifier); - break; - } - metrics::action_config_update(&action_identifier, "success"); - } - - log::debug!("Config forwarder stopped for {}", action_identifier); - }); -} - -/// Looks up matching flows for an event and requests their execution. -async fn handle_event( - action_identifier: &str, - event: ActionEvent, - kv: async_nats::jetstream::kv::Store, - client: async_nats::Client, -) { - let pattern = format!("{}.*.{}.*", event.event_type, event.project_id); - log::debug!( - "Handling action event event_type={} project_id={}", - event.event_type, - event.project_id - ); - - let flows = match get_flows(pattern.clone(), kv).await { - Ok(f) => f, - Err(err) => { - errors::record( - "flow_storage", - "action.event.find_flows", - &err, - format!( - "action.identifier={} event_type={} project_id={} pattern={}", - action_identifier, event.event_type, event.project_id, pattern - ), - ); - return; - } - }; - - let matched_count = flows.flows.len(); - log::info!( - "Matched flows for action event event_type={} project_id={} flow_count={}", - event.event_type, - event.project_id, - matched_count - ); - for flow in flows.flows { - let uuid = uuid::Uuid::new_v4().to_string(); - let flow_id = flow.flow_id; - let execution_flow: ExecutionFlow = convert_validation_flow(flow, event.payload.clone()); - let bytes = execution_flow.encode_to_vec(); - let topic = format!("execution.{}", uuid); - - log::info!( - "Requesting execution flow_id={} execution_id={} event_type={} project_id={}", - flow_id, - uuid, - event.event_type, - event.project_id - ); - - if let Err(err) = client.request(topic.clone(), bytes.into()).await { - errors::record( - "messaging", - "action.event.request_execution", - &err, - format!( - "action.identifier={} flow_id={} execution_id={} topic={}", - action_identifier, flow_id, uuid, topic - ), - ); - } - } -} - -/// Publishes execution results back to the original NATS reply subject. -async fn handle_result( - action_identifier: &str, - execution_result: ActionExecutionResponse, - client: async_nats::Client, - pending_replies: PendingReplies, -) { - let execution_id = execution_result.execution_identifier.clone(); - metrics::action_result(action_identifier, action_result_outcome(&execution_result)); - - let pending_reply = { - let mut pending = pending_replies.lock().await; - remove_pending_reply(&mut pending, &execution_id) - }; - - let Some(pending_reply) = pending_reply else { - metrics::action_failure(action_identifier, "result_unmatched"); - errors::record_message( - "protocol", - "action.result.match", - "No pending NATS reply subject found", - format!("action.identifier={action_identifier} execution_id={execution_id}"), - ); - return; - }; - metrics::action_execution_duration( - action_identifier, - pending_reply.started_at.elapsed().as_secs_f64(), - ); - - log::debug!( - "Publishing execution result execution_id={} reply_subject={}", - execution_id, - pending_reply.reply_subject - ); - - let payload = execution_result.encode_to_vec(); - if let Err(err) = client - .publish(pending_reply.reply_subject.clone(), payload.into()) - .await - { - metrics::action_failure(action_identifier, "result_publish"); - errors::record( - "messaging", - "action.result.publish", - &err, - format!( - "action.identifier={} execution_id={} reply_subject={}", - action_identifier, execution_id, pending_reply.reply_subject - ), - ); - return; - } - - if let Err(err) = client.flush().await { - metrics::action_failure(action_identifier, "result_flush"); - errors::record( - "messaging", - "action.result.flush", - &err, - format!("action.identifier={action_identifier} execution_id={execution_id}"), - ); - } -} - -#[tonic::async_trait] -impl ActionTransferService for AquilaActionTransferServiceServer { - type TransferStream = - Pin> + Send + 'static>>; - - #[tracing::instrument( - name = "aquila.action.transfer", - skip_all, - fields(rpc.system = "grpc", rpc.service = "ActionTransferService", rpc.method = "Transfer") - )] - async fn transfer( - &self, - request: tonic::Request>, - ) -> std::result::Result, tonic::Status> { - let token = extract_token(&request)?; - log::debug!("Action transfer stream opened"); - - let mut first_request = true; - let mut action_props: Option = None; - let mut stream = request.into_inner(); - - let actions = Arc::new(Mutex::new(self.actions.clone())); - let kv = self.kv.clone(); - let client = self.client.clone(); - let module_service = self.module_service.clone(); - let cfg_tx = self.action_config_tx.clone(); - let is_static = self.is_static; - let pending_replies: PendingReplies = Arc::new(Mutex::new(HashMap::new())); - - let (tx, rx) = - tokio::sync::mpsc::channel::>(32); - - let stream_span = tracing::info_span!( - "aquila.action.stream", - action.identifier = tracing::field::Empty - ); - tokio::spawn(async move { - let mut cfg_forwarder_started = false; - let mut connected_at = None; - let mut connected_identifier = None; - log::debug!("Action transfer stream started"); - - while let Some(next) = stream.next().await { - let transfer_request = match next { - Ok(tr) => tr, - Err(status) => { - log::warn!("Action transfer input stream failed status={:?}", status); - break; - } - }; - - let data = match transfer_request.data { - Some(d) => d, - None => { - log::warn!("Received empty action transfer request"); - continue; - } - }; - - if first_request { - first_request = false; - - match data { - tucana::aquila::action_transfer_request::Data::Logon(action_logon) => { - let identifier = match action_logon.module { - Some(ref m) => m.identifier.clone(), - None => { - log::warn!("Rejected action logon reason=missing_module"); - send_stream_error( - &tx, - Status::aborted("Please provide a module configuration."), - ) - .await; - break; - } - }; - - log::debug!("Received logon for action {}", identifier); - - let accepted = match handle_logon( - &token, - action_logon, - actions.clone(), - module_service.clone(), - client.clone(), - cfg_tx.clone(), - tx.clone(), - pending_replies.clone(), - &mut cfg_forwarder_started, - ) - .await - { - Ok(v) => v, - Err(status) => { - log::warn!( - "Action logon failed identifier={} code={:?} message={}", - identifier, - status.code(), - status.message() - ); - send_stream_error(&tx, status).await; - break; - } - }; - - action_props = Some(accepted); - tracing::Span::current() - .record("action.identifier", identifier.as_str()); - metrics::action_connection(&identifier, "accepted"); - metrics::action_active(&identifier, 1); - connected_at = Some(std::time::Instant::now()); - connected_identifier = Some(identifier); - } - _ => { - log::error!("Action stream protocol violation expected=logon"); - send_stream_error( - &tx, - Status::failed_precondition("first action stream message must be logon"), - ) - .await; - break; - } - } - - continue; - } - - let props = match action_props.clone() { - Some(p) => p, - None => { - log::error!("Missing action properties after logon"); - break; - } - }; - - let identifier = match props.module { - Some(ref m) => m.identifier.clone(), - None => { - log::error!("Logon state missing module"); - break; - } - }; - - if is_static { - let lock = actions.lock().await; - let configs = lock.get_action_configuration(&token, &identifier); - for conf in configs { - if let Err(err) = cfg_tx.send(conf) { - log::warn!("No action configuration receivers available: {:?}", err); - } - } - }; - - match data { - tucana::aquila::action_transfer_request::Data::Logon(_) => { - log::warn!( - "Action stream protocol violation identifier={} reason=duplicate_logon", - identifier - ); - send_stream_error( - &tx, - Status::failed_precondition("action stream logon was already accepted"), - ) - .await; - break; - } - tucana::aquila::action_transfer_request::Data::Event(event) => { - log::debug!("Received event action={}", identifier); - metrics::action_event(&identifier); - handle_event(&identifier, event, kv.clone(), client.clone()).await; - } - tucana::aquila::action_transfer_request::Data::Result(execution_result) => { - log::debug!( - "Received execution result execution_id={} action={}", - execution_result.execution_identifier, - identifier - ); - - handle_result( - &identifier, - execution_result, - client.clone(), - pending_replies.clone(), - ) - .await; - } - } - } - - if let Some(identifier) = connected_identifier { - metrics::action_active(&identifier, -1); - metrics::action_connection(&identifier, "closed"); - if let Some(connected_at) = connected_at { - metrics::action_connection_duration( - &identifier, - connected_at.elapsed().as_secs_f64(), - ); - } - } - log::debug!("Action transfer stream ended"); - } - .instrument(stream_span)); - - Ok(tonic::Response::new(Box::pin(ReceiverStream::new(rx)))) - } -} - -/// Forwards NATS execution requests to the connected action via gRPC and stores reply subjects. -async fn forward_nats_to_action( - action_identifier: String, - mut sub: Subscriber, - tx: tokio::sync::mpsc::Sender>, - pending_replies: PendingReplies, -) { - log::debug!("Waiting for incoming action execution request"); - - while let Some(msg) = sub.next().await { - let mut execution = match ActionExecutionRequest::decode(msg.payload.as_ref()) { - Ok(req) => req, - Err(err) => { - metrics::action_execution(&action_identifier, "invalid"); - metrics::action_failure(&action_identifier, "execution_decode"); - errors::record( - "protocol", - "action.execution.decode", - &err, - format!( - "action.identifier={} subject={} payload_bytes={}", - action_identifier, - msg.subject, - msg.payload.len() - ), - ); - continue; - } - }; - - let subject_execution_id = subject_execution_identifier(&msg.subject); - if execution.execution_identifier.is_empty() - && let Some(subject_execution_id) = subject_execution_id.as_ref() - { - log::warn!( - "Filled missing action execution identifier from NATS subject subject={} execution_id={}", - msg.subject, - subject_execution_id - ); - execution.execution_identifier = subject_execution_id.clone(); - } - - let execution_id = execution.execution_identifier.clone(); - - let Some(reply_subject) = msg.reply.clone() else { - metrics::action_execution(&action_identifier, "invalid"); - metrics::action_failure(&action_identifier, "missing_reply_subject"); - log::error!( - "Received request without NATS reply subject execution_id={}", - execution_id - ); - continue; - }; - - let keys = pending_reply_keys(&execution_id, subject_execution_id.as_deref()); - if keys.is_empty() { - metrics::action_execution(&action_identifier, "invalid"); - metrics::action_failure(&action_identifier, "missing_execution_identifier"); - log::error!( - "Cannot store NATS reply subject without execution identifier subject={} reply_subject={}", - msg.subject, - reply_subject - ); - continue; - } - - { - let mut pending = pending_replies.lock().await; - insert_pending_reply(&mut pending, reply_subject.clone(), keys.clone()); - } - - log::debug!( - "Stored reply subject reply_subject={} execution_id={} keys={:?}", - reply_subject, - execution_id, - keys - ); - - log::debug!( - "Forwarding execution request to action execution_id={} subject={}", - execution_id, - msg.subject - ); - - let resp = ActionTransferResponse { - data: Some(tucana::aquila::action_transfer_response::Data::Execution( - execution, - )), - }; - - if tx.send(Ok(resp)).await.is_err() { - metrics::action_execution(&action_identifier, "forward_failed"); - metrics::action_failure(&action_identifier, "execution_forward"); - log::debug!("Execution forwarder channel closed"); - - // cleanup, since the request can no longer be delivered to the action - let mut pending = pending_replies.lock().await; - remove_pending_reply(&mut pending, &execution_id); - - break; - } - metrics::action_execution(&action_identifier, "forwarded"); - } - - log::debug!("Execution forwarder stopped"); -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn module_configurations_apply_by_module_identifier() { - let configs = tucana::shared::ModuleConfigurations { - module_identifier: "gls-action".to_string(), - module_configurations: vec![tucana::shared::ModuleProjectConfigurations { - project_id: 1, - module_configurations: vec![tucana::shared::ModuleConfiguration { - identifier: "username".to_string(), - value: None, - }], - }], - }; - - assert!(applies_to_action(&configs, "gls-action")); - assert!(!applies_to_action(&configs, "another-action")); - } - - #[test] - fn overwrite_module_definition_sources_uses_action_source() { - let mut module = tucana::shared::Module { - flow_types: vec![tucana::shared::FlowType { - definition_source: Some("module.old".to_string()), - ..Default::default() - }], - runtime_flow_types: vec![tucana::shared::RuntimeFlowType { - definition_source: Some("module.old".to_string()), - ..Default::default() - }], - function_definitions: vec![tucana::shared::FunctionDefinition { - definition_source: "module.old".to_string(), - ..Default::default() - }], - runtime_function_definitions: vec![tucana::shared::RuntimeFunctionDefinition { - definition_source: "module.old".to_string(), - ..Default::default() - }], - definition_data_types: vec![tucana::shared::DefinitionDataType { - definition_source: "module.old".to_string(), - ..Default::default() - }], - ..Default::default() - }; - - overwrite_module_definition_sources(&mut module, "send-email"); - - assert_eq!( - module.flow_types[0].definition_source.as_deref(), - Some("action.send-email") - ); - assert_eq!( - module.runtime_flow_types[0].definition_source.as_deref(), - Some("action.send-email") - ); - assert_eq!( - module.function_definitions[0].definition_source, - "action.send-email" - ); - assert_eq!( - module.runtime_function_definitions[0].definition_source, - "action.send-email" - ); - assert_eq!( - module.definition_data_types[0].definition_source, - "action.send-email" - ); - } - - #[test] - fn pending_reply_keys_include_payload_and_subject_ids_once() { - assert_eq!( - pending_reply_keys("payload-id", Some("subject-id")), - vec!["payload-id".to_string(), "subject-id".to_string()] - ); - assert_eq!( - pending_reply_keys("same-id", Some("same-id")), - vec!["same-id".to_string()] - ); - assert_eq!( - pending_reply_keys("", Some("subject-id")), - vec!["subject-id".to_string()] - ); - } - - #[test] - fn remove_pending_reply_removes_all_aliases() { - let reply_subject = Subject::from("_INBOX.reply"); - let mut pending = HashMap::new(); - - insert_pending_reply( - &mut pending, - reply_subject.clone(), - vec!["payload-id".to_string(), "subject-id".to_string()], - ); - - let removed = remove_pending_reply(&mut pending, "subject-id") - .expect("pending reply should be found by alias"); - - assert_eq!(removed.reply_subject, reply_subject); - assert!(pending.is_empty()); - } - - #[test] - fn subject_execution_identifier_uses_last_subject_token() { - assert_eq!( - subject_execution_identifier(&Subject::from("action.example.execution-id")), - Some("execution-id".to_string()) - ); - } - - #[test] - fn action_result_outcome_distinguishes_success_error_and_missing() { - use tucana::shared::{Error, NodeExecutionResult, Value, node_execution_result}; - - let response = |result| ActionExecutionResponse { - execution_identifier: "execution-id".into(), - node_result: Some(NodeExecutionResult { - result, - ..Default::default() - }), - }; - - assert_eq!( - action_result_outcome(&response(Some(node_execution_result::Result::Success( - Value::default() - )))), - "success" - ); - assert_eq!( - action_result_outcome(&response(Some(node_execution_result::Result::Error( - Error::default() - )))), - "error" - ); - assert_eq!(action_result_outcome(&response(None)), "missing"); - assert_eq!( - action_result_outcome(&ActionExecutionResponse::default()), - "missing" - ); - } -} diff --git a/src/server/dynamic_server.rs b/src/server/dynamic_server.rs index afd6bd1..48eea54 100644 --- a/src/server/dynamic_server.rs +++ b/src/server/dynamic_server.rs @@ -1,3 +1,7 @@ +//! Assembles and serves the full gRPC service set used in dynamic mode: +//! action transfer, module registration, runtime execution results, and +//! runtime status, all gated behind the Sagittarius readiness interceptor. + use crate::{ configuration::{config::Config, service::ServiceConfiguration, state::AppReadiness}, sagittarius::{ @@ -6,7 +10,7 @@ use crate::{ test_execution_client_impl::SagittariusExecutionResponseSender, }, server::{ - action_transfer_service_server_impl::AquilaActionTransferServiceServer, + action_transfer::AquilaActionTransferServiceServer, create_readiness_interceptor, module_service_server_impl::AquilaModuleServiceServer, runtime_execution_service_server_impl::AquilaExecutionServiceServer, runtime_status_service_server_impl::AquilaRuntimeStatusServiceServer, @@ -88,6 +92,10 @@ impl AquilaDynamicServer { } } + /// Builds every service and blocks serving them until the listener + /// shuts down. Each service gets its own Sagittarius-backed client + /// (module updates, runtime status) so a slow or failing call on one + /// service can't stall another. pub async fn start(&self) -> Result<(), tonic::transport::Error> { let module_service = Arc::new(Mutex::new(SagittariusModuleServiceClient::new( self.channel.clone(), diff --git a/src/server/interceptor.rs b/src/server/interceptor.rs index 387bd70..03ffe07 100644 --- a/src/server/interceptor.rs +++ b/src/server/interceptor.rs @@ -1,7 +1,14 @@ +//! A tonic interceptor that rejects requests up front while a named +//! dependency isn't ready yet, instead of letting them into a handler that +//! would just fail partway through. + use crate::configuration::state::AppReadiness; use std::sync::Arc; use tonic::{Request, Status}; +/// Builds an interceptor that rejects every request with `Unavailable` +/// until `readiness` reports ready. `dependency_name` is only used for the +/// log line and error message. pub fn create_readiness_interceptor( readiness: Arc, dependency_name: &'static str, diff --git a/src/server/mod.rs b/src/server/mod.rs index e98eba2..a611870 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -1,4 +1,8 @@ -mod action_transfer_service_server_impl; +//! Aquila's gRPC surface: the services actions and runtimes connect to. +//! [`dynamic_server`] and [`static_server`] assemble the actual service set +//! for each run mode from the individual `*_service_server_impl` modules. + +mod action_transfer; mod interceptor; mod module_service_server_impl; mod runtime_execution_service_server_impl; diff --git a/src/server/module_service_server_impl.rs b/src/server/module_service_server_impl.rs index 4b5c766..7c46b31 100644 --- a/src/server/module_service_server_impl.rs +++ b/src/server/module_service_server_impl.rs @@ -1,3 +1,7 @@ +//! gRPC server for `ModuleService.Update`: authenticates the caller against +//! the first module's identifier, then relays the update to Sagittarius +//! unchanged. + use crate::{ authorization::authorization::extract_token, configuration::service::ServiceConfiguration, sagittarius::module_service_client_impl::SagittariusModuleServiceClient, @@ -44,6 +48,8 @@ impl ModuleService for AquilaModuleServiceServer { }; let modules_update_request = request.into_inner(); + // Every caller only ever sends modules for its own identifier, so + // authenticating against the first one covers the whole batch. let first_module_identifier = modules_update_request.modules.first(); let module_name = match first_module_identifier { diff --git a/src/server/runtime_execution_service_server_impl.rs b/src/server/runtime_execution_service_server_impl.rs index 05b0cc8..29255c2 100644 --- a/src/server/runtime_execution_service_server_impl.rs +++ b/src/server/runtime_execution_service_server_impl.rs @@ -1,3 +1,7 @@ +//! gRPC server for `ExecutionService.Update`: the endpoint the Taurus +//! runtime posts execution results to, which are then relayed onto the +//! Sagittarius execution stream via [`SagittariusExecutionResponseSender`]. + use crate::{ authorization::authorization::extract_token, configuration::service::ServiceConfiguration, sagittarius::test_execution_client_impl::SagittariusExecutionResponseSender, @@ -50,6 +54,9 @@ impl ExecutionService for AquilaExecutionServiceServer { } }; + // This endpoint is only ever called by the Taurus runtime, so the + // token is checked against that fixed identifier rather than one + // read from the request. if !self .service_configuration .has_runtime(&token, &String::from("taurus")) diff --git a/src/server/runtime_status_service_server_impl.rs b/src/server/runtime_status_service_server_impl.rs deleted file mode 100644 index fbcf6f6..0000000 --- a/src/server/runtime_status_service_server_impl.rs +++ /dev/null @@ -1,304 +0,0 @@ -use crate::{ - authorization::authorization::extract_token, configuration::service::ServiceConfiguration, - sagittarius::runtime_status_service_client_impl::SagittariusRuntimeStatusServiceClient, -}; -use std::{ - collections::HashMap, - sync::Arc, - time::{Duration, Instant, SystemTime, UNIX_EPOCH}, -}; -use tokio::sync::Mutex; -use tonic::Status; -use tucana::{ - aquila::{RuntimeStatusUpdateRequest, runtime_status_service_server::RuntimeStatusService}, - shared::{ModuleStatus, module_status}, -}; - -#[derive(Clone)] -struct RuntimeStatusSnapshot { - status: ModuleStatus, -} - -impl RuntimeStatusSnapshot { - fn from_update(request: &RuntimeStatusUpdateRequest) -> Option { - let status = request.status.as_ref()?; - - if status.identifier.is_empty() { - None - } else { - Some(Self { - status: status.clone(), - }) - } - } - - fn key(&self) -> String { - self.status.identifier.clone() - } - - fn identifier(&self) -> &str { - &self.status.identifier - } - - fn is_stopped(&self) -> bool { - self.status.status == module_status::StatusVariant::Stopped as i32 - } - - fn not_responding_update(&self) -> RuntimeStatusUpdateRequest { - let mut next_status = self.status.clone(); - next_status.status = module_status::StatusVariant::NotResponding as i32; - next_status.timestamp = epoch_seconds_now(); - - RuntimeStatusUpdateRequest { - status: Some(next_status), - } - } - - fn stopped_update(&self) -> RuntimeStatusUpdateRequest { - let mut next_status = self.status.clone(); - next_status.status = module_status::StatusVariant::Stopped as i32; - next_status.timestamp = epoch_seconds_now(); - - RuntimeStatusUpdateRequest { - status: Some(next_status), - } - } -} - -struct TrackedRuntime { - last_seen: Instant, - last_status: RuntimeStatusSnapshot, - not_responding_since: Option, -} - -pub struct AquilaRuntimeStatusServiceServer { - client: Arc>, - service_configuration: ServiceConfiguration, - tracked_runtimes: Arc>>, - not_responding_after: Duration, - stopped_after_not_responding: Duration, -} - -impl AquilaRuntimeStatusServiceServer { - pub fn new( - client: Arc>, - service_configuration: ServiceConfiguration, - not_responding_after: Duration, - stopped_after_not_responding: Duration, - monitor_interval: Duration, - ) -> Self { - let server = Self { - client, - service_configuration, - tracked_runtimes: Arc::new(Mutex::new(HashMap::new())), - not_responding_after, - stopped_after_not_responding, - }; - - server.spawn_timeout_monitor(monitor_interval); - server - } - - fn spawn_timeout_monitor(&self, monitor_interval: Duration) { - let tracked_runtimes = self.tracked_runtimes.clone(); - let _client = self.client.clone(); - let not_responding_after = self.not_responding_after; - let stopped_after_not_responding = self.stopped_after_not_responding; - - tokio::spawn(async move { - let mut interval = tokio::time::interval(monitor_interval); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); - - loop { - interval.tick().await; - - let timeout_updates = { - let mut tracked = tracked_runtimes.lock().await; - collect_timeout_updates( - &mut tracked, - Instant::now(), - not_responding_after, - stopped_after_not_responding, - ) - }; - - if timeout_updates.is_empty() { - continue; - } - - // Uncomment in #360 - // let mut client = _client.lock().await; - for _timeout_update in timeout_updates { - // let _ = client.update_runtime_status(timeout_update).await; - } - } - }); - } - - async fn track_runtime_update( - &self, - runtime_status_update_request: &RuntimeStatusUpdateRequest, - ) { - let Some(snapshot) = RuntimeStatusSnapshot::from_update(runtime_status_update_request) - else { - log::debug!( - "Skipping runtime heartbeat tracking because status payload is missing or identifier is empty." - ); - return; - }; - - let key = snapshot.key(); - let now = Instant::now(); - let mut tracked = self.tracked_runtimes.lock().await; - - if snapshot.is_stopped() { - tracked.remove(&key); - return; - } - - match tracked.get_mut(&key) { - Some(runtime) => { - if runtime.not_responding_since.is_some() { - log::info!( - "Runtime '{}' sent heartbeat again. Clearing NOT_RESPONDING flag.", - snapshot.identifier() - ); - } - runtime.last_seen = now; - runtime.last_status = snapshot; - runtime.not_responding_since = None; - } - None => { - tracked.insert( - key, - TrackedRuntime { - last_seen: now, - last_status: snapshot, - not_responding_since: None, - }, - ); - } - } - } -} - -fn collect_timeout_updates( - tracked: &mut HashMap, - now: Instant, - not_responding_after: Duration, - stopped_after_not_responding: Duration, -) -> Vec { - let mut timeout_updates = Vec::new(); - let mut runtimes_to_remove = Vec::new(); - - for (key, runtime) in tracked.iter_mut() { - let silence = now.duration_since(runtime.last_seen); - - match runtime.not_responding_since { - None if silence >= not_responding_after => { - log::warn!( - "Runtime '{}' has not sent status for {:?}. Marking as NOT_RESPONDING.", - runtime.last_status.identifier(), - silence - ); - runtime.not_responding_since = Some(now); - timeout_updates.push(runtime.last_status.not_responding_update()); - } - Some(since) if now.duration_since(since) >= stopped_after_not_responding => { - log::warn!( - "Runtime '{}' stayed NOT_RESPONDING for {:?}. Marking as STOPPED.", - runtime.last_status.identifier(), - now.duration_since(since) - ); - timeout_updates.push(runtime.last_status.stopped_update()); - runtimes_to_remove.push(key.clone()); - } - _ => {} - } - } - - for key in runtimes_to_remove { - tracked.remove(&key); - } - - timeout_updates -} - -fn epoch_seconds_now() -> i64 { - match SystemTime::now().duration_since(UNIX_EPOCH) { - Ok(duration) => duration.as_secs() as i64, - Err(error) => { - log::warn!("System time before UNIX_EPOCH: {:?}", error); - 0 - } - } -} - -#[tonic::async_trait] -impl RuntimeStatusService for AquilaRuntimeStatusServiceServer { - #[tracing::instrument( - name = "aquila.runtime_status.update", - skip_all, - fields(rpc.system = "grpc", rpc.service = "RuntimeStatusService", rpc.method = "Update") - )] - async fn update( - &self, - request: tonic::Request, - ) -> Result, tonic::Status> { - let token = match extract_token(&request) { - Ok(t) => t.to_string(), - Err(status) => { - log::warn!("Rejected runtime status update reason=missing_or_invalid_token"); - return Err(status); - } - }; - - let runtime_status_update_request = request.into_inner(); - - let runtime_identifier = match runtime_status_update_request.status.as_ref() { - Some(status) => status.identifier.clone(), - None => return Err(Status::invalid_argument("missing runtime status payload")), - }; - - if runtime_identifier.is_empty() { - return Err(Status::invalid_argument("runtime identifier is missing")); - } - - if !self - .service_configuration - .has_service(&token, &runtime_identifier) - { - log::warn!( - "Rejected runtime status update reason=token_not_registered runtime_identifier={}", - runtime_identifier - ); - return Err(Status::unauthenticated("token is not valid")); - } - self.track_runtime_update(&runtime_status_update_request) - .await; - - log::debug!( - "Received runtime status update runtime_identifier={}", - runtime_identifier - ); - - // Temporarily disabled: do not forward runtime status updates to Sagittarius. - // Re-enable with the timeout updates above when runtime usage reporting is restored. - // let mut client = self.client.lock().await; - // let response = client - // .update_runtime_status(runtime_status_update_request) - // .await; - let response = tucana::aquila::RuntimeStatusUpdateResponse { success: true }; - - log::debug!( - "Completed runtime status update success={}", - response.success - ); - - Ok(tonic::Response::new( - tucana::aquila::RuntimeStatusUpdateResponse { - success: response.success, - }, - )) - } -} diff --git a/src/server/runtime_status_service_server_impl/mod.rs b/src/server/runtime_status_service_server_impl/mod.rs new file mode 100644 index 0000000..cdf11a6 --- /dev/null +++ b/src/server/runtime_status_service_server_impl/mod.rs @@ -0,0 +1,128 @@ +//! gRPC server for runtime heartbeats (`RuntimeStatusService.Update`). +//! +//! Accepting a heartbeat and detecting a runtime's silence are two +//! different lifetimes: the RPC handler below only needs to authenticate +//! and record the heartbeat, while a background tick (spawned once at +//! construction, see [`monitor`]) is what actually notices when a runtime +//! stops sending them. Both share the same [`registry::TrackedRuntimeRegistry`]. + +mod monitor; +mod registry; + +use std::{sync::Arc, time::Duration}; + +use tokio::sync::Mutex; +use tonic::Status; +use tucana::aquila::runtime_status_service_server::RuntimeStatusService; + +use crate::{ + authorization::authorization::extract_token, configuration::service::ServiceConfiguration, + sagittarius::runtime_status_service_client_impl::SagittariusRuntimeStatusServiceClient, +}; + +use registry::TrackedRuntimeRegistry; + +pub struct AquilaRuntimeStatusServiceServer { + client: Arc>, + service_configuration: ServiceConfiguration, + tracked_runtimes: TrackedRuntimeRegistry, + not_responding_after: Duration, + stopped_after_not_responding: Duration, +} + +impl AquilaRuntimeStatusServiceServer { + pub fn new( + client: Arc>, + service_configuration: ServiceConfiguration, + not_responding_after: Duration, + stopped_after_not_responding: Duration, + monitor_interval: Duration, + ) -> Self { + let server = Self { + client, + service_configuration, + tracked_runtimes: TrackedRuntimeRegistry::default(), + not_responding_after, + stopped_after_not_responding, + }; + + monitor::spawn( + server.tracked_runtimes.clone(), + server.client.clone(), + server.not_responding_after, + server.stopped_after_not_responding, + monitor_interval, + ); + server + } +} + +#[tonic::async_trait] +impl RuntimeStatusService for AquilaRuntimeStatusServiceServer { + #[tracing::instrument( + name = "aquila.runtime_status.update", + skip_all, + fields(rpc.system = "grpc", rpc.service = "RuntimeStatusService", rpc.method = "Update") + )] + async fn update( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + let token = match extract_token(&request) { + Ok(t) => t.to_string(), + Err(status) => { + log::warn!("Rejected runtime status update reason=missing_or_invalid_token"); + return Err(status); + } + }; + + let runtime_status_update_request = request.into_inner(); + + let runtime_identifier = match runtime_status_update_request.status.as_ref() { + Some(status) => status.identifier.clone(), + None => return Err(Status::invalid_argument("missing runtime status payload")), + }; + + if runtime_identifier.is_empty() { + return Err(Status::invalid_argument("runtime identifier is missing")); + } + + if !self + .service_configuration + .has_service(&token, &runtime_identifier) + { + log::warn!( + "Rejected runtime status update reason=token_not_registered runtime_identifier={}", + runtime_identifier + ); + return Err(Status::unauthenticated("token is not valid")); + } + self.tracked_runtimes + .record_heartbeat(&runtime_status_update_request) + .await; + + log::debug!( + "Received runtime status update runtime_identifier={}", + runtime_identifier + ); + + // Temporarily disabled: do not forward runtime status updates to Sagittarius. + // Re-enable with the timeout updates above when runtime usage reporting is restored. + // let mut client = self.client.lock().await; + // let response = client + // .update_runtime_status(runtime_status_update_request) + // .await; + let response = tucana::aquila::RuntimeStatusUpdateResponse { success: true }; + + log::debug!( + "Completed runtime status update success={}", + response.success + ); + + Ok(tonic::Response::new( + tucana::aquila::RuntimeStatusUpdateResponse { + success: response.success, + }, + )) + } +} diff --git a/src/server/runtime_status_service_server_impl/monitor.rs b/src/server/runtime_status_service_server_impl/monitor.rs new file mode 100644 index 0000000..e52cda0 --- /dev/null +++ b/src/server/runtime_status_service_server_impl/monitor.rs @@ -0,0 +1,50 @@ +//! Background task that periodically checks tracked runtimes for timeouts +//! and would forward the resulting status transitions to Sagittarius. +//! +//! Forwarding is currently disabled pending #360 — the tick still runs and +//! computes the updates so the logic stays exercised, it just doesn't send +//! them anywhere yet. + +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use tokio::sync::Mutex; + +use crate::sagittarius::runtime_status_service_client_impl::SagittariusRuntimeStatusServiceClient; + +use super::registry::TrackedRuntimeRegistry; + +/// Spawns the periodic timeout check. Runs until the process exits; there is +/// no cancellation handle because the server itself never tears this down. +pub(super) fn spawn( + registry: TrackedRuntimeRegistry, + _client: Arc>, + not_responding_after: Duration, + stopped_after_not_responding: Duration, + monitor_interval: Duration, +) { + tokio::spawn(async move { + let mut interval = tokio::time::interval(monitor_interval); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + + loop { + interval.tick().await; + + let timeout_updates = registry + .collect_timeout_updates(Instant::now(), not_responding_after, stopped_after_not_responding) + .await; + + if timeout_updates.is_empty() { + continue; + } + + // Uncomment in #360 + // let mut client = _client.lock().await; + for _timeout_update in timeout_updates { + // let _ = client.update_runtime_status(timeout_update).await; + } + } + }); +} diff --git a/src/server/runtime_status_service_server_impl/registry.rs b/src/server/runtime_status_service_server_impl/registry.rs new file mode 100644 index 0000000..54a63ed --- /dev/null +++ b/src/server/runtime_status_service_server_impl/registry.rs @@ -0,0 +1,190 @@ +//! Tracks the last known status of every runtime that has sent a heartbeat, +//! and derives NOT_RESPONDING / STOPPED transitions from how long a runtime +//! has gone silent. This is pure bookkeeping — actually forwarding those +//! transitions to Sagittarius is [`super::monitor`]'s job. + +use std::{ + collections::HashMap, + sync::Arc, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use tokio::sync::Mutex; +use tucana::{aquila::RuntimeStatusUpdateRequest, shared::ModuleStatus, shared::module_status}; + +#[derive(Clone)] +pub(super) struct RuntimeStatusSnapshot { + status: ModuleStatus, +} + +impl RuntimeStatusSnapshot { + fn from_update(request: &RuntimeStatusUpdateRequest) -> Option { + let status = request.status.as_ref()?; + + if status.identifier.is_empty() { + None + } else { + Some(Self { + status: status.clone(), + }) + } + } + + fn key(&self) -> String { + self.status.identifier.clone() + } + + pub(super) fn identifier(&self) -> &str { + &self.status.identifier + } + + fn is_stopped(&self) -> bool { + self.status.status == module_status::StatusVariant::Stopped as i32 + } + + fn not_responding_update(&self) -> RuntimeStatusUpdateRequest { + let mut next_status = self.status.clone(); + next_status.status = module_status::StatusVariant::NotResponding as i32; + next_status.timestamp = epoch_seconds_now(); + + RuntimeStatusUpdateRequest { + status: Some(next_status), + } + } + + fn stopped_update(&self) -> RuntimeStatusUpdateRequest { + let mut next_status = self.status.clone(); + next_status.status = module_status::StatusVariant::Stopped as i32; + next_status.timestamp = epoch_seconds_now(); + + RuntimeStatusUpdateRequest { + status: Some(next_status), + } + } +} + +struct TrackedRuntime { + last_seen: Instant, + last_status: RuntimeStatusSnapshot, + /// `Some` once a timeout tick has flagged this runtime as NOT_RESPONDING; + /// cleared the moment another heartbeat arrives. + not_responding_since: Option, +} + +/// Shared, lock-protected map of runtime identifier to its last known status. +#[derive(Clone, Default)] +pub(super) struct TrackedRuntimeRegistry { + tracked: Arc>>, +} + +impl TrackedRuntimeRegistry { + /// Records a heartbeat, or drops tracking entirely if the runtime + /// reports itself STOPPED — a stopped runtime should not later be + /// timed out as if it had gone silent. + pub(super) async fn record_heartbeat(&self, request: &RuntimeStatusUpdateRequest) { + let Some(snapshot) = RuntimeStatusSnapshot::from_update(request) else { + log::debug!( + "Skipping runtime heartbeat tracking because status payload is missing or identifier is empty." + ); + return; + }; + + let key = snapshot.key(); + let now = Instant::now(); + let mut tracked = self.tracked.lock().await; + + if snapshot.is_stopped() { + tracked.remove(&key); + return; + } + + match tracked.get_mut(&key) { + Some(runtime) => { + if runtime.not_responding_since.is_some() { + log::info!( + "Runtime '{}' sent heartbeat again. Clearing NOT_RESPONDING flag.", + snapshot.identifier() + ); + } + runtime.last_seen = now; + runtime.last_status = snapshot; + runtime.not_responding_since = None; + } + None => { + tracked.insert( + key, + TrackedRuntime { + last_seen: now, + last_status: snapshot, + not_responding_since: None, + }, + ); + } + } + } + + /// Walks every tracked runtime and returns the status updates implied by + /// how long each has gone silent, advancing (and, for STOPPED, removing) + /// their tracked state in the process. + pub(super) async fn collect_timeout_updates( + &self, + now: Instant, + not_responding_after: Duration, + stopped_after_not_responding: Duration, + ) -> Vec { + let mut tracked = self.tracked.lock().await; + collect_timeout_updates(&mut tracked, now, not_responding_after, stopped_after_not_responding) + } +} + +fn collect_timeout_updates( + tracked: &mut HashMap, + now: Instant, + not_responding_after: Duration, + stopped_after_not_responding: Duration, +) -> Vec { + let mut timeout_updates = Vec::new(); + let mut runtimes_to_remove = Vec::new(); + + for (key, runtime) in tracked.iter_mut() { + let silence = now.duration_since(runtime.last_seen); + + match runtime.not_responding_since { + None if silence >= not_responding_after => { + log::warn!( + "Runtime '{}' has not sent status for {:?}. Marking as NOT_RESPONDING.", + runtime.last_status.identifier(), + silence + ); + runtime.not_responding_since = Some(now); + timeout_updates.push(runtime.last_status.not_responding_update()); + } + Some(since) if now.duration_since(since) >= stopped_after_not_responding => { + log::warn!( + "Runtime '{}' stayed NOT_RESPONDING for {:?}. Marking as STOPPED.", + runtime.last_status.identifier(), + now.duration_since(since) + ); + timeout_updates.push(runtime.last_status.stopped_update()); + runtimes_to_remove.push(key.clone()); + } + _ => {} + } + } + + for key in runtimes_to_remove { + tracked.remove(&key); + } + + timeout_updates +} + +fn epoch_seconds_now() -> i64 { + match SystemTime::now().duration_since(UNIX_EPOCH) { + Ok(duration) => duration.as_secs() as i64, + Err(error) => { + log::warn!("System time before UNIX_EPOCH: {:?}", error); + 0 + } + } +} diff --git a/src/server/static_server.rs b/src/server/static_server.rs index 7b12ad5..56a88e1 100644 --- a/src/server/static_server.rs +++ b/src/server/static_server.rs @@ -1,7 +1,11 @@ +//! Assembles and serves the gRPC service set used in static mode: only +//! action transfer, since there's no Sagittarius to register modules with +//! or report runtime status to. + use crate::{ configuration::{config::Config, service::ServiceConfiguration, state::AppReadiness}, server::{ - action_transfer_service_server_impl::AquilaActionTransferServiceServer, + action_transfer::AquilaActionTransferServiceServer, create_readiness_interceptor, }, }; diff --git a/src/startup/dynamic_mode.rs b/src/startup/dynamic_mode.rs index b936bc6..732b28a 100644 --- a/src/startup/dynamic_mode.rs +++ b/src/startup/dynamic_mode.rs @@ -1,3 +1,9 @@ +//! Dynamic mode wiring: the gRPC server and two independent, self-healing +//! Sagittarius streams (flow sync, test execution) all run as separate +//! tasks supervised by a single `select!` — if any one of them exits or +//! panics, the others are aborted and Aquila shuts down rather than +//! continuing in a partially working state. + use async_nats::Client; use crate::{ @@ -16,6 +22,9 @@ use crate::{ }; use std::{sync::Arc, time::Duration}; +/// Starts the gRPC server plus the flow-sync and test-execution stream +/// tasks, and blocks until one of them exits, panics, or a shutdown signal +/// arrives. pub async fn run( config: AquilaConfig, app_readiness: AppReadiness, @@ -79,6 +88,10 @@ pub async fn run( crate::configuration::env::Environment::Production => String::from("PRODUCTION"), }; + // Both stream tasks below share the same shape: connect, run the stream + // until it ends (which `logon`/`init_flow_stream` always eventually do, + // since Sagittarius connections aren't permanent), then reconnect with + // growing backoff. The task itself never exits on its own. let mut test_execution_task = tokio::spawn(async move { let mut backoff = Duration::from_millis(200); let max_backoff = Duration::from_secs(10); diff --git a/src/startup/mod.rs b/src/startup/mod.rs index 25ba8a2..383f9dd 100644 --- a/src/startup/mod.rs +++ b/src/startup/mod.rs @@ -1,3 +1,7 @@ +//! Bootstraps the NATS/JetStream connection shared by both run modes, then +//! hands off to [`static_mode`] or [`dynamic_mode`] depending on +//! [`AquilaConfig::is_static`]. + pub mod dynamic_mode; pub mod static_mode; @@ -7,6 +11,10 @@ use crate::configuration::{ use async_nats::jetstream::kv::Config; use std::sync::Arc; +/// Connects to NATS, ensures the flow KV bucket exists, and starts the +/// appropriate mode. Panics on any of those failures — none of them are +/// recoverable without operator intervention, so there's no useful degraded +/// mode to fall back to. pub async fn run( config: AquilaConfig, app_readiness: AppReadiness, diff --git a/src/startup/static_mode.rs b/src/startup/static_mode.rs index b862f4c..0034eca 100644 --- a/src/startup/static_mode.rs +++ b/src/startup/static_mode.rs @@ -1,3 +1,8 @@ +//! Static mode wiring: load a fixed flow export from disk into the KV +//! store once at startup, then serve the gRPC server with no ongoing +//! Sagittarius dependency. Readiness is set unconditionally since there's +//! nothing external left to wait on. + use crate::{ configuration::{config::Config, service::ServiceConfiguration, state::AppReadiness}, flow::get_flow_identifier, @@ -10,6 +15,8 @@ use serde_json::from_str; use std::{fs::File, io::Read, sync::Arc, sync::atomic::Ordering}; use tucana::shared::Flows; +/// Loads the fallback flow export and serves the static gRPC server until a +/// shutdown signal arrives. pub async fn run( config: Config, app_readiness: AppReadiness, @@ -86,6 +93,10 @@ pub async fn run( log::info!("Aquila shutdown complete"); } +/// Reads `path` as a JSON [`Flows`] export and stores each flow in the KV +/// store. Panics on any read/parse failure, since static mode has no flows +/// to serve without this file and continuing would just push the failure +/// downstream to the first action/runtime request. async fn init_flows_from_json( path: String, flow_store_client: Arc, diff --git a/src/telemetry/metrics.rs b/src/telemetry/metrics.rs index 7dd1d39..19b928d 100644 --- a/src/telemetry/metrics.rs +++ b/src/telemetry/metrics.rs @@ -1,3 +1,8 @@ +//! Aquila's OpenTelemetry metric instruments. Every recording function here +//! is a no-op until [`initialize`] has run (guarded by [`METRICS`] being +//! unset), so callers never need to check whether telemetry is enabled — +//! this module absorbs that instead of every call site doing so. + use std::sync::OnceLock; use opentelemetry::{ @@ -20,6 +25,10 @@ struct Metrics { action_failures: Counter, } +/// Registers every metric instrument against the global meter. Must be +/// called once during startup before any of the recording functions below +/// will actually emit anything; safe to call at most once (a second call is +/// silently dropped by [`OnceLock`]). pub fn initialize() { let meter = opentelemetry::global::meter(env!("CARGO_PKG_NAME")); let _ = METRICS.set(Metrics { diff --git a/src/telemetry/mod.rs b/src/telemetry/mod.rs index 98f9dc3..d76445b 100644 --- a/src/telemetry/mod.rs +++ b/src/telemetry/mod.rs @@ -1,3 +1,9 @@ +//! Aquila's telemetry surface: logging/tracing/error reporting setup comes +//! from the shared `code0_flow` crate, re-exported here so the rest of the +//! codebase depends on `crate::telemetry` rather than reaching into that +//! crate directly. [`metrics`] is Aquila-specific: the OpenTelemetry +//! instruments this service emits. + pub mod metrics; pub use code0_flow::flow_telemetry::{OpenTelemetry, Telemetry, TelemetrySettings, errors};