diff --git a/objectstore-metrics/src/lib.rs b/objectstore-metrics/src/lib.rs index af9d1209..8731d7a4 100644 --- a/objectstore-metrics/src/lib.rs +++ b/objectstore-metrics/src/lib.rs @@ -318,7 +318,7 @@ pub fn init(config: &MetricsConfig) -> Result<(), Error> { Ok(()) } -pub use mock::with_capturing_test_client; +pub use mock::{with_capturing_test_client, with_capturing_test_client_async}; // --------------------------------------------------------------------------- // Macros diff --git a/objectstore-metrics/src/mock.rs b/objectstore-metrics/src/mock.rs index 62b61d91..7d7ecf39 100644 --- a/objectstore-metrics/src/mock.rs +++ b/objectstore-metrics/src/mock.rs @@ -3,6 +3,7 @@ //! Provides [`with_capturing_test_client`], which installs a thread-local //! recorder that captures all emitted metrics as DogStatsD-format strings. +use std::future::Future; use std::sync::{Arc, Mutex}; use metrics::{Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit}; @@ -26,6 +27,30 @@ pub fn with_capturing_test_client(f: impl FnOnce()) -> Vec { recorder.consume() } +/// Awaits `future` with a thread-local mock recorder installed, then returns all captured +/// metrics as `"name:value|type|#key:value,key:value"` strings. +/// +/// The recorder stays installed across await points, so metrics emitted while the future is +/// suspended are captured as well. Since it is thread-local, the future must not migrate between +/// threads: run it on a current-thread runtime, as `#[tokio::test]` does by default. +/// +/// # Example +/// +/// ```ignore +/// let captured = objectstore_metrics::with_capturing_test_client_async(async { +/// objectstore_metrics::count!("test.counter"); +/// }) +/// .await; +/// assert!(captured.iter().any(|m| m.starts_with("test.counter:"))); +/// ``` +pub async fn with_capturing_test_client_async(future: impl Future) -> Vec { + let recorder = MockRecorder::default(); + let guard = metrics::set_default_local_recorder(&recorder); + future.await; + drop(guard); + recorder.consume() +} + /// A metrics recorder that formats and stores every operation as a string. #[derive(Clone, Default)] struct MockRecorder { diff --git a/objectstore-server/Cargo.toml b/objectstore-server/Cargo.toml index d6ddec62..9c68baa0 100644 --- a/objectstore-server/Cargo.toml +++ b/objectstore-server/Cargo.toml @@ -64,6 +64,7 @@ objectstore-options = { workspace = true, features = ["testing"] } objectstore-test = { workspace = true } stresstest = { workspace = true } tempfile = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } [[bin]] name = "objectstore" diff --git a/objectstore-server/src/web/metrics_body.rs b/objectstore-server/src/web/metrics_body.rs index 29f39554..15ce1f70 100644 --- a/objectstore-server/src/web/metrics_body.rs +++ b/objectstore-server/src/web/metrics_body.rs @@ -5,6 +5,7 @@ use std::task::{Context, Poll}; use axum::body::Body; use axum::http::{Method, StatusCode}; +use axum::response::Response; use bytes::Bytes; use http_body::{Body as HttpBody, Frame, SizeHint}; use pin_project_lite::pin_project; @@ -13,10 +14,11 @@ use tokio::time::Instant; use crate::extractors::downstream_service::DownstreamService; /// State of a response body. -#[derive(Clone, Copy)] +#[derive(Clone, Copy, Default)] enum BodyState { - /// The body is still streaming. - Streaming(StatusCode), + /// The body has not finished streaming: it was either never polled or interrupted mid-stream. + #[default] + Pending, /// The body was streamed to completion. Completed(StatusCode), /// An error was encountered while streaming the body. @@ -26,15 +28,15 @@ enum BodyState { /// Tracks request timing and emits `server.requests.duration` when dropped. /// /// [`MetricsBody`] owns this guard so request duration spans full response-body streaming. -pub(crate) struct EmitMetricsGuard { +pub struct EmitMetricsGuard { route: String, method: Method, start: Instant, - body_state: Option, + body_state: BodyState, } impl EmitMetricsGuard { - pub(crate) fn new(route: &str, method: &Method, service: DownstreamService) -> Self { + pub fn new(route: &str, method: &Method, service: DownstreamService) -> Self { objectstore_metrics::count!( "server.requests", route = route.to_owned(), @@ -46,32 +48,30 @@ impl EmitMetricsGuard { route: route.to_owned(), method: method.clone(), start: Instant::now(), - body_state: None, + body_state: BodyState::Pending, } } - fn mark_streaming(&mut self, status: StatusCode) { - self.body_state = Some(BodyState::Streaming(status)); - } - - fn mark_completed(&mut self) { - if let Some(BodyState::Streaming(status)) = self.body_state { - self.body_state = Some(BodyState::Completed(status)); + /// Records that the body ended with `status`, unless it already errored. + fn complete(&mut self, status: StatusCode) { + if matches!(self.body_state, BodyState::Pending) { + self.body_state = BodyState::Completed(status); } } fn mark_errored(&mut self) { - self.body_state = Some(BodyState::Errored); + self.body_state = BodyState::Errored; } } impl Drop for EmitMetricsGuard { fn drop(&mut self) { let state = match self.body_state { - Some(BodyState::Completed(status)) => status.as_u16(), - Some(BodyState::Streaming(_)) | None => 499, - Some(BodyState::Errored) => 500, + BodyState::Pending => 499, + BodyState::Completed(status) => status.as_u16(), + BodyState::Errored => 500, }; + objectstore_metrics::record!( "server.requests.duration" = self.start.elapsed(), route = self.route.clone(), @@ -87,25 +87,45 @@ pin_project! { /// /// The guard emits the request-duration metric on drop, so keeping it inside /// the body defers that emission until streaming completes. + /// + /// The body checks these conditions where hyper stops polling: + /// 1. The body returns end-of-stream (`None`). + /// 2. An explicit `content-length` was specified and that many bytes were written. + /// 3. For a buffered body, hyper yields a single frame. + /// 4. For a body with trailers, hyper yields the trailers frame and then drops the body. pub struct MetricsBody { - guard: EmitMetricsGuard, #[pin] inner: Body, + guard: EmitMetricsGuard, + status: StatusCode, + remaining: Option, } } impl MetricsBody { - /// Creates a new [`MetricsBody`] that keeps `guard` alive while polling `inner`. - /// - /// `status` is the response status from headers. - pub fn new(mut guard: EmitMetricsGuard, status: StatusCode, inner: Body) -> Self { - guard.mark_streaming(status); - // An empty response body reports end-of-stream immediately and is never polled by hyper, - // so mark it completed up front. - if inner.is_end_stream() { - guard.mark_completed(); + /// Wraps a response body with a [`MetricsBody`] that keeps `guard` alive until the body ends. + pub fn wrap_response(response: Response, mut guard: EmitMetricsGuard) -> Response { + let status = response.status(); + let content_length = response + .headers() + .get(axum::http::header::CONTENT_LENGTH) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()); + + if content_length == Some(0) || response.body().is_end_stream() { + // Fast-path: the body won't be polled, so we complete the guard immediately + guard.complete(status); + response + } else { + response.map(|inner| { + Body::new(Self { + guard, + inner, + status, + remaining: content_length, + }) + }) } - Self { guard, inner } } } @@ -120,22 +140,24 @@ impl http_body::Body for MetricsBody { let mut this = self.project(); let poll = this.inner.as_mut().poll_frame(cx); match &poll { - // End-of-stream for a streamed body. - Poll::Ready(None) => this.guard.mark_completed(), - // End-of-stream for a buffered body (yields a single frame and reports end-of-stream - // immediately, so hyper doesn't attempt to poll it again). - Poll::Ready(Some(Ok(frame))) if frame.is_trailers() || this.inner.is_end_stream() => { - this.guard.mark_completed() + Poll::Ready(None) => this.guard.complete(*this.status), + Poll::Ready(Some(Ok(frame))) => { + if let Some(remaining) = this.remaining.as_mut() { + let frame_len = frame.data_ref().map_or(0, |data| data.len() as u64); + *remaining = remaining.saturating_sub(frame_len); + } + + if *this.remaining == Some(0) || this.inner.is_end_stream() || frame.is_trailers() { + this.guard.complete(*this.status); + } } Poll::Ready(Some(Err(_))) => this.guard.mark_errored(), - _ => {} + Poll::Pending => {} } poll } fn size_hint(&self) -> SizeHint { - // Delegating [`size_hint`](http_body::Body::size_hint) preserves `Content-Length` on buffered - // responses instead of forcing chunked encoding. self.inner.size_hint() } @@ -147,143 +169,195 @@ impl http_body::Body for MetricsBody { #[cfg(test)] mod tests { use std::convert::Infallible; + use std::io; use std::pin::Pin; - use std::sync::Arc; use std::task::{Context, Poll}; + use std::time::Duration; use axum::Router; use axum::body::{self, Body, Bytes}; - use axum::http::{HeaderMap, Request, StatusCode}; + use axum::handler::Handler; + use axum::http::{HeaderMap, Request, header}; use axum::middleware::from_fn; use axum::routing::get; + use futures::StreamExt; use http_body::Frame; use tower::ServiceExt; use crate::web::middleware::emit_request_metrics; - fn make_request(uri: &str) -> Request { - Request::builder().uri(uri).body(Body::empty()).unwrap() + /// How the client consumes the response body. + enum Client { + /// Reads the body to end-of-stream. + ReadToEnd, + /// Reads a single chunk and then disconnects, like hyper, which stops polling once + /// `content-length` bytes have been written. + ReadOneChunk, + /// Disconnects without reading the body. + Disconnect, } - /// Runs a request whose handler returns `body`, driving it to completion on a - /// current-thread runtime, and returns the captured `server.requests.duration` metric - /// string. If `consume_body` is false, the response body is dropped without being read, - /// simulating a client that disconnects mid-stream. - fn capture_duration_metric(body: Body, consume_body: bool) -> String { - let captured = objectstore_metrics::with_capturing_test_client(|| { - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - - rt.block_on(async { - let shared = Arc::new(std::sync::Mutex::new(Some(body))); - let app = Router::new() - .route( - "/v1/test/_/key", - get(move || { - let shared = shared.clone(); - async move { shared.lock().unwrap().take().unwrap() } - }), - ) - .layer(from_fn(emit_request_metrics)); - - let resp = app.oneshot(make_request("/v1/test/_/key")).await.unwrap(); - assert_eq!(resp.status(), StatusCode::OK); - - if consume_body { - // Drain the body to end-of-stream (or until it errors); the result is - // ignored so a server-side stream error does not panic the test. - let _ = body::to_bytes(resp.into_body(), usize::MAX).await; - } else { - // Drop the response (and its body) without reading it: the wrapping - // `MetricsBody` never reaches end-of-stream, mirroring a client disconnect. - drop(resp); + impl Client { + async fn consume(self, response: axum::response::Response) { + // Read results are ignored: a failing stream must be tracked, not panic the test. + match self { + Client::ReadToEnd => { + let _ = body::to_bytes(response.into_body(), usize::MAX).await; + } + Client::ReadOneChunk => { + let _ = response.into_body().into_data_stream().next().await; } - }); - }); + Client::Disconnect => drop(response), + } + } + } - captured - .into_iter() - .find(|m| m.starts_with("server.requests.duration:")) - .expect("duration metric not captured") + /// Serves `handler` behind [`emit_request_metrics`] and returns the status and duration + /// tracked in `server.requests.duration`. + async fn track_request(handler: H, client: Client) -> (u16, Duration) + where + H: Handler, + T: 'static, + { + let app = Router::new() + .route("/", get(handler)) + .layer(from_fn(emit_request_metrics)); + + let captured = objectstore_metrics::with_capturing_test_client_async(async move { + let request = Request::get("/").body(Body::empty()).unwrap(); + let response = app.oneshot(request).await.unwrap(); + client.consume(response).await; + }) + .await; + + // The metric is formatted as `server.requests.duration:|d|#`. + let metric = captured + .iter() + .find_map(|m| m.strip_prefix("server.requests.duration:")) + .expect("duration metric not captured"); + let (seconds, tags) = metric + .split_once("|d|#") + .expect("malformed duration metric"); + let (_, status) = tags.rsplit_once("status:").expect("status tag"); + + let status = status.parse().expect("numeric status"); + let duration = Duration::from_secs_f64(seconds.parse().expect("numeric duration")); + (status, duration) } - /// A body streamed to completion reports the real response status. - #[test] - fn completed_stream_reports_real_status() { - let stream = async_stream::stream! { - yield Ok::<_, std::io::Error>(Bytes::from_static(b"hello")); - }; - let metric = capture_duration_metric(Body::from_stream(stream), true); - assert!(metric.contains("status:200"), "unexpected metric: {metric}"); + /// A stream that yields one chunk and then ends. + fn ending_stream() -> Body { + Body::from_stream(async_stream::stream! { + yield Ok::<_, io::Error>(Bytes::from_static(b"hello")); + }) } - /// A body that needs no polling is completed when it is wrapped, because hyper may not poll - /// it before dropping it. - #[test] - fn empty_body_reports_real_status() { - let metric = capture_duration_metric(Body::empty(), false); - assert!(metric.contains("status:200"), "unexpected metric: {metric}"); + /// A stream that yields one chunk and then never ends. + fn stalling_stream() -> Body { + Body::from_stream(async_stream::stream! { + yield Ok::<_, io::Error>(Bytes::from_static(b"hello")); + std::future::pending::<()>().await; + }) } - /// A buffered body completes after its final data frame, without requiring a subsequent poll - /// that returns `None`. - #[test] - fn buffered_body_reports_real_status() { - let metric = capture_duration_metric(Body::from("hello"), true); - assert!(metric.contains("status:200"), "unexpected metric: {metric}"); + /// A stream that yields one chunk after `delay` and then ends. + fn slow_stream(delay: Duration) -> Body { + Body::from_stream(async_stream::stream! { + tokio::time::sleep(delay).await; + yield Ok::<_, io::Error>(Bytes::from_static(b"hello")); + }) } - /// Hyper treats trailers as terminal, so completion must be recorded when their frame is - /// yielded rather than waiting for an unobserved following `None`. - #[test] - fn trailers_report_real_status() { - struct TrailersBody(bool); - - impl http_body::Body for TrailersBody { - type Data = Bytes; - type Error = Infallible; - - fn poll_frame( - mut self: Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll, Self::Error>>> { - Poll::Ready(if self.0 { - None - } else { - self.0 = true; - Some(Ok(Frame::trailers(HeaderMap::new()))) - }) - } + /// A stream that yields one chunk and then fails. + fn failing_stream() -> Body { + Body::from_stream(async_stream::stream! { + yield Ok(Bytes::from_static(b"hello")); + yield Err(io::Error::other("boom")); + }) + } + + /// A body that yields a trailers frame before its end-of-stream. + struct TrailersBody(bool); + + impl http_body::Body for TrailersBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + Poll::Ready(if self.0 { + None + } else { + self.0 = true; + Some(Ok(Frame::trailers(HeaderMap::new()))) + }) } + } - let metric = capture_duration_metric(Body::new(TrailersBody(false)), true); - assert!(metric.contains("status:200"), "unexpected metric: {metric}"); + #[tokio::test] + async fn completed_stream_reports_real_status() { + let (status, _) = track_request(|| async { ending_stream() }, Client::ReadToEnd).await; + assert_eq!(status, 200); } - /// A body dropped before end-of-stream (client disconnect) reports `499`, overriding the - /// `200` status that was sent in the headers. - #[test] - fn interrupted_stream_reports_499() { - // A stream that yields one chunk then pends forever, so it never reaches end-of-stream. - let stream = async_stream::stream! { - yield Ok::<_, std::io::Error>(Bytes::from_static(b"hello")); - std::future::pending::<()>().await; - }; - let metric = capture_duration_metric(Body::from_stream(stream), false); - assert!(metric.contains("status:499"), "unexpected metric: {metric}"); + /// Hyper may drop an empty body without ever polling it. + #[tokio::test] + async fn empty_body_reports_real_status() { + let (status, _) = track_request(|| async { Body::empty() }, Client::Disconnect).await; + assert_eq!(status, 200); } - /// A body stream that errors server-side reports `500`, overriding the `200` status that - /// was sent in the headers. - #[test] - fn errored_stream_reports_500() { - let stream = async_stream::stream! { - yield Ok(Bytes::from_static(b"hello")); - yield Err(std::io::Error::other("boom")); - }; - let metric = capture_duration_metric(Body::from_stream(stream), true); - assert!(metric.contains("status:500"), "unexpected metric: {metric}"); + /// A buffered body reports end-of-stream with its final frame, so no `None` poll follows. + #[tokio::test] + async fn buffered_body_reports_real_status() { + let (status, _) = track_request(|| async { Body::from("hello") }, Client::ReadToEnd).await; + assert_eq!(status, 200); + } + + /// Hyper treats a trailers frame as terminal and never polls for the following `None`. + #[tokio::test] + async fn trailers_report_real_status() { + let handler = || async { Body::new(TrailersBody(false)) }; + let (status, _) = track_request(handler, Client::ReadToEnd).await; + assert_eq!(status, 200); + } + + /// Hyper stops polling once `content-length` bytes have been written, so completion is + /// derived from the byte count rather than a following end-of-stream poll. + #[tokio::test] + async fn content_length_stream_reports_real_status() { + let handler = || async { ([(header::CONTENT_LENGTH, "5")], stalling_stream()) }; + let (status, _) = track_request(handler, Client::ReadOneChunk).await; + assert_eq!(status, 200); + } + + #[tokio::test] + async fn interrupted_stream_reports_499() { + let (status, _) = track_request(|| async { stalling_stream() }, Client::Disconnect).await; + assert_eq!(status, 499); + } + + /// A client that disconnects after only part of `content-length` was written still reports + /// `499`: an incomplete byte count must not complete the request. + #[tokio::test] + async fn interrupted_content_length_stream_reports_499() { + let handler = || async { ([(header::CONTENT_LENGTH, "10")], stalling_stream()) }; + let (status, _) = track_request(handler, Client::ReadOneChunk).await; + assert_eq!(status, 499); + } + + #[tokio::test] + async fn errored_stream_reports_500() { + let (status, _) = track_request(|| async { failing_stream() }, Client::ReadToEnd).await; + assert_eq!(status, 500); + } + + #[tokio::test(start_paused = true)] + async fn duration_covers_body_streaming() { + let handler = || async { slow_stream(Duration::from_secs(5)) }; + let (_, duration) = track_request(handler, Client::ReadToEnd).await; + assert_eq!(duration, Duration::from_secs(5)); } } diff --git a/objectstore-server/src/web/middleware.rs b/objectstore-server/src/web/middleware.rs index 1d037612..82f93888 100644 --- a/objectstore-server/src/web/middleware.rs +++ b/objectstore-server/src/web/middleware.rs @@ -10,7 +10,7 @@ use axum::response::{IntoResponse, Response}; use objectstore_log::tracing; use tower_http::set_header::SetResponseHeaderLayer; -use crate::endpoints::is_internal_route; +use crate::endpoints; use crate::extractors::downstream_service::DownstreamService; use crate::web::RequestCounter; use crate::web::metrics_body::{EmitMetricsGuard, MetricsBody}; @@ -23,7 +23,7 @@ const SERVER: &str = concat!("objectstore/", env!("CARGO_PKG_VERSION")); /// maximum. /// /// Use with [`from_fn_with_state`](axum::middleware::from_fn_with_state), passing a -/// [`RequestCounter`]. Internal routes (see [`is_internal_route`]) are excluded. +/// [`RequestCounter`]. Internal routes are excluded. pub async fn limit_web_concurrency( State(counter): State, mut request: Request, @@ -32,7 +32,7 @@ pub async fn limit_web_concurrency( let matched_path = request.extract_parts::().await; let route = matched_path.as_ref().map_or("unknown", |m| m.as_str()); - if !is_internal_route(route) && counter.count() >= counter.limit() { + if !endpoints::is_internal_route(route) && counter.count() >= counter.limit() { let service = request.extract_parts::().await.unwrap(); objectstore_metrics::count!("web.concurrency.rejected", service = service.to_string()); objectstore_log::warn!("Request rejected: web concurrency limit reached"); @@ -103,13 +103,13 @@ pub async fn bind_sentry_body(request: Request, next: Next) -> Response { /// the response body (see [`MetricsBody`]) so it is dropped only once the body has finished /// streaming, not when the handler produces the response headers. /// -/// Internal routes (see [`is_internal_route`]) are excluded from metrics. +/// Internal routes are excluded from metrics. pub async fn emit_request_metrics(mut request: Request, next: Next) -> Response { let matched_path = request.extract_parts::().await; let route = matched_path.as_ref().map_or("unknown", |m| m.as_str()); let service = request.extract_parts::().await.unwrap(); - let should_emit = !is_internal_route(route); + let should_emit = !endpoints::is_internal_route(route); let guard = should_emit.then(|| EmitMetricsGuard::new(route, request.method(), service)); let response = next.run(request).await; @@ -117,10 +117,7 @@ pub async fn emit_request_metrics(mut request: Request, next: Next) -> Response // Move the guard into the response body so the duration metric is emitted only when the // body has finished streaming. The header status is applied on successful completion. match guard { - Some(guard) => { - let status = response.status(); - response.map(|body| Body::new(MetricsBody::new(guard, status, body))) - } + Some(guard) => MetricsBody::wrap_response(response, guard), None => response, } } @@ -214,63 +211,4 @@ mod tests { resume.notify_one(); blocking.await.unwrap().unwrap(); } - - /// The request-duration metric must be emitted only once the response body has finished - /// streaming, not when the handler produces the response headers. - /// - /// This runs on a current-thread runtime inside [`with_capturing_test_client`] so all - /// metric emissions happen on the capturing thread. A sentinel `test.marker` metric is - /// emitted after the response headers are received but before the body is consumed; the - /// duration metric must appear *after* that sentinel. - #[test] - fn duration_measured_end_to_end() { - use axum::body::{self, Bytes}; - use axum::middleware::from_fn; - - let captured = objectstore_metrics::with_capturing_test_client(|| { - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - - rt.block_on(async { - let app = Router::new() - .route( - "/v1/test/_/key", - get(|| async { - let stream = async_stream::stream! { - yield Ok::<_, std::io::Error>(Bytes::from_static(b"hello")); - }; - Body::from_stream(stream).into_response() - }), - ) - .layer(from_fn(emit_request_metrics)); - - let resp = app.oneshot(make_request("/v1/test/_/key")).await.unwrap(); - assert_eq!(resp.status(), StatusCode::OK); - - // Headers received. The duration metric must not have been emitted yet. - objectstore_metrics::count!("test.marker"); - - // Consuming the body drops the wrapping `MetricsBody` and its guard, which - // emits the duration metric. - let bytes = body::to_bytes(resp.into_body(), usize::MAX).await.unwrap(); - assert_eq!(&bytes[..], b"hello"); - }); - }); - - let marker = captured - .iter() - .position(|m| m.starts_with("test.marker:")) - .expect("sentinel marker not captured"); - let duration = captured - .iter() - .position(|m| m.starts_with("server.requests.duration:")) - .expect("duration metric not captured"); - - assert!( - duration > marker, - "duration metric emitted before body was consumed: {captured:?}" - ); - } }