diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index 8c13f3975c2..0a52272153e 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -42,6 +42,7 @@ use spacetimedb_client_api_messages::name::{ self, DatabaseName, DomainName, MigrationPolicy, PrePublishAutoMigrateResult, PrePublishManualMigrateResult, PrePublishResult, PrettyPrintStyle, PublishOp, PublishResult, }; +use spacetimedb_datastore::db_metrics::DB_METRICS; use spacetimedb_lib::db::raw_def::v10::RawModuleDefV10; use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9; use spacetimedb_lib::{http as st_http, ConnectionId}; @@ -1621,7 +1622,10 @@ where S: NodeDelegate + ControlStateDelegate + Authorization + Clone + 'static, { pub fn into_router(self, ctx: S) -> axum::Router { - let db_router = axum::Router::::new() + let egress_metrics_middleware = + axum::middleware::from_fn_with_state(ctx.clone(), http_response_egress_metrics_middleware::); + + let counted_db_router = axum::Router::::new() .route("/", self.db_put) .route("/", self.db_get) .route("/", self.db_delete) @@ -1629,7 +1633,6 @@ where .route("/names", self.names_post) .route("/names", self.names_put) .route("/identity", self.identity_get) - .route("/subscribe", self.subscribe_get) .route("/call/:reducer", self.call_reducer_procedure_post) .route("/schema", self.schema_get) .route("/logs", self.logs_get) @@ -1639,7 +1642,13 @@ where .route("/pre_publish", self.pre_publish) .route("/reset", self.db_reset) .route("/lock", self.lock_post) - .route("/unlock", self.unlock_post); + .route("/unlock", self.unlock_post) + .route_layer(egress_metrics_middleware.clone()); + + // Add the subscribe route after `egress_metrics_middleware` + // so that its egress bytes don't get counted into `http_response_size_bytes`; + // we have different metrics tracking WebSocket message size. + let db_router = counted_db_router.route("/subscribe", self.subscribe_get); let authed_root_router = axum::Router::new().route( "/", @@ -1661,7 +1670,8 @@ where let http_route_router = axum::Router::::new() .route("/:name_or_identity/route", self.http_route_root) .route("/:name_or_identity/route/", self.http_route_root_slash) - .route("/:name_or_identity/route/*path", self.http_route); + .route("/:name_or_identity/route/*path", self.http_route) + .route_layer(egress_metrics_middleware); axum::Router::new() .merge(authed_root_router) @@ -1670,6 +1680,68 @@ where } } +/// Middleware which counts response bytes in the metric `spacetime_http_response_size_bytes_total`. +/// +/// This middleware is intended to be supplied to all HTTP routes which apply to a particular database, +/// *except* the WebSocket connection route (confusingly named `subscribe`), whose egress is measured separately. +async fn http_response_egress_metrics_middleware( + State(worker_ctx): State, + Path(DatabaseParam { name_or_identity }): Path, + request: Request, + next: axum::middleware::Next, +) -> axum::response::Response +where + S: ControlStateDelegate + Clone + Send + Sync + 'static, +{ + let Ok(Ok(database_identity)) = name_or_identity.try_resolve(&worker_ctx).await else { + // The provided name doesn't map to an `Identity`. + // Run the route unchanged (as opposed to returning an error from the middleware) + // to preserve the error-handling behavior of the route. + return next.run(request).await; + }; + + if !matches!( + worker_ctx_find_database(&worker_ctx, &database_identity).await, + Ok(Some(_)) + ) { + // We have what appears to be an `Identity`, but it doesn't name any database. + // Don't create a metrics label for it, + // and run the route unchanged (as opposed to returning an error from the middleware) + // to preserve the error-handling behavior of the route. + return next.run(request).await; + } + + let response = next.run(request).await; + let (parts, body) = response.into_parts(); + + // Count the number of bytes used by the headers. + // For guest-defined routes bound to HTTP handlers, these may be arbitrarily large and are worth billing for; + // for built-in routes they will be small and it doesn't really matter one way or another whether we do or don't bill. + // N.b. headers installed by other middleware may or may not be counted here, + // depending on the order in which the middleware applies. + let header_bytes: usize = parts + .headers + .iter() + .map(|(name, value)| name.as_str().len() + value.as_bytes().len()) + .sum(); + + let counter = DB_METRICS + .http_response_size_bytes + .with_label_values(&database_identity); + counter.inc_by(header_bytes as u64); + + // `/logs?follow=true` can stream indefinitely. + // Counting frames as they are emitted preserves streaming behavior and avoids buffering the response. + let body = body.map_frame(move |frame| { + if let Some(data) = frame.data_ref() { + counter.inc_by(data.len() as u64); + } + frame + }); + + axum::response::Response::from_parts(parts, Body::new(body)) +} + #[cfg(test)] mod tests { use super::*; @@ -1690,10 +1762,13 @@ mod tests { use spacetimedb_client_api_messages::name::{ DomainName, InsertDomainResult, RegisterTldResult, SetDomainsResult, Tld, }; + use spacetimedb_lib::Hash; use spacetimedb_paths::server::ModuleLogsDir; use spacetimedb_paths::FromPathUnchecked; use spacetimedb_schema::auto_migrate::{MigrationPolicy, PrettyPrintStyle}; + use std::collections::HashMap; use tower::util::ServiceExt; + #[derive(Clone, Default)] struct DummyValidator; @@ -1710,8 +1785,14 @@ mod tests { } impl TokenSigner for DummyJwtProvider { - fn sign(&self, _claims: &T) -> Result { - Err(JwtError::from(JwtErrorKind::InvalidSignature)) + fn sign(&self, claims: &T) -> Result { + use base64::{engine::general_purpose, Engine}; + + let payload = serde_json::to_vec(claims).map_err(|_| JwtError::from(JwtErrorKind::InvalidSignature))?; + Ok(format!( + "test.{}.signature", + general_purpose::URL_SAFE_NO_PAD.encode(payload) + )) } } @@ -1736,6 +1817,8 @@ mod tests { jwt: DummyJwtProvider, client_actor_index: std::sync::Arc, module_logs_dir: ModuleLogsDir, + databases: std::sync::Arc>, + dns: std::sync::Arc>, } impl DummyState { @@ -1746,8 +1829,77 @@ mod tests { }, client_actor_index: std::sync::Arc::new(ClientActorIndex::new()), module_logs_dir: ModuleLogsDir::from_path_unchecked(std::env::temp_dir()), + databases: std::sync::Arc::new(HashMap::new()), + dns: std::sync::Arc::new(HashMap::new()), } } + + fn with_database(mut self, database_identity: Identity) -> Self { + let mut databases = HashMap::new(); + databases.insert(database_identity, test_database(database_identity)); + self.databases = std::sync::Arc::new(databases); + self + } + + fn with_dns(mut self, name: &str, database_identity: Identity) -> Self { + let mut dns = HashMap::new(); + dns.insert(name.to_owned(), database_identity); + self.dns = std::sync::Arc::new(dns); + self + } + } + + fn test_identity(byte: u8) -> Identity { + Identity::from_byte_array([byte; 32]) + } + + fn test_database(database_identity: Identity) -> Database { + Database { + id: u64::from(database_identity.to_byte_array()[0]), + database_identity, + owner_identity: test_identity(254), + host_type: HostType::Wasm, + initial_program: Hash::from_byte_array([0; 32]), + bootstrap_generation: 0, + } + } + + fn http_response_size_metric(database_identity: Identity) -> u64 { + DB_METRICS + .http_response_size_bytes + .with_label_values(&database_identity) + .get() + } + + fn collected_http_response_size_metric(database_identity: Identity) -> Option { + let db_label = database_identity.to_hex(); + for metric_family in prometheus::core::Collector::collect(&DB_METRICS.http_response_size_bytes) { + if metric_family.name() != "spacetime_http_response_size_bytes_total" { + continue; + } + + for metric in metric_family.get_metric() { + let has_db_label = metric + .get_label() + .iter() + .any(|label| label.name() == "db" && label.value() == db_label.as_str()); + if has_db_label { + return Some(metric.get_counter().value() as u64); + } + } + } + + None + } + + fn remove_http_response_size_metric(database_identity: Identity) { + let _ = DB_METRICS + .http_response_size_bytes + .remove_label_values(&database_identity); + } + + fn text_plain_header_bytes() -> u64 { + "content-type".len() as u64 + "text/plain; charset=utf-8".len() as u64 } impl HasWebSocketOptions for DummyState { @@ -1817,8 +1969,8 @@ mod tests { async fn get_database_by_id(&self, _id: u64) -> anyhow::Result> { Ok(None) } - async fn get_database_by_identity(&self, _database_identity: &Identity) -> anyhow::Result> { - Ok(None) + async fn get_database_by_identity(&self, database_identity: &Identity) -> anyhow::Result> { + Ok(self.databases.get(database_identity).cloned()) } async fn get_databases(&self) -> anyhow::Result> { Ok(Vec::new()) @@ -1835,8 +1987,8 @@ mod tests { async fn get_energy_balance(&self, _identity: &Identity) -> anyhow::Result> { Ok(None) } - async fn lookup_database_identity(&self, _domain: &str) -> anyhow::Result> { - Ok(None) + async fn lookup_database_identity(&self, domain: &str) -> anyhow::Result> { + Ok(self.dns.get(domain).copied()) } async fn reverse_lookup(&self, _database_identity: &Identity) -> anyhow::Result> { Ok(Vec::new()) @@ -1966,4 +2118,192 @@ mod tests { // - `name_or_identity.resolve(worker_ctx)` -> `NameOrIdentity::resolve` assert_eq!(body, "`not-a-database` not found"); } + + #[tokio::test] + async fn http_response_egress_metric_counts_database_routes() { + let database_identity = test_identity(11); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new().with_database(database_identity); + let app = DatabaseRoutes:: { + db_get: axum::routing::get(|| async { ([(http::header::CONTENT_TYPE, "text/plain")], "hello") }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri(format!("/{database_identity}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "hello"); + assert_eq!( + http_response_size_metric(database_identity), + "content-type".len() as u64 + "text/plain".len() as u64 + "hello".len() as u64 + ); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn http_response_egress_metric_counts_user_http_route_headers_and_body() { + let database_identity = test_identity(12); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new() + .with_database(database_identity) + .with_dns("metric-test", database_identity); + let app = DatabaseRoutes:: { + http_route: axum::routing::any(|| async { ([("x-large-test-header", "abcdef")], "route-body") }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri("/metric-test/route/health") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "route-body"); + assert_eq!( + http_response_size_metric(database_identity), + "x-large-test-header".len() as u64 + + "abcdef".len() as u64 + + text_plain_header_bytes() + + "route-body".len() as u64 + ); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn http_response_egress_metric_skips_unresolved_and_nonexistent_databases() { + let resolved_but_missing_identity = test_identity(13); + let arbitrary_identity = test_identity(14); + remove_http_response_size_metric(resolved_but_missing_identity); + remove_http_response_size_metric(arbitrary_identity); + + let state = DummyState::new().with_dns("missing-db", resolved_but_missing_identity); + let app = DatabaseRoutes:: { + db_get: axum::routing::get(|| async { "not counted" }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let _ = app + .clone() + .oneshot(Request::builder().uri("/unresolved-name").body(Body::empty()).unwrap()) + .await + .unwrap() + .into_body() + .collect() + .await + .unwrap(); + let _ = app + .clone() + .oneshot(Request::builder().uri("/missing-db").body(Body::empty()).unwrap()) + .await + .unwrap() + .into_body() + .collect() + .await + .unwrap(); + let _ = app + .oneshot( + Request::builder() + .uri(format!("/{arbitrary_identity}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap() + .into_body() + .collect() + .await + .unwrap(); + + assert_eq!(collected_http_response_size_metric(resolved_but_missing_identity), None); + assert_eq!(collected_http_response_size_metric(arbitrary_identity), None); + + remove_http_response_size_metric(resolved_but_missing_identity); + remove_http_response_size_metric(arbitrary_identity); + } + + #[tokio::test] + async fn http_response_egress_metric_does_not_count_subscribe() { + let database_identity = test_identity(15); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new().with_database(database_identity); + let app = DatabaseRoutes:: { + subscribe_get: axum::routing::get(|| async { "subscribe response" }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri(format!("/{database_identity}/subscribe")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "subscribe response"); + assert_eq!(http_response_size_metric(database_identity), 0); + + remove_http_response_size_metric(database_identity); + } + + #[tokio::test] + async fn http_response_egress_metric_counts_error_responses_for_existing_database() { + let database_identity = test_identity(16); + remove_http_response_size_metric(database_identity); + + let state = DummyState::new().with_database(database_identity); + let app = DatabaseRoutes:: { + db_get: axum::routing::get(|| async { (StatusCode::BAD_REQUEST, "bad request body") }), + ..Default::default() + } + .into_router(state.clone()) + .with_state(state); + + let response = app + .oneshot( + Request::builder() + .uri(format!("/{database_identity}")) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = response.into_body().collect().await.unwrap().to_bytes(); + + assert_eq!(body, "bad request body"); + assert_eq!( + http_response_size_metric(database_identity), + text_plain_header_bytes() + "bad request body".len() as u64 + ); + + remove_http_response_size_metric(database_identity); + } } diff --git a/crates/core/src/host/host_controller.rs b/crates/core/src/host/host_controller.rs index 60ffb1ba042..fbd717b51f9 100644 --- a/crates/core/src/host/host_controller.rs +++ b/crates/core/src/host/host_controller.rs @@ -1578,8 +1578,10 @@ pub async fn extract_schema(program_bytes: Box<[u8]>, host_type: HostType) -> an .await } -// Remove all gauges associated with a database. -// This is useful if a database is being deleted. +/// Removes metrics associated with a database when the database is deleted. +/// +/// Despite the historical function name, this cleans up per-database metric +/// series even when they are not literally `Gauge`s or `IntGauge`s. pub fn remove_database_gauges<'a, I>(db: &Identity, table_names: I) where I: IntoIterator, @@ -1609,4 +1611,5 @@ where V8HeapMetrics::remove_all_metric_label_values_for_database(db); let _ = WORKER_METRICS.v8_request_queue_length.remove_label_values(db); + let _ = DB_METRICS.http_response_size_bytes.remove_label_values(db); } diff --git a/crates/datastore/src/db_metrics/mod.rs b/crates/datastore/src/db_metrics/mod.rs index b8366034bac..2c7375362b0 100644 --- a/crates/datastore/src/db_metrics/mod.rs +++ b/crates/datastore/src/db_metrics/mod.rs @@ -225,6 +225,13 @@ An individual HTTP response's size in bytes is the sum of the sizes of the heade #[labels(db: Identity)] pub procedure_http_response_size_bytes: IntCounterVec, + #[name = spacetime_http_response_size_bytes_total] + #[help = "Total logical bytes sent in HTTP responses for routes scoped to a specific database. + +An individual HTTP response's size in bytes is the sum of the sizes of the header names, header values and body."] + #[labels(db: Identity)] + pub http_response_size_bytes: IntCounterVec, + #[name = spacetime_procedure_num_http_requests] #[help = "Number of HTTP requests performed by procedures running in databases.