// HTTP load generation is intentionally independent from gateway internals. use std::collections::BTreeMap; use std::error::Error as StdError; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use bytes::Bytes; use http::{HeaderMap, HeaderName, HeaderValue}; use reqwest::{Client, Method}; use crate::runtime::{BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot}; const MAX_ERROR_SAMPLES: usize = 32; const MAX_STATUS_SAMPLES: usize = 32; const MAX_STATUS_SAMPLE_BODY_CHARS: usize = 512; const FIRST_BODY_BACKGROUND_DRAIN_CHUNKS_ENV: &str = "AETHER_TESTKIT_FIRST_BODY_BACKGROUND_DRAIN_CHUNKS"; const FIRST_BODY_BACKGROUND_DRAIN_MS_ENV: &str = "AETHER_TESTKIT_FIRST_BODY_BACKGROUND_DRAIN_MS"; const FIRST_BODY_AUTO_CLIENT_SHARDS_MAX_ENV: &str = "AETHER_TESTKIT_FIRST_BODY_AUTO_CLIENT_SHARDS_MAX"; const DEFAULT_FIRST_BODY_BACKGROUND_DRAIN_CHUNKS: usize = 64; const DEFAULT_FIRST_BODY_BACKGROUND_DRAIN_MS: u64 = 100; const DEFAULT_FIRST_BODY_AUTO_CLIENT_SHARDS_MAX: usize = 512; const MAX_SSE_CONTROL_LINE_BYTES: usize = 4 * 1024; const SSE_COMPLETION_EVENT_NAMES: &[&str] = &[ "response.completed", "response.done", "response.incomplete", "message.completed", "message_stop", ]; const SSE_ERROR_EVENT_NAMES: &[&str] = &["error", "response.failed"]; #[derive(Debug, Clone, Copy, Default, serde::Serialize, PartialEq, Eq)] pub enum HttpLoadProbeResponseMode { #[default] HeadersOnly, FirstBodyByte, FullBody, } /// Optional checks applied while consuming a probe response body. /// /// The default load-probe API keeps its historical behavior. Callers that /// exercise an SSE endpoint can opt into protocol-level completion checking /// without changing the request configuration shared by existing probes. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct HttpLoadProbeOptions { pub require_sse_done: bool, } #[derive(Debug, Clone)] pub struct HttpLoadProbeConfig { pub url: String, pub warmup_url: Option, pub method: Method, pub headers: BTreeMap, pub header_sets: Vec>, pub body: Option>, pub total_requests: usize, pub concurrency: usize, /// Number of successful warmup requests. For HTTP/2, one request per /// client shard is enough to establish that shard's multiplexed socket; /// this value is intentionally not a promise about TCP connection count. pub warmup_connections: usize, pub timeout: Duration, pub connect_timeout: Option, pub response_mode: HttpLoadProbeResponseMode, pub client_shards: usize, pub pool_max_idle_per_host: Option, pub start_ramp: Duration, pub http1_only: bool, pub http2_prior_knowledge: bool, pub first_body_hold: Duration, } impl Default for HttpLoadProbeConfig { fn default() -> Self { Self { url: String::new(), warmup_url: None, method: Method::GET, headers: BTreeMap::new(), header_sets: Vec::new(), body: None, total_requests: 100, concurrency: 10, warmup_connections: 0, timeout: Duration::from_secs(30), connect_timeout: None, response_mode: HttpLoadProbeResponseMode::HeadersOnly, client_shards: 1, pool_max_idle_per_host: None, start_ramp: Duration::ZERO, http1_only: false, http2_prior_knowledge: false, first_body_hold: Duration::ZERO, } } } impl HttpLoadProbeConfig { pub fn validate(&self) -> Result<(), String> { if self.url.trim().is_empty() { return Err("load probe url cannot be empty".to_string()); } if self.total_requests == 0 { return Err("load probe total_requests must be positive".to_string()); } if self.concurrency == 0 { return Err("load probe concurrency must be positive".to_string()); } if self.timeout.is_zero() { return Err("load probe timeout must be positive".to_string()); } if matches!(self.connect_timeout, Some(timeout) if timeout.is_zero()) { return Err("load probe connect_timeout must be positive when set".to_string()); } if self.client_shards == 0 { return Err("load probe client_shards must be positive".to_string()); } for (index, headers) in self.header_sets.iter().enumerate() { if headers.is_empty() { return Err(format!("load probe header_sets[{index}] cannot be empty")); } } if self.http1_only && self.http2_prior_knowledge { return Err( "load probe cannot enable both http1_only and http2_prior_knowledge".to_string(), ); } Ok(()) } } #[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)] pub struct HttpLoadProbeErrorSample { pub request_index: usize, pub url: String, pub phase: String, pub kind: String, pub elapsed_ms: u64, pub message: String, #[serde(skip_serializing_if = "Option::is_none")] pub source: Option, } #[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)] pub struct HttpLoadProbeStatusSample { pub request_index: usize, pub url: String, pub status: u16, pub elapsed_ms: u64, pub body: String, } #[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)] pub struct HttpLoadProbeResult { pub url: String, pub method: String, pub response_mode: HttpLoadProbeResponseMode, pub require_sse_done: bool, pub total_requests: usize, pub concurrency: usize, /// Highest number of requests concurrently owned by the load probe. /// /// This is a client-side measurement. It proves how many request tasks /// remained in flight (including response-body draining), but it does not /// by itself prove that the gateway admitted the same number. pub max_in_flight_requests: usize, /// Number of successful warmup requests. For HTTP/2, one request per /// client shard is enough to establish that shard's multiplexed socket; /// this value is intentionally not a promise about TCP connection count. pub warmup_connections: usize, #[serde(skip_serializing_if = "Option::is_none")] pub warmup_url: Option, pub client_shards: usize, pub start_ramp_ms: u64, #[serde(skip_serializing_if = "Option::is_none")] pub pool_max_idle_per_host: Option, pub http1_only: bool, pub http2_prior_knowledge: bool, pub timeout_ms: u64, #[serde(skip_serializing_if = "Option::is_none")] pub connect_timeout_ms: Option, pub first_body_hold_ms: u64, pub duration_ms: u64, pub throughput_rps: u64, pub p99_ms: u64, pub completed_requests: usize, pub failed_requests: usize, pub p50_ms: u64, pub p95_ms: u64, pub max_ms: u64, pub mean_ms: u64, #[serde(skip_serializing_if = "Option::is_none")] pub headers_p50_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub headers_p95_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub headers_p99_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub first_body_p50_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub first_body_p95_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub first_body_p99_ms: Option, pub runtime: BenchmarkRuntimeSnapshot, pub status_counts: BTreeMap, pub error_counts: BTreeMap, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub error_samples: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub non_success_status_samples: Vec, } #[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)] pub struct MultiUrlHttpLoadProbeResult { pub target_urls: Vec, pub target_request_counts: BTreeMap, pub method: String, pub response_mode: HttpLoadProbeResponseMode, pub require_sse_done: bool, pub total_requests: usize, pub concurrency: usize, /// Highest number of requests concurrently owned by the load probe. /// /// This is a client-side measurement. It proves how many request tasks /// remained in flight (including response-body draining), but it does not /// by itself prove that the gateway admitted the same number. pub max_in_flight_requests: usize, pub warmup_connections: usize, #[serde(skip_serializing_if = "Option::is_none")] pub warmup_url: Option, pub client_shards: usize, pub start_ramp_ms: u64, #[serde(skip_serializing_if = "Option::is_none")] pub pool_max_idle_per_host: Option, pub http1_only: bool, pub http2_prior_knowledge: bool, pub timeout_ms: u64, #[serde(skip_serializing_if = "Option::is_none")] pub connect_timeout_ms: Option, pub first_body_hold_ms: u64, pub duration_ms: u64, pub throughput_rps: u64, pub p99_ms: u64, pub completed_requests: usize, pub failed_requests: usize, pub p50_ms: u64, pub p95_ms: u64, pub max_ms: u64, pub mean_ms: u64, #[serde(skip_serializing_if = "Option::is_none")] pub headers_p50_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub headers_p95_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub headers_p99_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub first_body_p50_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub first_body_p95_ms: Option, #[serde(skip_serializing_if = "Option::is_none")] pub first_body_p99_ms: Option, pub runtime: BenchmarkRuntimeSnapshot, pub status_counts: BTreeMap, pub error_counts: BTreeMap, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub error_samples: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub non_success_status_samples: Vec, } #[derive(Debug, Default)] struct ProbeWorkerStats { latencies_ms: Vec, header_latencies_ms: Vec, first_body_latencies_ms: Vec, status_counts: BTreeMap, error_counts: BTreeMap, error_samples: Vec, non_success_status_samples: Vec, target_request_counts: BTreeMap, failed_requests: usize, completed_requests: usize, } pub async fn run_http_load_probe( config: &HttpLoadProbeConfig, ) -> Result { run_http_load_probe_with_options(config, HttpLoadProbeOptions::default()).await } pub async fn run_http_load_probe_with_options( config: &HttpLoadProbeConfig, options: HttpLoadProbeOptions, ) -> Result { config.validate()?; validate_probe_options(config, options)?; run_http_load_probe_against_urls(config, std::slice::from_ref(&config.url), options) .await .map(|result| HttpLoadProbeResult { url: result .target_urls .into_iter() .next() .unwrap_or_else(|| config.url.clone()), method: result.method, response_mode: result.response_mode, require_sse_done: result.require_sse_done, total_requests: result.total_requests, concurrency: result.concurrency, max_in_flight_requests: result.max_in_flight_requests, warmup_connections: result.warmup_connections, warmup_url: result.warmup_url, client_shards: result.client_shards, start_ramp_ms: result.start_ramp_ms, pool_max_idle_per_host: result.pool_max_idle_per_host, http1_only: result.http1_only, http2_prior_knowledge: result.http2_prior_knowledge, timeout_ms: result.timeout_ms, connect_timeout_ms: result.connect_timeout_ms, first_body_hold_ms: result.first_body_hold_ms, duration_ms: result.duration_ms, throughput_rps: result.throughput_rps, p99_ms: result.p99_ms, completed_requests: result.completed_requests, failed_requests: result.failed_requests, p50_ms: result.p50_ms, p95_ms: result.p95_ms, max_ms: result.max_ms, mean_ms: result.mean_ms, headers_p50_ms: result.headers_p50_ms, headers_p95_ms: result.headers_p95_ms, headers_p99_ms: result.headers_p99_ms, first_body_p50_ms: result.first_body_p50_ms, first_body_p95_ms: result.first_body_p95_ms, first_body_p99_ms: result.first_body_p99_ms, runtime: result.runtime, status_counts: result.status_counts, error_counts: result.error_counts, error_samples: result.error_samples, non_success_status_samples: result.non_success_status_samples, }) } pub async fn run_multi_url_http_load_probe( config: &HttpLoadProbeConfig, urls: &[String], ) -> Result { run_multi_url_http_load_probe_with_options(config, urls, HttpLoadProbeOptions::default()).await } pub async fn run_multi_url_http_load_probe_with_options( config: &HttpLoadProbeConfig, urls: &[String], options: HttpLoadProbeOptions, ) -> Result { config.validate()?; validate_probe_options(config, options)?; if urls.is_empty() { return Err("multi-url load probe requires at least one target url".to_string()); } run_http_load_probe_against_urls(config, urls, options).await } fn validate_probe_options( config: &HttpLoadProbeConfig, options: HttpLoadProbeOptions, ) -> Result<(), String> { if options.require_sse_done && config.response_mode != HttpLoadProbeResponseMode::FullBody { return Err( "load probe --require-sse-done requires --response-mode full so the body can be consumed" .to_string(), ); } Ok(()) } async fn run_http_load_probe_against_urls( config: &HttpLoadProbeConfig, urls: &[String], options: HttpLoadProbeOptions, ) -> Result { let effective_client_shards = effective_probe_client_shards(config); let clients = Arc::new(build_probe_clients(config, effective_client_shards)?); let target_urls = Arc::new(urls.to_vec()); let total_requests = config.total_requests; let request_headers = build_request_header_sets(config)?; let request_body = config.body.clone().map(Bytes::from); let response_mode = config.response_mode; let require_sse_done = options.require_sse_done; let first_body_hold = config.first_body_hold; let start_ramp = config.start_ramp; warmup_probe_connections(config, Arc::clone(&clients)).await?; let mut runtime_sampler = BenchmarkRuntimeSampler::new(); let started_at = Instant::now(); let next_request = Arc::new(AtomicUsize::new(0)); let in_flight_requests = Arc::new(AtomicUsize::new(0)); let max_in_flight_requests = Arc::new(AtomicUsize::new(0)); let mut workers = tokio::task::JoinSet::new(); for worker_index in 0..config.concurrency { let client = clients[worker_index % clients.len()].clone(); let next_request = Arc::clone(&next_request); let in_flight_requests = Arc::clone(&in_flight_requests); let max_in_flight_requests = Arc::clone(&max_in_flight_requests); let method = config.method.clone(); let urls = Arc::clone(&target_urls); let request_headers = Arc::clone(&request_headers); let request_body = request_body.clone(); let start_delay = worker_start_delay(start_ramp, worker_index, config.concurrency); workers.spawn(async move { let mut stats = ProbeWorkerStats::default(); if !start_delay.is_zero() { tokio::time::sleep(start_delay).await; } loop { let current = next_request.fetch_add(1, Ordering::AcqRel); if current >= total_requests { break; } let current_in_flight = in_flight_requests.fetch_add(1, Ordering::AcqRel) + 1; max_in_flight_requests.fetch_max(current_in_flight, Ordering::AcqRel); let started_at = Instant::now(); let url = urls[current % urls.len()].clone(); let mut request = client.request(method.clone(), &url); let headers = &request_headers[current % request_headers.len()]; for (name, value) in headers.iter() { request = request.header(name, value); } if let Some(body) = request_body.as_ref() { request = request.body(body.clone()); } match request.send().await { Ok(response) => { let headers_latency_ms = started_at.elapsed().as_millis() as u64; let status = response.status().as_u16(); // Count the HTTP response before consuming its body. A stream can // terminate at the protocol layer while still having a valid 2xx status. *stats.status_counts.entry(status).or_insert(0) += 1; let body_result = observe_response_body( response, response_mode, started_at, first_body_hold, require_sse_done && (200..300).contains(&status), ) .await; match body_result { Err(error) => { stats.failed_requests += 1; record_load_error( &mut stats.error_counts, &mut stats.error_samples, current, &url, started_at.elapsed().as_millis() as u64, error, ); } Ok(observation) => { if !(200..300).contains(&status) { record_non_success_status_sample( &mut stats.non_success_status_samples, current, &url, status, started_at.elapsed().as_millis() as u64, observation.body_sample.as_deref().unwrap_or_default(), ); } *stats.target_request_counts.entry(url).or_insert(0) += 1; if let Some(first_body_latency_ms) = observation.first_body_latency_ms { stats.first_body_latencies_ms.push(first_body_latency_ms); } } } let latency_ms = started_at.elapsed().as_millis() as u64; stats.latencies_ms.push(latency_ms); stats.header_latencies_ms.push(headers_latency_ms); stats.completed_requests += 1; } Err(err) => { let latency_ms = started_at.elapsed().as_millis() as u64; stats.latencies_ms.push(latency_ms); stats.failed_requests += 1; record_load_error( &mut stats.error_counts, &mut stats.error_samples, current, &url, latency_ms, classify_reqwest_error("send", &err), ); stats.completed_requests += 1; } } in_flight_requests.fetch_sub(1, Ordering::AcqRel); } stats }); } let mut aggregate = ProbeWorkerStats::default(); while let Some(result) = workers.join_next().await { let worker = result.map_err(|err| format!("load probe worker task failed: {err}"))?; merge_probe_worker_stats(&mut aggregate, worker); } let ProbeWorkerStats { mut latencies_ms, header_latencies_ms: mut header_latencies, first_body_latencies_ms: mut first_body_latencies, status_counts, error_counts, error_samples, non_success_status_samples, target_request_counts, failed_requests, completed_requests, } = aggregate; let mut latencies = std::mem::take(&mut latencies_ms); latencies.sort_unstable(); header_latencies.sort_unstable(); first_body_latencies.sort_unstable(); let (p50_ms, p95_ms, p99_ms, max_ms, mean_ms) = summarize_latencies(&latencies); let (headers_p50_ms, headers_p95_ms, headers_p99_ms, _, _) = summarize_latencies(&header_latencies); let (first_body_p50_ms, first_body_p95_ms, first_body_p99_ms, _, _) = summarize_latencies(&first_body_latencies); let duration_ms = started_at.elapsed().as_millis() as u64; let throughput_rps = if duration_ms == 0 { completed_requests as u64 } else { ((completed_requests as u64) * 1_000) / duration_ms.max(1) }; Ok(MultiUrlHttpLoadProbeResult { target_urls: urls.to_vec(), target_request_counts, method: config.method.as_str().to_string(), response_mode: config.response_mode, require_sse_done, total_requests: config.total_requests, concurrency: config.concurrency, max_in_flight_requests: max_in_flight_requests.load(Ordering::Acquire), warmup_connections: config.warmup_connections, warmup_url: config.warmup_url.clone(), client_shards: effective_client_shards, start_ramp_ms: config.start_ramp.as_millis() as u64, pool_max_idle_per_host: config.pool_max_idle_per_host, http1_only: config.http1_only, http2_prior_knowledge: config.http2_prior_knowledge, timeout_ms: config.timeout.as_millis() as u64, connect_timeout_ms: config .connect_timeout .map(|timeout| timeout.as_millis() as u64), first_body_hold_ms: config.first_body_hold.as_millis() as u64, duration_ms, throughput_rps, p99_ms, completed_requests, failed_requests, p50_ms, p95_ms, max_ms, mean_ms, headers_p50_ms: (!header_latencies.is_empty()).then_some(headers_p50_ms), headers_p95_ms: (!header_latencies.is_empty()).then_some(headers_p95_ms), headers_p99_ms: (!header_latencies.is_empty()).then_some(headers_p99_ms), first_body_p50_ms: (!first_body_latencies.is_empty()).then_some(first_body_p50_ms), first_body_p95_ms: (!first_body_latencies.is_empty()).then_some(first_body_p95_ms), first_body_p99_ms: (!first_body_latencies.is_empty()).then_some(first_body_p99_ms), runtime: runtime_sampler.snapshot(), status_counts, error_counts, error_samples, non_success_status_samples, }) } fn effective_probe_client_shards(config: &HttpLoadProbeConfig) -> usize { let configured = config.client_shards.max(1); if configured != 1 || config.http1_only || !config.http2_prior_knowledge || config.response_mode != HttpLoadProbeResponseMode::FirstBodyByte { return configured; } let max_auto = env_usize( FIRST_BODY_AUTO_CLIENT_SHARDS_MAX_ENV, DEFAULT_FIRST_BODY_AUTO_CLIENT_SHARDS_MAX, ) .max(1); config.concurrency.max(1).min(max_auto) } fn build_probe_clients( config: &HttpLoadProbeConfig, client_shards: usize, ) -> Result, String> { let mut clients = Vec::with_capacity(client_shards); for _ in 0..client_shards { let mut builder = Client::builder().timeout(config.timeout); if let Some(connect_timeout) = config.connect_timeout { builder = builder.connect_timeout(connect_timeout); } if let Some(pool_max_idle_per_host) = config.pool_max_idle_per_host { builder = builder.pool_max_idle_per_host(pool_max_idle_per_host); } if config.http1_only { builder = builder.http1_only(); } if config.http2_prior_knowledge { builder = builder.http2_prior_knowledge(); } clients.push( builder .build() .map_err(|err| format!("failed to build load probe http client: {err}"))?, ); } Ok(clients) } fn worker_start_delay(start_ramp: Duration, worker_index: usize, concurrency: usize) -> Duration { if start_ramp.is_zero() || concurrency <= 1 || worker_index == 0 { return Duration::ZERO; } let ramp_nanos = start_ramp.as_nanos(); let offset_nanos = ramp_nanos .saturating_mul(worker_index as u128) .checked_div((concurrency - 1) as u128) .unwrap_or_default(); Duration::from_nanos(offset_nanos.min(u64::MAX as u128) as u64) } async fn warmup_probe_connections( config: &HttpLoadProbeConfig, clients: Arc>, ) -> Result<(), String> { if config.warmup_connections == 0 { return Ok(()); } let warmup_url = config.warmup_url.as_deref().unwrap_or(config.url.as_str()); let mut workers = tokio::task::JoinSet::new(); let used_client_count = clients.len().min(config.concurrency); for (client_index, requests) in warmup_request_distribution(used_client_count, config.warmup_connections) .into_iter() .enumerate() { let client = clients[client_index].clone(); let warmup_url = warmup_url.to_string(); workers.spawn(async move { for _ in 0..requests { let response = client .get(&warmup_url) .send() .await .map_err(|err| format!("warmup request failed: {err}"))?; let response = response .error_for_status() .map_err(|err| format!("warmup request returned error status: {err}"))?; response .bytes() .await .map_err(|err| format!("warmup response body failed: {err}"))?; } Ok::<(), String>(()) }); } while let Some(result) = workers.join_next().await { result .map_err(|err| format!("warmup worker task failed: {err}"))? .map_err(|err| format!("failed to warm load probe connections: {err}"))?; } Ok(()) } fn warmup_request_distribution(client_count: usize, total_requests: usize) -> Vec { if client_count == 0 || total_requests == 0 { return Vec::new(); } let shard_count = client_count.min(total_requests); let requests_per_shard = total_requests / shard_count; let shards_with_extra_request = total_requests % shard_count; (0..shard_count) .map(|index| requests_per_shard + usize::from(index < shards_with_extra_request)) .collect() } #[derive(Debug, Clone, PartialEq, Eq)] struct ClassifiedLoadError { key: String, phase: String, kind: String, message: String, source: Option, } impl ClassifiedLoadError { fn static_body(kind: &str, message: &str) -> Self { Self { key: format!("body:{kind}"), phase: "body".to_string(), kind: kind.to_string(), message: message.to_string(), source: None, } } } fn merge_probe_worker_stats(target: &mut ProbeWorkerStats, mut source: ProbeWorkerStats) { target.latencies_ms.append(&mut source.latencies_ms); target .header_latencies_ms .append(&mut source.header_latencies_ms); target .first_body_latencies_ms .append(&mut source.first_body_latencies_ms); merge_counts(&mut target.status_counts, source.status_counts); merge_counts(&mut target.error_counts, source.error_counts); merge_counts( &mut target.target_request_counts, source.target_request_counts, ); append_bounded( &mut target.error_samples, source.error_samples, MAX_ERROR_SAMPLES, ); append_bounded( &mut target.non_success_status_samples, source.non_success_status_samples, MAX_STATUS_SAMPLES, ); target.failed_requests = target .failed_requests .saturating_add(source.failed_requests); target.completed_requests = target .completed_requests .saturating_add(source.completed_requests); } fn merge_counts(target: &mut BTreeMap, source: BTreeMap) { for (key, count) in source { let current = target.entry(key).or_insert(0); *current = current.saturating_add(count); } } fn append_bounded(target: &mut Vec, source: Vec, limit: usize) { let remaining = limit.saturating_sub(target.len()); target.extend(source.into_iter().take(remaining)); } fn record_load_error( error_counts: &mut BTreeMap, error_samples: &mut Vec, request_index: usize, url: &str, elapsed_ms: u64, error: ClassifiedLoadError, ) { *error_counts.entry(error.key).or_insert(0) += 1; if error_samples.len() < MAX_ERROR_SAMPLES { error_samples.push(HttpLoadProbeErrorSample { request_index, url: url.to_string(), phase: error.phase, kind: error.kind, elapsed_ms, message: error.message, source: error.source, }); } } fn record_non_success_status_sample( non_success_status_samples: &mut Vec, request_index: usize, url: &str, status: u16, elapsed_ms: u64, body: &str, ) { if non_success_status_samples.len() < MAX_STATUS_SAMPLES { non_success_status_samples.push(HttpLoadProbeStatusSample { request_index, url: url.to_string(), status, elapsed_ms, body: compact_error_text(body, MAX_STATUS_SAMPLE_BODY_CHARS), }); } } fn classify_reqwest_error(phase: &str, err: &reqwest::Error) -> ClassifiedLoadError { let kind = if err.is_timeout() && err.is_connect() { "connect_timeout" } else if err.is_timeout() && err.is_body() { "body_timeout" } else if err.is_timeout() { "timeout" } else if err.is_connect() { "connect" } else if err.is_body() { "body" } else if err.is_request() { "request" } else if err.is_decode() { "decode" } else if err.is_redirect() { "redirect" } else { "other" }; ClassifiedLoadError { key: format!("{phase}:{kind}"), phase: phase.to_string(), kind: kind.to_string(), message: compact_error_text(err.to_string(), 240), source: error_source_chain(err, 240), } } fn classify_incomplete_sse(body_error: Option<&reqwest::Error>) -> ClassifiedLoadError { let mut incomplete = ClassifiedLoadError::static_body( "sse_incomplete", "SSE response body ended without [DONE] or a recognized completion event", ); if let Some(body_error) = body_error { let body_error = classify_reqwest_error("body", body_error); incomplete.message = format!("{}: {}", incomplete.message, body_error.message); incomplete.source = body_error.source; } incomplete } fn classify_sse_error() -> ClassifiedLoadError { ClassifiedLoadError::static_body( "sse_error", "SSE response body contained an in-band error event", ) } fn error_source_chain(err: &(dyn StdError + 'static), max_chars: usize) -> Option { let mut sources = Vec::new(); let mut next = err.source(); while let Some(source) = next { sources.push(compact_error_text(source.to_string(), max_chars)); if sources.len() >= 4 { break; } next = source.source(); } (!sources.is_empty()).then(|| compact_error_text(sources.join(" | "), max_chars)) } fn compact_error_text(value: impl AsRef, max_chars: usize) -> String { let mut compact = value .as_ref() .split_whitespace() .collect::>() .join(" "); if compact.chars().count() > max_chars { compact = compact.chars().take(max_chars.saturating_sub(1)).collect(); compact.push_str("..."); } compact } async fn observe_response_body( mut response: reqwest::Response, response_mode: HttpLoadProbeResponseMode, started_at: Instant, first_body_hold: Duration, require_sse_done: bool, ) -> Result { match response_mode { HttpLoadProbeResponseMode::HeadersOnly => Ok(BodyObservation::default()), HttpLoadProbeResponseMode::FirstBodyByte => { let first = response .chunk() .await .map_err(|err| classify_reqwest_error("body", &err))? .ok_or_else(|| { ClassifiedLoadError::static_body( "empty_body", "response body ended before the first chunk", ) })?; let first_body_latency_ms = started_at.elapsed().as_millis() as u64; let body_sample = compact_bytes_sample(&first, MAX_STATUS_SAMPLE_BODY_CHARS); if !first_body_hold.is_zero() { tokio::time::sleep(first_body_hold).await; } drain_first_body_response_tail(response).await?; Ok(BodyObservation { first_body_latency_ms: Some(first_body_latency_ms), body_sample: Some(body_sample), }) } HttpLoadProbeResponseMode::FullBody => { let mut first_body_latency_ms = None; let mut body_sample = Vec::new(); let mut sse_completion = SseCompletionDetector::default(); loop { let chunk = match response.chunk().await { Ok(Some(chunk)) => chunk, Ok(None) => break, Err(err) => { if require_sse_done { if sse_completion.has_error() { return Err(classify_sse_error()); } if first_body_latency_ms.is_some() && !sse_completion.is_complete() { return Err(classify_incomplete_sse(Some(&err))); } } return Err(classify_reqwest_error("body", &err)); } }; if first_body_latency_ms.is_none() { first_body_latency_ms = Some(started_at.elapsed().as_millis() as u64); } if require_sse_done { sse_completion.observe(&chunk); } append_body_sample(&mut body_sample, &chunk, MAX_STATUS_SAMPLE_BODY_CHARS); } sse_completion.finish(); if first_body_latency_ms.is_none() { return Err(ClassifiedLoadError::static_body( "empty_body", "response body ended before the first chunk", )); } if require_sse_done { if sse_completion.has_error() { return Err(classify_sse_error()); } if !sse_completion.is_complete() { return Err(classify_incomplete_sse(None)); } } Ok(BodyObservation { first_body_latency_ms, body_sample: Some(String::from_utf8_lossy(&body_sample).into_owned()), }) } } } #[derive(Debug, Default)] struct SseCompletionDetector { complete: bool, error: bool, pending_complete: bool, pending_error: bool, line: Vec, line_overflowed: bool, } impl SseCompletionDetector { fn observe(&mut self, chunk: &[u8]) { if self.error { return; } for &byte in chunk { if byte == b'\n' { if !self.line_overflowed { if trim_ascii_whitespace(&self.line).is_empty() { self.commit_event(); } else { self.pending_error |= sse_line_is_error(&self.line); self.pending_complete |= sse_line_is_completion(&self.line); } } self.line.clear(); self.line_overflowed = false; if self.error { return; } } else if !self.line_overflowed { if self.line.len() < MAX_SSE_CONTROL_LINE_BYTES { self.line.push(byte); } else { self.line.clear(); self.line_overflowed = true; } } } } fn finish(&mut self) { // An SSE event is dispatched only after its terminating blank line. Do not commit a // partial final line or an event that was truncated between its last field and separator. self.line.clear(); self.line_overflowed = false; self.pending_complete = false; self.pending_error = false; } fn commit_event(&mut self) { if self.pending_error { self.error = true; } else if self.pending_complete { self.complete = true; } self.pending_complete = false; self.pending_error = false; } fn is_complete(&self) -> bool { self.complete } fn has_error(&self) -> bool { self.error } } fn sse_line_is_error(line: &[u8]) -> bool { let line = trim_ascii_whitespace(line); let Some(separator) = line.iter().position(|byte| *byte == b':') else { return false; }; let field = trim_ascii_whitespace(&line[..separator]); let value = trim_ascii_whitespace(&line[separator + 1..]); if field == b"event" { return std::str::from_utf8(value) .ok() .is_some_and(is_sse_error_event_name); } if field != b"data" || (!contains_bytes(value, b"\"error\"") && !SSE_ERROR_EVENT_NAMES .iter() .any(|name| contains_bytes(value, name.as_bytes()))) { return false; } serde_json::from_slice::(value) .ok() .is_some_and(|payload| { payload.as_object().is_some_and(|payload| { payload.get("error").is_some_and(|error| !error.is_null()) || payload .get("type") .and_then(serde_json::Value::as_str) .is_some_and(is_sse_error_event_name) }) }) } fn is_sse_error_event_name(name: &str) -> bool { SSE_ERROR_EVENT_NAMES.contains(&name) } fn sse_line_is_completion(line: &[u8]) -> bool { let line = trim_ascii_whitespace(line); let Some(separator) = line.iter().position(|byte| *byte == b':') else { return false; }; let field = trim_ascii_whitespace(&line[..separator]); let value = trim_ascii_whitespace(&line[separator + 1..]); if field == b"event" { return std::str::from_utf8(value) .ok() .is_some_and(is_sse_completion_event_name); } if field != b"data" { return false; } if value == b"[DONE]" { return true; } if !SSE_COMPLETION_EVENT_NAMES .iter() .any(|name| contains_bytes(value, name.as_bytes())) { return false; } serde_json::from_slice::(value) .ok() .and_then(|payload| { payload .get("type") .and_then(serde_json::Value::as_str) .map(is_sse_completion_event_name) }) .unwrap_or(false) } fn is_sse_completion_event_name(name: &str) -> bool { SSE_COMPLETION_EVENT_NAMES.contains(&name) } fn contains_bytes(haystack: &[u8], needle: &[u8]) -> bool { !needle.is_empty() && haystack .windows(needle.len()) .any(|window| window == needle) } fn trim_ascii_whitespace(mut value: &[u8]) -> &[u8] { while value.first().is_some_and(u8::is_ascii_whitespace) { value = &value[1..]; } while value.last().is_some_and(u8::is_ascii_whitespace) { value = &value[..value.len() - 1]; } value } fn compact_bytes_sample(bytes: &[u8], max_chars: usize) -> String { compact_error_text(String::from_utf8_lossy(bytes), max_chars) } fn append_body_sample(target: &mut Vec, chunk: &[u8], max_chars: usize) { if target.len() >= max_chars { return; } let remaining = max_chars.saturating_sub(target.len()); target.extend_from_slice(&chunk[..chunk.len().min(remaining)]); } async fn drain_first_body_response_tail( mut response: reqwest::Response, ) -> Result<(), ClassifiedLoadError> { let max_chunks = env_usize( FIRST_BODY_BACKGROUND_DRAIN_CHUNKS_ENV, DEFAULT_FIRST_BODY_BACKGROUND_DRAIN_CHUNKS, ); if max_chunks == 0 { return Ok(()); } let timeout = Duration::from_millis(env_u64( FIRST_BODY_BACKGROUND_DRAIN_MS_ENV, DEFAULT_FIRST_BODY_BACKGROUND_DRAIN_MS, )); let drain = async { for _ in 0..max_chunks { let Some(chunk) = response .chunk() .await .map_err(|err| classify_reqwest_error("body", &err))? else { break; }; drop(chunk); } Ok(()) }; if timeout.is_zero() { return drain.await; } match tokio::time::timeout(timeout, drain).await { Ok(result) => result, Err(_) => Ok(()), } } fn env_usize(key: &str, default_value: usize) -> usize { std::env::var(key) .ok() .and_then(|value| value.trim().parse::().ok()) .unwrap_or(default_value) } fn env_u64(key: &str, default_value: u64) -> u64 { std::env::var(key) .ok() .and_then(|value| value.trim().parse::().ok()) .unwrap_or(default_value) } #[derive(Debug, Clone, Default, PartialEq, Eq)] struct BodyObservation { first_body_latency_ms: Option, body_sample: Option, } fn build_request_header_sets(config: &HttpLoadProbeConfig) -> Result>, String> { let raw_sets = if config.header_sets.is_empty() { vec![config.headers.clone()] } else { config.header_sets.clone() }; let mut sets = Vec::with_capacity(raw_sets.len().max(1)); for headers in raw_sets { sets.push(build_headers(&headers)?); } if sets.is_empty() { sets.push(HeaderMap::new()); } Ok(Arc::new(sets)) } fn build_headers(headers: &BTreeMap) -> Result { let mut result = HeaderMap::new(); for (name, value) in headers { let name = HeaderName::try_from(name.as_str()) .map_err(|err| format!("invalid load probe header name `{name}`: {err}"))?; let value = HeaderValue::from_str(value) .map_err(|err| format!("invalid load probe header value for `{name}`: {err}"))?; result.insert(name, value); } Ok(result) } fn summarize_latencies(latencies: &[u64]) -> (u64, u64, u64, u64, u64) { if latencies.is_empty() { return (0, 0, 0, 0, 0); } let max_ms = *latencies.last().unwrap_or(&0); let mean_ms = latencies.iter().sum::() / latencies.len() as u64; let p50_ms = percentile(latencies, 50); let p95_ms = percentile(latencies, 95); let p99_ms = percentile(latencies, 99); (p50_ms, p95_ms, p99_ms, max_ms, mean_ms) } fn percentile(latencies: &[u64], percentile: u8) -> u64 { if latencies.is_empty() { return 0; } let last_index = latencies.len() - 1; let rank = ((last_index as f64) * (percentile as f64 / 100.0)).round() as usize; latencies[rank.min(last_index)] } #[cfg(test)] mod tests { use super::{ build_headers, build_request_header_sets, run_http_load_probe, run_http_load_probe_with_options, summarize_latencies, validate_probe_options, warmup_request_distribution, worker_start_delay, HttpLoadProbeConfig, HttpLoadProbeOptions, HttpLoadProbeResponseMode, SseCompletionDetector, }; use reqwest::Method; use std::collections::BTreeMap; use std::time::Duration; use tokio::io::{AsyncReadExt, AsyncWriteExt}; #[test] fn validates_probe_config() { assert!(HttpLoadProbeConfig { url: String::new(), ..HttpLoadProbeConfig::default() } .validate() .is_err()); assert!(HttpLoadProbeConfig { total_requests: 0, ..HttpLoadProbeConfig::default() } .validate() .is_err()); assert!(HttpLoadProbeConfig { concurrency: 0, ..HttpLoadProbeConfig::default() } .validate() .is_err()); assert!(HttpLoadProbeConfig { timeout: Duration::ZERO, ..HttpLoadProbeConfig::default() } .validate() .is_err()); assert!(HttpLoadProbeConfig { connect_timeout: Some(Duration::ZERO), ..HttpLoadProbeConfig::default() } .validate() .is_err()); assert!(HttpLoadProbeConfig { client_shards: 0, ..HttpLoadProbeConfig::default() } .validate() .is_err()); assert!(HttpLoadProbeConfig { http1_only: true, http2_prior_knowledge: true, ..HttpLoadProbeConfig::default() } .validate() .is_err()); } #[test] fn summarizes_latency_distribution() { let (p50_ms, p95_ms, p99_ms, max_ms, mean_ms) = summarize_latencies(&[10, 20, 30, 40, 50, 60, 70, 80, 90, 100]); assert_eq!(p50_ms, 60); assert_eq!(p95_ms, 100); assert_eq!(p99_ms, 100); assert_eq!(max_ms, 100); assert_eq!(mean_ms, 55); } #[test] fn default_probe_config_is_reasonable() { let config = HttpLoadProbeConfig::default(); assert_eq!(config.method, Method::GET); assert!(config.warmup_url.is_none()); assert!(config.headers.is_empty()); assert!(config.header_sets.is_empty()); assert!(config.body.is_none()); assert_eq!(config.total_requests, 100); assert_eq!(config.concurrency, 10); assert_eq!(config.warmup_connections, 0); assert_eq!(config.timeout, Duration::from_secs(30)); assert_eq!(config.connect_timeout, None); assert_eq!(config.response_mode, HttpLoadProbeResponseMode::HeadersOnly); assert_eq!(config.client_shards, 1); assert_eq!(config.pool_max_idle_per_host, None); assert_eq!(config.start_ramp, Duration::ZERO); assert!(!config.http1_only); assert!(!config.http2_prior_knowledge); assert_eq!(config.first_body_hold, Duration::ZERO); } #[test] fn sse_completion_check_requires_full_body_mode() { let options = HttpLoadProbeOptions { require_sse_done: true, }; assert!(validate_probe_options(&HttpLoadProbeConfig::default(), options).is_err()); let config = HttpLoadProbeConfig { response_mode: HttpLoadProbeResponseMode::FullBody, ..HttpLoadProbeConfig::default() }; assert!(validate_probe_options(&config, options).is_ok()); } #[test] fn detects_sse_completion_markers_across_chunks() { let mut done = SseCompletionDetector::default(); done.observe(b"data: [DO"); assert!(!done.is_complete()); done.observe(b"NE]\n\n"); assert!(done.is_complete()); let mut explicit_event = SseCompletionDetector::default(); explicit_event.observe(b"event: response.comp"); assert!(!explicit_event.is_complete()); explicit_event.observe(b"leted\ndata: {}\n\n"); assert!(explicit_event.is_complete()); let mut explicit_data_type = SseCompletionDetector::default(); explicit_data_type.observe(b"data: {\"type\":\"message_stop\"}\n\n"); assert!(explicit_data_type.is_complete()); let mut incomplete_response = SseCompletionDetector::default(); incomplete_response.observe(b"event: response.incomplete\ndata: {}\n\n"); assert!(incomplete_response.is_complete()); assert!(!incomplete_response.has_error()); let mut truncated_event = SseCompletionDetector::default(); truncated_event.observe(b"event: response.completed\n"); assert!(!truncated_event.is_complete()); truncated_event.observe(b"\n"); assert!(truncated_event.is_complete()); let mut content_only = SseCompletionDetector::default(); content_only.observe(b"data: {\"delta\":{\"content\":\"[DONE] response.completed\"}}\n\n"); assert!(!content_only.is_complete()); } #[test] fn detects_in_band_sse_errors_across_chunks() { let mut data_error = SseCompletionDetector::default(); data_error.observe(b"data: {\"err"); assert!(!data_error.has_error()); data_error.observe( b"or\":{\"type\":\"execution_runtime_stream_read_error\"}}\n\ndata: [DONE]\n\n", ); assert!(data_error.has_error()); let mut event_error = SseCompletionDetector::default(); event_error.observe(b"event:err"); assert!(!event_error.has_error()); event_error.observe(b"or\ndata: {\"message\":\"upstream failed\"}\n\ndata: [DONE]\n\n"); assert!(event_error.has_error()); let mut response_failed = SseCompletionDetector::default(); response_failed .observe(b"data: [DONE]\n\ndata: {\"type\":\"response.failed\",\"response\":{}}\n\n"); assert!(response_failed.has_error()); let mut null_error = SseCompletionDetector::default(); null_error.observe(b"data: {\"error\":null}\n\ndata: [DONE]\n\n"); assert!(!null_error.has_error()); assert!(null_error.is_complete()); let mut nested_error = SseCompletionDetector::default(); nested_error .observe(b"data: {\"delta\":{\"error\":\"quoted output\"}}\n\ndata: [DONE]\n\n"); assert!(!nested_error.has_error()); assert!(nested_error.is_complete()); } #[tokio::test] async fn missing_sse_completion_is_failed_without_losing_http_status() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await .expect("test listener should bind"); let address = listener .local_addr() .expect("test listener address should resolve"); let server = tokio::spawn(async move { let (mut stream, _) = listener .accept() .await .expect("test connection should arrive"); let mut request = [0_u8; 1024]; let _ = stream .read(&mut request) .await .expect("test request should read"); let body = "data: {\"partial\":true}\n\n"; let response = format!( "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len() ); stream .write_all(response.as_bytes()) .await .expect("test response should write"); }); let config = HttpLoadProbeConfig { url: format!("http://{address}/stream"), total_requests: 1, concurrency: 1, response_mode: HttpLoadProbeResponseMode::FullBody, ..HttpLoadProbeConfig::default() }; let result = run_http_load_probe_with_options( &config, HttpLoadProbeOptions { require_sse_done: true, }, ) .await .expect("load probe should complete"); server.await.expect("test server should stop"); assert!(result.require_sse_done); assert_eq!(result.completed_requests, 1); assert_eq!(result.max_in_flight_requests, 1); assert_eq!(result.failed_requests, 1); assert_eq!(result.status_counts.get(&200), Some(&1)); assert_eq!(result.error_counts.get("body:sse_incomplete"), Some(&1)); assert_eq!(result.error_samples[0].phase, "body"); assert_eq!(result.error_samples[0].kind, "sse_incomplete"); } #[tokio::test] async fn in_band_sse_error_is_failed_even_when_done_follows() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await .expect("test listener should bind"); let address = listener .local_addr() .expect("test listener address should resolve"); let server = tokio::spawn(async move { let (mut stream, _) = listener .accept() .await .expect("test connection should arrive"); let mut request = [0_u8; 1024]; let _ = stream .read(&mut request) .await .expect("test request should read"); let body = concat!( "data: {\"error\":{\"type\":\"execution_runtime_stream_read_error\",", "\"code\":502}}\n\n", "data: [DONE]\n\n", ); let response = format!( "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len() ); stream .write_all(response.as_bytes()) .await .expect("test response should write"); }); let config = HttpLoadProbeConfig { url: format!("http://{address}/stream"), total_requests: 1, concurrency: 1, response_mode: HttpLoadProbeResponseMode::FullBody, ..HttpLoadProbeConfig::default() }; let result = run_http_load_probe_with_options( &config, HttpLoadProbeOptions { require_sse_done: true, }, ) .await .expect("load probe should complete"); server.await.expect("test server should stop"); assert_eq!(result.completed_requests, 1); assert_eq!(result.failed_requests, 1); assert_eq!(result.status_counts.get(&200), Some(&1)); assert_eq!(result.error_counts.get("body:sse_error"), Some(&1)); assert_eq!(result.error_samples[0].phase, "body"); assert_eq!(result.error_samples[0].kind, "sse_error"); } #[tokio::test] async fn result_serializes_effective_request_timeout_ms() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await .expect("test listener should bind"); let address = listener .local_addr() .expect("test listener address should resolve"); let server = tokio::spawn(async move { let (mut stream, _) = listener .accept() .await .expect("test connection should arrive"); let mut request = [0_u8; 1024]; let _ = stream .read(&mut request) .await .expect("test request should read"); stream .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") .await .expect("test response should write"); }); let config = HttpLoadProbeConfig { url: format!("http://{address}/health"), total_requests: 1, concurrency: 1, timeout: Duration::from_secs(120), ..HttpLoadProbeConfig::default() }; let result = run_http_load_probe(&config) .await .expect("load probe should complete"); server.await.expect("test server should stop"); assert_eq!(result.timeout_ms, 120_000); let serialized = serde_json::to_value(result).expect("result should serialize"); assert_eq!( serialized .get("timeout_ms") .and_then(|value| value.as_u64()), Some(120_000) ); } #[test] fn validates_probe_headers() { let mut headers = BTreeMap::new(); headers.insert("x-aether-test".to_string(), "ok".to_string()); let built = build_headers(&headers).expect("headers should build"); assert_eq!( built .get("x-aether-test") .and_then(|value| value.to_str().ok()), Some("ok") ); let invalid = BTreeMap::from([("bad header".to_string(), "ok".to_string())]); assert!(build_headers(&invalid).is_err()); } #[test] fn validates_header_sets() { let mut config = HttpLoadProbeConfig { url: "http://127.0.0.1/".to_string(), header_sets: vec![BTreeMap::new()], ..HttpLoadProbeConfig::default() }; assert!(config.validate().is_err()); config.header_sets = vec![BTreeMap::from([( "authorization".to_string(), "Bearer test".to_string(), )])]; assert!(config.validate().is_ok()); let sets = build_request_header_sets(&config).expect("header set should build"); assert_eq!(sets.len(), 1); assert_eq!( sets[0] .get("authorization") .and_then(|value| value.to_str().ok()), Some("Bearer test") ); } #[test] fn spreads_worker_start_delay_across_ramp() { assert_eq!( worker_start_delay(Duration::from_millis(900), 0, 10), Duration::ZERO ); assert_eq!( worker_start_delay(Duration::from_millis(900), 5, 10), Duration::from_millis(500) ); assert_eq!( worker_start_delay(Duration::from_millis(900), 9, 10), Duration::from_millis(900) ); assert_eq!( worker_start_delay(Duration::from_millis(900), 3, 1), Duration::ZERO ); } #[test] fn distributes_warmup_requests_across_each_used_client_shard() { assert_eq!(warmup_request_distribution(0, 10), Vec::::new()); assert_eq!(warmup_request_distribution(4, 0), Vec::::new()); assert_eq!(warmup_request_distribution(4, 4), vec![1, 1, 1, 1]); assert_eq!(warmup_request_distribution(4, 6), vec![2, 2, 1, 1]); assert_eq!(warmup_request_distribution(4, 10), vec![3, 3, 2, 2]); assert_eq!(warmup_request_distribution(8, 3), vec![1, 1, 1]); } }