diff --git a/aquila.yml b/aquila.yml index 8762d82..2e5db1a 100644 --- a/aquila.yml +++ b/aquila.yml @@ -47,6 +47,9 @@ grpc: # Expose the standard gRPC health service. health_service: false + # Maximum number of post-logon messages handled concurrently per action stream. + action_transfer_concurrency_limit: 256 + runtime_status: # Frequency of Aquila's own heartbeat to the Sagittarius gateway, in minutes. heartbeat_interval_minutes: 5 diff --git a/docs/installation.mdx b/docs/installation.mdx index 5ab807d..88a1d2f 100644 --- a/docs/installation.mdx +++ b/docs/installation.mdx @@ -101,6 +101,7 @@ provide a more specific runtime filter, for example `RUST_LOG=aquila::server=tra | `grpc.host` | Aquila gRPC bind host. | | `grpc.port` | Aquila gRPC bind port. | | `grpc.health_service` | Enables the gRPC health service. | +| `grpc.action_transfer_concurrency_limit` | Maximum concurrent post-logon handlers per action stream (default: `256`). | | `runtime_status.not_responding_after_secs` | Heartbeat timeout before `not_responding`. | | `runtime_status.stopped_after_not_responding_secs` | Additional timeout before `stopped`. | | `runtime_status.monitor_interval_secs` | Heartbeat monitor interval. | diff --git a/src/configuration/config/display.rs b/src/configuration/config/display.rs index 012f21c..e110202 100644 --- a/src/configuration/config/display.rs +++ b/src/configuration/config/display.rs @@ -50,6 +50,11 @@ impl fmt::Display for Config { " Health service: {}", self.grpc.health_service )?; + writeln!( + formatter, + " Action transfer concurrency: {}", + self.grpc.action_transfer_concurrency_limit + )?; writeln!(formatter, " Static mode")?; writeln!(formatter, " Flow path: {}", self.static_config.flow_path)?; writeln!(formatter, " Dynamic mode")?; @@ -93,6 +98,7 @@ mod tests { assert!(output.starts_with("Aquila configuration\n")); assert!(output.contains(" Environment: development")); assert!(output.contains(" Address: 127.0.0.1:8081")); + assert!(output.contains(" Action transfer concurrency: 256")); assert!(output.contains(" Request timeout: 5s")); assert!(output.contains(" Backend token: [FILTERED]")); assert!(!output.contains("super-secret")); diff --git a/src/configuration/config/mod.rs b/src/configuration/config/mod.rs index 61375e4..1be082a 100644 --- a/src/configuration/config/mod.rs +++ b/src/configuration/config/mod.rs @@ -74,6 +74,9 @@ pub struct Grpc { pub host: String, pub port: u16, pub health_service: bool, + /// Maximum number of post-logon messages processed concurrently for one + /// action transfer stream. + pub action_transfer_concurrency_limit: usize, } #[derive(Clone, Debug, Deserialize, Serialize)] @@ -138,6 +141,7 @@ impl Default for Grpc { host: "127.0.0.1".into(), port: 8081, health_service: false, + action_transfer_concurrency_limit: 256, } } } @@ -261,4 +265,25 @@ mod tests { fn opentelemetry_default_service_name_is_aquila() { assert_eq!(Config::default().opentelemetry.service_name, "aquila"); } + + #[test] + fn action_transfer_concurrency_limit_is_configurable() { + let config: Config = ConfigLoader::builder() + .add_source( + ConfigLoader::try_from(&Config::default()) + .expect("default configuration should serialize"), + ) + .set_override("grpc.action_transfer_concurrency_limit", 17) + .expect("concurrency override should apply") + .build() + .expect("configuration should build") + .try_deserialize() + .expect("configuration should deserialize"); + + assert_eq!( + Config::default().grpc.action_transfer_concurrency_limit, + 256 + ); + assert_eq!(config.grpc.action_transfer_concurrency_limit, 17); + } } diff --git a/src/server/action_transfer/mod.rs b/src/server/action_transfer/mod.rs index bb3669d..066c086 100644 --- a/src/server/action_transfer/mod.rs +++ b/src/server/action_transfer/mod.rs @@ -19,11 +19,14 @@ mod shard_registry; pub use flow_execution_registry::ActionFlowExecutionRegistry; pub use shard_registry::{ActionShardRegistry, ShardAssignment}; -use std::{pin::Pin, sync::Arc}; +use std::{future::Future, pin::Pin, sync::Arc}; use futures::StreamExt; use futures_core::Stream; -use tokio::sync::Mutex; +use tokio::{ + sync::{Mutex, Semaphore}, + task::JoinSet, +}; use tokio_stream::wrappers::ReceiverStream; use tonic::Status; use tracing::Instrument; @@ -38,7 +41,9 @@ use crate::{ }; use logon::{extract_token, handle_logon}; -use nats_bridge::{handle_flow_execution, handle_result, handle_sub_flow_execution, send_stream_error}; +use nats_bridge::{ + handle_flow_execution, handle_result, handle_sub_flow_execution, send_stream_error, +}; use pending_replies::PendingReplyStore; /// Every dependency an action's connection needs, bundled into one @@ -69,6 +74,68 @@ pub(super) struct ActionTransferContext { pub(super) shard_registry: ActionShardRegistry, /// Whether Aquila is running in static mode, which changes how config updates are sourced. pub(super) is_static: bool, + /// Per-connection limit for post-logon message handlers. Parsing stays on + /// the stream task; handler work runs in this bounded task set. + pub(super) concurrency_limit: usize, +} + +/// Owns every post-logon handler spawned for one action stream. +/// +/// Permits are acquired inside the spawned tasks so a slow handler never +/// prevents the stream task from parsing later messages or noticing EOF. The +/// join set gives the connection a single place to reap completed work and to +/// cancel everything that is still running when the stream closes. +struct PostLogonTaskSet { + semaphore: Arc, + tasks: JoinSet<()>, +} + +impl PostLogonTaskSet { + fn new(concurrency_limit: usize) -> Self { + assert!( + concurrency_limit > 0, + "grpc.action_transfer_concurrency_limit must be at least 1" + ); + + Self { + semaphore: Arc::new(Semaphore::new(concurrency_limit)), + tasks: JoinSet::new(), + } + } + + fn spawn(&mut self, future: F) + where + F: Future + Send + 'static, + { + self.reap_finished(); + + let semaphore = self.semaphore.clone(); + self.tasks.spawn(async move { + let Ok(_permit) = semaphore.acquire_owned().await else { + return; + }; + future.await; + }); + } + + fn reap_finished(&mut self) { + while let Some(result) = self.tasks.try_join_next() { + if let Err(error) = result { + log::warn!("Action transfer message handler failed: {error}"); + } + } + } + + async fn shutdown(&mut self) { + self.tasks.abort_all(); + while let Some(result) = self.tasks.join_next().await { + if let Err(error) = result + && !error.is_cancelled() + { + log::warn!("Action transfer message handler failed during shutdown: {error}"); + } + } + } } /// Implements the `ActionTransfer` gRPC service that a connected action @@ -126,6 +193,7 @@ impl ActionTransferService for AquilaActionTransferServiceServer { action.identifier = tracing::field::Empty ); tokio::spawn(async move { + let mut post_logon_tasks = PostLogonTaskSet::new(context.concurrency_limit); let mut cfg_forwarder_started = false; let mut flow_forwarder_started = false; let mut connected_at = None; @@ -268,13 +336,17 @@ impl ActionTransferService for AquilaActionTransferServiceServer { identifier ); - handle_result( - &identifier, - execution_result, - context.client.clone(), - pending_replies.clone(), - ) - .await; + let client = context.client.clone(); + let pending_replies = pending_replies.clone(); + post_logon_tasks.spawn(async move { + handle_result( + &identifier, + execution_result, + client, + pending_replies, + ) + .await; + }); } tucana::aquila::action_transfer_request::Data::SubFlowExecution(request) => { log::debug!( @@ -283,13 +355,11 @@ impl ActionTransferService for AquilaActionTransferServiceServer { request.execution_identifier ); - handle_sub_flow_execution( - &identifier, - request, - context.client.clone(), - tx.clone(), - ) - .await; + let client = context.client.clone(); + let tx = tx.clone(); + post_logon_tasks.spawn(async move { + handle_sub_flow_execution(&identifier, request, client, tx).await; + }); } tucana::aquila::action_transfer_request::Data::FlowExecution(request) => { log::debug!( @@ -299,19 +369,30 @@ impl ActionTransferService for AquilaActionTransferServiceServer { request.flow_id ); - handle_flow_execution( - &identifier, - request, - context.kv.clone(), - context.client.clone(), - context.flow_execution_registry.clone(), - tx.clone(), - ) - .await; + let kv = context.kv.clone(); + let client = context.client.clone(); + let registry = context.flow_execution_registry.clone(); + let tx = tx.clone(); + post_logon_tasks.spawn(async move { + handle_flow_execution( + &identifier, + request, + kv, + client, + registry, + tx, + ) + .await; + }); } } } + // Handler futures may be blocked on NATS or waiting for a permit. + // Once the stream is gone none can produce a useful response, so + // cancel and join them before releasing per-connection state. + post_logon_tasks.shutdown().await; + if let Some(identifier) = connected_identifier { metrics::action_active(&identifier, -1); metrics::action_connection(&identifier, "closed"); @@ -349,3 +430,157 @@ async fn log_unclaimed_shards(context: &ActionTransferContext, identifier: &str, ); } } + +#[cfg(test)] +mod tests { + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + + use tokio::sync::{Notify, Semaphore, mpsc}; + + use super::PostLogonTaskSet; + + #[tokio::test] + async fn parallel_flow_starts_are_not_serialized() { + let mut tasks = PostLogonTaskSet::new(2); + let (started_tx, mut started_rx) = mpsc::unbounded_channel(); + let release = Arc::new(Semaphore::new(0)); + + for execution_id in ["flow-1", "flow-2"] { + let started_tx = started_tx.clone(); + let release = release.clone(); + tasks.spawn(async move { + started_tx.send(execution_id).unwrap(); + let _permit = release.acquire().await.unwrap(); + }); + } + + let first = tokio::time::timeout(std::time::Duration::from_secs(1), started_rx.recv()) + .await + .expect("first flow should start") + .unwrap(); + let second = tokio::time::timeout(std::time::Duration::from_secs(1), started_rx.recv()) + .await + .expect("second flow should start while the first is running") + .unwrap(); + + assert_ne!(first, second); + release.add_permits(2); + tasks.shutdown().await; + } + + #[tokio::test] + async fn slow_subflow_does_not_block_an_unrelated_flow() { + let mut tasks = PostLogonTaskSet::new(2); + let slow_subflow = Arc::new(Notify::new()); + let (subflow_started_tx, subflow_started_rx) = tokio::sync::oneshot::channel(); + let (flow_started_tx, flow_started_rx) = tokio::sync::oneshot::channel(); + + let slow_subflow_task = slow_subflow.clone(); + tasks.spawn(async move { + subflow_started_tx.send(()).unwrap(); + slow_subflow_task.notified().await; + }); + subflow_started_rx.await.unwrap(); + + tasks.spawn(async move { + flow_started_tx.send(()).unwrap(); + }); + + tokio::time::timeout(std::time::Duration::from_secs(1), flow_started_rx) + .await + .expect("flow should start while the subflow is still waiting") + .unwrap(); + + slow_subflow.notify_one(); + tasks.shutdown().await; + } + + #[tokio::test] + async fn configured_concurrency_limit_is_never_exceeded() { + const LIMIT: usize = 3; + const TASK_COUNT: usize = 12; + + let mut tasks = PostLogonTaskSet::new(LIMIT); + let active = Arc::new(AtomicUsize::new(0)); + let maximum = Arc::new(AtomicUsize::new(0)); + let release = Arc::new(Semaphore::new(0)); + let (started_tx, mut started_rx) = mpsc::unbounded_channel(); + let (finished_tx, mut finished_rx) = mpsc::unbounded_channel(); + + for _ in 0..TASK_COUNT { + let active = active.clone(); + let maximum = maximum.clone(); + let release = release.clone(); + let started_tx = started_tx.clone(); + let finished_tx = finished_tx.clone(); + tasks.spawn(async move { + let now_active = active.fetch_add(1, Ordering::SeqCst) + 1; + maximum.fetch_max(now_active, Ordering::SeqCst); + started_tx.send(()).unwrap(); + + let _release = release.acquire().await.unwrap(); + active.fetch_sub(1, Ordering::SeqCst); + finished_tx.send(()).unwrap(); + }); + } + + for _ in 0..LIMIT { + tokio::time::timeout(std::time::Duration::from_secs(1), started_rx.recv()) + .await + .expect("task up to the limit should start") + .unwrap(); + } + assert!( + tokio::time::timeout(std::time::Duration::from_millis(50), started_rx.recv()) + .await + .is_err(), + "a task beyond the limit started before a permit was released" + ); + + release.add_permits(TASK_COUNT); + for _ in 0..TASK_COUNT { + tokio::time::timeout(std::time::Duration::from_secs(1), finished_rx.recv()) + .await + .expect("all tasks should finish") + .unwrap(); + } + + assert_eq!(maximum.load(Ordering::SeqCst), LIMIT); + tasks.shutdown().await; + } + + #[tokio::test] + async fn stream_shutdown_cancels_outstanding_tasks() { + struct NotifyOnDrop(Option>); + + impl Drop for NotifyOnDrop { + fn drop(&mut self) { + if let Some(tx) = self.0.take() { + let _ = tx.send(()); + } + } + } + + let mut tasks = PostLogonTaskSet::new(1); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (dropped_tx, dropped_rx) = tokio::sync::oneshot::channel(); + + tasks.spawn(async move { + let _drop_notification = NotifyOnDrop(Some(dropped_tx)); + started_tx.send(()).unwrap(); + std::future::pending::<()>().await; + }); + started_rx.await.unwrap(); + + tasks.shutdown().await; + + tokio::time::timeout(std::time::Duration::from_secs(1), dropped_rx) + .await + .expect("shutdown should drop the outstanding handler future") + .unwrap(); + assert!(tasks.tasks.is_empty()); + } +} diff --git a/src/server/dynamic_server.rs b/src/server/dynamic_server.rs index 74398b9..c4a6dbb 100644 --- a/src/server/dynamic_server.rs +++ b/src/server/dynamic_server.rs @@ -60,6 +60,7 @@ pub struct AquilaDynamicServer { action_flow_tx: tokio::sync::broadcast::Sender, flow_execution_registry: ActionFlowExecutionRegistry, execution_response_sender: SagittariusExecutionResponseSender, + action_transfer_concurrency_limit: usize, sagittarius_unary_rpc_timeout: Duration, } @@ -99,6 +100,7 @@ impl AquilaDynamicServer { action_flow_tx, flow_execution_registry: ActionFlowExecutionRegistry::new(), execution_response_sender, + action_transfer_concurrency_limit: config.grpc.action_transfer_concurrency_limit, sagittarius_unary_rpc_timeout: Duration::from_secs( config.dynamic_config.backend_unary_timeout_secs, ), @@ -152,6 +154,7 @@ impl AquilaDynamicServer { flow_execution_registry: self.flow_execution_registry.clone(), shard_registry: ActionShardRegistry::new(), is_static: false, + concurrency_limit: self.action_transfer_concurrency_limit, }); info!("Starting dynamic gRPC Server..."); diff --git a/src/server/static_server.rs b/src/server/static_server.rs index 47c5242..8b8bc2b 100644 --- a/src/server/static_server.rs +++ b/src/server/static_server.rs @@ -28,6 +28,7 @@ pub struct AquilaStaticServer { kv_store: Arc, action_config_tx: tokio::sync::broadcast::Sender, action_flow_tx: tokio::sync::broadcast::Sender, + action_transfer_concurrency_limit: usize, } impl AquilaStaticServer { @@ -58,6 +59,7 @@ impl AquilaStaticServer { kv_store, action_config_tx, action_flow_tx, + action_transfer_concurrency_limit: config.grpc.action_transfer_concurrency_limit, } } @@ -76,6 +78,7 @@ impl AquilaStaticServer { flow_execution_registry: ActionFlowExecutionRegistry::new(), shard_registry: ActionShardRegistry::new(), is_static: true, + concurrency_limit: self.action_transfer_concurrency_limit, }); info!("Starting static gRPC Server...");