From ed10f57aaee84db4d273b6dab24fdd89597d29ea Mon Sep 17 00:00:00 2001 From: Dmitry Porokh Date: Fri, 31 Jul 2026 14:42:47 -0700 Subject: [PATCH] feat(api): add bounded admission control --- Cargo.lock | 2 + book/src/configuration/configurability.md | 1 + crates/api-core/Cargo.toml | 3 + crates/api-core/src/admission.rs | 650 ++++++++++++++++++ crates/api-core/src/cfg/README.md | 10 + crates/api-core/src/cfg/file.rs | 163 +++++ crates/api-core/src/cfg/load.rs | 1 + crates/api-core/src/handlers/scout_stream.rs | 141 +++- crates/api-core/src/lib.rs | 1 + crates/api-core/src/listener.rs | 14 + .../src/test_support/default_config.rs | 1 + docs/observability/core_metrics.md | 6 + 12 files changed, 989 insertions(+), 4 deletions(-) create mode 100644 crates/api-core/src/admission.rs diff --git a/Cargo.lock b/Cargo.lock index a798547c73..05b91d8286 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1262,6 +1262,8 @@ dependencies = [ "pkcs1 0.7.5", "prometheus", "prometheus-text-parser", + "prost", + "prost-types", "rand 0.10.1", "rcgen", "regex", diff --git a/book/src/configuration/configurability.md b/book/src/configuration/configurability.md index 592fe9f4fa..6784882193 100644 --- a/book/src/configuration/configurability.md +++ b/book/src/configuration/configurability.md @@ -685,6 +685,7 @@ These don't fit any sub-section but show up in production tuning: | Field | Default | When to touch | |-------|---------|---------------| | `max_database_connections` | `1000` | Drop when running multiple `nico-api` replicas to avoid saturating Postgres `max_connections`. | +| `api_admission_control` | enabled, `64` executing, `1024` pending, `5s` timeout | Tune the bounded shared budget for gRPC and admin business requests after scale testing; disable only as a rollback escape hatch. | | `max_find_by_ids` | `100` | Increase if scripts paginate batch lookups; raise the API-side limit to match the client. | | `compute_allocation_enforcement` | `WarnOnly` | Switch to `Enforce` once tenant compute pools are sized correctly — flips over-allocation from a warning to a refusal. | | `bmc_session_lockout_threshold` | `3` | Number of consecutive 401/403s from a BMC before NICo stops session-token logins for that BMC. Raise on environments with flaky BMC firmware. | diff --git a/crates/api-core/Cargo.toml b/crates/api-core/Cargo.toml index 6c19d80453..abad1cfc56 100644 --- a/crates/api-core/Cargo.toml +++ b/crates/api-core/Cargo.toml @@ -244,9 +244,12 @@ ctor = { workspace = true } data-encoding = { workspace = true } figment = { workspace = true, features = ["env", "test", "toml"] } mockito = { workspace = true } +prost = { workspace = true } +prost-types = { workspace = true } rcgen = { workspace = true } strum = { workspace = true } tempfile = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } [lints] workspace = true diff --git a/crates/api-core/src/admission.rs b/crates/api-core/src/admission.rs new file mode 100644 index 0000000000..519efcd0c3 --- /dev/null +++ b/crates/api-core/src/admission.rs @@ -0,0 +1,650 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +//! Shared admission control for gRPC and admin HTTP business requests. +//! +//! The middleware owns transport classification and response mapping. The +//! controller owns admission policy and returns an RAII work permit. Keeping +//! that boundary explicit lets fair scheduling replace the global semaphore +//! policy without moving handler futures out of their requesting tasks. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use axum::body::Body; +use axum::extract::{Request, State}; +use axum::http::{Response, StatusCode, header}; +use axum::middleware::Next; +use axum::response::IntoResponse; +use opentelemetry::metrics::{Meter, ObservableGauge}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError}; +use tokio_util::sync::CancellationToken; + +use crate::cfg::file::ApiAdmissionControlConfig; +use crate::logging::log_limiter::LogLimiter; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RequestTransport { + Grpc, + Http, +} + +impl RequestTransport { + fn classify(path: &str) -> Option { + // Keep these route prefixes in sync with the gRPC services and admin + // routes mounted by the listener. The descriptor-backed test below + // makes adding a gRPC service without updating this policy fail in CI; + // the table test documents the intentional admin bypasses. + if path.starts_with("/forge.Forge/") { + return Some(Self::Grpc); + } + + if !is_path_or_child(path, "/admin") + || is_path_or_child(path, "/admin/static") + || is_path_or_child(path, "/admin/auth-callback") + || is_path_or_child(path, "/admin/logs") + { + return None; + } + + Some(Self::Http) + } + + fn overloaded_response(self) -> Response { + match self { + Self::Grpc => tonic::Status::resource_exhausted("API admission capacity exhausted") + .into_http::(), + Self::Http => ( + StatusCode::SERVICE_UNAVAILABLE, + [(header::RETRY_AFTER, "1")], + "API admission capacity exhausted", + ) + .into_response(), + } + } +} + +fn is_path_or_child(path: &str, root: &str) -> bool { + path == root + || path + .strip_prefix(root) + .is_some_and(|suffix| suffix.starts_with('/')) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, carbide_instrument::LabelValue)] +enum RejectionReason { + QueueFull, + QueueTimeout, + ControllerUnavailable, + ShuttingDown, +} + +#[derive(carbide_instrument::Event)] +#[event( + event_name = "api_admission_request_admitted", + metric_name = "carbide_api_admission_admitted_total", + component = "nico-api", + log = off, + metric = counter, + describe = "Number of API requests admitted for execution" +)] +struct RequestAdmitted; + +#[derive(carbide_instrument::Event)] +#[event( + event_name = "api_admission_request_rejected", + metric_name = "carbide_api_admission_rejected_total", + component = "nico-api", + log = off, + metric = counter, + describe = "Number of API requests rejected before handler execution" +)] +struct RequestRejected { + #[label] + reason: RejectionReason, +} + +#[derive(carbide_instrument::Event)] +#[event( + event_name = "api_admission_pending_wait_finished", + metric_name = "carbide_api_admission_pending_wait_duration_seconds", + component = "nico-api", + log = off, + metric = histogram, + describe = "Duration API requests spent waiting for admission" +)] +struct PendingWaitFinished { + #[observation] + duration: Duration, +} + +#[derive(carbide_instrument::Event)] +#[event( + event_name = "api_admission_handler_execution_finished", + metric_name = "carbide_api_admission_handler_execution_duration_seconds", + component = "nico-api", + log = off, + metric = histogram, + describe = "Duration of admitted API request handler execution" +)] +struct HandlerExecutionFinished { + #[observation] + duration: Duration, +} + +pub(crate) struct AdmissionController { + work_slots: Arc, + pending_slots: Arc, + pending_timeout: Duration, + shutdown: CancellationToken, + rejection_log_limiter: LogLimiter, + _work_in_flight_gauge: ObservableGauge, + _pending_requests_gauge: ObservableGauge, +} + +impl AdmissionController { + pub(crate) fn new( + config: &ApiAdmissionControlConfig, + meter: &Meter, + shutdown: CancellationToken, + ) -> eyre::Result> { + config.validate()?; + + let work_slots = Arc::new(Semaphore::new(config.max_work_in_flight)); + let pending_slots = Arc::new(Semaphore::new(config.max_pending)); + let work_in_flight_gauge = register_occupancy_gauge( + meter, + "carbide_api_admission_work_in_flight", + "Number of API requests currently holding an execution slot", + config.max_work_in_flight, + Arc::clone(&work_slots), + ); + let pending_requests_gauge = register_occupancy_gauge( + meter, + "carbide_api_admission_pending_requests", + "Number of API requests currently waiting for an execution slot", + config.max_pending, + Arc::clone(&pending_slots), + ); + + Ok(Arc::new(Self { + work_slots, + pending_slots, + pending_timeout: config.pending_timeout, + shutdown, + rejection_log_limiter: LogLimiter::default(), + _work_in_flight_gauge: work_in_flight_gauge, + _pending_requests_gauge: pending_requests_gauge, + })) + } + + async fn acquire(&self) -> Result { + if self.shutdown.is_cancelled() { + return self.reject(RejectionReason::ShuttingDown); + } + + match Arc::clone(&self.work_slots).try_acquire_owned() { + Ok(permit) => return Ok(self.admit(permit)), + Err(TryAcquireError::Closed) => { + return self.reject(RejectionReason::ControllerUnavailable); + } + Err(TryAcquireError::NoPermits) => {} + } + + let pending_permit = match Arc::clone(&self.pending_slots).try_acquire_owned() { + Ok(permit) => permit, + Err(TryAcquireError::NoPermits) => { + return self.reject(RejectionReason::QueueFull); + } + Err(TryAcquireError::Closed) => { + return self.reject(RejectionReason::ControllerUnavailable); + } + }; + + let pending_wait = PendingWait { + _permit: pending_permit, + started: Instant::now(), + }; + let acquisition = tokio::time::timeout( + self.pending_timeout, + Arc::clone(&self.work_slots).acquire_owned(), + ); + let result = tokio::select! { + biased; + () = self.shutdown.cancelled() => Err(RejectionReason::ShuttingDown), + result = acquisition => match result { + Ok(Ok(permit)) => Ok(permit), + Ok(Err(_)) => Err(RejectionReason::ControllerUnavailable), + Err(_) => Err(RejectionReason::QueueTimeout), + }, + }; + drop(pending_wait); + + match result { + Ok(permit) => Ok(self.admit(permit)), + Err(reason) => self.reject(reason), + } + } + + fn admit(&self, permit: OwnedSemaphorePermit) -> WorkPermit { + carbide_instrument::emit(RequestAdmitted); + WorkPermit { + _permit: permit, + started: Instant::now(), + } + } + + fn reject(&self, reason: RejectionReason) -> Result { + carbide_instrument::emit(RequestRejected { reason }); + if self + .rejection_log_limiter + .should_log(&reason, "API request rejected by admission control") + { + tracing::warn!(?reason, "API request rejected by admission control"); + } + Err(reason) + } + + #[cfg(test)] + fn occupancy(&self) -> (usize, usize) { + ( + self.work_slots.available_permits(), + self.pending_slots.available_permits(), + ) + } +} + +fn register_occupancy_gauge( + meter: &Meter, + name: &'static str, + description: &'static str, + capacity: usize, + semaphore: Arc, +) -> ObservableGauge { + meter + .u64_observable_gauge(name) + .with_description(description) + .with_callback(move |observer| { + let occupied = capacity.saturating_sub(semaphore.available_permits()); + observer.observe(occupied as u64, &[]); + }) + .build() +} + +#[derive(Debug)] +struct WorkPermit { + _permit: OwnedSemaphorePermit, + started: Instant, +} + +struct PendingWait { + _permit: OwnedSemaphorePermit, + started: Instant, +} + +impl Drop for PendingWait { + fn drop(&mut self) { + carbide_instrument::emit(PendingWaitFinished { + duration: self.started.elapsed(), + }); + } +} + +impl Drop for WorkPermit { + fn drop(&mut self) { + carbide_instrument::emit(HandlerExecutionFinished { + duration: self.started.elapsed(), + }); + } +} + +pub(crate) async fn enforce( + State(controller): State>, + request: Request, + next: Next, +) -> Response { + let Some(transport) = RequestTransport::classify(request.uri().path()) else { + return next.run(request).await; + }; + + let permit = match controller.acquire().await { + Ok(permit) => permit, + Err(_) => return transport.overloaded_response(), + }; + let response = next.run(request).await; + drop(permit); + response +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use axum::Router; + use axum::http::Request; + use axum::routing::get; + use carbide_instrument::testing::MetricsCapture; + use prost::Message; + use prost_types::FileDescriptorSet; + use tokio::sync::Notify; + use tower::ServiceExt; + + use super::*; + + static ADMISSION_TEST_SERIAL: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + + fn controller( + max_work_in_flight: usize, + max_pending: usize, + pending_timeout: Duration, + shutdown: CancellationToken, + ) -> Arc { + AdmissionController::new( + &ApiAdmissionControlConfig { + enabled: true, + max_work_in_flight, + max_pending, + pending_timeout, + }, + &opentelemetry::global::meter("api-admission-tests"), + shutdown, + ) + .expect("test admission config is valid") + } + + #[test] + fn request_classification_covers_business_and_infrastructure_routes() { + let cases = [ + ("/forge.Forge/FindMachines", Some(RequestTransport::Grpc)), + ("/admin", Some(RequestTransport::Http)), + ("/admin/machine", Some(RequestTransport::Http)), + ("/", None), + ( + "/grpc.reflection.v1alpha.ServerReflection/ServerReflectionInfo", + None, + ), + ("/admin/static/carbide.css", None), + ("/admin/auth-callback", None), + ("/admin/logs/api/stream", None), + ("/administrator", None), + ("/admin/staticity", Some(RequestTransport::Http)), + ("/unrecognized", None), + ]; + + for (path, expected) in cases { + assert_eq!(RequestTransport::classify(path), expected, "path: {path}"); + } + } + + #[test] + fn every_forge_grpc_route_is_classified_for_admission() { + let descriptor = FileDescriptorSet::decode(::rpc::REFLECTION_API_SERVICE_DESCRIPTOR) + .expect("API service descriptor is valid"); + let mut route_count = 0; + + for file in descriptor.file { + let package = file.package.unwrap_or_default(); + // The reflection descriptor also contains protocols used by API + // clients. Only the forge package is mounted by this listener. + if package != "forge" { + continue; + } + + for service in file.service { + let service_name = service.name.clone().expect("service has a name"); + let qualified_service = if package.is_empty() { + service_name + } else { + format!("{package}.{service_name}") + }; + + for method in service.method { + let method = method.name.expect("method has a name"); + let path = format!("/{qualified_service}/{method}"); + route_count += 1; + assert_eq!( + RequestTransport::classify(&path), + Some(RequestTransport::Grpc), + "gRPC route {path} bypasses admission; update the route policy" + ); + } + } + } + + assert!(route_count > 0, "API descriptor contains no gRPC routes"); + } + + #[test] + fn overload_responses_match_the_request_transport() { + let grpc = RequestTransport::Grpc.overloaded_response(); + assert_eq!(grpc.status(), StatusCode::OK); + assert_eq!(grpc.headers().get("grpc-status").unwrap(), "8"); + + let http = RequestTransport::Http.overloaded_response(); + assert_eq!(http.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(http.headers().get(header::RETRY_AFTER).unwrap(), "1"); + } + + #[tokio::test] + async fn pending_capacity_is_bounded_and_cancellation_releases_it() { + let _serial = ADMISSION_TEST_SERIAL.lock().await; + let controller = controller(1, 1, Duration::from_secs(1), CancellationToken::new()); + let executing = controller.acquire().await.expect("first work is admitted"); + assert_eq!(controller.occupancy(), (0, 1)); + + { + let pending = controller.acquire(); + tokio::pin!(pending); + assert!( + tokio::time::timeout(Duration::from_millis(10), &mut pending) + .await + .is_err(), + "request should remain pending" + ); + assert_eq!(controller.occupancy(), (0, 0)); + assert_eq!( + controller + .acquire() + .await + .expect_err("pending queue is full"), + RejectionReason::QueueFull + ); + } + + assert_eq!(controller.occupancy(), (0, 1)); + drop(executing); + assert_eq!(controller.occupancy(), (1, 1)); + } + + #[tokio::test] + async fn timed_out_request_is_removed_and_metrics_record_the_outcomes() { + let _serial = ADMISSION_TEST_SERIAL.lock().await; + let metrics = MetricsCapture::start(); + let controller = controller(1, 1, Duration::from_millis(10), CancellationToken::new()); + let executing = controller.acquire().await.expect("first work is admitted"); + let rejection = controller + .acquire() + .await + .expect_err("second request should time out"); + assert_eq!(rejection, RejectionReason::QueueTimeout); + assert_eq!(controller.occupancy(), (0, 1)); + drop(executing); + + assert_eq!( + metrics.counter_delta("carbide_api_admission_admitted_total", &[]), + 1.0 + ); + assert_eq!( + metrics.counter_delta( + "carbide_api_admission_rejected_total", + &[("reason", "queue_timeout")], + ), + 1.0 + ); + assert_eq!( + metrics + .histogram_count_delta("carbide_api_admission_pending_wait_duration_seconds", &[],), + 1 + ); + assert_eq!( + metrics.histogram_count_delta( + "carbide_api_admission_handler_execution_duration_seconds", + &[], + ), + 1 + ); + } + + #[tokio::test] + async fn shutdown_and_closed_controller_are_rejected() { + let _serial = ADMISSION_TEST_SERIAL.lock().await; + let shutdown = CancellationToken::new(); + let shutting_down_controller = controller(1, 1, Duration::from_secs(1), shutdown.clone()); + shutdown.cancel(); + assert_eq!( + shutting_down_controller + .acquire() + .await + .expect_err("shutdown rejects work"), + RejectionReason::ShuttingDown + ); + + let controller = controller(1, 1, Duration::from_secs(1), CancellationToken::new()); + controller.work_slots.close(); + assert_eq!( + controller + .acquire() + .await + .expect_err("closed controller rejects work"), + RejectionReason::ControllerUnavailable + ); + } + + #[tokio::test] + async fn grpc_and_admin_routes_share_capacity_and_bypasses_remain_available() { + let _serial = ADMISSION_TEST_SERIAL.lock().await; + let controller = controller(1, 1, Duration::from_secs(1), CancellationToken::new()); + let handler_calls = Arc::new(AtomicUsize::new(0)); + let handler_started = Arc::new(Notify::new()); + let release_handler = Arc::new(Notify::new()); + let blocking_handler = { + let handler_calls = Arc::clone(&handler_calls); + let handler_started = Arc::clone(&handler_started); + let release_handler = Arc::clone(&release_handler); + move || { + let handler_calls = Arc::clone(&handler_calls); + let handler_started = Arc::clone(&handler_started); + let release_handler = Arc::clone(&release_handler); + async move { + handler_calls.fetch_add(1, Ordering::SeqCst); + handler_started.notify_one(); + release_handler.notified().await; + "business response" + } + } + }; + let immediate_handler = { + let handler_calls = Arc::clone(&handler_calls); + move || { + let handler_calls = Arc::clone(&handler_calls); + async move { + handler_calls.fetch_add(1, Ordering::SeqCst); + "business response" + } + } + }; + let router = Router::new() + .route("/forge.Forge/Test", get(immediate_handler)) + .route("/admin/business", get(blocking_handler)) + .route("/admin/static/test.css", get(|| async { "static" })) + .layer(axum::middleware::from_fn_with_state( + Arc::clone(&controller), + enforce, + )); + + let executing_router = router.clone(); + let executing = tokio::spawn(async move { + executing_router + .oneshot( + Request::builder() + .uri("/admin/business") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap() + }); + handler_started.notified().await; + assert_eq!(handler_calls.load(Ordering::SeqCst), 1); + + let pending_router = router.clone(); + let pending = tokio::spawn(async move { + pending_router + .oneshot( + Request::builder() + .uri("/forge.Forge/Test") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap() + }); + for _ in 0..100 { + if controller.occupancy() == (0, 0) { + break; + } + tokio::task::yield_now().await; + } + assert_eq!(controller.occupancy(), (0, 0)); + + let bypass = router + .clone() + .oneshot( + Request::builder() + .uri("/admin/static/test.css") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(bypass.status(), StatusCode::OK); + + let rejected = router + .clone() + .oneshot( + Request::builder() + .uri("/admin/business") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(rejected.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(rejected.headers().get(header::RETRY_AFTER).unwrap(), "1"); + assert_eq!(handler_calls.load(Ordering::SeqCst), 1); + + release_handler.notify_one(); + let executing = executing.await.unwrap(); + let admitted_grpc = pending.await.unwrap(); + assert_eq!(executing.status(), StatusCode::OK); + assert_eq!(admitted_grpc.status(), StatusCode::OK); + assert_eq!(handler_calls.load(Ordering::SeqCst), 2); + // Capacity is released when handlers return, even while their response + // bodies remain alive and unconsumed. + assert_eq!(controller.occupancy(), (1, 1)); + } +} diff --git a/crates/api-core/src/cfg/README.md b/crates/api-core/src/cfg/README.md index 33ca562972..af552db0d8 100644 --- a/crates/api-core/src/cfg/README.md +++ b/crates/api-core/src/cfg/README.md @@ -20,6 +20,7 @@ applicable. | `database_pool_acquire_timeout` | `Duration` | `30s` | `server` | How long a caller may wait for a connection from the pool before the attempt fails (sqlx's own default); trips on a stalled database or a saturated pool alike. Must be greater than zero (startup rejects `0`). | | `database_pool_idle_timeout` | `Duration` | `10m` | `server` | Idle time after which the pool closes a connection, keeping the pool's own reaping well inside the Postgres server's 60-minute idle-session reaper. Must be greater than zero (startup rejects `0`). | | `database_pool_max_lifetime` | `Duration` | `30m` | `server` | Maximum age of a pooled connection before it is recycled, so the pool re-balances onto the current primary after a database failover. Must be greater than zero (startup rejects `0`). | +| `api_admission_control` | `ApiAdmissionControlConfig` | *(see below)* | `server` | Shared execution and pending-request limits for gRPC and admin HTTP business requests. | | `ib_config` | `Option` | — | `hardware` | InfiniBand fabric configuration (see [IBFabricConfig](#ibfabricconfig)). | | `asn` | `u32` | **required** | `networking` | Autonomous System Number, fixed per environment. Used by nico-dpu-agent for `frr.conf` BGP routing. | | `dhcp_servers` | `Vec` | `[]` | `networking` | DHCP server addresses announced to DPUs during network provisioning. | @@ -251,6 +252,15 @@ available for topology-specific flows. ## Sub-Structs +### `ApiAdmissionControlConfig` + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `enabled` | `bool` | `true` | Enable bounded API admission. Set to `false` to restore unrestricted request admission. | +| `max_work_in_flight` | `usize` | `64` | Maximum business requests executing concurrently. When enabled, must be greater than zero and no greater than `tokio::sync::Semaphore::MAX_PERMITS`. | +| `max_pending` | `usize` | `1024` | Maximum business requests waiting for execution. When enabled, must be greater than zero and no greater than `tokio::sync::Semaphore::MAX_PERMITS`. | +| `pending_timeout` | `Duration` | `5s` | Maximum time a pending request may wait for execution. Must be greater than zero when admission control is enabled. | + ### `TlsConfig` | Field | Type | Default | Description | diff --git a/crates/api-core/src/cfg/file.rs b/crates/api-core/src/cfg/file.rs index f4035a566c..66de2d694f 100644 --- a/crates/api-core/src/cfg/file.rs +++ b/crates/api-core/src/cfg/file.rs @@ -148,6 +148,11 @@ pub struct CarbideConfig { )] pub database_pool_max_lifetime: std::time::Duration, + /// Bounds the number of API requests that may execute or wait for + /// execution. The limits are shared by gRPC and admin HTTP traffic. + #[serde(default)] + pub api_admission_control: ApiAdmissionControlConfig, + /// InfiniBand fabric configuration, used by the IB /// fabric manager for partition and UFM management. pub ib_config: Option, @@ -789,6 +794,69 @@ pub struct CarbideConfig { pub certificates: CertificatesConfig, } +/// Global admission limits for business requests handled by nico-api. +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] +pub struct ApiAdmissionControlConfig { + /// Whether admission control is active. + #[serde(default = "default_to_true")] + pub enabled: bool, + + /// Maximum number of requests executing business handlers concurrently. + #[serde(default = "default_api_admission_max_work_in_flight")] + pub max_work_in_flight: usize, + + /// Maximum number of requests waiting for an execution slot. + #[serde(default = "default_api_admission_max_pending")] + pub max_pending: usize, + + /// Maximum time a pending request may wait for execution. + #[serde( + default = "default_api_admission_pending_timeout", + deserialize_with = "deserialize_duration", + serialize_with = "as_std_duration" + )] + pub pending_timeout: std::time::Duration, +} + +impl Default for ApiAdmissionControlConfig { + fn default() -> Self { + Self { + enabled: true, + max_work_in_flight: default_api_admission_max_work_in_flight(), + max_pending: default_api_admission_max_pending(), + pending_timeout: default_api_admission_pending_timeout(), + } + } +} + +impl ApiAdmissionControlConfig { + /// Reject invalid bounds before the API listener starts. + pub fn validate(&self) -> eyre::Result<()> { + if !self.enabled { + return Ok(()); + } + + for (field, value) in [ + ("max_work_in_flight", self.max_work_in_flight), + ("max_pending", self.max_pending), + ] { + if value == 0 { + eyre::bail!("api_admission_control.{field} must be greater than zero"); + } + if value > tokio::sync::Semaphore::MAX_PERMITS { + eyre::bail!( + "api_admission_control.{field} must not exceed {}", + tokio::sync::Semaphore::MAX_PERMITS + ); + } + } + if self.pending_timeout.is_zero() { + eyre::bail!("api_admission_control.pending_timeout must be greater than zero"); + } + Ok(()) + } +} + /// `[certificates]` config section: selects the backend that vends machine and /// service certificates, independently of where credentials are stored. #[derive(Clone, Debug, Default, Deserialize, Serialize)] @@ -2612,6 +2680,18 @@ pub const fn default_database_pool_max_lifetime() -> std::time::Duration { std::time::Duration::from_secs(30 * 60) } +const fn default_api_admission_max_work_in_flight() -> usize { + 64 +} + +const fn default_api_admission_max_pending() -> usize { + 1024 +} + +const fn default_api_admission_pending_timeout() -> std::time::Duration { + std::time::Duration::from_secs(5) +} + pub const fn default_bmc_session_lockout_threshold() -> u32 { 3 } @@ -3868,6 +3948,80 @@ mod tests { assert!(config.enabled); } + #[test] + fn api_admission_control_only_validates_bounds_when_enabled() { + type ZeroOut = fn(&mut ApiAdmissionControlConfig); + let cases: [(&str, ZeroOut); 3] = [ + ("max_work_in_flight", |config| config.max_work_in_flight = 0), + ("max_pending", |config| config.max_pending = 0), + ("pending_timeout", |config| { + config.pending_timeout = std::time::Duration::ZERO + }), + ]; + + let disabled = ApiAdmissionControlConfig { + enabled: false, + max_work_in_flight: 0, + max_pending: 0, + pending_timeout: std::time::Duration::ZERO, + }; + disabled + .validate() + .expect("disabled admission control ignores its bounds"); + + for (field, zero_out) in cases { + let mut config = ApiAdmissionControlConfig::default(); + zero_out(&mut config); + let error = config + .validate() + .expect_err("zero admission values must be rejected"); + assert!( + error.to_string().contains(field), + "error must name {field}, got: {error}" + ); + } + } + + fn assert_api_admission_semaphore_bound( + field: &str, + set_value: fn(&mut ApiAdmissionControlConfig, usize), + ) { + let mut config = ApiAdmissionControlConfig::default(); + set_value(&mut config, tokio::sync::Semaphore::MAX_PERMITS); + config + .validate() + .expect("Tokio's semaphore maximum must be accepted"); + + set_value(&mut config, tokio::sync::Semaphore::MAX_PERMITS + 1); + let error = config + .validate() + .expect_err("values above Tokio's semaphore maximum must be rejected"); + assert!( + error.to_string().contains(field), + "error must name {field}, got: {error}" + ); + assert!( + error + .to_string() + .contains(&tokio::sync::Semaphore::MAX_PERMITS.to_string()), + "error must name the maximum, got: {error}" + ); + } + + #[test] + fn api_admission_control_validates_max_work_in_flight_upper_bound() { + assert_api_admission_semaphore_bound("max_work_in_flight", |config, value| { + config.max_work_in_flight = value; + }); + } + + #[test] + fn api_admission_control_validates_max_pending_upper_bound() { + assert_api_admission_semaphore_bound("max_pending", |config, value| { + config.max_pending = value; + }); + } + #[test] fn periodic_state_republish_rejects_zero_interval() { for enabled in [true, false] { @@ -3994,6 +4148,15 @@ mod tests { config.database_pool_max_lifetime, std::time::Duration::from_secs(30 * 60) ); + assert_eq!( + config.api_admission_control, + ApiAdmissionControlConfig { + enabled: true, + max_work_in_flight: 64, + max_pending: 1024, + pending_timeout: std::time::Duration::from_secs(5), + } + ); assert!(config.dhcp_servers.is_empty()); assert!(!config.allow_insecure_discovery); assert!(config.route_servers.is_empty()); diff --git a/crates/api-core/src/cfg/load.rs b/crates/api-core/src/cfg/load.rs index 6f0b71224e..22534424b7 100644 --- a/crates/api-core/src/cfg/load.rs +++ b/crates/api-core/src/cfg/load.rs @@ -152,6 +152,7 @@ pub fn parse_carbide_config( // Validate that admin-UI tool entries have unique names. config.validate_web_ui_sidebar_tools()?; + config.api_admission_control.validate()?; if let Some(config) = &config.dsx_exchange_event_bus { config.periodic_state_republish.validate()?; diff --git a/crates/api-core/src/handlers/scout_stream.rs b/crates/api-core/src/handlers/scout_stream.rs index 0fc4b4aebd..40176aa8ca 100644 --- a/crates/api-core/src/handlers/scout_stream.rs +++ b/crates/api-core/src/handlers/scout_stream.rs @@ -15,6 +15,9 @@ * limitations under the License. */ +use std::future::Future; +use std::time::Duration; + use ::rpc::protos::forge as rpc; use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; @@ -24,6 +27,13 @@ use crate::CarbideError; use crate::api::{Api, ScoutStreamType, log_request_data}; use crate::handlers::utils::convert_and_log_machine_id; +// Keep this hard-coded unless an operational need for tuning is demonstrated. +// Scout sends Init immediately after opening the RPC, so ten seconds is generous +// for a healthy connection and serves as a protocol safety bound rather than a +// deployment-specific policy. Making it configurable prematurely would add +// configuration surface and allow this resource-leak protection to be weakened. +const SCOUT_STREAM_INIT_TIMEOUT: Duration = Duration::from_secs(10); + // scout_stream handles the bidirectional streaming connection from scout agents. // scout agents call scout_stream and send an Init message, and then carbide-api // will send down "request" messages to connected agent(s) to either instruct them @@ -37,10 +47,7 @@ pub(crate) async fn scout_stream( let mut stream = request.into_inner(); - let init_message = stream - .message() - .await? - .ok_or_else(|| CarbideError::InvalidArgument("invalid message received".to_string()))?; + let init_message = receive_initial_message(stream.message(), SCOUT_STREAM_INIT_TIMEOUT).await?; // As part of "constructing" the new scout stream, we expect // an Init message as the first thing from the client (in this @@ -98,6 +105,23 @@ pub(crate) async fn scout_stream( Ok(Response::new(Box::pin(ReceiverStream::new(server_rx)))) } +async fn receive_initial_message( + message: F, + timeout: Duration, +) -> Result +where + F: Future, Status>>, +{ + match tokio::time::timeout(timeout, message).await { + Ok(result) => result?.ok_or_else(|| { + CarbideError::InvalidArgument("invalid message received".to_string()).into() + }), + Err(_) => Err(Status::deadline_exceeded( + "timed out waiting for initial ScoutStream init message", + )), + } +} + pub async fn show_connections( api: &Api, request: Request, @@ -217,3 +241,112 @@ fn format_system_time(time: std::time::SystemTime) -> String { Err(_) => "unknown".to_string(), } } + +#[cfg(test)] +mod tests { + use std::convert::Infallible; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::task::Poll; + + use ::rpc::forge::forge_server::ForgeServer; + use axum::Router; + use axum::body::Body; + use axum::http::Request as AxumRequest; + use axum::routing::post; + use futures::stream; + use tokio::sync::Notify; + use tokio_util::sync::CancellationToken; + use tower::ServiceExt; + + use super::*; + use crate::admission::{AdmissionController, enforce as enforce_admission}; + use crate::cfg::file::ApiAdmissionControlConfig; + use crate::tests::common::api_fixtures::create_test_env; + + #[crate::sqlx_test] + async fn stalled_initial_message_times_out_and_releases_admission_capacity(pool: sqlx::PgPool) { + let env = create_test_env(pool).await; + let controller = AdmissionController::new( + &ApiAdmissionControlConfig { + enabled: true, + max_work_in_flight: 1, + max_pending: 1, + pending_timeout: SCOUT_STREAM_INIT_TIMEOUT + Duration::from_secs(1), + }, + &opentelemetry::global::meter("scout-stream-init-timeout-test"), + CancellationToken::new(), + ) + .expect("test admission config is valid"); + let probe_calls = Arc::new(AtomicUsize::new(0)); + let probe_handler = { + let probe_calls = Arc::clone(&probe_calls); + move || { + let probe_calls = Arc::clone(&probe_calls); + async move { + probe_calls.fetch_add(1, Ordering::SeqCst); + "probe response" + } + } + }; + let router = Router::new() + .route_service( + "/forge.Forge/{*rpc}", + ForgeServer::from_arc(Arc::clone(&env.api)), + ) + .route("/admin/probe", post(probe_handler)) + .layer(axum::middleware::from_fn_with_state( + controller, + enforce_admission, + )); + + tokio::time::pause(); + let stalled_message_polled = Arc::new(Notify::new()); + let pending_body = { + let stalled_message_polled = Arc::clone(&stalled_message_polled); + Body::from_stream(stream::poll_fn(move |_| { + stalled_message_polled.notify_one(); + Poll::>>::Pending + })) + }; + let stalled_router = router.clone(); + let stalled_request = tokio::spawn(async move { + stalled_router + .oneshot( + AxumRequest::post("/forge.Forge/ScoutStream") + .header(axum::http::header::CONTENT_TYPE, "application/grpc") + .header("te", "trailers") + .body(pending_body) + .unwrap(), + ) + .await + .unwrap() + }); + stalled_message_polled.notified().await; + + let probe_router = router.clone(); + let probe_request = tokio::spawn(async move { + probe_router + .oneshot( + AxumRequest::post("/admin/probe") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap() + }); + tokio::task::yield_now().await; + assert_eq!(probe_calls.load(Ordering::SeqCst), 0); + + tokio::time::advance(SCOUT_STREAM_INIT_TIMEOUT + Duration::from_millis(1)).await; + + let stalled_response = stalled_request.await.unwrap(); + assert_eq!(stalled_response.status(), axum::http::StatusCode::OK); + assert_eq!(stalled_response.headers().get("grpc-status").unwrap(), "4"); + + let probe_response = probe_request.await.unwrap(); + assert_eq!(probe_response.status(), axum::http::StatusCode::OK); + assert_eq!(probe_calls.load(Ordering::SeqCst), 1); + tokio::time::resume(); + } +} diff --git a/crates/api-core/src/lib.rs b/crates/api-core/src/lib.rs index 612527ae0d..be5d1d9c29 100644 --- a/crates/api-core/src/lib.rs +++ b/crates/api-core/src/lib.rs @@ -41,6 +41,7 @@ // `cfg::file` config types). // Anything that doesn't need to cross a crate boundary should stay private. +mod admission; mod api; mod attestation; mod auth; diff --git a/crates/api-core/src/listener.rs b/crates/api-core/src/listener.rs index c1af02b038..34a85b748c 100644 --- a/crates/api-core/src/listener.rs +++ b/crates/api-core/src/listener.rs @@ -38,6 +38,7 @@ use tower_http::add_extension::AddExtensionLayer; use tower_http::auth::AsyncRequireAuthorizationLayer; use tower_http::normalize_path::NormalizePath; +use crate::admission::{AdmissionController, enforce as enforce_admission}; use crate::api::Api; use crate::auth; use crate::auth::Authorization; @@ -377,6 +378,19 @@ pub async fn start( None => router, }; + let admission_config = &api_service.runtime_config.api_admission_control; + admission_config.validate()?; + let router = if admission_config.enabled { + let controller = AdmissionController::new(admission_config, &meter, cancel_token.clone())?; + router.layer(axum::middleware::from_fn_with_state( + controller, + enforce_admission, + )) + } else { + tracing::info!("API admission control disabled"); + router + }; + let app = tower::ServiceBuilder::new() .layer(LogLayer::new(meter.clone())) .layer(cert_description_layer) diff --git a/crates/api-core/src/test_support/default_config.rs b/crates/api-core/src/test_support/default_config.rs index 413f719c18..e872d67186 100644 --- a/crates/api-core/src/test_support/default_config.rs +++ b/crates/api-core/src/test_support/default_config.rs @@ -109,6 +109,7 @@ pub fn get() -> CarbideConfig { database_pool_acquire_timeout: default_database_pool_acquire_timeout(), database_pool_idle_timeout: default_database_pool_idle_timeout(), database_pool_max_lifetime: default_database_pool_max_lifetime(), + api_admission_control: Default::default(), compute_allocation_enforcement: Default::default(), asn: 0, datacenter_asn: 0, diff --git a/docs/observability/core_metrics.md b/docs/observability/core_metrics.md index 5895211e5b..12592e66b9 100644 --- a/docs/observability/core_metrics.md +++ b/docs/observability/core_metrics.md @@ -5,6 +5,12 @@ This file contains a list of metrics exported by NVIDIA Infra Controller (NICo). + + + + + +
NameTypeDescription
carbide_active_host_firmware_update_countgaugeNumber of host machines in the system currently working on updating their firmware.
carbide_api_admission_admitted_totalcounterNumber of API requests admitted for execution
carbide_api_admission_handler_execution_duration_secondshistogramDuration of admitted API request handler execution
carbide_api_admission_pending_requestsgaugeNumber of API requests currently waiting for an execution slot
carbide_api_admission_pending_wait_duration_secondshistogramDuration API requests spent waiting for admission
carbide_api_admission_rejected_totalcounterNumber of API requests rejected before handler execution
carbide_api_admission_work_in_flightgaugeNumber of API requests currently holding an execution slot
carbide_api_db_queries_totalcounterNumber of database queries that occurred inside a span
carbide_api_db_span_query_time_millisecondshistogramTotal time the request spent inside a span on database transactions
carbide_api_grpc_server_duration_millisecondshistogramProcessing time for a request on the carbide API server