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,