runtime: add execution concurrency limits

2026-05-01 12:11:53 +00:00
parent 3d2d97c1e6
commit 8ac55ebcc2
13 changed files with 448 additions and 15 deletions
@@ -1,4 +1,5 @@
 use std::collections::BTreeMap;
+use std::sync::Arc;

 use crank_adapter_graphql::{GraphqlAdapter, GraphqlRequest};
 use crank_adapter_grpc::{GrpcAdapter, GrpcRequest, GrpcWindowRequest};
@@ -7,10 +8,11 @@ use crank_adapter_soap::{SoapAdapter, SoapRequest};
 use crank_adapter_websocket::{WebsocketAdapter, WebsocketWindowRequest};
 use crank_core::{ExecutionMode, Target, TransportBehavior};
 use serde_json::{Map, Value, json};
+use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 use tracing::debug;

 use crate::{
-    AdapterResponse, PreparedRequest, ResolvedAuth, RuntimeError, RuntimeOperation,
+    AdapterResponse, PreparedRequest, ResolvedAuth, RuntimeError, RuntimeLimits, RuntimeOperation,
    RuntimeRequestContext, WindowExecutionResult,
 };

@@ -21,6 +23,11 @@ pub struct RuntimeExecutor {
    rest_adapter: RestAdapter,
    soap_adapter: SoapAdapter,
    websocket_adapter: WebsocketAdapter,
+    limits: RuntimeLimits,
+    unary_limiter: Arc<Semaphore>,
+    window_limiter: Arc<Semaphore>,
+    session_limiter: Arc<Semaphore>,
+    job_limiter: Arc<Semaphore>,
 }

 impl Default for RuntimeExecutor {
@@ -31,12 +38,21 @@ impl Default for RuntimeExecutor {

 impl RuntimeExecutor {
    pub fn new() -> Self {
+        Self::with_limits(RuntimeLimits::default())
+    }
+
+    pub fn with_limits(limits: RuntimeLimits) -> Self {
        Self {
            graphql_adapter: GraphqlAdapter::new(),
            grpc_adapter: GrpcAdapter::new(),
            rest_adapter: RestAdapter::new(),
            soap_adapter: SoapAdapter::new(),
            websocket_adapter: WebsocketAdapter::new(),
+            unary_limiter: Arc::new(Semaphore::new(limits.max_concurrent_unary)),
+            window_limiter: Arc::new(Semaphore::new(limits.max_concurrent_window)),
+            session_limiter: Arc::new(Semaphore::new(limits.max_concurrent_sessions)),
+            job_limiter: Arc::new(Semaphore::new(limits.max_concurrent_jobs)),
+            limits,
        }
    }

@@ -77,6 +93,7 @@ impl RuntimeExecutor {
        request_context: Option<&RuntimeRequestContext>,
    ) -> Result<Value, RuntimeError> {
        log_runtime_event("unary.execute", operation, request_context);
+        let _permit = self.acquire_unary_permit(operation)?;
        let prepared_request = self.prepare_request(operation, input)?;
        let prepared_request = apply_resolved_auth(prepared_request, resolved_auth);
        self.execute_prepared(operation, prepared_request, request_context)
@@ -120,6 +137,7 @@ impl RuntimeExecutor {
        request_context: Option<&RuntimeRequestContext>,
    ) -> Result<WindowExecutionResult, RuntimeError> {
        log_runtime_event("window.execute", operation, request_context);
+        let _permit = self.acquire_window_permit()?;
        let Some(streaming) = operation.execution_config.streaming.as_ref() else {
            return Err(RuntimeError::MissingStreamingConfig {
                operation_id: operation.operation_id.as_str().to_owned(),
@@ -189,6 +207,7 @@ impl RuntimeExecutor {
        request_context: Option<&RuntimeRequestContext>,
    ) -> Result<WindowExecutionResult, RuntimeError> {
        log_runtime_event("session.seed", operation, request_context);
+        let _permit = self.acquire_session_permit()?;
        let Some(streaming) = operation.execution_config.streaming.as_ref() else {
            return Err(RuntimeError::MissingStreamingConfig {
                operation_id: operation.operation_id.as_str().to_owned(),
@@ -482,6 +501,48 @@ impl RuntimeExecutor {
            }
        }
    }
+
+    fn acquire_unary_permit(
+        &self,
+        operation: &RuntimeOperation,
+    ) -> Result<OwnedSemaphorePermit, RuntimeError> {
+        let (kind, limit, limiter) = if operation
+            .execution_config
+            .streaming
+            .as_ref()
+            .is_some_and(|streaming| streaming.mode == ExecutionMode::AsyncJob)
+        {
+            (
+                "async_job",
+                self.limits.max_concurrent_jobs,
+                Arc::clone(&self.job_limiter),
+            )
+        } else {
+            (
+                "unary",
+                self.limits.max_concurrent_unary,
+                Arc::clone(&self.unary_limiter),
+            )
+        };
+
+        try_acquire_limit(limiter, kind, limit)
+    }
+
+    fn acquire_window_permit(&self) -> Result<OwnedSemaphorePermit, RuntimeError> {
+        try_acquire_limit(
+            Arc::clone(&self.window_limiter),
+            "window",
+            self.limits.max_concurrent_window,
+        )
+    }
+
+    fn acquire_session_permit(&self) -> Result<OwnedSemaphorePermit, RuntimeError> {
+        try_acquire_limit(
+            Arc::clone(&self.session_limiter),
+            "session",
+            self.limits.max_concurrent_sessions,
+        )
+    }
 }

 impl PreparedRequest {
@@ -556,6 +617,16 @@ fn runtime_context_headers(
        .unwrap_or_default()
 }

+fn try_acquire_limit(
+    limiter: Arc<Semaphore>,
+    kind: &'static str,
+    limit: usize,
+) -> Result<OwnedSemaphorePermit, RuntimeError> {
+    limiter
+        .try_acquire_owned()
+        .map_err(|_| RuntimeError::ConcurrencyLimitExceeded { kind, limit })
+}
+
 fn log_runtime_event(
    stage: &'static str,
    operation: &RuntimeOperation,
@@ -665,7 +736,8 @@ mod tests {
    use tracing_subscriber::{filter::LevelFilter, fmt::MakeWriter, prelude::*};

    use crate::{
-        PreparedRequest, RuntimeError, RuntimeExecutor, RuntimeOperation, RuntimeRequestContext,
+        PreparedRequest, RuntimeError, RuntimeExecutor, RuntimeLimits, RuntimeOperation,
+        RuntimeRequestContext,
    };

    fn timestamp(value: &str) -> OffsetDateTime {
@@ -788,6 +860,96 @@ mod tests {
        assert!(logs.contains("op_rest_runtime"));
    }

+    #[tokio::test]
+    async fn rejects_unary_execution_when_unary_limit_is_exhausted() {
+        let executor = RuntimeExecutor::with_limits(RuntimeLimits {
+            max_concurrent_unary: 0,
+            ..RuntimeLimits::default()
+        });
+        let operation = test_rest_operation("http://example.invalid", false, false);
+
+        let error = executor
+            .execute(&operation, &json!({ "email": "user@example.com" }))
+            .await
+            .unwrap_err();
+
+        assert!(matches!(
+            error,
+            RuntimeError::ConcurrencyLimitExceeded {
+                kind: "unary",
+                limit: 0
+            }
+        ));
+    }
+
+    #[tokio::test]
+    async fn rejects_window_execution_when_window_limit_is_exhausted() {
+        let executor = RuntimeExecutor::with_limits(RuntimeLimits {
+            max_concurrent_window: 0,
+            ..RuntimeLimits::default()
+        });
+        let operation = test_window_snapshot_operation(
+            "http://example.invalid",
+            AggregationMode::RawItems,
+            None,
+            None,
+        );
+
+        let error = executor
+            .execute_window(&operation, &json!({}))
+            .await
+            .unwrap_err();
+
+        assert!(matches!(
+            error,
+            RuntimeError::ConcurrencyLimitExceeded {
+                kind: "window",
+                limit: 0
+            }
+        ));
+    }
+
+    #[tokio::test]
+    async fn rejects_session_seed_when_session_limit_is_exhausted() {
+        let executor = RuntimeExecutor::with_limits(RuntimeLimits {
+            max_concurrent_sessions: 0,
+            ..RuntimeLimits::default()
+        });
+        let operation = test_session_seed_operation("http://example.invalid");
+
+        let error = executor
+            .execute_session_seed(&operation, &json!({}))
+            .await
+            .unwrap_err();
+
+        assert!(matches!(
+            error,
+            RuntimeError::ConcurrencyLimitExceeded {
+                kind: "session",
+                limit: 0
+            }
+        ));
+    }
+
+    #[tokio::test]
+    async fn rejects_async_job_execution_when_job_limit_is_exhausted() {
+        let executor = RuntimeExecutor::with_limits(RuntimeLimits {
+            max_concurrent_jobs: 0,
+            ..RuntimeLimits::default()
+        });
+        let operation = test_async_job_operation("http://example.invalid");
+
+        let error = executor.execute(&operation, &json!({})).await.unwrap_err();
+
+        assert!(matches!(
+            error,
+            RuntimeError::ConcurrencyLimitExceeded {
+                kind: "async_job",
+                limit: 0
+            }
+        ));
+    }
+
    #[tokio::test]
    async fn executes_graphql_operation_end_to_end() {
        let endpoint = spawn_graphql_server().await;
@@ -2106,6 +2268,49 @@ mod tests {
        operation
    }

+    fn test_session_seed_operation(base_url: &str) -> RuntimeOperation {
+        let mut operation =
+            test_window_snapshot_operation(base_url, AggregationMode::RawItems, Some(10), None);
+
+        operation
+            .execution_config
+            .streaming
+            .as_mut()
+            .expect("streaming config")
+            .mode = ExecutionMode::Session;
+
+        operation
+    }
+
+    fn test_async_job_operation(base_url: &str) -> RuntimeOperation {
+        let mut operation = test_rest_operation(base_url, false, false);
+        operation.execution_config.streaming = Some(StreamingConfig {
+            mode: ExecutionMode::AsyncJob,
+            transport_behavior: TransportBehavior::RequestResponse,
+            window_duration_ms: None,
+            poll_interval_ms: Some(5_000),
+            upstream_timeout_ms: Some(operation.execution_config.timeout_ms),
+            idle_timeout_ms: None,
+            max_session_lifetime_ms: Some(300_000),
+            max_items: None,
+            max_bytes: None,
+            aggregation_mode: AggregationMode::RawItems,
+            summary_path: None,
+            items_path: None,
+            cursor_path: None,
+            status_path: None,
+            done_path: None,
+            redacted_paths: Vec::new(),
+            truncate_item_fields: false,
+            max_field_length: None,
+            drop_duplicates: false,
+            sampling_rate: None,
+            tool_family: ToolFamilyConfig::default(),
+        });
+
+        operation
+    }
+
    fn object_schema(field_name: &str, kind: SchemaKind) -> Schema {
        Schema {
            kind: SchemaKind::Object,