From 73e7a2ded55c98d572ae6dcaa13fab8fc841d090 Mon Sep 17 00:00:00 2001 From: Rob Lyon Date: Tue, 21 Jul 2026 08:54:44 -0700 Subject: [PATCH] refactor: share the metrics helpers, error type, and prefetch observation The agent kept its own copy of the telemetry registration helpers: agent_register /agent_counter/agent_gauge/agent_histogram had the same bodies as their telemetry counterparts, transfer_buckets was character-identical, saturating_i64 differed only in its integer type, and both encode methods were the same line. The telemetry helpers are now pub(crate) and the copies are gone, with full metric names passed at the call site instead of formatting a flywheel_agent_ prefix nineteen times at startup. AgentMetrics::new is a struct literal, flywheel_requests_total goes through int_counter rather than being hand-rolled, and a new int_counter_vec helper covers the two IntCounterVecs. ChannelStoreError differed from MetadataError only by IncompatibleStore versus AlreadyExists, which forced four duplicated helper pairs in rocksdb.rs. MetadataError gains AlreadyExists and ChannelStoreError is deleted along with its re-export and the four channel_* helpers. The agent and shard prefetch observations are now one PrefetchObservation parameterized by its recorder. complete takes &mut self, which removes the Option from the unfold state and both expects from the per-chunk path. --- src/agent/mod.rs | 415 +++++++++++--------------------- src/cache/service.rs | 4 +- src/channel/mod.rs | 2 +- src/channel/registry.rs | 12 - src/channel/service.rs | 8 +- src/storage/metadata/mod.rs | 2 + src/storage/metadata/rocksdb.rs | 99 +++----- src/telemetry.rs | 210 +++++++++++++--- src/transport/http/mod.rs | 105 +------- 9 files changed, 345 insertions(+), 512 deletions(-) diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 6e85545..6a11b6b 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -22,7 +22,10 @@ use crate::{ REQUEST_PREFETCH_CONCURRENCY_HEADER, REQUEST_PURPOSE_HEADER, REQUEST_PURPOSE_PREFETCH, REQUEST_SESSION_HEADER, }, - telemetry::{MakeRequestUlid, REQUEST_ID_HEADER}, + telemetry::{ + MakeRequestUlid, PrefetchObservation, PrefetchRecorder, REQUEST_ID_HEADER, encode_registry, + histogram, int_counter, int_gauge, observe_prefetch_body, saturating_i64, transfer_buckets, + }, }; use axum::{ Json, Router, @@ -33,11 +36,7 @@ use axum::{ routing::get, }; use discovery::{Resolver, RingState, RingStatus}; -use futures_util::StreamExt; -use prometheus::{ - Encoder, Histogram, HistogramOpts, IntCounter, IntGauge, Opts, Registry, TextEncoder, - core::Collector, -}; +use prometheus::{Histogram, IntCounter, IntGauge, Registry}; use ring::RingMember; use serde::Serialize; use std::{ @@ -83,7 +82,7 @@ struct AgentState { ring: RingState, resolver: Arc, client: reqwest::Client, - metrics: AgentMetrics, + metrics: Arc, } impl Agent { @@ -114,7 +113,7 @@ impl Agent { ring, resolver, client, - metrics: AgentMetrics::new()?, + metrics: Arc::new(AgentMetrics::new()?), }), }) } @@ -347,7 +346,12 @@ async fn forward(State(state): State>, request: Request) -> Resp .get(REQUEST_PREFETCH_CONCURRENCY_HEADER) .and_then(|value| value.to_str().ok()) .and_then(|value| value.parse::().ok()); - AgentPrefetchObservation::start(Arc::clone(&state), session, configured_concurrency) + PrefetchObservation::start( + Arc::clone(&state.metrics), + None, + session, + configured_concurrency, + ) }); let started = Instant::now(); let response = forward_request(&state, request, prefetch).await; @@ -426,116 +430,6 @@ async fn forward_request(state: &Arc, request: Request, prefetch: bo } } -struct AgentPrefetchObservation { - state: Arc, - started: Instant, - session: Option, - configured_concurrency: Option, - status: u16, - bytes: u64, - completed: bool, -} - -impl AgentPrefetchObservation { - fn start( - state: Arc, - session: Option, - configured_concurrency: Option, - ) -> Self { - state.metrics.prefetch_requests.inc(); - state.metrics.prefetch_in_flight.inc(); - let in_flight = state.metrics.prefetch_in_flight.get(); - tracing::debug!( - session_id = session.as_deref().unwrap_or(""), - configured_concurrency, - in_flight, - "agent prefetch request started" - ); - Self { - state, - started: Instant::now(), - session, - configured_concurrency, - status: 0, - bytes: 0, - completed: false, - } - } - - fn set_status(&mut self, status: u16) { - self.status = status; - } - - fn add_bytes(&mut self, bytes: usize) { - self.bytes = self.bytes.saturating_add(bytes as u64); - } - - fn complete(mut self) { - self.completed = true; - } -} - -impl Drop for AgentPrefetchObservation { - fn drop(&mut self) { - let duration = self.started.elapsed(); - self.state.metrics.prefetch_in_flight.dec(); - self.state - .metrics - .prefetch_response_bytes - .inc_by(self.bytes); - self.state - .metrics - .prefetch_transfer_duration - .observe(duration.as_secs_f64()); - if self.completed { - self.state.metrics.prefetch_completed.inc(); - } else { - self.state.metrics.prefetch_cancelled.inc(); - } - tracing::debug!( - session_id = self.session.as_deref().unwrap_or(""), - configured_concurrency = self.configured_concurrency, - status = self.status, - bytes = self.bytes, - duration_ms = duration.as_millis() as u64, - completed = self.completed, - "agent prefetch request finished" - ); - } -} - -fn observe_prefetch_body( - response: Response, - mut observation: AgentPrefetchObservation, -) -> Response { - observation.set_status(response.status().as_u16()); - let (parts, body) = response.into_parts(); - let stream = Box::pin(body.into_data_stream()); - let observed = futures_util::stream::unfold( - (stream, Some(observation)), - |(mut stream, mut observation)| async move { - match stream.next().await { - Some(Ok(bytes)) => { - observation - .as_mut() - .expect("prefetch observation exists while streaming") - .add_bytes(bytes.len()); - Some((Ok(bytes), (stream, observation))) - } - Some(Err(error)) => Some((Err(error), (stream, observation))), - None => { - observation - .take() - .expect("prefetch observation exists at completion") - .complete(); - None - } - } - }, - ); - Response::from_parts(parts, Body::from_stream(observed)) -} - /// Counts a failure toward ejection only when it is attributable to the backend /// (connect refused or timed out). A client aborting its own upload mid-stream /// also fails `send()`, and that must not eject a healthy shard — a genuinely @@ -677,171 +571,138 @@ struct AgentMetrics { impl AgentMetrics { fn new() -> prometheus::Result { let registry = Registry::new(); - let requests = agent_counter( - ®istry, - "requests_total", - "Requests received by the routing agent.", - )?; - let forwarded = agent_counter( - ®istry, - "forwarded_total", - "Requests forwarded to a shard.", - )?; - let forward_duration = agent_histogram( - ®istry, - "forward_duration_seconds", - "Time from agent request receipt until response headers are ready.", - request_buckets(), - )?; - let forward_failures = agent_counter( - ®istry, - "forward_failures_total", - "Shard connection failures and send timeouts.", - )?; - let synthesized_misses = agent_counter( - ®istry, - "synthesized_misses_total", - "Build-cache misses synthesized without a shard response.", - )?; - let synthesized_bypasses = agent_counter( - ®istry, - "synthesized_write_bypasses_total", - "Build-cache write successes synthesized without storing the body.", - )?; - let unavailable = agent_counter( - ®istry, - "unavailable_total", - "Non-build-cache requests that could not reach an owner.", - )?; - let ejections = agent_counter( - ®istry, - "ejections_total", - "Shard ejections from the routing ring.", - )?; - let prefetch_requests = agent_counter( - ®istry, - "prefetch_requests_total", - "Prefetch requests received by the routing agent.", - )?; - let prefetch_hits = agent_counter( - ®istry, - "prefetch_hits_total", - "Prefetch requests answered successfully by a shard.", - )?; - let prefetch_misses = agent_counter( - ®istry, - "prefetch_misses_total", - "Prefetch requests answered with a cache miss.", - )?; - let prefetch_unavailable = agent_counter( - ®istry, - "prefetch_unavailable_total", - "Prefetch requests unavailable because no shard produced a cache response.", - )?; - let prefetch_response_bytes = agent_counter( - ®istry, - "prefetch_response_bytes_total", - "Prefetch response bytes actually streamed through the agent.", - )?; - let prefetch_in_flight = agent_gauge( - ®istry, - "prefetch_in_flight", - "Prefetch response bodies currently streaming through the agent.", - )?; - let prefetch_completed = agent_counter( - ®istry, - "prefetch_completed_total", - "Prefetch response bodies streamed to completion.", - )?; - let prefetch_cancelled = agent_counter( - ®istry, - "prefetch_cancelled_total", - "Prefetch response bodies dropped before completion.", - )?; - let prefetch_transfer_duration = agent_histogram( - ®istry, - "prefetch_transfer_duration_seconds", - "Prefetch response body lifetime through the agent.", - transfer_buckets(), - )?; - let ring_members = agent_gauge( - ®istry, - "ring_members", - "Members in the current routing ring.", - )?; - let ring_ejected = agent_gauge( - ®istry, - "ring_ejected", - "Members currently ejected from the routing ring.", - )?; - Ok(Self { + requests: int_counter( + ®istry, + "flywheel_agent_requests_total", + "Requests received by the routing agent.", + )?, + forwarded: int_counter( + ®istry, + "flywheel_agent_forwarded_total", + "Requests forwarded to a shard.", + )?, + forward_duration: histogram( + ®istry, + "flywheel_agent_forward_duration_seconds", + "Time from agent request receipt until response headers are ready.", + request_buckets(), + )?, + forward_failures: int_counter( + ®istry, + "flywheel_agent_forward_failures_total", + "Shard connection failures and send timeouts.", + )?, + synthesized_misses: int_counter( + ®istry, + "flywheel_agent_synthesized_misses_total", + "Build-cache misses synthesized without a shard response.", + )?, + synthesized_bypasses: int_counter( + ®istry, + "flywheel_agent_synthesized_write_bypasses_total", + "Build-cache write successes synthesized without storing the body.", + )?, + unavailable: int_counter( + ®istry, + "flywheel_agent_unavailable_total", + "Non-build-cache requests that could not reach an owner.", + )?, + ejections: int_counter( + ®istry, + "flywheel_agent_ejections_total", + "Shard ejections from the routing ring.", + )?, + prefetch_requests: int_counter( + ®istry, + "flywheel_agent_prefetch_requests_total", + "Prefetch requests received by the routing agent.", + )?, + prefetch_hits: int_counter( + ®istry, + "flywheel_agent_prefetch_hits_total", + "Prefetch requests answered successfully by a shard.", + )?, + prefetch_misses: int_counter( + ®istry, + "flywheel_agent_prefetch_misses_total", + "Prefetch requests answered with a cache miss.", + )?, + prefetch_unavailable: int_counter( + ®istry, + "flywheel_agent_prefetch_unavailable_total", + "Prefetch requests unavailable because no shard produced a cache response.", + )?, + prefetch_response_bytes: int_counter( + ®istry, + "flywheel_agent_prefetch_response_bytes_total", + "Prefetch response bytes actually streamed through the agent.", + )?, + prefetch_in_flight: int_gauge( + ®istry, + "flywheel_agent_prefetch_in_flight", + "Prefetch response bodies currently streaming through the agent.", + )?, + prefetch_completed: int_counter( + ®istry, + "flywheel_agent_prefetch_completed_total", + "Prefetch response bodies streamed to completion.", + )?, + prefetch_cancelled: int_counter( + ®istry, + "flywheel_agent_prefetch_cancelled_total", + "Prefetch response bodies dropped before completion.", + )?, + prefetch_transfer_duration: histogram( + ®istry, + "flywheel_agent_prefetch_transfer_duration_seconds", + "Prefetch response body lifetime through the agent.", + transfer_buckets(), + )?, + ring_members: int_gauge( + ®istry, + "flywheel_agent_ring_members", + "Members in the current routing ring.", + )?, + ring_ejected: int_gauge( + ®istry, + "flywheel_agent_ring_ejected", + "Members currently ejected from the routing ring.", + )?, registry, - requests, - forwarded, - forward_duration, - forward_failures, - synthesized_misses, - synthesized_bypasses, - unavailable, - ejections, - prefetch_requests, - prefetch_hits, - prefetch_misses, - prefetch_unavailable, - prefetch_response_bytes, - prefetch_in_flight, - prefetch_completed, - prefetch_cancelled, - prefetch_transfer_duration, - ring_members, - ring_ejected, }) } fn encode(&self, members: usize, ejected: usize) -> prometheus::Result> { - self.ring_members.set(saturating_i64(members)); - self.ring_ejected.set(saturating_i64(ejected)); - let mut body = Vec::new(); - TextEncoder::new().encode(&self.registry.gather(), &mut body)?; - Ok(body) + self.ring_members.set(saturating_i64(members as u64)); + self.ring_ejected.set(saturating_i64(ejected as u64)); + encode_registry(&self.registry) } } -fn agent_register(registry: &Registry, metric: T) -> prometheus::Result -where - T: Collector + Clone + 'static, -{ - registry.register(Box::new(metric.clone()))?; - Ok(metric) -} - -fn agent_counter(registry: &Registry, suffix: &str, help: &str) -> prometheus::Result { - agent_register( - registry, - IntCounter::with_opts(Opts::new(format!("flywheel_agent_{suffix}"), help))?, - ) -} +impl PrefetchRecorder for AgentMetrics { + fn prefetch_started(&self) -> i64 { + self.prefetch_requests.inc(); + self.prefetch_in_flight.inc(); + self.prefetch_in_flight.get() + } -fn agent_gauge(registry: &Registry, suffix: &str, help: &str) -> prometheus::Result { - agent_register( - registry, - IntGauge::with_opts(Opts::new(format!("flywheel_agent_{suffix}"), help))?, - ) -} + /// The agent classifies a prefetch where it learns the outcome — a forwarded + /// response, an empty ring, or a send failure — so there is nothing left to + /// record once the body is on its way back. + fn prefetch_response(&self, _status: u16) {} -fn agent_histogram( - registry: &Registry, - suffix: &str, - help: &str, - buckets: Vec, -) -> prometheus::Result { - agent_register( - registry, - Histogram::with_opts( - HistogramOpts::new(format!("flywheel_agent_{suffix}"), help).buckets(buckets), - )?, - ) + fn prefetch_finished(&self, duration: Duration, bytes: u64, completed: bool) { + self.prefetch_in_flight.dec(); + self.prefetch_response_bytes.inc_by(bytes); + self.prefetch_transfer_duration + .observe(duration.as_secs_f64()); + if completed { + self.prefetch_completed.inc(); + } else { + self.prefetch_cancelled.inc(); + } + } } fn request_buckets() -> Vec { @@ -849,11 +710,3 @@ fn request_buckets() -> Vec { 0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, ] } - -fn transfer_buckets() -> Vec { - vec![0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0] -} - -fn saturating_i64(value: usize) -> i64 { - i64::try_from(value).unwrap_or(i64::MAX) -} diff --git a/src/cache/service.rs b/src/cache/service.rs index ce3e2cc..c7e4176 100644 --- a/src/cache/service.rs +++ b/src/cache/service.rs @@ -4,7 +4,7 @@ use crate::{ cache::recent_use::RecentUse, cache::space::SpaceLedger, cache::stripes::Stripes, - channel::{ChannelGates, ChannelId, ChannelStoreError, Lifecycle}, + channel::{ChannelGates, ChannelId, Lifecycle}, clock::Clock, storage::{ local::{ArtifactFiles, FilePublication, LocalError, StageOutcome, StagedArtifact}, @@ -534,8 +534,6 @@ pub enum CacheError { Local(#[from] LocalError), #[error(transparent)] Metadata(#[from] MetadataError), - #[error(transparent)] - ChannelStore(#[from] ChannelStoreError), #[error("channel does not exist")] MissingChannel, #[error("channel is being deleted")] diff --git a/src/channel/mod.rs b/src/channel/mod.rs index 18283ce..47850ad 100644 --- a/src/channel/mod.rs +++ b/src/channel/mod.rs @@ -7,6 +7,6 @@ mod token; pub use crate::storage::records::RecordError; pub use identity::{ChannelId, ChannelIdError}; pub use policy::{Access, Lifecycle}; -pub use registry::{ChannelRecord, ChannelStoreError}; +pub use registry::ChannelRecord; pub use service::{ChannelError, ChannelGates, ChannelLease, ChannelService, IssuedChannel}; pub use token::{ChannelToken, TokenDigest}; diff --git a/src/channel/registry.rs b/src/channel/registry.rs index 8f02479..6f6ce91 100644 --- a/src/channel/registry.rs +++ b/src/channel/registry.rs @@ -9,15 +9,3 @@ pub struct ChannelRecord { pub state: Lifecycle, pub created_at: u64, } - -#[derive(Debug, thiserror::Error)] -pub enum ChannelStoreError { - #[error("channel registry failed: {0}")] - Store(String), - #[error("channel already exists")] - AlreadyExists, - #[error("durable record failed: {0}")] - Record(#[from] crate::storage::records::RecordError), - #[error("channel registry task failed: {0}")] - Task(#[from] tokio::task::JoinError), -} diff --git a/src/channel/service.rs b/src/channel/service.rs index 71da6fc..1f80d71 100644 --- a/src/channel/service.rs +++ b/src/channel/service.rs @@ -1,7 +1,7 @@ -use super::{Access, ChannelId, ChannelRecord, ChannelStoreError, ChannelToken, Lifecycle}; +use super::{Access, ChannelId, ChannelRecord, ChannelToken, Lifecycle}; use crate::storage::{ local::{ArtifactFiles, LocalError}, - metadata::RocksMetadata, + metadata::{MetadataError, RocksMetadata}, }; use dashmap::DashMap; use std::{ @@ -80,7 +80,7 @@ impl ChannelService { }; match self.store.create_channel(record.clone()).await { Ok(()) => record, - Err(ChannelStoreError::AlreadyExists) => self + Err(MetadataError::AlreadyExists) => self .store .channel(ChannelId::DEFAULT) .await? @@ -238,7 +238,7 @@ pub enum ChannelError { #[error("the persisted default channel violates its invariants")] InvalidDefault, #[error(transparent)] - Store(#[from] ChannelStoreError), + Store(#[from] MetadataError), #[error(transparent)] Local(#[from] LocalError), } diff --git a/src/storage/metadata/mod.rs b/src/storage/metadata/mod.rs index 79daf19..d7f053a 100644 --- a/src/storage/metadata/mod.rs +++ b/src/storage/metadata/mod.rs @@ -48,6 +48,8 @@ pub enum MetadataError { Store(String), #[error("durable store format is incompatible: {0}")] IncompatibleStore(String), + #[error("channel already exists")] + AlreadyExists, #[error("durable record failed: {0}")] Record(#[from] crate::storage::records::RecordError), #[error("metadata task failed: {0}")] diff --git a/src/storage/metadata/rocksdb.rs b/src/storage/metadata/rocksdb.rs index b5b57f1..bc39b1f 100644 --- a/src/storage/metadata/rocksdb.rs +++ b/src/storage/metadata/rocksdb.rs @@ -1,7 +1,7 @@ use super::{Candidate, Durability, Evicted, MetadataError, ReferenceRecord}; use crate::{ artifact::{ArtifactId, ArtifactMetadata, Digest}, - channel::{ChannelId, ChannelRecord, ChannelStoreError}, + channel::{ChannelId, ChannelRecord}, storage::records::{decode_record, encode_record}, }; use rocksdb::{ColumnFamilyDescriptor, DB, IteratorMode, Options, WriteBatch, WriteOptions}; @@ -326,26 +326,21 @@ impl RocksMetadata { } impl RocksMetadata { - pub(crate) async fn create_channel( - &self, - channel: ChannelRecord, - ) -> Result<(), ChannelStoreError> { + pub(crate) async fn create_channel(&self, channel: ChannelRecord) -> Result<(), MetadataError> { let database = Arc::clone(&self.database); - channel_blocking(move || { - let family = database - .cf_handle(CHANNELS) - .ok_or_else(channel_missing_cf)?; + blocking(move || { + let family = database.cf_handle(CHANNELS).ok_or_else(missing_cf)?; let key = channel.id.as_key(); if database .get_cf(&family, key) - .map_err(channel_store_error)? + .map_err(store_error)? .is_some() { - return Err(ChannelStoreError::AlreadyExists); + return Err(MetadataError::AlreadyExists); } let mut batch = WriteBatch::default(); batch.put_cf(&family, key, encode_record(&channel)?); - channel_write_sync(&database, batch) + write_sync(&database, batch) }) .await } @@ -353,81 +348,67 @@ impl RocksMetadata { pub(crate) async fn channel( &self, id: ChannelId, - ) -> Result, ChannelStoreError> { + ) -> Result, MetadataError> { let database = Arc::clone(&self.database); - channel_blocking(move || { - let family = database - .cf_handle(CHANNELS) - .ok_or_else(channel_missing_cf)?; + blocking(move || { + let family = database.cf_handle(CHANNELS).ok_or_else(missing_cf)?; database .get_cf(&family, id.as_key()) - .map_err(channel_store_error)? - .map(|bytes| decode_record(&bytes).map_err(ChannelStoreError::from)) + .map_err(store_error)? + .map(|bytes| decode_record(&bytes).map_err(MetadataError::from)) .transpose() }) .await } - pub(crate) async fn store_channel( - &self, - channel: ChannelRecord, - ) -> Result<(), ChannelStoreError> { + pub(crate) async fn store_channel(&self, channel: ChannelRecord) -> Result<(), MetadataError> { let database = Arc::clone(&self.database); - channel_blocking(move || { - let family = database - .cf_handle(CHANNELS) - .ok_or_else(channel_missing_cf)?; + blocking(move || { + let family = database.cf_handle(CHANNELS).ok_or_else(missing_cf)?; let mut batch = WriteBatch::default(); batch.put_cf(&family, channel.id.as_key(), encode_record(&channel)?); - channel_write_sync(&database, batch) + write_sync(&database, batch) }) .await } - pub(crate) async fn channels(&self) -> Result, ChannelStoreError> { + pub(crate) async fn channels(&self) -> Result, MetadataError> { let database = Arc::clone(&self.database); - channel_blocking(move || { - let family = database - .cf_handle(CHANNELS) - .ok_or_else(channel_missing_cf)?; + blocking(move || { + let family = database.cf_handle(CHANNELS).ok_or_else(missing_cf)?; database .iterator_cf(&family, IteratorMode::Start) .map(|item| { - let (_, bytes) = item.map_err(channel_store_error)?; - decode_record::(&bytes).map_err(ChannelStoreError::from) + let (_, bytes) = item.map_err(store_error)?; + decode_record::(&bytes).map_err(MetadataError::from) }) .collect() }) .await } - pub(crate) async fn delete_channel_data(&self, id: ChannelId) -> Result<(), ChannelStoreError> { + pub(crate) async fn delete_channel_data(&self, id: ChannelId) -> Result<(), MetadataError> { let database = Arc::clone(&self.database); - channel_blocking(move || { + blocking(move || { let prefix = id.as_key(); let end = prefix_end(prefix); let mut batch = WriteBatch::default(); for name in [ARTIFACTS, REFERENCES, EVICTION] { - let family = database.cf_handle(name).ok_or_else(channel_missing_cf)?; + let family = database.cf_handle(name).ok_or_else(missing_cf)?; batch.delete_range_cf(&family, prefix, end); } - channel_write_sync(&database, batch) + write_sync(&database, batch) }) .await } - pub(crate) async fn finish_channel_deletion( - &self, - id: ChannelId, - ) -> Result<(), ChannelStoreError> { + pub(crate) async fn finish_channel_deletion(&self, id: ChannelId) -> Result<(), MetadataError> { let database = Arc::clone(&self.database); - channel_blocking(move || { - let channels = database - .cf_handle(CHANNELS) - .ok_or_else(channel_missing_cf)?; + blocking(move || { + let channels = database.cf_handle(CHANNELS).ok_or_else(missing_cf)?; let mut batch = WriteBatch::default(); batch.delete_cf(&channels, id.as_key()); - channel_write_sync(&database, batch) + write_sync(&database, batch) }) .await } @@ -507,28 +488,6 @@ fn store_error(error: impl std::fmt::Display) -> MetadataError { MetadataError::Store(error.to_string()) } -async fn channel_blocking( - operation: impl FnOnce() -> Result + Send + 'static, -) -> Result { - tokio::task::spawn_blocking(operation).await? -} - -fn channel_write_sync(database: &DB, batch: WriteBatch) -> Result<(), ChannelStoreError> { - let mut options = WriteOptions::default(); - options.set_sync(true); - database - .write_opt(batch, &options) - .map_err(channel_store_error) -} - -fn channel_missing_cf() -> ChannelStoreError { - ChannelStoreError::Store("missing RocksDB column family".to_owned()) -} - -fn channel_store_error(error: impl std::fmt::Display) -> ChannelStoreError { - ChannelStoreError::Store(error.to_string()) -} - /// The exclusive upper bound of a channel's key range. A channel key is canonical /// Crockford base32 text, so its last byte is an ASCII digit or capital letter and /// incrementing it can never carry. diff --git a/src/telemetry.rs b/src/telemetry.rs index 7dd71bb..5232163 100644 --- a/src/telemetry.rs +++ b/src/telemetry.rs @@ -1,10 +1,15 @@ use crate::cache::SpaceSnapshot; +use axum::{body::Body, response::Response}; +use futures_util::StreamExt; use http::{HeaderName, HeaderValue, Request, StatusCode}; use prometheus::{ Encoder, Histogram, HistogramOpts, HistogramVec, IntCounter, IntCounterVec, IntGauge, Opts, Registry, TextEncoder, core::Collector, }; -use std::time::Duration; +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; use tower_http::request_id::{MakeRequestId, RequestId}; pub(crate) const REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id"); @@ -51,12 +56,10 @@ pub struct Metrics { impl Metrics { pub fn new(configured_foreground_limit: usize) -> prometheus::Result { let registry = Registry::new(); - let requests = register( + let requests = int_counter( ®istry, - IntCounter::with_opts(Opts::new( - "flywheel_requests_total", - "HTTP requests received by this replica.", - ))?, + "flywheel_requests_total", + "HTTP requests received by this replica.", )?; let http_request_duration = register( ®istry, @@ -134,30 +137,22 @@ impl Metrics { "flywheel_prefetch_requests_total", "Prefetch object requests received by this shard.", )?; - let prefetch_responses = register( + let prefetch_responses = int_counter_vec( ®istry, - IntCounterVec::new( - Opts::new( - "flywheel_prefetch_responses_total", - "Prefetch responses classified at response headers.", - ), - &["outcome"], - )?, + "flywheel_prefetch_responses_total", + "Prefetch responses classified at response headers.", + &["outcome"], )?; let prefetch_in_flight = int_gauge( ®istry, "flywheel_prefetch_in_flight", "Prefetch response bodies currently streaming from this shard.", )?; - let prefetch_transfers = register( + let prefetch_transfers = int_counter_vec( ®istry, - IntCounterVec::new( - Opts::new( - "flywheel_prefetch_transfers_total", - "Prefetch response bodies by terminal outcome.", - ), - &["outcome"], - )?, + "flywheel_prefetch_transfers_total", + "Prefetch response bodies by terminal outcome.", + &["outcome"], )?; let prefetch_response_bytes = int_counter( ®istry, @@ -274,13 +269,36 @@ impl Metrics { self.foreground_rejections.inc(); } - pub(crate) fn prefetch_started(&self) -> i64 { + pub fn record_space(&self, space: SpaceSnapshot) { + self.free_observed_bytes + .set(saturating_i64(space.free_observed)); + self.reserved_bytes.set(saturating_i64(space.reserved)); + self.committed_since_bytes + .set(saturating_i64(space.committed_since)); + } + + pub fn encode(&self) -> prometheus::Result> { + encode_registry(&self.registry) + } +} + +/// How a prefetch body is counted. The shard and the routing agent keep separate +/// registries, so an observation records against whichever handle started it. +pub(crate) trait PrefetchRecorder { + /// Counts a started prefetch and returns the resulting in-flight count. + fn prefetch_started(&self) -> i64; + fn prefetch_response(&self, status: u16); + fn prefetch_finished(&self, duration: Duration, bytes: u64, completed: bool); +} + +impl PrefetchRecorder for Metrics { + fn prefetch_started(&self) -> i64 { self.prefetch_requests.inc(); self.prefetch_in_flight.inc(); self.prefetch_in_flight.get() } - pub(crate) fn prefetch_response(&self, status: u16) { + fn prefetch_response(&self, status: u16) { let outcome = if (200..300).contains(&status) { "hit" } else if status == 404 { @@ -291,7 +309,7 @@ impl Metrics { self.prefetch_responses.with_label_values(&[outcome]).inc(); } - pub(crate) fn prefetch_finished(&self, duration: Duration, bytes: u64, completed: bool) { + fn prefetch_finished(&self, duration: Duration, bytes: u64, completed: bool) { self.prefetch_in_flight.dec(); self.prefetch_response_bytes.inc_by(bytes); self.prefetch_transfer_duration @@ -300,22 +318,117 @@ impl Metrics { .with_label_values(&[if completed { "completed" } else { "cancelled" }]) .inc(); } +} - pub fn record_space(&self, space: SpaceSnapshot) { - self.free_observed_bytes - .set(saturating_i64(space.free_observed)); - self.reserved_bytes.set(saturating_i64(space.reserved)); - self.committed_since_bytes - .set(saturating_i64(space.committed_since)); +/// One prefetch response body, from its headers until the body is dropped. The +/// terminal counters are recorded in `Drop` so a client that abandons the transfer +/// is counted exactly like one that finishes it. +pub(crate) struct PrefetchObservation { + recorder: Arc, + started: Instant, + /// The matched route, where the caller has one; the agent forwards by ring + /// position and never matches a route. + route: Option, + session: Option, + configured_concurrency: Option, + status: u16, + bytes: u64, + completed: bool, +} + +impl PrefetchObservation { + pub(crate) fn start( + recorder: Arc, + route: Option, + session: Option, + configured_concurrency: Option, + ) -> Self { + let in_flight = recorder.prefetch_started(); + tracing::debug!( + route = route.as_deref().unwrap_or(""), + session_id = session.as_deref().unwrap_or(""), + configured_concurrency, + in_flight, + "prefetch request started" + ); + Self { + recorder, + started: Instant::now(), + route, + session, + configured_concurrency, + status: 0, + bytes: 0, + completed: false, + } } - pub fn encode(&self) -> prometheus::Result> { - let mut body = Vec::new(); - TextEncoder::new().encode(&self.registry.gather(), &mut body)?; - Ok(body) + fn set_status(&mut self, status: u16) { + self.status = status; + self.recorder.prefetch_response(status); + } + + fn add_bytes(&mut self, bytes: usize) { + self.bytes = self.bytes.saturating_add(bytes as u64); } } +impl Drop for PrefetchObservation { + fn drop(&mut self) { + let duration = self.started.elapsed(); + self.recorder + .prefetch_finished(duration, self.bytes, self.completed); + tracing::debug!( + route = self.route.as_deref().unwrap_or(""), + session_id = self.session.as_deref().unwrap_or(""), + configured_concurrency = self.configured_concurrency, + status = self.status, + bytes = self.bytes, + duration_ms = duration.as_millis() as u64, + completed = self.completed, + "prefetch request finished" + ); + } +} + +/// Wraps a prefetch response body so the observation follows the bytes the client +/// actually receives. The observation lives in the stream state and is dropped with +/// it, whether the stream ends or the client goes away. +pub(crate) fn observe_prefetch_body( + response: Response, + mut observation: PrefetchObservation, +) -> Response +where + R: PrefetchRecorder + Send + Sync + 'static, +{ + observation.set_status(response.status().as_u16()); + let (parts, body) = response.into_parts(); + let stream = Box::pin(body.into_data_stream()); + let observed = futures_util::stream::unfold( + (stream, observation), + |(mut stream, mut observation)| async move { + match stream.next().await { + Some(Ok(bytes)) => { + observation.add_bytes(bytes.len()); + Some((Ok(bytes), (stream, observation))) + } + Some(Err(error)) => Some((Err(error), (stream, observation))), + None => { + observation.completed = true; + None + } + } + }, + ); + Response::from_parts(parts, Body::from_stream(observed)) +} + +pub(crate) fn encode_registry(registry: &Registry) -> prometheus::Result> { + let mut body = Vec::new(); + TextEncoder::new().encode(®istry.gather(), &mut body)?; + Ok(body) +} + fn register(registry: &Registry, metric: T) -> prometheus::Result where T: Collector + Clone + 'static, @@ -324,15 +437,32 @@ where Ok(metric) } -fn int_counter(registry: &Registry, name: &str, help: &str) -> prometheus::Result { +pub(crate) fn int_counter( + registry: &Registry, + name: &str, + help: &str, +) -> prometheus::Result { register(registry, IntCounter::with_opts(Opts::new(name, help))?) } -fn int_gauge(registry: &Registry, name: &str, help: &str) -> prometheus::Result { +fn int_counter_vec( + registry: &Registry, + name: &str, + help: &str, + labels: &[&str], +) -> prometheus::Result { + register(registry, IntCounterVec::new(Opts::new(name, help), labels)?) +} + +pub(crate) fn int_gauge( + registry: &Registry, + name: &str, + help: &str, +) -> prometheus::Result { register(registry, IntGauge::with_opts(Opts::new(name, help))?) } -fn histogram( +pub(crate) fn histogram( registry: &Registry, name: &str, help: &str, @@ -344,10 +474,10 @@ fn histogram( ) } -fn transfer_buckets() -> Vec { +pub(crate) fn transfer_buckets() -> Vec { vec![0.01, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0] } -fn saturating_i64(value: u64) -> i64 { +pub(crate) fn saturating_i64(value: u64) -> i64 { i64::try_from(value).unwrap_or(i64::MAX) } diff --git a/src/transport/http/mod.rs b/src/transport/http/mod.rs index 087f4c6..47e8a57 100644 --- a/src/transport/http/mod.rs +++ b/src/transport/http/mod.rs @@ -12,7 +12,9 @@ use crate::{ }, proxy::ProxyService, reference::Reference, - telemetry::{MakeRequestUlid, Metrics, REQUEST_ID_HEADER}, + telemetry::{ + MakeRequestUlid, Metrics, PrefetchObservation, REQUEST_ID_HEADER, observe_prefetch_body, + }, }; use async_compression::tokio::bufread::ZstdDecoder; use axum::{ @@ -335,7 +337,7 @@ async fn count_request( .and_then(|value| value.parse::().ok()); PrefetchObservation::start( Arc::clone(&metrics), - route.clone(), + Some(route.clone()), session, configured_concurrency, ) @@ -355,105 +357,6 @@ async fn count_request( } } -struct PrefetchObservation { - metrics: Arc, - started: Instant, - route: String, - session: Option, - configured_concurrency: Option, - status: u16, - bytes: u64, - completed: bool, -} - -impl PrefetchObservation { - fn start( - metrics: Arc, - route: String, - session: Option, - configured_concurrency: Option, - ) -> Self { - let in_flight = metrics.prefetch_started(); - tracing::debug!( - %route, - session_id = session.as_deref().unwrap_or(""), - configured_concurrency, - in_flight, - "prefetch request started" - ); - Self { - metrics, - started: Instant::now(), - route, - session, - configured_concurrency, - status: 0, - bytes: 0, - completed: false, - } - } - - fn set_status(&mut self, status: u16) { - self.status = status; - self.metrics.prefetch_response(status); - } - - fn add_bytes(&mut self, bytes: usize) { - self.bytes = self.bytes.saturating_add(bytes as u64); - } - - fn complete(mut self) { - self.completed = true; - } -} - -impl Drop for PrefetchObservation { - fn drop(&mut self) { - let duration = self.started.elapsed(); - self.metrics - .prefetch_finished(duration, self.bytes, self.completed); - tracing::debug!( - route = %self.route, - session_id = self.session.as_deref().unwrap_or(""), - configured_concurrency = self.configured_concurrency, - status = self.status, - bytes = self.bytes, - duration_ms = duration.as_millis() as u64, - completed = self.completed, - "prefetch request finished" - ); - } -} - -fn observe_prefetch_body(response: Response, mut observation: PrefetchObservation) -> Response { - observation.set_status(response.status().as_u16()); - let (parts, body) = response.into_parts(); - let stream = Box::pin(body.into_data_stream()); - let observed = futures_util::stream::unfold( - (stream, Some(observation)), - |(mut stream, mut observation)| async move { - match stream.next().await { - Some(Ok(bytes)) => { - observation - .as_mut() - .expect("prefetch observation exists while streaming") - .add_bytes(bytes.len()); - Some((Ok(bytes), (stream, observation))) - } - Some(Err(error)) => Some((Err(error), (stream, observation))), - None => { - observation - .take() - .expect("prefetch observation exists at completion") - .complete(); - None - } - } - }, - ); - Response::from_parts(parts, Body::from_stream(observed)) -} - async fn put_artifact( State(state): State>, Path(path): Path,