fix: harden typed metrics review findings

crates/crank-observability/src/instrumentation.rs

@@ -10,7 +10,7 @@ use crank_metrics::{
 };
 use metrics::Unit;
 
-use crate::{MetricKind, MetricUnit, metric_schema};
+use crate::{MetricKind, MetricUnit};
 
 pub async fn record_http_request(request: Request, next: Next) -> Response {
     let route = request
@@ -24,11 +24,7 @@ pub async fn record_http_request(request: Request, next: Next) -> Response {
     let _inflight = InFlightGuard::http();
 
     let response = next.run(request).await;
-    let exemplar = response
-        .headers()
-        .get("x-trace-id")
-        .and_then(|value| value.to_str().ok())
-        .and_then(ExemplarTraceId::parse);
+    let exemplar = response.extensions().get::<ExemplarTraceId>().copied();
     crank_metrics::record_http_request_with_exemplar(
         route,
         method,
@@ -46,8 +42,8 @@ pub fn record_db_pool_connections(total: u32, idle: usize) {
     crank_metrics::set_db_pool_connections(DbPoolState::Used, total.saturating_sub(idle));
 }
 
-pub(crate) fn register_metric_schema() {
-    for definition in metric_schema() {
+pub(crate) fn register_metric_schema(service: crank_metrics::MetricService) {
+    for definition in crank_metrics::metric_schema_for_service(service) {
         let unit = match definition.unit {
             MetricUnit::Count => Unit::Count,
             MetricUnit::Seconds => Unit::Seconds,
@@ -65,7 +61,7 @@ pub(crate) fn register_metric_schema() {
         }
     }
 
-    crank_metrics::initialize_gauges();
+    crank_metrics::initialize_gauges(service);
 }
 
 #[cfg(test)]

crates/crank-observability/src/lifecycle.rs

@@ -47,8 +47,10 @@ impl ObservabilityLifecycle {
         build_subscriber_with_tracer(config, io::stdout, Some(tracer))?
             .try_init()
             .map_err(|_| ObservabilityInitError::SubscriberAlreadyInitialized)?;
+        let metrics_service = crank_metrics::MetricService::parse(identity.service())
+            .ok_or(MetricsSurfaceError::RecorderConfiguration)?;
         let metrics_handle = install_prometheus_recorder(&identity)?;
-        register_metric_schema();
+        register_metric_schema(metrics_service);
         let sentry_guard = init_sentry(&identity, redaction_limits, sentry_config)?;
 
         Ok(Self {

crates/crank-observability/src/prometheus.rs

@@ -1,4 +1,11 @@
-use std::net::SocketAddr;
+use std::{
+    collections::HashMap,
+    net::SocketAddr,
+    sync::{
+        Arc,
+        atomic::{AtomicUsize, Ordering},
+    },
+};
 
 use axum::{
     Router,
@@ -23,6 +30,8 @@ use crank_metrics::{ExemplarObservation, MAX_EXPOSITION_BYTES, MetricService, ex
 const PROMETHEUS_CONTENT_TYPE: &str = "text/plain; version=0.0.4; charset=utf-8";
 const OPENMETRICS_CONTENT_TYPE: &str = "application/openmetrics-text; version=1.0.0; charset=utf-8";
 const EXPOSITION_TOO_LARGE: &str = "metrics exposition exceeds configured bound\n";
+const TOO_MANY_SCRAPES: &str = "metrics scrape concurrency limit exceeded\n";
+const MAX_CONCURRENT_SCRAPES: usize = 2;
 
 #[derive(Clone)]
 pub struct MetricsConfig {
@@ -93,6 +102,7 @@ struct MetricsState {
     handle: PrometheusHandle,
     token_digest: Option<[u8; 32]>,
     requires_authentication: bool,
+    active_scrapes: Arc<AtomicUsize>,
 }
 
 pub struct MetricsSurface {
@@ -108,6 +118,7 @@ impl MetricsSurface {
                 handle,
                 token_digest: config.token_digest,
                 requires_authentication: config.requires_authentication(),
+                active_scrapes: Arc::new(AtomicUsize::new(0)),
             },
             config,
             _recorder: None,
@@ -206,20 +217,23 @@ fn prometheus_builder(
 }
 
 async fn render_metrics(State(state): State<MetricsState>, headers: HeaderMap) -> Response {
+    let Some(_permit) = ScrapePermit::try_acquire(&state.active_scrapes) else {
+        return (StatusCode::SERVICE_UNAVAILABLE, TOO_MANY_SCRAPES).into_response();
+    };
     let legacy = state.handle.render();
     let openmetrics = headers
         .get(header::ACCEPT)
         .and_then(|value| value.to_str().ok())
-        .is_some_and(|value| {
-            value
-                .split(',')
-                .any(|part| part.trim().starts_with("application/openmetrics-text"))
-        });
+        .is_some_and(prefers_openmetrics);
     let body = if openmetrics {
         render_openmetrics(&legacy, &exemplar_snapshot())
     } else {
         legacy
     };
+    exposition_response(body, openmetrics)
+}
+
+fn exposition_response(body: String, openmetrics: bool) -> Response {
     if body.len() > MAX_EXPOSITION_BYTES {
         return (StatusCode::SERVICE_UNAVAILABLE, EXPOSITION_TOO_LARGE).into_response();
     }
@@ -235,14 +249,15 @@ async fn render_metrics(State(state): State<MetricsState>, headers: HeaderMap) -
     response
 }
 
-fn render_openmetrics(legacy: &str, exemplars: &[ExemplarObservation]) -> String {
+fn render_openmetrics(legacy: &str, exemplars: &[Arc<ExemplarObservation>]) -> String {
+    let index: HashMap<String, &ExemplarObservation> = exemplars
+        .iter()
+        .map(|exemplar| (exemplar_key(exemplar), exemplar.as_ref()))
+        .collect();
     let mut output = String::with_capacity(legacy.len() + exemplars.len().saturating_mul(96) + 6);
     for line in legacy.lines() {
         output.push_str(line);
-        if let Some(exemplar) = exemplars
-            .iter()
-            .find(|candidate| line_matches_exemplar(line, candidate))
-        {
+        if let Some(exemplar) = sample_key(line).and_then(|key| index.get(&key).copied()) {
             output.push_str(" # {trace_id=\"");
             output.push_str(exemplar.trace_id.as_str());
             output.push_str("\"} ");
@@ -254,19 +269,79 @@ fn render_openmetrics(legacy: &str, exemplars: &[ExemplarObservation]) -> String
     output
 }
 
-fn line_matches_exemplar(line: &str, exemplar: &ExemplarObservation) -> bool {
-    if !line.starts_with(exemplar.metric) || !line[exemplar.metric.len()..].starts_with("_bucket{")
-    {
-        return false;
-    }
-    let expected_bound = exemplar
+fn exemplar_key(exemplar: &ExemplarObservation) -> String {
+    let mut labels = exemplar.labels.clone();
+    labels.sort_unstable();
+    let bound = exemplar
         .bucket_upper_bound
         .map_or_else(|| "+Inf".to_owned(), |value| value.to_string());
-    line.contains(&format!("le=\"{expected_bound}\""))
-        && exemplar
-            .labels
-            .iter()
-            .all(|(key, value)| line.contains(&format!("{key}=\"{value}\"")))
+    format!("{}|{:?}|{bound}", exemplar.metric, labels)
+}
+
+fn sample_key(line: &str) -> Option<String> {
+    let (head, _) = line.split_once("} ")?;
+    let (metric, raw_labels) = head.split_once("_bucket{")?;
+    let mut labels = Vec::new();
+    let mut bound = None;
+    for item in raw_labels.split(',') {
+        let (name, value) = item.split_once("=\"")?;
+        let value = value.strip_suffix('"')?;
+        match name {
+            "le" => bound = Some(value),
+            "service" | "version" | "environment" => {}
+            _ => labels.push((name, value)),
+        }
+    }
+    labels.sort_unstable();
+    Some(format!("{metric}|{labels:?}|{}", bound?))
+}
+
+fn prefers_openmetrics(value: &str) -> bool {
+    let mut open_q = 0.0_f32;
+    let mut legacy_q = 0.0_f32;
+    for range in value.split(',') {
+        let mut parts = range.trim().split(';');
+        let media = parts.next().unwrap_or_default().trim();
+        let mut q = 1.0_f32;
+        let mut supported_version = true;
+        for parameter in parts {
+            let Some((name, raw)) = parameter.trim().split_once('=') else {
+                continue;
+            };
+            match name.trim() {
+                "q" => q = raw.trim().parse().unwrap_or(0.0),
+                "version" if media == "application/openmetrics-text" => {
+                    supported_version = raw.trim().trim_matches('"') == "1.0.0";
+                }
+                _ => {}
+            }
+        }
+        if media == "application/openmetrics-text" && supported_version {
+            open_q = open_q.max(q.clamp(0.0, 1.0));
+        } else if media == "text/plain" || media == "*/*" {
+            legacy_q = legacy_q.max(q.clamp(0.0, 1.0));
+        }
+    }
+    open_q > 0.0 && open_q >= legacy_q
+}
+
+struct ScrapePermit(Arc<AtomicUsize>);
+
+impl ScrapePermit {
+    fn try_acquire(active: &Arc<AtomicUsize>) -> Option<Self> {
+        active
+            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| {
+                (count < MAX_CONCURRENT_SCRAPES).then_some(count + 1)
+            })
+            .ok()
+            .map(|_| Self(Arc::clone(active)))
+    }
+}
+
+impl Drop for ScrapePermit {
+    fn drop(&mut self) {
+        self.0.fetch_sub(1, Ordering::Release);
+    }
 }
 
 async fn metrics_health() -> impl IntoResponse {
@@ -313,9 +388,17 @@ fn token_digest(token: &[u8]) -> [u8; 32] {
 
 #[cfg(test)]
 mod rendering_tests {
+    use std::{
+        sync::{Arc, atomic::AtomicUsize},
+        time::{Duration, Instant},
+    };
+
     use crank_metrics::{ExemplarObservation, ExemplarTraceId};
 
-    use super::render_openmetrics;
+    use super::{
+        MAX_CONCURRENT_SCRAPES, ScrapePermit, exposition_response, prefers_openmetrics,
+        render_openmetrics,
+    };
 
     #[test]
     fn openmetrics_adds_bounded_exemplar_without_changing_aggregate() {
@@ -327,9 +410,66 @@ mod rendering_tests {
             value: 0.007,
             trace_id: ExemplarTraceId::parse("0123456789abcdef0123456789abcdef").unwrap(),
         };
-        let rendered = render_openmetrics(legacy, &[exemplar]);
+        let rendered = render_openmetrics(legacy, &[Arc::new(exemplar)]);
         assert!(rendered.contains("# {trace_id=\"0123456789abcdef0123456789abcdef\"} 0.007"));
         assert!(rendered.ends_with("# EOF\n"));
         assert_eq!(rendered.matches(" 1").count(), legacy.matches(" 1").count());
     }
+
+    #[test]
+    fn accept_negotiation_is_exact_and_honors_quality_and_version() {
+        assert!(prefers_openmetrics(
+            "application/openmetrics-text; version=1.0.0"
+        ));
+        assert!(prefers_openmetrics("application/openmetrics-text"));
+        assert!(!prefers_openmetrics("application/openmetrics-text; q=0"));
+        assert!(!prefers_openmetrics("application/openmetrics-textual"));
+        assert!(!prefers_openmetrics(
+            "application/openmetrics-text; version=0.0.1"
+        ));
+        assert!(!prefers_openmetrics(
+            "application/openmetrics-text;q=0.2,text/plain;q=0.8"
+        ));
+    }
+
+    #[test]
+    fn output_and_concurrent_scrapes_fail_closed_at_their_bounds() {
+        assert_eq!(
+            exposition_response("x".repeat(crank_metrics::MAX_EXPOSITION_BYTES + 1), false)
+                .status(),
+            axum::http::StatusCode::SERVICE_UNAVAILABLE
+        );
+        let active = Arc::new(AtomicUsize::new(0));
+        let permits = (0..MAX_CONCURRENT_SCRAPES)
+            .map(|_| ScrapePermit::try_acquire(&active).unwrap())
+            .collect::<Vec<_>>();
+        assert!(ScrapePermit::try_acquire(&active).is_none());
+        drop(permits);
+        assert!(ScrapePermit::try_acquire(&active).is_some());
+    }
+
+    #[test]
+    fn maximum_fixture_renders_in_linear_bounded_time() {
+        let trace = ExemplarTraceId::parse("0123456789abcdef0123456789abcdef").unwrap();
+        let mut legacy = String::new();
+        let mut exemplars = Vec::new();
+        for index in 0..5_000 {
+            let value: &'static str = Box::leak(format!("route-{index}").into_boxed_str());
+            legacy.push_str(&format!(
+                "crank_http_request_duration_seconds_bucket{{route=\"{value}\",method=\"GET\",le=\"0.01\"}} 1\n"
+            ));
+            exemplars.push(Arc::new(ExemplarObservation {
+                metric: "crank_http_request_duration_seconds",
+                labels: vec![("route", value), ("method", "GET")],
+                bucket_upper_bound: Some(0.01),
+                value: 0.007,
+                trace_id: trace,
+            }));
+        }
+        let started = Instant::now();
+        let rendered = render_openmetrics(&legacy, &exemplars);
+        assert_eq!(rendered.matches("# {trace_id=").count(), 5_000);
+        assert!(started.elapsed() < Duration::from_secs(3));
+        assert!(rendered.len() < crank_metrics::MAX_EXPOSITION_BYTES);
+    }
 }

crates/crank-observability/tests/http_metrics.rs

@@ -5,6 +5,7 @@ use axum::{
     body::{Body, to_bytes},
     http::{Request, StatusCode, header},
     middleware,
+    response::IntoResponse,
     routing::get,
 };
 use crank_observability::{
@@ -32,10 +33,12 @@ async fn http_metrics_use_matched_routes_and_closed_labels() {
         .route(
             "/api/admin/workspaces/{workspace_id}/operations/{operation_id}",
             get(|| async {
-                (
-                    StatusCode::NO_CONTENT,
-                    [("x-trace-id", "0123456789abcdef0123456789abcdef")],
-                )
+                let mut response = StatusCode::NO_CONTENT.into_response();
+                response.extensions_mut().insert(
+                    crank_metrics::ExemplarTraceId::parse("0123456789abcdef0123456789abcdef")
+                        .unwrap(),
+                );
+                response
             }),
         )
         .layer(middleware::from_fn(record_http_request));

crates/crank-observability/tests/prometheus.rs

@@ -28,6 +28,35 @@ fn loopback_is_allowed_without_a_token() {
     assert!(!config.requires_authentication());
 }
 
+#[test]
+fn disabled_config_remains_disabled_without_auth_side_effects() {
+    let config = MetricsConfig::new(
+        false,
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 9464),
+        None,
+    )
+    .expect("disabled listener does not require a token");
+    assert!(!config.enabled());
+}
+
+#[tokio::test]
+async fn bind_collision_returns_a_typed_safe_error() {
+    let listener = tokio::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0))
+        .await
+        .unwrap();
+    let address = listener.local_addr().unwrap();
+    let config = MetricsConfig::new(true, address, None).unwrap();
+    let error = match MetricsSurface::for_test(config, identity())
+        .unwrap()
+        .bind()
+        .await
+    {
+        Ok(_) => panic!("occupied port must fail closed"),
+        Err(error) => error,
+    };
+    assert_eq!(error.to_string(), "failed to bind metrics listener");
+}
+
 #[tokio::test]
 async fn default_service_ports_bind_and_serve_real_metrics_listeners() {
     for (service, port) in [("admin-api", 9464), ("mcp-server", 9465)] {