1050 lines
36 KiB
Rust
1050 lines
36 KiB
Rust
use std::{
|
|
pin::Pin,
|
|
sync::{
|
|
atomic::{AtomicU64, Ordering},
|
|
Arc,
|
|
},
|
|
task::{Context, Poll},
|
|
time::{Duration, Instant},
|
|
};
|
|
|
|
use axum::{
|
|
body::Body,
|
|
extract::{Request, State},
|
|
http::{header, HeaderValue, StatusCode},
|
|
middleware::Next,
|
|
response::{IntoResponse, Response},
|
|
Json,
|
|
};
|
|
use bytes::Bytes;
|
|
use http_body::Frame;
|
|
use rand::Rng;
|
|
use serde_json::json;
|
|
use subtle::ConstantTimeEq;
|
|
use tokio::sync::{mpsc, oneshot};
|
|
use tower::{Layer, Service};
|
|
use tower_http::trace::{MakeSpan, OnRequest, OnResponse, TraceLayer};
|
|
use tracing::{debug, error, field::Empty, info, info_span, warn, Span};
|
|
|
|
pub use crate::core::token_bucket::TokenBucket;
|
|
use crate::{
|
|
observability::{
|
|
inflight_tracker::InFlightRequestTracker,
|
|
metrics::{method_to_static_str, metrics_labels, Metrics},
|
|
},
|
|
routers::error::extract_error_code_from_response,
|
|
server::AppState,
|
|
wasm::{
|
|
module::{MiddlewareAttachPoint, WasmModuleAttachPoint},
|
|
spec::{
|
|
apply_modify_action_to_headers, build_wasm_headers_from_axum_headers,
|
|
smg::gateway::middleware_types::{
|
|
Action, Request as WasmRequest, Response as WasmResponse,
|
|
},
|
|
},
|
|
types::WasmComponentInput,
|
|
},
|
|
};
|
|
|
|
/// A body wrapper that holds a token and returns it when the body is fully consumed or dropped.
|
|
/// This ensures that for streaming responses, the token is only returned after the entire
|
|
/// stream has been sent to the client.
|
|
pub struct TokenGuardBody {
|
|
inner: Body,
|
|
/// The token bucket to return tokens to. Uses Option so we can take() on drop.
|
|
token_bucket: Option<Arc<TokenBucket>>,
|
|
/// Number of tokens to return.
|
|
tokens: f64,
|
|
}
|
|
|
|
impl TokenGuardBody {
|
|
/// Create a new TokenGuardBody that will return tokens when dropped.
|
|
pub fn new(inner: Body, token_bucket: Arc<TokenBucket>, tokens: f64) -> Self {
|
|
Self {
|
|
inner,
|
|
token_bucket: Some(token_bucket),
|
|
tokens,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Drop for TokenGuardBody {
|
|
fn drop(&mut self) {
|
|
if let Some(bucket) = self.token_bucket.take() {
|
|
debug!(
|
|
"TokenGuardBody: stream ended, returning {} tokens to bucket",
|
|
self.tokens
|
|
);
|
|
// Use lock-free sync return - no runtime needed, guaranteed token return
|
|
bucket.return_tokens_sync(self.tokens);
|
|
}
|
|
}
|
|
}
|
|
|
|
impl http_body::Body for TokenGuardBody {
|
|
type Data = Bytes;
|
|
type Error = axum::Error;
|
|
|
|
fn poll_frame(
|
|
self: Pin<&mut Self>,
|
|
cx: &mut Context<'_>,
|
|
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
|
|
// SAFETY: We never move the inner body, and Body is Unpin
|
|
// (it's a type alias for UnsyncBoxBody which is Unpin)
|
|
let this = self.get_mut();
|
|
Pin::new(&mut this.inner).poll_frame(cx)
|
|
}
|
|
|
|
fn is_end_stream(&self) -> bool {
|
|
self.inner.is_end_stream()
|
|
}
|
|
|
|
fn size_hint(&self) -> http_body::SizeHint {
|
|
self.inner.size_hint()
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct AuthConfig {
|
|
pub api_key: Option<String>,
|
|
}
|
|
|
|
/// Middleware to validate Bearer token against configured API key
|
|
/// Only active when router has an API key configured
|
|
pub async fn auth_middleware(
|
|
State(auth_config): State<AuthConfig>,
|
|
request: Request<Body>,
|
|
next: Next,
|
|
) -> Result<Response, StatusCode> {
|
|
if let Some(expected_key) = &auth_config.api_key {
|
|
// Extract Authorization header
|
|
let auth_header = request
|
|
.headers()
|
|
.get(header::AUTHORIZATION)
|
|
.and_then(|h| h.to_str().ok());
|
|
|
|
match auth_header {
|
|
Some(header_value) if header_value.starts_with("Bearer ") => {
|
|
let token = &header_value[7..]; // Skip "Bearer "
|
|
// Use constant-time comparison to prevent timing attacks
|
|
let token_bytes = token.as_bytes();
|
|
let expected_bytes = expected_key.as_bytes();
|
|
|
|
// Check if lengths match first (this is not constant-time but necessary)
|
|
if token_bytes.len() != expected_bytes.len() {
|
|
return Err(StatusCode::UNAUTHORIZED);
|
|
}
|
|
|
|
// Constant-time comparison of the actual values
|
|
if token_bytes.ct_eq(expected_bytes).unwrap_u8() != 1 {
|
|
return Err(StatusCode::UNAUTHORIZED);
|
|
}
|
|
}
|
|
_ => return Err(StatusCode::UNAUTHORIZED),
|
|
}
|
|
}
|
|
|
|
Ok(next.run(request).await)
|
|
}
|
|
|
|
/// Alphanumeric characters for request ID generation (as bytes for O(1) indexing)
|
|
const REQUEST_ID_CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
|
|
|
|
/// Generate OpenAI-compatible request ID based on endpoint.
|
|
fn generate_request_id(path: &str) -> String {
|
|
let prefix = if path.contains("/chat/completions") {
|
|
"chatcmpl-"
|
|
} else if path.contains("/completions") {
|
|
"cmpl-"
|
|
} else if path.contains("/generate") {
|
|
"gnt-"
|
|
} else if path.contains("/responses") {
|
|
"resp-"
|
|
} else {
|
|
"req-"
|
|
};
|
|
|
|
// Generate a random string similar to OpenAI's format
|
|
// Use byte array indexing (O(1)) instead of chars().nth() (O(n))
|
|
let mut rng = rand::rng();
|
|
let random_part: String = (0..24)
|
|
.map(|_| {
|
|
let idx = rng.random_range(0..REQUEST_ID_CHARS.len());
|
|
REQUEST_ID_CHARS[idx] as char
|
|
})
|
|
.collect();
|
|
|
|
format!("{}{}", prefix, random_part)
|
|
}
|
|
|
|
/// Extension type for storing request ID
|
|
#[derive(Clone, Debug)]
|
|
pub struct RequestId(pub String);
|
|
|
|
/// Tower Layer for request ID middleware
|
|
#[derive(Clone)]
|
|
pub struct RequestIdLayer {
|
|
headers: Arc<Vec<String>>,
|
|
}
|
|
|
|
impl RequestIdLayer {
|
|
pub fn new(headers: Vec<String>) -> Self {
|
|
Self {
|
|
headers: Arc::new(headers),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<S> Layer<S> for RequestIdLayer {
|
|
type Service = RequestIdMiddleware<S>;
|
|
|
|
fn layer(&self, inner: S) -> Self::Service {
|
|
RequestIdMiddleware {
|
|
inner,
|
|
headers: self.headers.clone(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Tower Service for request ID middleware
|
|
#[derive(Clone)]
|
|
pub struct RequestIdMiddleware<S> {
|
|
inner: S,
|
|
headers: Arc<Vec<String>>,
|
|
}
|
|
|
|
impl<S> Service<Request> for RequestIdMiddleware<S>
|
|
where
|
|
S: Service<Request, Response = Response> + Send + 'static,
|
|
S::Future: Send + 'static,
|
|
{
|
|
type Response = S::Response;
|
|
type Error = S::Error;
|
|
type Future =
|
|
Pin<Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>>;
|
|
|
|
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
self.inner.poll_ready(cx)
|
|
}
|
|
|
|
fn call(&mut self, mut req: Request) -> Self::Future {
|
|
let headers = self.headers.clone();
|
|
|
|
// Extract request ID from headers or generate new one
|
|
let mut request_id = None;
|
|
|
|
for header_name in headers.iter() {
|
|
if let Some(header_value) = req.headers().get(header_name) {
|
|
if let Ok(value) = header_value.to_str() {
|
|
request_id = Some(value.to_string());
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
let request_id = request_id.unwrap_or_else(|| generate_request_id(req.uri().path()));
|
|
|
|
// Insert request ID into request extensions for other middleware/handlers to use
|
|
req.extensions_mut().insert(RequestId(request_id.clone()));
|
|
|
|
// Call the inner service
|
|
let future = self.inner.call(req);
|
|
|
|
Box::pin(async move {
|
|
let mut response = future.await?;
|
|
|
|
// Add request ID to response headers
|
|
response.headers_mut().insert(
|
|
"x-request-id",
|
|
HeaderValue::from_str(&request_id)
|
|
.unwrap_or_else(|_| HeaderValue::from_static("invalid-request-id")),
|
|
);
|
|
|
|
Ok(response)
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Custom span maker that includes request ID
|
|
#[derive(Clone, Debug)]
|
|
pub struct RequestSpan;
|
|
|
|
impl<B> MakeSpan<B> for RequestSpan {
|
|
fn make_span(&mut self, request: &Request<B>) -> Span {
|
|
// Don't try to extract request ID here - it won't be available yet
|
|
// The RequestIdLayer runs after TraceLayer creates the span
|
|
info_span!(
|
|
target: "smg::otel-trace",
|
|
"http_request",
|
|
method = %request.method(),
|
|
uri = %request.uri(),
|
|
version = ?request.version(),
|
|
request_id = Empty, // Will be set later
|
|
status_code = Empty,
|
|
latency = Empty,
|
|
error = Empty,
|
|
module = "smg"
|
|
)
|
|
}
|
|
}
|
|
|
|
/// Custom on_request handler
|
|
#[derive(Clone, Debug)]
|
|
pub struct RequestLogger;
|
|
|
|
impl<B> OnRequest<B> for RequestLogger {
|
|
fn on_request(&mut self, request: &Request<B>, span: &Span) {
|
|
let _enter = span.enter();
|
|
|
|
// Try to get the request ID from extensions
|
|
// This will work if RequestIdLayer has already run
|
|
if let Some(request_id) = request.extensions().get::<RequestId>() {
|
|
span.record("request_id", request_id.0.as_str());
|
|
}
|
|
|
|
let method = method_to_static_str(request.method().as_str());
|
|
let path = normalize_path_for_metrics(request.uri().path());
|
|
Metrics::record_http_request(method, &path);
|
|
|
|
// Log the request start
|
|
info!(
|
|
target: "smg::request",
|
|
"started processing request"
|
|
);
|
|
}
|
|
}
|
|
|
|
/// Custom on_response handler
|
|
#[derive(Clone, Debug)]
|
|
pub struct ResponseLogger {
|
|
_start_time: Instant,
|
|
}
|
|
|
|
impl Default for ResponseLogger {
|
|
fn default() -> Self {
|
|
Self {
|
|
_start_time: Instant::now(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<B> OnResponse<B> for ResponseLogger {
|
|
fn on_response(self, response: &Response<B>, latency: Duration, span: &Span) {
|
|
let status = response.status();
|
|
let status_code = status.as_u16();
|
|
|
|
let error_code = extract_error_code_from_response(response);
|
|
|
|
// Layer 1: HTTP metrics
|
|
Metrics::record_http_response(status_code, error_code);
|
|
|
|
// Record these in the span for structured logging/observability tools
|
|
span.record("status_code", status_code);
|
|
// Use microseconds as integer to avoid format! string allocation
|
|
span.record("latency", latency.as_micros() as u64);
|
|
|
|
// Log the response completion
|
|
let _enter = span.enter();
|
|
if status.is_server_error() {
|
|
error!(
|
|
target: "smg::response",
|
|
"request failed with server error"
|
|
);
|
|
} else if status.is_client_error() {
|
|
warn!(
|
|
target: "smg::response",
|
|
"request failed with client error"
|
|
);
|
|
} else {
|
|
info!(
|
|
target: "smg::response",
|
|
"finished processing request"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Create a configured TraceLayer for HTTP logging
|
|
/// Note: Actual request/response logging with request IDs is done in RequestIdService
|
|
pub fn create_logging_layer() -> TraceLayer<
|
|
tower_http::classify::SharedClassifier<tower_http::classify::ServerErrorsAsFailures>,
|
|
RequestSpan,
|
|
RequestLogger,
|
|
ResponseLogger,
|
|
> {
|
|
TraceLayer::new_for_http()
|
|
.make_span_with(RequestSpan)
|
|
.on_request(RequestLogger)
|
|
.on_response(ResponseLogger::default())
|
|
}
|
|
|
|
/// Request queue entry
|
|
pub struct QueuedRequest {
|
|
/// Time when the request was queued
|
|
queued_at: Instant,
|
|
/// Channel to send the permit back when acquired
|
|
permit_tx: oneshot::Sender<Result<(), StatusCode>>,
|
|
}
|
|
|
|
/// Queue metrics for monitoring
|
|
#[derive(Debug, Default)]
|
|
pub struct QueueMetrics {
|
|
pub total_queued: AtomicU64,
|
|
pub current_queued: AtomicU64,
|
|
pub total_timeout: AtomicU64,
|
|
pub total_rejected: AtomicU64,
|
|
}
|
|
|
|
/// Queue processor that handles queued requests
|
|
pub struct QueueProcessor {
|
|
token_bucket: Arc<TokenBucket>,
|
|
queue_rx: mpsc::Receiver<QueuedRequest>,
|
|
queue_timeout: Duration,
|
|
}
|
|
|
|
impl QueueProcessor {
|
|
pub fn new(
|
|
token_bucket: Arc<TokenBucket>,
|
|
queue_rx: mpsc::Receiver<QueuedRequest>,
|
|
queue_timeout: Duration,
|
|
) -> Self {
|
|
Self {
|
|
token_bucket,
|
|
queue_rx,
|
|
queue_timeout,
|
|
}
|
|
}
|
|
|
|
pub async fn run(mut self) {
|
|
debug!("Starting concurrency queue processor");
|
|
|
|
// Process requests in a single task to reduce overhead
|
|
while let Some(queued) = self.queue_rx.recv().await {
|
|
// Check timeout immediately
|
|
let elapsed = queued.queued_at.elapsed();
|
|
if elapsed >= self.queue_timeout {
|
|
warn!("Request already timed out in queue");
|
|
let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT));
|
|
continue;
|
|
}
|
|
|
|
let remaining_timeout = self.queue_timeout - elapsed;
|
|
|
|
// Try to acquire token for this request
|
|
if self.token_bucket.try_acquire(1.0).await.is_ok() {
|
|
// Got token immediately
|
|
debug!("Queue: acquired token immediately for queued request");
|
|
let _ = queued.permit_tx.send(Ok(()));
|
|
} else {
|
|
// Need to wait for token
|
|
let token_bucket = self.token_bucket.clone();
|
|
|
|
// Spawn task only when we actually need to wait
|
|
tokio::spawn(async move {
|
|
if token_bucket
|
|
.acquire_timeout(1.0, remaining_timeout)
|
|
.await
|
|
.is_ok()
|
|
{
|
|
debug!("Queue: acquired token after waiting");
|
|
let _ = queued.permit_tx.send(Ok(()));
|
|
} else {
|
|
warn!("Queue: request timed out waiting for token");
|
|
let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT));
|
|
}
|
|
});
|
|
}
|
|
}
|
|
|
|
warn!("Concurrency queue processor shutting down");
|
|
}
|
|
}
|
|
|
|
/// State for the concurrency limiter
|
|
pub struct ConcurrencyLimiter {
|
|
pub queue_tx: Option<mpsc::Sender<QueuedRequest>>,
|
|
}
|
|
|
|
impl ConcurrencyLimiter {
|
|
/// Create new concurrency limiter with optional queue
|
|
pub fn new(
|
|
token_bucket: Option<Arc<TokenBucket>>,
|
|
queue_size: usize,
|
|
queue_timeout: Duration,
|
|
) -> (Self, Option<QueueProcessor>) {
|
|
match (token_bucket, queue_size) {
|
|
(None, _) => (Self { queue_tx: None }, None),
|
|
(Some(bucket), size) if size > 0 => {
|
|
let (queue_tx, queue_rx) = mpsc::channel(size);
|
|
let processor = QueueProcessor::new(bucket, queue_rx, queue_timeout);
|
|
(
|
|
Self {
|
|
queue_tx: Some(queue_tx),
|
|
},
|
|
Some(processor),
|
|
)
|
|
}
|
|
(Some(_), _) => (Self { queue_tx: None }, None),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Middleware function for concurrency limiting with optional queuing
|
|
pub async fn concurrency_limit_middleware(
|
|
State(app_state): State<Arc<AppState>>,
|
|
request: Request<Body>,
|
|
next: Next,
|
|
) -> Response {
|
|
// Check mesh global rate limit first if mesh is enabled
|
|
// If mesh is not enabled, this check is skipped and local rate limiting is used
|
|
if let Some(sync_manager) = &app_state.mesh_sync_manager {
|
|
let (is_exceeded, current_count, limit) = sync_manager.check_global_rate_limit();
|
|
if is_exceeded {
|
|
debug!(
|
|
"Global rate limit exceeded: {}/{} req/s",
|
|
current_count, limit
|
|
);
|
|
return (
|
|
StatusCode::TOO_MANY_REQUESTS,
|
|
Json(json!({
|
|
"error": "Rate limit exceeded",
|
|
"current_count": current_count,
|
|
"limit": limit
|
|
})),
|
|
)
|
|
.into_response();
|
|
}
|
|
}
|
|
|
|
let token_bucket = match &app_state.context.rate_limiter {
|
|
Some(bucket) => bucket.clone(),
|
|
None => {
|
|
// Rate limiting disabled, pass through immediately
|
|
return next.run(request).await;
|
|
}
|
|
};
|
|
|
|
// Static counter for embeddings queue size
|
|
static EMBEDDINGS_QUEUE_SIZE: AtomicU64 = AtomicU64::new(0);
|
|
|
|
// Identify if this is an embeddings request based on path
|
|
let is_embeddings = request.uri().path().contains("/v1/embeddings");
|
|
|
|
// Try to acquire token immediately
|
|
if token_bucket.try_acquire(1.0).await.is_ok() {
|
|
debug!("Acquired token immediately");
|
|
Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED);
|
|
let response = next.run(request).await;
|
|
|
|
// Wrap the response body with TokenGuardBody to return token when stream ends
|
|
// This ensures that for streaming responses, the token is only returned
|
|
// after the entire stream has been sent to the client.
|
|
let (parts, body) = response.into_parts();
|
|
let guarded_body = TokenGuardBody::new(body, token_bucket, 1.0);
|
|
Response::from_parts(parts, Body::new(guarded_body))
|
|
} else {
|
|
// No tokens available, try to queue if enabled
|
|
if let Some(queue_tx) = &app_state.concurrency_queue_tx {
|
|
debug!("No tokens available, attempting to queue request");
|
|
|
|
// Create a channel for the token response
|
|
let (permit_tx, permit_rx) = oneshot::channel();
|
|
|
|
let queued = QueuedRequest {
|
|
queued_at: Instant::now(),
|
|
permit_tx,
|
|
};
|
|
|
|
// Try to send to queue
|
|
match queue_tx.try_send(queued) {
|
|
Ok(_) => {
|
|
// On successful enqueue, update embeddings queue counter if applicable
|
|
if is_embeddings {
|
|
EMBEDDINGS_QUEUE_SIZE.fetch_add(1, Ordering::Relaxed);
|
|
}
|
|
|
|
// Wait for token from queue processor
|
|
match permit_rx.await {
|
|
Ok(Ok(())) => {
|
|
debug!("Acquired token from queue");
|
|
Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_ALLOWED);
|
|
// Dequeue for embeddings
|
|
if is_embeddings {
|
|
EMBEDDINGS_QUEUE_SIZE.fetch_sub(1, Ordering::Relaxed);
|
|
}
|
|
|
|
let response = next.run(request).await;
|
|
|
|
// Wrap the response body with TokenGuardBody to return token when stream ends
|
|
let (parts, body) = response.into_parts();
|
|
let guarded_body = TokenGuardBody::new(body, token_bucket, 1.0);
|
|
Response::from_parts(parts, Body::new(guarded_body))
|
|
}
|
|
Ok(Err(status)) => {
|
|
warn!("Queue returned error status: {}", status);
|
|
Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED);
|
|
// Dequeue for embeddings on error
|
|
if is_embeddings {
|
|
EMBEDDINGS_QUEUE_SIZE.fetch_sub(1, Ordering::Relaxed);
|
|
}
|
|
status.into_response()
|
|
}
|
|
Err(_) => {
|
|
error!("Queue response channel closed");
|
|
Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED);
|
|
// Dequeue for embeddings on channel error
|
|
if is_embeddings {
|
|
EMBEDDINGS_QUEUE_SIZE.fetch_sub(1, Ordering::Relaxed);
|
|
}
|
|
StatusCode::INTERNAL_SERVER_ERROR.into_response()
|
|
}
|
|
}
|
|
}
|
|
Err(_) => {
|
|
warn!("Request queue is full, returning 429");
|
|
Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED);
|
|
StatusCode::TOO_MANY_REQUESTS.into_response()
|
|
}
|
|
}
|
|
} else {
|
|
warn!("No tokens available and queuing is disabled, returning 429");
|
|
Metrics::record_http_rate_limit(metrics_labels::RATE_LIMIT_REJECTED);
|
|
StatusCode::TOO_MANY_REQUESTS.into_response()
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// HTTP Metrics Layer (Layer 1: SMG metrics)
|
|
// ============================================================================
|
|
|
|
/// Global counter for active HTTP connections (handlers currently executing)
|
|
static ACTIVE_HTTP_CONNECTIONS: AtomicU64 = AtomicU64::new(0);
|
|
|
|
/// Tower Layer for HTTP metrics collection (SMG Layer 1 metrics)
|
|
#[derive(Clone)]
|
|
pub struct HttpMetricsLayer {
|
|
tracker: Arc<InFlightRequestTracker>,
|
|
}
|
|
|
|
impl HttpMetricsLayer {
|
|
pub fn new(tracker: Arc<InFlightRequestTracker>) -> Self {
|
|
Self { tracker }
|
|
}
|
|
}
|
|
|
|
impl<S> Layer<S> for HttpMetricsLayer {
|
|
type Service = HttpMetricsMiddleware<S>;
|
|
|
|
fn layer(&self, inner: S) -> Self::Service {
|
|
HttpMetricsMiddleware {
|
|
inner,
|
|
in_flight_request_tracker: self.tracker.clone(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Tower Service for HTTP metrics collection
|
|
#[derive(Clone)]
|
|
pub struct HttpMetricsMiddleware<S> {
|
|
inner: S,
|
|
in_flight_request_tracker: Arc<InFlightRequestTracker>,
|
|
}
|
|
|
|
impl<S> Service<Request> for HttpMetricsMiddleware<S>
|
|
where
|
|
S: Service<Request, Response = Response> + Send + Clone + 'static,
|
|
S::Future: Send + 'static,
|
|
{
|
|
type Response = S::Response;
|
|
type Error = S::Error;
|
|
type Future =
|
|
Pin<Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>>;
|
|
|
|
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
|
self.inner.poll_ready(cx)
|
|
}
|
|
|
|
fn call(&mut self, req: Request) -> Self::Future {
|
|
// Convert method to static string to avoid allocation
|
|
let method = method_to_static_str(req.method().as_str());
|
|
let path = normalize_path_for_metrics(req.uri().path());
|
|
let start = Instant::now();
|
|
|
|
let mut inner = self.inner.clone();
|
|
let in_flight_request_tracker = self.in_flight_request_tracker.clone();
|
|
|
|
Box::pin(async move {
|
|
// Increment inside async block - ensures no leak if future is dropped before polling
|
|
let active = ACTIVE_HTTP_CONNECTIONS.fetch_add(1, Ordering::Relaxed) + 1;
|
|
Metrics::set_http_connections_active(active as usize);
|
|
|
|
let guard = in_flight_request_tracker.track();
|
|
|
|
// Capture result before decrementing to ensure decrement happens on error too
|
|
let result = inner.call(req).await;
|
|
|
|
drop(guard);
|
|
|
|
// Always decrement, regardless of success or failure
|
|
let active = ACTIVE_HTTP_CONNECTIONS.fetch_sub(1, Ordering::Relaxed) - 1;
|
|
Metrics::set_http_connections_active(active as usize);
|
|
|
|
let response = result?;
|
|
|
|
let duration = start.elapsed();
|
|
Metrics::record_http_duration(method, &path, duration);
|
|
|
|
Ok(response)
|
|
})
|
|
}
|
|
}
|
|
|
|
/// Normalize path for metrics to avoid high cardinality.
|
|
/// Replaces dynamic segments (IDs, UUIDs) with `{id}` placeholder.
|
|
/// Only allocates when normalization is needed; uses single-pass with byte offsets.
|
|
fn normalize_path_for_metrics(path: &str) -> String {
|
|
let bytes = path.as_bytes();
|
|
let mut segment_start = 0;
|
|
let mut segment_idx = 0;
|
|
let mut result: Option<String> = None;
|
|
|
|
for (pos, &b) in bytes.iter().enumerate() {
|
|
if b == b'/' || pos == bytes.len() - 1 {
|
|
// Determine segment end (include last char if not a slash)
|
|
let segment_end = if b == b'/' { pos } else { pos + 1 };
|
|
let segment = &path[segment_start..segment_end];
|
|
|
|
// Check segments after index 2 for dynamic IDs
|
|
if segment_idx > 2 && !segment.is_empty() && is_dynamic_id(segment) {
|
|
// Initialize result with everything before this segment
|
|
let result = result.get_or_insert_with(|| {
|
|
let mut s = String::with_capacity(path.len());
|
|
s.push_str(&path[..segment_start]);
|
|
s
|
|
});
|
|
result.push_str("{id}");
|
|
} else if let Some(ref mut r) = result {
|
|
// Already normalizing, append this segment as-is
|
|
r.push_str(segment);
|
|
}
|
|
|
|
// Add slash after segment (except at end)
|
|
if b == b'/' {
|
|
if let Some(ref mut r) = result {
|
|
r.push('/');
|
|
}
|
|
segment_start = pos + 1;
|
|
segment_idx += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
result.unwrap_or_else(|| path.to_owned())
|
|
}
|
|
|
|
/// Check if segment looks like a dynamic ID (prefixed ID, UUID, or numeric).
|
|
#[inline]
|
|
fn is_dynamic_id(s: &str) -> bool {
|
|
// Prefixed IDs: resp_xxx, chatcmpl_xxx (len > 10 with underscore)
|
|
if s.len() > 10 && s.contains('_') {
|
|
return true;
|
|
}
|
|
// UUIDs: 32+ hex chars with dashes
|
|
if s.len() >= 32 && s.bytes().all(|b| b.is_ascii_hexdigit() || b == b'-') {
|
|
return true;
|
|
}
|
|
// Numeric IDs
|
|
!s.is_empty() && s.bytes().all(|b| b.is_ascii_digit())
|
|
}
|
|
|
|
pub async fn wasm_middleware(
|
|
State(app_state): State<Arc<AppState>>,
|
|
request: Request<Body>,
|
|
next: Next,
|
|
) -> Result<Response, StatusCode> {
|
|
// Check if WASM is enabled
|
|
if !app_state.context.router_config.enable_wasm {
|
|
return Ok(next.run(request).await);
|
|
}
|
|
|
|
// Get WASM manager
|
|
let wasm_manager = match &app_state.context.wasm_manager {
|
|
Some(manager) => manager,
|
|
None => {
|
|
return Ok(next.run(request).await);
|
|
}
|
|
};
|
|
|
|
// Get request ID from extensions or generate one
|
|
let request_id = request
|
|
.extensions()
|
|
.get::<RequestId>()
|
|
.map(|r| r.0.clone())
|
|
.unwrap_or_else(|| generate_request_id(request.uri().path()));
|
|
|
|
// ===== OnRequest Phase =====
|
|
let on_request_attach_point =
|
|
WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnRequest);
|
|
|
|
let modules_on_request =
|
|
match wasm_manager.get_modules_by_attach_point(on_request_attach_point.clone()) {
|
|
Ok(modules) => modules,
|
|
Err(e) => {
|
|
error!("Failed to get WASM modules for OnRequest: {}", e);
|
|
return Ok(next.run(request).await);
|
|
}
|
|
};
|
|
|
|
let response = if modules_on_request.is_empty() {
|
|
next.run(request).await
|
|
} else {
|
|
// Extract request body once before processing modules
|
|
let method = request.method().clone();
|
|
let uri = request.uri().clone();
|
|
let mut headers = request.headers().clone();
|
|
let max_body_size = wasm_manager.get_max_body_size();
|
|
let body_bytes = match axum::body::to_bytes(request.into_body(), max_body_size).await {
|
|
Ok(bytes) => bytes.to_vec(),
|
|
Err(e) => {
|
|
error!("Failed to read request body: {}", e);
|
|
// Create a minimal request with empty body for error recovery
|
|
let error_request = Request::builder()
|
|
.uri(uri)
|
|
.body(Body::empty())
|
|
.unwrap_or_else(|_| Request::new(Body::empty()));
|
|
return Ok(next.run(error_request).await);
|
|
}
|
|
};
|
|
|
|
// Process each OnRequest module
|
|
let mut modified_body = body_bytes;
|
|
|
|
// Pre-compute strings once before the loop to avoid repeated allocations
|
|
let method_str = method.to_string();
|
|
let path_str = uri.path().to_string();
|
|
let query_str = uri.query().unwrap_or("").to_string();
|
|
|
|
for module in modules_on_request {
|
|
// Build WebAssembly request from collected data
|
|
let wasm_headers = build_wasm_headers_from_axum_headers(&headers);
|
|
let wasm_request = WasmRequest {
|
|
method: method_str.clone(),
|
|
path: path_str.clone(),
|
|
query: query_str.clone(),
|
|
headers: wasm_headers,
|
|
body: modified_body.clone(),
|
|
request_id: request_id.clone(),
|
|
now_epoch_ms: std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap_or_else(|_| Duration::from_millis(0))
|
|
.as_millis() as u64,
|
|
};
|
|
|
|
// Execute WASM component
|
|
let action = match wasm_manager
|
|
.execute_module_for_attach_point(
|
|
&module,
|
|
on_request_attach_point.clone(),
|
|
WasmComponentInput::MiddlewareRequest(wasm_request),
|
|
)
|
|
.await
|
|
{
|
|
Some(action) => action,
|
|
None => continue, // Continue to next module on error
|
|
};
|
|
|
|
// Process action
|
|
match action {
|
|
Action::Continue => {}
|
|
Action::Reject(status) => {
|
|
return Err(StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_REQUEST));
|
|
}
|
|
Action::Modify(modify) => {
|
|
// Apply modifications to headers and body
|
|
apply_modify_action_to_headers(&mut headers, &modify);
|
|
if let Some(body_bytes) = modify.body_replace {
|
|
modified_body = body_bytes;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Reconstruct request with modifications
|
|
let mut final_request = Request::builder()
|
|
.method(method)
|
|
.uri(uri)
|
|
.body(Body::from(modified_body))
|
|
.unwrap_or_else(|_| Request::new(Body::empty()));
|
|
*final_request.headers_mut() = headers;
|
|
|
|
// Continue with request processing
|
|
next.run(final_request).await
|
|
};
|
|
|
|
// ===== OnResponse Phase =====
|
|
let on_response_attach_point =
|
|
WasmModuleAttachPoint::Middleware(MiddlewareAttachPoint::OnResponse);
|
|
|
|
let modules_on_response =
|
|
match wasm_manager.get_modules_by_attach_point(on_response_attach_point.clone()) {
|
|
Ok(modules) => modules,
|
|
Err(e) => {
|
|
error!("Failed to get WASM modules for OnResponse: {}", e);
|
|
return Ok(response);
|
|
}
|
|
};
|
|
if modules_on_response.is_empty() {
|
|
return Ok(response);
|
|
}
|
|
// Extract response data once before processing modules
|
|
let mut status = response.status();
|
|
let mut headers = response.headers().clone();
|
|
let max_body_size = wasm_manager.get_max_body_size();
|
|
let mut body_bytes = match axum::body::to_bytes(response.into_body(), max_body_size).await {
|
|
Ok(bytes) => bytes.to_vec(),
|
|
Err(e) => {
|
|
error!("Failed to read response body: {}", e);
|
|
// Create a minimal response with empty body for error recovery
|
|
let error_response = Response::builder()
|
|
.status(status)
|
|
.body(Body::empty())
|
|
.unwrap_or_else(|_| Response::new(Body::empty()));
|
|
return Ok(error_response);
|
|
}
|
|
};
|
|
|
|
// Process each OnResponse module
|
|
for module in modules_on_response {
|
|
// Build WebAssembly response from collected data
|
|
let wasm_headers = build_wasm_headers_from_axum_headers(&headers);
|
|
let wasm_response = WasmResponse {
|
|
status: status.as_u16(),
|
|
headers: wasm_headers,
|
|
body: body_bytes.clone(),
|
|
};
|
|
|
|
// Execute WASM component
|
|
let action = match wasm_manager
|
|
.execute_module_for_attach_point(
|
|
&module,
|
|
on_response_attach_point.clone(),
|
|
WasmComponentInput::MiddlewareResponse(wasm_response),
|
|
)
|
|
.await
|
|
{
|
|
Some(action) => action,
|
|
None => continue, // Continue to next module on error
|
|
};
|
|
|
|
// Process action - apply modifications incrementally
|
|
match action {
|
|
Action::Continue => {
|
|
// Continue to next module
|
|
}
|
|
Action::Reject(status_code) => {
|
|
// Override response status
|
|
status = StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_REQUEST);
|
|
// Return immediately with current state
|
|
let final_response = Response::builder()
|
|
.status(status)
|
|
.body(Body::from(body_bytes))
|
|
.unwrap_or_else(|_| Response::new(Body::empty()));
|
|
let mut final_response = final_response;
|
|
*final_response.headers_mut() = headers;
|
|
return Ok(final_response);
|
|
}
|
|
Action::Modify(modify) => {
|
|
// Apply status modification
|
|
if let Some(new_status) = modify.status {
|
|
status = StatusCode::from_u16(new_status).unwrap_or(status);
|
|
}
|
|
// Apply headers modifications
|
|
apply_modify_action_to_headers(&mut headers, &modify);
|
|
// Apply body_replace
|
|
if let Some(new_body) = modify.body_replace {
|
|
body_bytes = new_body;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Reconstruct final response with all modifications
|
|
let final_response = Response::builder()
|
|
.status(status)
|
|
.body(Body::from(body_bytes))
|
|
.unwrap_or_else(|_| Response::new(Body::empty()));
|
|
let mut final_response = final_response;
|
|
*final_response.headers_mut() = headers;
|
|
Ok(final_response)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_normalize_path_no_ids() {
|
|
// Common API paths should pass through unchanged
|
|
assert_eq!(
|
|
normalize_path_for_metrics("/v1/chat/completions"),
|
|
"/v1/chat/completions"
|
|
);
|
|
assert_eq!(
|
|
normalize_path_for_metrics("/v1/completions"),
|
|
"/v1/completions"
|
|
);
|
|
assert_eq!(normalize_path_for_metrics("/v1/models"), "/v1/models");
|
|
assert_eq!(normalize_path_for_metrics("/health"), "/health");
|
|
}
|
|
|
|
#[test]
|
|
fn test_normalize_path_with_prefixed_id() {
|
|
// Prefixed IDs (resp_xxx, chatcmpl_xxx) should be normalized
|
|
assert_eq!(
|
|
normalize_path_for_metrics("/v1/responses/resp_abc123def456"),
|
|
"/v1/responses/{id}"
|
|
);
|
|
assert_eq!(
|
|
normalize_path_for_metrics("/v1/chat/completions/chatcmpl_abc123xyz"),
|
|
"/v1/chat/completions/{id}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_normalize_path_with_uuid() {
|
|
assert_eq!(
|
|
normalize_path_for_metrics("/v1/responses/550e8400-e29b-41d4-a716-446655440000"),
|
|
"/v1/responses/{id}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_normalize_path_with_numeric_id() {
|
|
assert_eq!(
|
|
normalize_path_for_metrics("/v1/workers/12345"),
|
|
"/v1/workers/{id}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_is_dynamic_id() {
|
|
// Prefixed IDs
|
|
assert!(is_dynamic_id("resp_abc123def"));
|
|
assert!(is_dynamic_id("chatcmpl_xyz789"));
|
|
assert!(!is_dynamic_id("short_id")); // Too short
|
|
|
|
// UUIDs
|
|
assert!(is_dynamic_id("550e8400-e29b-41d4-a716-446655440000"));
|
|
assert!(is_dynamic_id("550e8400e29b41d4a716446655440000")); // No dashes
|
|
|
|
// Numeric
|
|
assert!(is_dynamic_id("12345"));
|
|
assert!(!is_dynamic_id("")); // Empty
|
|
|
|
// Regular words
|
|
assert!(!is_dynamic_id("completions"));
|
|
assert!(!is_dynamic_id("chat"));
|
|
}
|
|
}
|