From 0b573ffe588b37408d2426e3a5018101867e5daa Mon Sep 17 00:00:00 2001 From: Makisuo Date: Wed, 13 May 2026 01:52:09 +0200 Subject: [PATCH 1/8] Add native Rust telemetry ingest pipeline --- apps/ingest/Cargo.lock | 43 + apps/ingest/Cargo.toml | 6 +- apps/ingest/src/main.rs | 883 +++++-- apps/ingest/src/telemetry.rs | 2123 +++++++++++++++++ .../0011_org_ingest_sampling_policies.sql | 8 + packages/db/src/schema/index.ts | 1 + .../schema/org-ingest-sampling-policies.ts | 15 + 7 files changed, 2941 insertions(+), 138 deletions(-) create mode 100644 apps/ingest/src/telemetry.rs create mode 100644 packages/db/drizzle/0011_org_ingest_sampling_policies.sql create mode 100644 packages/db/src/schema/org-ingest-sampling-policies.ts diff --git a/apps/ingest/Cargo.lock b/apps/ingest/Cargo.lock index b96e810bc3..55b11cc143 100644 --- a/apps/ingest/Cargo.lock +++ b/apps/ingest/Cargo.lock @@ -334,6 +334,20 @@ dependencies = [ "typenum", ] +[[package]] +name = "dashmap" +version = "6.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + [[package]] name = "digest" version = "0.10.7" @@ -574,6 +588,12 @@ dependencies = [ "tracing", ] +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + [[package]] name = "hashbrown" version = "0.15.5" @@ -589,6 +609,12 @@ version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "hmac" version = "0.12.1" @@ -975,12 +1001,15 @@ dependencies = [ "axum", "base64", "chrono", + "crc32fast", + "dashmap", "dotenvy", "flate2", "hmac", "metrics", "metrics-exporter-prometheus", "moka", + "num_cpus", "opentelemetry", "opentelemetry-otlp", "opentelemetry-proto", @@ -992,6 +1021,7 @@ dependencies = [ "sha2", "tikv-jemallocator", "tokio", + "tonic", "tower-http", "tracing", "tracing-opentelemetry", @@ -1133,6 +1163,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -2106,8 +2146,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f32a6f80051a4111560201420c7885d0082ba9efe2ab61875c587bb6b18b9a0" dependencies = [ "async-trait", + "axum", "base64", "bytes", + "h2", "http", "http-body", "http-body-util", @@ -2116,6 +2158,7 @@ dependencies = [ "hyper-util", "percent-encoding", "pin-project", + "socket2", "sync_wrapper", "tokio", "tokio-stream", diff --git a/apps/ingest/Cargo.toml b/apps/ingest/Cargo.toml index 6b7c2651ab..74f4c3f7d7 100644 --- a/apps/ingest/Cargo.toml +++ b/apps/ingest/Cargo.toml @@ -25,18 +25,22 @@ opentelemetry-otlp = { version = "0.31", default-features = false, features = [ "trace", ] } opentelemetry-proto = { version = "0.31.0", features = [ - "gen-tonic-messages", + "gen-tonic", "trace", "logs", "metrics", "with-serde", ] } prost = "0.14.3" +crc32fast = "1.5" +dashmap = "6.1" +num_cpus = "1.17" reqwest = { version = "0.13.2", default-features = false, features = ["rustls", "http2", "json"] } serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0.145" sha2 = "0.10.9" tokio = { version = "1.48.0", features = ["macros", "rt-multi-thread", "signal"] } +tonic = { version = "0.14", features = ["transport"] } tracing = "0.1.41" tracing-opentelemetry = "0.32" tracing-subscriber = { version = "0.3.20", features = ["fmt", "env-filter"] } diff --git a/apps/ingest/src/main.rs b/apps/ingest/src/main.rs index fde7968ea1..de1a4e5800 100644 --- a/apps/ingest/src/main.rs +++ b/apps/ingest/src/main.rs @@ -2,8 +2,11 @@ static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; mod autumn; +mod telemetry; use std::io::{Read, Write}; +use std::path::PathBuf; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -21,6 +24,7 @@ use axum::Router; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use base64::Engine; use chrono::DateTime; +use dashmap::DashMap; use flate2::read::GzDecoder; use flate2::write::GzEncoder; use flate2::Compression; @@ -30,7 +34,6 @@ use moka::future::Cache; use opentelemetry::trace::TracerProvider as _; use opentelemetry::KeyValue as OtelKeyValue; use opentelemetry_otlp::{Protocol, SpanExporter, WithExportConfig}; -use tracing::Instrument; use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest; use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; @@ -41,15 +44,17 @@ use opentelemetry_sdk::runtime::Tokio as OtelTokio; use opentelemetry_sdk::trace::span_processor_with_async_runtime::BatchSpanProcessor; use opentelemetry_sdk::trace::{BatchConfigBuilder, SdkTracerProvider}; use opentelemetry_sdk::Resource as OtelResource; -use tracing_subscriber::layer::SubscriberExt; -use tracing_subscriber::util::SubscriberInitExt; use prost::Message; use reqwest::Client; use serde::{Deserialize, Serialize}; use serde_json::{Map as JsonMap, Value as JsonValue}; use sha2::Sha256; +use telemetry::{PipelineError, SamplingPolicy, TelemetryPipeline, TinybirdConfig}; use tower_http::cors::{Any, CorsLayer}; +use tracing::Instrument; use tracing::{debug, error, info, warn, Span}; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::util::SubscriberInitExt; const INGEST_SOURCE: &str = "maple-ingest-gateway"; const CLOUDFLARE_LOGPUSH_SOURCE: &str = "cloudflare-logpush"; @@ -79,10 +84,14 @@ type HmacSha256 = Hmac; #[derive(Clone)] struct AppConfig { port: u16, + otlp_grpc_port: Option, forward_endpoint: String, forward_self_managed_endpoint: Option, forward_timeout: Duration, + write_mode: WriteMode, + tinybird: TinybirdConfig, max_request_body_bytes: usize, + org_max_in_flight: u64, require_tls: bool, key_store_backend: KeyStoreBackend, lookup_hmac_key: String, @@ -91,12 +100,44 @@ struct AppConfig { autumn_flush_interval_secs: u64, } +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum WriteMode { + Tinybird, + Forward, + Dual, +} + +impl WriteMode { + fn from_env() -> Result { + let raw = std::env::var("INGEST_WRITE_MODE") + .unwrap_or_else(|_| "tinybird".to_string()) + .trim() + .to_ascii_lowercase(); + match raw.as_str() { + "tinybird" | "native" => Ok(Self::Tinybird), + "forward" | "collector" => Ok(Self::Forward), + "dual" | "dual_write" => Ok(Self::Dual), + _ => Err("INGEST_WRITE_MODE must be tinybird, forward, or dual".to_string()), + } + } + + fn uses_tinybird(self) -> bool { + matches!(self, Self::Tinybird | Self::Dual) + } + + fn uses_forward(self) -> bool { + matches!(self, Self::Forward | Self::Dual) + } +} + #[derive(Clone)] enum KeyStoreBackend { // No-DB local backend: every well-formed ingest key resolves to a single // override org. Selected for single-tenant local dev so contributors don't // need CF D1 credentials to boot the service. - Static { org_id: String }, + Static { + org_id: String, + }, // Cloudflare D1 REST backend used in multi-tenant / production deploys. D1 { cf_account_id: String, @@ -114,6 +155,11 @@ impl AppConfig { .or_else(|| std::env::var("PORT").ok()), 3474, )?; + let otlp_grpc_port = parse_optional_u16( + "INGEST_OTLP_GRPC_PORT", + std::env::var("INGEST_OTLP_GRPC_PORT").ok(), + )?; + let write_mode = WriteMode::from_env()?; let forward_endpoint = std::env::var("INGEST_FORWARD_OTLP_ENDPOINT") .unwrap_or_else(|_| "http://127.0.0.1:4318".to_string()) @@ -139,11 +185,94 @@ impl AppConfig { 10_000, )?; + let tinybird = TinybirdConfig { + endpoint: std::env::var("TINYBIRD_HOST") + .unwrap_or_default() + .trim() + .trim_end_matches('/') + .to_string(), + token: std::env::var("TINYBIRD_TOKEN") + .unwrap_or_default() + .trim() + .to_string(), + queue_dir: PathBuf::from( + std::env::var("INGEST_QUEUE_DIR") + .unwrap_or_else(|_| "/var/lib/maple-ingest/wal".to_string()), + ), + queue_max_bytes: parse_u64( + "INGEST_QUEUE_MAX_BYTES", + std::env::var("INGEST_QUEUE_MAX_BYTES").ok(), + 20 * 1024 * 1024 * 1024, + )?, + org_queue_max_bytes: parse_u64( + "INGEST_ORG_QUEUE_MAX_BYTES", + std::env::var("INGEST_ORG_QUEUE_MAX_BYTES").ok(), + 1024 * 1024 * 1024, + )?, + queue_channel_capacity: parse_usize( + "INGEST_QUEUE_CHANNEL_CAPACITY", + std::env::var("INGEST_QUEUE_CHANNEL_CAPACITY").ok(), + 100_000, + )?, + wal_shards: parse_usize( + "INGEST_WAL_SHARDS", + std::env::var("INGEST_WAL_SHARDS").ok(), + (num_cpus::get().max(1) * 2).max(2), + )?, + batch_max_rows: parse_usize( + "INGEST_BATCH_MAX_ROWS", + std::env::var("INGEST_BATCH_MAX_ROWS").ok(), + 5_000, + )?, + batch_max_bytes: parse_usize( + "INGEST_BATCH_MAX_BYTES", + std::env::var("INGEST_BATCH_MAX_BYTES").ok(), + 4 * 1024 * 1024, + )?, + batch_max_wait: Duration::from_millis(parse_u64( + "INGEST_BATCH_MAX_WAIT_MS", + std::env::var("INGEST_BATCH_MAX_WAIT_MS").ok(), + 100, + )?), + export_concurrency_per_shard: parse_usize( + "INGEST_TINYBIRD_CONCURRENCY_PER_SHARD", + std::env::var("INGEST_TINYBIRD_CONCURRENCY_PER_SHARD").ok(), + 1, + )?, + datasource_traces: std::env::var("INGEST_TINYBIRD_DATASOURCE_TRACES") + .unwrap_or_else(|_| "traces".to_string()), + datasource_logs: std::env::var("INGEST_TINYBIRD_DATASOURCE_LOGS") + .unwrap_or_else(|_| "logs".to_string()), + datasource_metrics_sum: std::env::var("INGEST_TINYBIRD_DATASOURCE_METRICS_SUM") + .unwrap_or_else(|_| "metrics_sum".to_string()), + datasource_metrics_gauge: std::env::var("INGEST_TINYBIRD_DATASOURCE_METRICS_GAUGE") + .unwrap_or_else(|_| "metrics_gauge".to_string()), + datasource_metrics_histogram: std::env::var( + "INGEST_TINYBIRD_DATASOURCE_METRICS_HISTOGRAM", + ) + .unwrap_or_else(|_| "metrics_histogram".to_string()), + datasource_metrics_exponential_histogram: std::env::var( + "INGEST_TINYBIRD_DATASOURCE_METRICS_EXPONENTIAL_HISTOGRAM", + ) + .unwrap_or_else(|_| "metrics_exponential_histogram".to_string()), + }; + if write_mode.uses_tinybird() { + tinybird.validate()?; + } + let max_request_body_bytes = parse_usize( "INGEST_MAX_REQUEST_BODY_BYTES", std::env::var("INGEST_MAX_REQUEST_BODY_BYTES").ok(), 20 * 1024 * 1024, )?; + let org_max_in_flight = parse_u64( + "INGEST_ORG_MAX_IN_FLIGHT", + std::env::var("INGEST_ORG_MAX_IN_FLIGHT").ok(), + 1_000, + )?; + if org_max_in_flight == 0 { + return Err("INGEST_ORG_MAX_IN_FLIGHT must be greater than 0".to_string()); + } let require_tls = parse_bool( "INGEST_REQUIRE_TLS", @@ -199,10 +328,14 @@ impl AppConfig { Ok(Self { port, + otlp_grpc_port, forward_endpoint, forward_self_managed_endpoint, forward_timeout: Duration::from_millis(forward_timeout_ms), + write_mode, + tinybird, max_request_body_bytes, + org_max_in_flight, require_tls, key_store_backend, lookup_hmac_key, @@ -296,6 +429,11 @@ struct CloudflareConnectorResolver { cache: Cache, } +struct SamplingPolicyResolver { + store: Arc, + cache: Cache, +} + /// Database-agnostic surface used by the two resolvers. Implementations: /// `LibsqlKeyStore` (local dev / legacy) and `D1KeyStore` (Cloudflare D1 REST /// in production, where the API service writes ingest-key rows). Both back @@ -314,11 +452,13 @@ trait KeyStore: Send + Sync { secret_hash: &str, ) -> Result, String>; - async fn record_connector_success( + async fn fetch_sampling_policy( &self, - connector_id: &str, - now_ms: i64, - ) -> Result<(), String>; + org_id: &str, + ) -> Result, String>; + + async fn record_connector_success(&self, connector_id: &str, now_ms: i64) + -> Result<(), String>; async fn record_connector_failure( &self, @@ -343,10 +483,20 @@ struct ConnectorRow { self_managed: bool, } +#[derive(Clone, Debug)] +struct SamplingPolicyRow { + trace_sample_ratio: f64, + always_keep_error_spans: bool, + always_keep_slow_spans_ms: Option, +} + struct AppState { config: AppConfig, http_client: Client, + telemetry_pipeline: Option, resolver: IngestKeyResolver, + org_inflight_limiter: OrgInFlightLimiter, + sampling_resolver: SamplingPolicyResolver, cloudflare_resolver: CloudflareConnectorResolver, metrics_handle: metrics_exporter_prometheus::PrometheusHandle, autumn_tracker: Option, @@ -412,9 +562,34 @@ impl Signal { } } -struct EnrichResult { - payload: Vec, - item_count: usize, +enum DecodedPayload { + Traces(ExportTraceServiceRequest), + Logs(ExportLogsServiceRequest), + Metrics(ExportMetricsServiceRequest), +} + +impl DecodedPayload { + fn item_count(&self) -> usize { + match self { + Self::Traces(request) => count_trace_items(request), + Self::Logs(request) => count_log_items(request), + Self::Metrics(request) => count_metric_items(request), + } + } + + fn encode(&self, payload_format: PayloadFormat) -> Result, ApiError> { + match (self, payload_format) { + (Self::Traces(request), PayloadFormat::Protobuf) => Ok(request.encode_to_vec()), + (Self::Logs(request), PayloadFormat::Protobuf) => Ok(request.encode_to_vec()), + (Self::Metrics(request), PayloadFormat::Protobuf) => Ok(request.encode_to_vec()), + (Self::Traces(request), PayloadFormat::Json) => serde_json::to_vec(request) + .map_err(|_| ApiError::service_unavailable("Failed to serialize traces payload")), + (Self::Logs(request), PayloadFormat::Json) => serde_json::to_vec(request) + .map_err(|_| ApiError::service_unavailable("Failed to serialize logs payload")), + (Self::Metrics(request), PayloadFormat::Json) => serde_json::to_vec(request) + .map_err(|_| ApiError::service_unavailable("Failed to serialize metrics payload")), + } + } } struct InFlightGuard; @@ -425,6 +600,73 @@ impl Drop for InFlightGuard { } } +#[derive(Clone)] +struct OrgInFlightLimiter { + max_per_org: u64, + counts: Arc>>, +} + +struct OrgInFlightPermit { + org_id: String, + counter: Arc, +} + +impl OrgInFlightLimiter { + fn new(max_per_org: u64) -> Self { + Self { + max_per_org, + counts: Arc::new(DashMap::new()), + } + } + + fn try_acquire(&self, org_id: &str) -> Option { + let counter = self + .counts + .entry(org_id.to_string()) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .clone(); + + loop { + let current = counter.load(Ordering::Relaxed); + if current >= self.max_per_org { + counter!( + "ingest_org_throttled_total", + "org_id" => org_id.to_string(), + "reason" => "in_flight" + ) + .increment(1); + return None; + } + if counter + .compare_exchange(current, current + 1, Ordering::AcqRel, Ordering::Relaxed) + .is_ok() + { + gauge!( + "ingest_org_requests_in_flight", + "org_id" => org_id.to_string() + ) + .set((current + 1) as f64); + return Some(OrgInFlightPermit { + org_id: org_id.to_string(), + counter, + }); + } + } + } +} + +impl Drop for OrgInFlightPermit { + fn drop(&mut self) { + let current = self.counter.fetch_sub(1, Ordering::AcqRel); + let next = current.saturating_sub(1); + gauge!( + "ingest_org_requests_in_flight", + "org_id" => self.org_id.clone() + ) + .set(next as f64); + } +} + #[derive(Serialize)] struct ErrorBody { error: String, @@ -460,6 +702,10 @@ impl ApiError { Self::new(StatusCode::PAYLOAD_TOO_LARGE, message) } + fn too_many_requests(message: impl Into) -> Self { + Self::new(StatusCode::TOO_MANY_REQUESTS, message) + } + fn service_unavailable(message: impl Into) -> Self { Self::new(StatusCode::SERVICE_UNAVAILABLE, message) } @@ -547,7 +793,9 @@ fn init_tracing(forward_endpoint: &str, bind_port: u16) -> Option exporter, Err(error) => { - eprintln!("Failed to build OTLP span exporter: {error}; falling back to stdout-only tracing"); + eprintln!( + "Failed to build OTLP span exporter: {error}; falling back to stdout-only tracing" + ); tracing_subscriber::registry() .with(env_filter) .with(fmt_layer) @@ -615,8 +863,10 @@ async fn main() { let http_client = match Client::builder() .timeout(config.forward_timeout) - .pool_max_idle_per_host(5) + .pool_max_idle_per_host(256) .pool_idle_timeout(Duration::from_secs(30)) + .http2_keep_alive_interval(Duration::from_secs(20)) + .http2_keep_alive_timeout(Duration::from_secs(5)) .build() { Ok(client) => client, @@ -626,6 +876,18 @@ async fn main() { } }; + let telemetry_pipeline = if config.write_mode.uses_tinybird() { + match TelemetryPipeline::new(config.tinybird.clone(), http_client.clone()).await { + Ok(pipeline) => Some(pipeline), + Err(error) => { + eprintln!("Telemetry pipeline init error: {error}"); + std::process::exit(1); + } + } + } else { + None + }; + // Cloudflare D1 REST backend — the API writes ingest-key rows to D1, so // ingest reads them from the same place. We run a probe query before // accepting traffic; if anything is wrong (auth, schema, network) the @@ -655,6 +917,10 @@ async fn main() { .time_to_live(Duration::from_secs(60)) .max_capacity(1_000) .build(); + let sampling_policy_cache = Cache::builder() + .time_to_live(Duration::from_secs(30)) + .max_capacity(10_000) + .build(); let state = Arc::new(AppState { resolver: IngestKeyResolver { @@ -662,11 +928,17 @@ async fn main() { lookup_hmac_key: config.lookup_hmac_key.clone(), cache: ingest_key_cache, }, + org_inflight_limiter: OrgInFlightLimiter::new(config.org_max_in_flight), + sampling_resolver: SamplingPolicyResolver { + store: Arc::clone(&store), + cache: sampling_policy_cache, + }, cloudflare_resolver: CloudflareConnectorResolver { store: Arc::clone(&store), lookup_hmac_key: config.lookup_hmac_key.clone(), cache: cloudflare_connector_cache, }, + telemetry_pipeline, http_client, config: config.clone(), metrics_handle: prometheus_handle, @@ -683,6 +955,7 @@ async fn main() { HeaderName::from_static("x-maple-ingest-key"), ]); + let grpc_state = Arc::clone(&state); let app = Router::new() .route("/health", get(health)) .route("/metrics", get(serve_metrics)) @@ -740,6 +1013,10 @@ async fn main() { ); } + if let Some(grpc_port) = config.otlp_grpc_port { + tokio::spawn(run_grpc_server(grpc_state, grpc_port)); + } + let serve_result = axum::serve(listener, app) .with_graceful_shutdown(shutdown_signal()) .await; @@ -779,6 +1056,203 @@ async fn shutdown_signal() { } } +async fn run_grpc_server(state: Arc, port: u16) { + use opentelemetry_proto::tonic::collector::logs::v1::logs_service_server::LogsServiceServer; + use opentelemetry_proto::tonic::collector::metrics::v1::metrics_service_server::MetricsServiceServer; + use opentelemetry_proto::tonic::collector::trace::v1::trace_service_server::TraceServiceServer; + + let addr = ([0, 0, 0, 0], port).into(); + let server = tonic::transport::Server::builder() + .add_service(TraceServiceServer::new(GrpcTraceService { + state: Arc::clone(&state), + })) + .add_service(LogsServiceServer::new(GrpcLogsService { + state: Arc::clone(&state), + })) + .add_service(MetricsServiceServer::new(GrpcMetricsService { state })); + + info!(port, "Maple OTLP gRPC server listening"); + if let Err(error) = server.serve_with_shutdown(addr, shutdown_signal()).await { + error!(error = %error, "OTLP gRPC server failed"); + } +} + +#[derive(Clone)] +struct GrpcTraceService { + state: Arc, +} + +#[derive(Clone)] +struct GrpcLogsService { + state: Arc, +} + +#[derive(Clone)] +struct GrpcMetricsService { + state: Arc, +} + +#[tonic::async_trait] +impl opentelemetry_proto::tonic::collector::trace::v1::trace_service_server::TraceService + for GrpcTraceService +{ + async fn export( + &self, + request: tonic::Request, + ) -> Result< + tonic::Response< + opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceResponse, + >, + tonic::Status, + > { + let resolved = resolve_grpc_ingest_key(&self.state, request.metadata()).await?; + let mut inner = request.into_inner(); + enrich_trace_request(&mut inner, &resolved); + accept_grpc_decoded( + &self.state, + Signal::Traces, + DecodedPayload::Traces(inner), + &resolved, + ) + .await?; + Ok(tonic::Response::new( + opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceResponse { + partial_success: None, + }, + )) + } +} + +#[tonic::async_trait] +impl opentelemetry_proto::tonic::collector::logs::v1::logs_service_server::LogsService + for GrpcLogsService +{ + async fn export( + &self, + request: tonic::Request, + ) -> Result< + tonic::Response, + tonic::Status, + > { + let resolved = resolve_grpc_ingest_key(&self.state, request.metadata()).await?; + let mut inner = request.into_inner(); + enrich_logs_request(&mut inner, &resolved); + accept_grpc_decoded( + &self.state, + Signal::Logs, + DecodedPayload::Logs(inner), + &resolved, + ) + .await?; + Ok(tonic::Response::new( + opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceResponse { + partial_success: None, + }, + )) + } +} + +#[tonic::async_trait] +impl opentelemetry_proto::tonic::collector::metrics::v1::metrics_service_server::MetricsService + for GrpcMetricsService +{ + async fn export( + &self, + request: tonic::Request, + ) -> Result< + tonic::Response< + opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceResponse, + >, + tonic::Status, + > { + let resolved = resolve_grpc_ingest_key(&self.state, request.metadata()).await?; + let mut inner = request.into_inner(); + enrich_metrics_request(&mut inner, &resolved); + accept_grpc_decoded( + &self.state, + Signal::Metrics, + DecodedPayload::Metrics(inner), + &resolved, + ) + .await?; + Ok(tonic::Response::new( + opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceResponse { + partial_success: None, + }, + )) + } +} + +async fn accept_grpc_decoded( + state: &AppState, + signal: Signal, + decoded: DecodedPayload, + resolved: &ResolvedIngestKey, +) -> Result<(), tonic::Status> { + let _org_inflight_permit = state + .org_inflight_limiter + .try_acquire(&resolved.org_id) + .ok_or_else(|| tonic::Status::resource_exhausted("Per-org ingest limit exceeded"))?; + process_decoded_payload( + state, + signal, + PayloadFormat::Protobuf, + None, + &decoded, + resolved, + ) + .await + .map(|_| ()) + .map_err(|error| { + if error.status == StatusCode::TOO_MANY_REQUESTS { + tonic::Status::resource_exhausted(error.message) + } else { + tonic::Status::unavailable(error.message) + } + }) +} + +async fn resolve_grpc_ingest_key( + state: &AppState, + metadata: &tonic::metadata::MetadataMap, +) -> Result { + let token = metadata + .get("authorization") + .and_then(|value| value.to_str().ok()) + .and_then(|value| { + if value.len() > 7 && value[..7].eq_ignore_ascii_case("Bearer ") { + Some(value[7..].trim().to_string()) + } else { + None + } + }) + .or_else(|| { + metadata + .get("x-maple-ingest-key") + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + }) + .ok_or_else(|| tonic::Status::unauthenticated("Missing ingest key"))?; + + if is_sentinel_token(&token) { + return Ok(ResolvedIngestKey { + org_id: SENTINEL_ORG_ID.to_string(), + key_type: IngestKeyType::Public, + key_id: "sentinel".to_string(), + self_managed: false, + }); + } + + state + .resolver + .resolve_ingest_key(&token) + .await + .map_err(|_| tonic::Status::unavailable("Ingest authentication unavailable"))? + .ok_or_else(|| tonic::Status::unauthenticated("Invalid ingest key")) +} + async fn health() -> &'static str { "OK" } @@ -937,9 +1411,10 @@ async fn handle_cloudflare_logpush( ); let span_handle = span.clone(); - let result = handle_cloudflare_logpush_inner(&state, &connector_id, secret.as_deref(), &headers, body) - .instrument(span) - .await; + let result = + handle_cloudflare_logpush_inner(&state, &connector_id, secret.as_deref(), &headers, body) + .instrument(span) + .await; let duration = start.elapsed(); match result { @@ -1045,6 +1520,19 @@ async fn handle_signal_inner( resolve_ms = key_resolve_start.elapsed().as_millis() as u64, "Authenticated" ); + let _org_inflight_permit = state + .org_inflight_limiter + .try_acquire(&resolved_key.org_id) + .ok_or_else(|| { + warn!( + org_id = %resolved_key.org_id, + "Per-org in-flight ingest limit exceeded" + ); + ( + ApiError::too_many_requests("Per-org ingest limit exceeded"), + "throttle", + ) + })?; // --- Payload validation --- if body.len() > state.config.max_request_body_bytes { @@ -1100,64 +1588,46 @@ async fn handle_signal_inner( .record(decoded_payload.len() as f64); // --- Enrich --- - let enrich_result = enrich_payload(signal, payload_format, &decoded_payload, &resolved_key) - .map_err(|e| { - warn!( - format = payload_format.label(), - signal = signal.path(), - org_id = resolved_key.org_id.as_str(), - key_type = resolved_key.key_type.as_str(), - decoded_bytes = decoded_payload.len(), - reason = %e.message, - "Invalid OTLP payload" - ); - (e, "enrich") - })?; + let decoded = + decode_and_enrich_payload(signal, payload_format, &decoded_payload, &resolved_key) + .map_err(|e| { + warn!( + format = payload_format.label(), + signal = signal.path(), + org_id = resolved_key.org_id.as_str(), + key_type = resolved_key.key_type.as_str(), + decoded_bytes = decoded_payload.len(), + reason = %e.message, + "Invalid OTLP payload" + ); + (e, "enrich") + })?; + let item_count = decoded.item_count(); - Span::current().record("maple.ingest.item_count", enrich_result.item_count); - debug!(item_count = enrich_result.item_count, "Payload enriched"); + Span::current().record("maple.ingest.item_count", item_count); + debug!(item_count, "Payload enriched"); counter!( "ingest_items_total", "signal" => signal.path() ) - .increment(enrich_result.item_count as u64); + .increment(item_count as u64); let decoded_bytes = decoded_payload.len(); - // --- Encode & Forward --- - let outbound_body = encode_payload(&enrich_result.payload, content_encoding.as_deref()) - .map_err(|e| (e, "encode"))?; - - let outbound_bytes = outbound_body.len(); - let forward_span = tracing::info_span!( - "forward", - otel.name = "POST", - otel.kind = "client", - otel.status_code = tracing::field::Empty, - "http.request.method" = "POST", - "http.request.body.size" = outbound_bytes, - "http.response.status_code" = tracing::field::Empty, - "url.full" = tracing::field::Empty, - "server.address" = tracing::field::Empty, - "error.type" = tracing::field::Empty, - "maple.signal" = signal.path(), - "maple.ingest.upstream_pool" = tracing::field::Empty, - ); - let response = forward_to_collector( + let response = process_decoded_payload( state, signal, - payload_format.content_type(), + payload_format, content_encoding.as_deref(), - outbound_body, + &decoded, &resolved_key, ) - .instrument(forward_span) .await .map_err(|e| (e, "forward"))?; Ok(( response, - enrich_result.item_count, + item_count, resolved_key.org_id.clone(), decoded_bytes, )) @@ -1208,6 +1678,20 @@ async fn handle_cloudflare_logpush_inner( key_id = %resolved.secret_key_id, "Authenticated Cloudflare Logpush connector" ); + let _org_inflight_permit = state + .org_inflight_limiter + .try_acquire(&resolved.org_id) + .ok_or_else(|| { + warn!( + org_id = %resolved.org_id, + connector_id = %resolved.connector_id, + "Per-org in-flight ingest limit exceeded" + ); + ( + ApiError::too_many_requests("Per-org ingest limit exceeded"), + "throttle", + ) + })?; if body.len() > state.config.max_request_body_bytes { warn!( @@ -1292,36 +1776,21 @@ async fn handle_cloudflare_logpush_inner( ) .increment(item_count as u64); - let outbound = request.encode_to_vec(); - let outbound_bytes = outbound.len(); - let forward_span = tracing::info_span!( - "forward", - otel.name = "POST", - otel.kind = "client", - otel.status_code = tracing::field::Empty, - "http.request.method" = "POST", - "http.request.body.size" = outbound_bytes, - "http.response.status_code" = tracing::field::Empty, - "url.full" = tracing::field::Empty, - "server.address" = tracing::field::Empty, - "error.type" = tracing::field::Empty, - "maple.signal" = Signal::Logs.path(), - "maple.ingest.upstream_pool" = tracing::field::Empty, - ); - let response = match forward_to_collector( + let resolved_key = ResolvedIngestKey { + org_id: resolved.org_id.clone(), + key_type: IngestKeyType::Connector, + key_id: resolved.secret_key_id.clone(), + self_managed: resolved.self_managed, + }; + let decoded = DecodedPayload::Logs(request); + let response = match process_decoded_payload( state, Signal::Logs, - "application/x-protobuf", + PayloadFormat::Protobuf, None, - outbound, - &ResolvedIngestKey { - org_id: resolved.org_id.clone(), - key_type: IngestKeyType::Connector, - key_id: resolved.secret_key_id.clone(), - self_managed: resolved.self_managed, - }, + &decoded, + &resolved_key, ) - .instrument(forward_span) .await { Ok(response) => response, @@ -1722,79 +2191,48 @@ fn encode_payload(payload: &[u8], content_encoding: Option<&str>) -> Result Result { +) -> Result { match (signal, payload_format) { (Signal::Traces, PayloadFormat::Protobuf) => { let mut request = ExportTraceServiceRequest::decode(payload) .map_err(|_| ApiError::bad_request("Invalid OTLP traces protobuf payload"))?; enrich_trace_request(&mut request, resolved_key); - let item_count = count_trace_items(&request); - Ok(EnrichResult { - payload: request.encode_to_vec(), - item_count, - }) + Ok(DecodedPayload::Traces(request)) } (Signal::Logs, PayloadFormat::Protobuf) => { let mut request = ExportLogsServiceRequest::decode(payload) .map_err(|_| ApiError::bad_request("Invalid OTLP logs protobuf payload"))?; enrich_logs_request(&mut request, resolved_key); - let item_count = count_log_items(&request); - Ok(EnrichResult { - payload: request.encode_to_vec(), - item_count, - }) + Ok(DecodedPayload::Logs(request)) } (Signal::Metrics, PayloadFormat::Protobuf) => { let mut request = ExportMetricsServiceRequest::decode(payload) .map_err(|_| ApiError::bad_request("Invalid OTLP metrics protobuf payload"))?; enrich_metrics_request(&mut request, resolved_key); - let item_count = count_metric_items(&request); - Ok(EnrichResult { - payload: request.encode_to_vec(), - item_count, - }) + Ok(DecodedPayload::Metrics(request)) } (Signal::Traces, PayloadFormat::Json) => { let mut request: ExportTraceServiceRequest = serde_json::from_slice(payload) .map_err(|_| ApiError::bad_request("Invalid OTLP traces JSON payload"))?; enrich_trace_request(&mut request, resolved_key); - let item_count = count_trace_items(&request); - let payload = serde_json::to_vec(&request) - .map_err(|_| ApiError::service_unavailable("Failed to serialize traces payload"))?; - Ok(EnrichResult { - payload, - item_count, - }) + Ok(DecodedPayload::Traces(request)) } (Signal::Logs, PayloadFormat::Json) => { let mut request: ExportLogsServiceRequest = serde_json::from_slice(payload) .map_err(|_| ApiError::bad_request("Invalid OTLP logs JSON payload"))?; enrich_logs_request(&mut request, resolved_key); - let item_count = count_log_items(&request); - let payload = serde_json::to_vec(&request) - .map_err(|_| ApiError::service_unavailable("Failed to serialize logs payload"))?; - Ok(EnrichResult { - payload, - item_count, - }) + Ok(DecodedPayload::Logs(request)) } (Signal::Metrics, PayloadFormat::Json) => { let mut request: ExportMetricsServiceRequest = serde_json::from_slice(payload) .map_err(|_| ApiError::bad_request("Invalid OTLP metrics JSON payload"))?; enrich_metrics_request(&mut request, resolved_key); - let item_count = count_metric_items(&request); - let payload = serde_json::to_vec(&request).map_err(|_| { - ApiError::service_unavailable("Failed to serialize metrics payload") - })?; - Ok(EnrichResult { - payload, - item_count, - }) + Ok(DecodedPayload::Metrics(request)) } } } @@ -2040,6 +2478,111 @@ async fn forward_to_collector( .map_err(|_| ApiError::service_unavailable("Telemetry backend unavailable")) } +async fn process_decoded_payload( + state: &AppState, + signal: Signal, + payload_format: PayloadFormat, + content_encoding: Option<&str>, + decoded: &DecodedPayload, + resolved_key: &ResolvedIngestKey, +) -> Result { + if state.config.write_mode.uses_tinybird() { + let pipeline = state + .telemetry_pipeline + .as_ref() + .ok_or_else(|| ApiError::service_unavailable("Telemetry pipeline is not configured"))?; + let native_start = Instant::now(); + let stats = match decoded { + DecodedPayload::Traces(request) => { + let policy = state + .sampling_resolver + .resolve_policy(&resolved_key.org_id) + .await; + pipeline + .accept_traces(&resolved_key.org_id, request, &policy) + .await + } + DecodedPayload::Logs(request) => { + pipeline.accept_logs(&resolved_key.org_id, request).await + } + DecodedPayload::Metrics(request) => { + pipeline.accept_metrics(&resolved_key.org_id, request).await + } + } + .map_err(|error| { + let api_error = match &error { + PipelineError::Throttled(_) => { + ApiError::too_many_requests("Per-org ingest queue limit exceeded") + } + PipelineError::Backpressure(_) => { + ApiError::service_unavailable("Telemetry backend unavailable") + } + PipelineError::QueueUnavailable(_) | PipelineError::Encode(_) => { + ApiError::service_unavailable("Telemetry backend unavailable") + } + }; + error!( + error = %error, + signal = signal.path(), + org_id = %resolved_key.org_id, + "Native telemetry pipeline rejected payload" + ); + api_error + })?; + histogram!( + "ingest_native_accept_duration_seconds", + "signal" => signal.path() + ) + .record(native_start.elapsed().as_secs_f64()); + counter!( + "ingest_native_rows_total", + "signal" => signal.path() + ) + .increment(stats.rows as u64); + if stats.dropped > 0 { + counter!( + "ingest_native_sampled_dropped_total", + "signal" => signal.path() + ) + .increment(stats.dropped as u64); + } + Span::current().record("maple.ingest.native_rows", stats.rows as u64); + Span::current().record("maple.ingest.sampled_dropped", stats.dropped as u64); + } + + if state.config.write_mode.uses_forward() { + let outbound_payload = decoded.encode(payload_format)?; + let outbound_body = encode_payload(&outbound_payload, content_encoding)?; + let outbound_bytes = outbound_body.len(); + let forward_span = tracing::info_span!( + "forward", + otel.name = "POST", + otel.kind = "client", + otel.status_code = tracing::field::Empty, + "http.request.method" = "POST", + "http.request.body.size" = outbound_bytes, + "http.response.status_code" = tracing::field::Empty, + "url.full" = tracing::field::Empty, + "server.address" = tracing::field::Empty, + "error.type" = tracing::field::Empty, + "maple.signal" = signal.path(), + "maple.ingest.upstream_pool" = tracing::field::Empty, + ); + return forward_to_collector( + state, + signal, + payload_format.content_type(), + content_encoding, + outbound_body, + resolved_key, + ) + .instrument(forward_span) + .await; + } + + Ok(StatusCode::OK.into_response()) +} + impl IngestKeyResolver { async fn resolve_ingest_key(&self, raw_key: &str) -> Result, String> { if let Some(cached) = self.cache.get(raw_key).await { @@ -2123,15 +2666,38 @@ impl CloudflareConnectorResolver { async fn record_failure(&self, connector_id: &str, error_message: &str) -> Result<(), String> { self.store - .record_connector_failure( - connector_id, - error_message, - current_time_millis() as i64, - ) + .record_connector_failure(connector_id, error_message, current_time_millis() as i64) .await } } +impl SamplingPolicyResolver { + async fn resolve_policy(&self, org_id: &str) -> SamplingPolicy { + if let Some(policy) = self.cache.get(org_id).await { + return policy; + } + + let policy = match self.store.fetch_sampling_policy(org_id).await { + Ok(Some(row)) => SamplingPolicy { + trace_sample_ratio: row.trace_sample_ratio, + always_keep_error_spans: row.always_keep_error_spans, + always_keep_slow_spans_ms: row.always_keep_slow_spans_ms, + }, + Ok(None) => SamplingPolicy::default(), + Err(error) => { + warn!( + org_id, + error = %error, + "Sampling policy lookup failed; using unsampled default" + ); + SamplingPolicy::default() + } + }; + self.cache.insert(org_id.to_string(), policy.clone()).await; + policy + } +} + /// Cloudflare D1 REST-backed KeyStore. Hits /// `POST /accounts/{acct}/d1/database/{db}/query` for every cache miss. /// The HMAC-fingerprint canary, the 60s in-process cache, and SQL strings @@ -2144,12 +2710,7 @@ struct D1KeyStore { } impl D1KeyStore { - fn new( - http: reqwest::Client, - account_id: &str, - database_id: &str, - api_token: String, - ) -> Self { + fn new(http: reqwest::Client, account_id: &str, database_id: &str, api_token: String) -> Self { Self { http, endpoint: format!( @@ -2332,6 +2893,30 @@ impl KeyStore for D1KeyStore { })) } + async fn fetch_sampling_policy( + &self, + org_id: &str, + ) -> Result, String> { + let sql = "SELECT trace_sample_ratio, always_keep_error_spans, always_keep_slow_spans_ms \ + FROM org_ingest_sampling_policies WHERE org_id = ? LIMIT 1"; + let rows = self + .query(sql, vec![serde_json::Value::String(org_id.to_string())]) + .await?; + let Some(row) = rows.into_iter().next() else { + return Ok(None); + }; + Ok(Some(SamplingPolicyRow { + trace_sample_ratio: row + .get("trace_sample_ratio") + .and_then(serde_json::Value::as_f64) + .unwrap_or(1.0), + always_keep_error_spans: d1_truthy(&row, "always_keep_error_spans"), + always_keep_slow_spans_ms: row + .get("always_keep_slow_spans_ms") + .and_then(serde_json::Value::as_u64), + })) + } + async fn record_connector_success( &self, connector_id: &str, @@ -2394,6 +2979,13 @@ impl KeyStore for StaticKeyStore { Ok(None) } + async fn fetch_sampling_policy( + &self, + _org_id: &str, + ) -> Result, String> { + Ok(None) + } + async fn record_connector_success( &self, _connector_id: &str, @@ -2519,6 +3111,20 @@ fn parse_u16(name: &str, raw: Option, default: u16) -> Result) -> Result, String> { + let Some(raw) = raw else { + return Ok(None); + }; + let value = raw.trim(); + if value.is_empty() || value == "0" { + return Ok(None); + } + value + .parse::() + .map(Some) + .map_err(|_| format!("{name} must be a valid u16")) +} + fn parse_u64(name: &str, raw: Option, default: u64) -> Result { let Some(raw) = raw else { return Ok(default); @@ -2762,11 +3368,8 @@ mod tests { #[test] fn self_managed_goes_to_self_managed_pool_when_configured() { - let (endpoint, pool) = select_forward_endpoint( - "http://shared:4318", - Some("http://self-managed:4318"), - true, - ); + let (endpoint, pool) = + select_forward_endpoint("http://shared:4318", Some("http://self-managed:4318"), true); assert_eq!(endpoint, "http://self-managed:4318"); assert_eq!(pool, "self_managed"); } @@ -2820,6 +3423,12 @@ mod tests { ) -> Result, String> { Ok(None) } + async fn fetch_sampling_policy( + &self, + _org_id: &str, + ) -> Result, String> { + Ok(None) + } async fn record_connector_success( &self, _connector_id: &str, diff --git a/apps/ingest/src/telemetry.rs b/apps/ingest/src/telemetry.rs new file mode 100644 index 0000000000..af2b4b7cb2 --- /dev/null +++ b/apps/ingest/src/telemetry.rs @@ -0,0 +1,2123 @@ +use std::collections::BTreeMap; +use std::fs::{File, OpenOptions}; +use std::hash::{Hash, Hasher}; +use std::io::{Read, Seek, SeekFrom, Write}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use crc32fast::Hasher as Crc32; +use dashmap::DashMap; +use flate2::write::GzEncoder; +use flate2::Compression; +use metrics::{counter, gauge, histogram}; +use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; +use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest; +use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest; +use opentelemetry_proto::tonic::common::v1::{any_value, AnyValue, KeyValue}; +use opentelemetry_proto::tonic::logs::v1::LogRecord; +use opentelemetry_proto::tonic::metrics::v1::{ + metric, number_data_point, Exemplar, NumberDataPoint, +}; +use opentelemetry_proto::tonic::trace::v1::{span, status, Span}; +use reqwest::Client; +use serde::Serialize; +use serde_json::{json, Map, Value}; +use tokio::sync::mpsc; +use tokio::time::sleep; +use tracing::{error, info, warn}; + +const WAL_MAGIC: &[u8; 4] = b"MTW1"; +const WAL_V1_HEADER_LEN: usize = 20; +const WAL_V2_HEADER_LEN: usize = 22; + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub enum TelemetrySignal { + Traces, + Logs, + Metrics, +} + +#[derive(Clone, Debug)] +pub struct TinybirdConfig { + pub endpoint: String, + pub token: String, + pub queue_dir: PathBuf, + pub queue_max_bytes: u64, + pub org_queue_max_bytes: u64, + pub queue_channel_capacity: usize, + pub wal_shards: usize, + pub batch_max_rows: usize, + pub batch_max_bytes: usize, + pub batch_max_wait: Duration, + pub export_concurrency_per_shard: usize, + pub datasource_traces: String, + pub datasource_logs: String, + pub datasource_metrics_sum: String, + pub datasource_metrics_gauge: String, + pub datasource_metrics_histogram: String, + pub datasource_metrics_exponential_histogram: String, +} + +impl TinybirdConfig { + pub fn validate(&self) -> Result<(), String> { + if self.endpoint.is_empty() { + return Err( + "TINYBIRD_HOST is required when INGEST_WRITE_MODE uses tinybird".to_string(), + ); + } + if self.token.is_empty() { + return Err( + "TINYBIRD_TOKEN is required when INGEST_WRITE_MODE uses tinybird".to_string(), + ); + } + if self.wal_shards == 0 { + return Err("INGEST_WAL_SHARDS must be greater than 0".to_string()); + } + if self.batch_max_rows == 0 || self.batch_max_bytes == 0 { + return Err( + "INGEST_BATCH_MAX_ROWS and INGEST_BATCH_MAX_BYTES must be greater than 0" + .to_string(), + ); + } + if self.queue_max_bytes == 0 { + return Err("INGEST_QUEUE_MAX_BYTES must be greater than 0".to_string()); + } + if self.org_queue_max_bytes == 0 { + return Err("INGEST_ORG_QUEUE_MAX_BYTES must be greater than 0".to_string()); + } + if self.export_concurrency_per_shard == 0 { + return Err("INGEST_TINYBIRD_CONCURRENCY_PER_SHARD must be greater than 0".to_string()); + } + Ok(()) + } +} + +#[derive(Clone, Debug)] +pub struct SamplingPolicy { + pub trace_sample_ratio: f64, + pub always_keep_error_spans: bool, + pub always_keep_slow_spans_ms: Option, +} + +impl Default for SamplingPolicy { + fn default() -> Self { + Self { + trace_sample_ratio: 1.0, + always_keep_error_spans: true, + always_keep_slow_spans_ms: None, + } + } +} + +impl SamplingPolicy { + fn clamped_ratio(&self) -> f64 { + if !self.trace_sample_ratio.is_finite() { + return 1.0; + } + self.trace_sample_ratio.clamp(0.000001, 1.0) + } +} + +#[derive(Debug)] +pub enum PipelineError { + Backpressure(&'static str), + Throttled(&'static str), + QueueUnavailable(String), + Encode(String), +} + +impl std::fmt::Display for PipelineError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Backpressure(message) => f.write_str(message), + Self::Throttled(message) => f.write_str(message), + Self::QueueUnavailable(message) => f.write_str(message), + Self::Encode(message) => f.write_str(message), + } + } +} + +impl std::error::Error for PipelineError {} + +#[derive(Clone, Debug, Default)] +pub struct AcceptStats { + pub rows: usize, + pub dropped: usize, +} + +#[derive(Clone)] +pub struct TelemetryPipeline { + inner: Arc, +} + +struct PipelineInner { + cfg: Arc, + wal: Arc, + shard_senders: Vec>, + org_queue_bytes: Arc>>, +} + +#[derive(Clone, Debug)] +struct QueuedFrame { + shard: usize, + start: u64, + end: u64, + org_id: String, + queued_bytes: u64, + signal: TelemetrySignal, + datasource: String, + row_count: usize, + payload: Vec, +} + +#[derive(Debug)] +struct EncodedFrame { + routing_key: u64, + org_id: String, + signal: TelemetrySignal, + datasource: String, + row_count: usize, + payload: Vec, +} + +impl TelemetryPipeline { + pub async fn new(cfg: TinybirdConfig, http: Client) -> Result { + cfg.validate()?; + std::fs::create_dir_all(&cfg.queue_dir) + .map_err(|error| format!("create ingest WAL dir: {error}"))?; + let cfg = Arc::new(cfg); + let wal = Arc::new(ShardedWal::open(&cfg)?); + let org_queue_bytes = Arc::new(DashMap::new()); + let mut shard_senders = Vec::with_capacity(cfg.wal_shards); + + for shard in 0..cfg.wal_shards { + let (sender, receiver) = mpsc::channel(cfg.queue_channel_capacity); + shard_senders.push(sender); + let worker = ExportWorker { + shard, + cfg: Arc::clone(&cfg), + wal: Arc::clone(&wal), + org_queue_bytes: Arc::clone(&org_queue_bytes), + http: http.clone(), + receiver, + }; + tokio::spawn(worker.run()); + } + + let pipeline = Self { + inner: Arc::new(PipelineInner { + cfg, + wal, + shard_senders, + org_queue_bytes, + }), + }; + pipeline.replay_committed_frames().await; + Ok(pipeline) + } + + pub async fn accept_traces( + &self, + org_id: &str, + request: &ExportTraceServiceRequest, + sampling_policy: &SamplingPolicy, + ) -> Result { + let (frames, stats) = encode_traces(&self.inner.cfg, org_id, request, sampling_policy)?; + self.commit_frames(frames).await?; + Ok(stats) + } + + pub async fn accept_logs( + &self, + org_id: &str, + request: &ExportLogsServiceRequest, + ) -> Result { + let (frames, stats) = encode_logs(&self.inner.cfg, org_id, request)?; + self.commit_frames(frames).await?; + Ok(stats) + } + + pub async fn accept_metrics( + &self, + org_id: &str, + request: &ExportMetricsServiceRequest, + ) -> Result { + let (frames, stats) = encode_metrics(&self.inner.cfg, org_id, request)?; + self.commit_frames(frames).await?; + Ok(stats) + } + + async fn commit_frames(&self, frames: Vec) -> Result<(), PipelineError> { + if frames.is_empty() { + return Ok(()); + } + + for frame in frames { + let shard = (frame.routing_key as usize) % self.inner.shard_senders.len(); + let sender = self.inner.shard_senders[shard].clone(); + let permit = sender + .try_reserve_owned() + .map_err(|_| PipelineError::Backpressure("Telemetry queue is full"))?; + let queued_bytes = frame.payload.len() as u64; + self.reserve_org_queue_bytes(&frame.org_id, queued_bytes)?; + let (start, end) = self + .inner + .wal + .append(shard, &frame) + .await + .map_err(|error| { + self.release_org_queue_bytes(&frame.org_id, queued_bytes); + PipelineError::QueueUnavailable(error) + })?; + permit.send(QueuedFrame { + shard, + start, + end, + org_id: frame.org_id, + queued_bytes, + signal: frame.signal, + datasource: frame.datasource, + row_count: frame.row_count, + payload: frame.payload, + }); + } + + Ok(()) + } + + fn reserve_org_queue_bytes(&self, org_id: &str, bytes: u64) -> Result<(), PipelineError> { + if org_id.is_empty() || bytes == 0 { + return Ok(()); + } + + let counter = self + .inner + .org_queue_bytes + .entry(org_id.to_string()) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .clone(); + loop { + let current = counter.load(Ordering::Relaxed); + if current.saturating_add(bytes) > self.inner.cfg.org_queue_max_bytes { + counter!( + "ingest_org_throttled_total", + "org_id" => org_id.to_string(), + "reason" => "queue_bytes" + ) + .increment(1); + return Err(PipelineError::Throttled( + "Telemetry org queue byte limit exceeded", + )); + } + if counter + .compare_exchange( + current, + current + bytes, + Ordering::AcqRel, + Ordering::Relaxed, + ) + .is_ok() + { + gauge!( + "ingest_org_queue_bytes", + "org_id" => org_id.to_string() + ) + .set((current + bytes) as f64); + return Ok(()); + } + } + } + + fn release_org_queue_bytes(&self, org_id: &str, bytes: u64) { + release_org_queue_bytes(&self.inner.org_queue_bytes, org_id, bytes); + } + + async fn replay_committed_frames(&self) { + for shard in 0..self.inner.shard_senders.len() { + let frames = match self.inner.wal.replay(shard).await { + Ok(frames) => frames, + Err(error) => { + error!(shard, error = %error, "Failed to replay ingest WAL shard"); + continue; + } + }; + if frames.is_empty() { + continue; + } + info!( + shard, + frames = frames.len(), + "Replaying committed ingest WAL frames" + ); + for frame in frames { + add_org_queue_bytes( + &self.inner.org_queue_bytes, + &frame.org_id, + frame.queued_bytes, + ); + if self.inner.shard_senders[shard].send(frame).await.is_err() { + error!(shard, "Ingest WAL replay worker stopped"); + break; + } + } + } + } +} + +struct ShardedWal { + shards: Vec>, +} + +struct WalShard { + path: PathBuf, + cursor_path: PathBuf, + max_bytes: u64, + file: Mutex, +} + +impl ShardedWal { + fn open(cfg: &TinybirdConfig) -> Result { + let mut shards = Vec::with_capacity(cfg.wal_shards); + let max_bytes_per_shard = (cfg.queue_max_bytes / cfg.wal_shards as u64).max(1); + for shard in 0..cfg.wal_shards { + let path = cfg.queue_dir.join(format!("shard-{shard:03}.wal")); + let cursor_path = cfg.queue_dir.join(format!("shard-{shard:03}.cursor")); + let file = OpenOptions::new() + .create(true) + .read(true) + .append(true) + .open(&path) + .map_err(|error| format!("open ingest WAL {path:?}: {error}"))?; + shards.push(Arc::new(WalShard { + path, + cursor_path, + max_bytes: max_bytes_per_shard, + file: Mutex::new(file), + })); + } + Ok(Self { shards }) + } + + async fn append(&self, shard: usize, frame: &EncodedFrame) -> Result<(u64, u64), String> { + let shard_ref = Arc::clone( + self.shards + .get(shard) + .ok_or_else(|| format!("invalid WAL shard {shard}"))?, + ); + let encoded = encode_wal_frame(frame)?; + tokio::task::spawn_blocking(move || { + let mut file = shard_ref + .file + .lock() + .map_err(|_| "WAL shard mutex poisoned".to_string())?; + let start = file + .seek(SeekFrom::End(0)) + .map_err(|error| format!("seek WAL: {error}"))?; + if start.saturating_add(encoded.len() as u64) > shard_ref.max_bytes { + return Err("Telemetry WAL shard is full".to_string()); + } + file.write_all(&encoded) + .map_err(|error| format!("write WAL: {error}"))?; + file.sync_data() + .map_err(|error| format!("sync WAL: {error}"))?; + let end = start + encoded.len() as u64; + histogram!("ingest_wal_commit_bytes", "shard" => shard.to_string()) + .record(encoded.len() as f64); + Ok((start, end)) + }) + .await + .map_err(|error| format!("join WAL append: {error}"))? + } + + async fn replay(&self, shard: usize) -> Result, String> { + let shard_ref = Arc::clone( + self.shards + .get(shard) + .ok_or_else(|| format!("invalid WAL shard {shard}"))?, + ); + tokio::task::spawn_blocking(move || replay_shard(shard, &shard_ref)) + .await + .map_err(|error| format!("join WAL replay: {error}"))? + } + + async fn mark_exported(&self, shard: usize, offset: u64) -> Result<(), String> { + let shard_ref = Arc::clone( + self.shards + .get(shard) + .ok_or_else(|| format!("invalid WAL shard {shard}"))?, + ); + tokio::task::spawn_blocking(move || { + let mut file = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(&shard_ref.cursor_path) + .map_err(|error| format!("open WAL cursor: {error}"))?; + file.write_all(offset.to_string().as_bytes()) + .map_err(|error| format!("write WAL cursor: {error}"))?; + file.sync_data() + .map_err(|error| format!("sync WAL cursor: {error}")) + }) + .await + .map_err(|error| format!("join WAL cursor: {error}"))? + } +} + +fn read_cursor(path: &Path) -> u64 { + std::fs::read_to_string(path) + .ok() + .and_then(|raw| raw.trim().parse::().ok()) + .unwrap_or(0) +} + +fn replay_shard(shard: usize, shard_ref: &WalShard) -> Result, String> { + let mut file = OpenOptions::new() + .read(true) + .open(&shard_ref.path) + .map_err(|error| format!("open WAL replay: {error}"))?; + let mut offset = read_cursor(&shard_ref.cursor_path); + file.seek(SeekFrom::Start(offset)) + .map_err(|error| format!("seek WAL replay: {error}"))?; + + let mut frames = Vec::new(); + loop { + let start = offset; + let Some(frame) = read_wal_frame(&mut file, start)? else { + break; + }; + frames.push(QueuedFrame { + shard, + start, + end: frame.end, + org_id: frame.org_id, + queued_bytes: frame.payload.len() as u64, + signal: frame.signal, + datasource: frame.datasource, + row_count: frame.row_count, + payload: frame.payload, + }); + offset = frame.end; + } + Ok(frames) +} + +fn encode_wal_frame(frame: &EncodedFrame) -> Result, String> { + let datasource = frame.datasource.as_bytes(); + let org_id = frame.org_id.as_bytes(); + if datasource.len() > u16::MAX as usize { + return Err("datasource name too long".to_string()); + } + if org_id.len() > u16::MAX as usize { + return Err("org id too long".to_string()); + } + if frame.payload.len() > u32::MAX as usize { + return Err("WAL payload too large".to_string()); + } + let mut crc = Crc32::new(); + crc.update(&[signal_tag(frame.signal)]); + crc.update(org_id); + crc.update(datasource); + crc.update(&frame.payload); + let checksum = crc.finalize(); + + let mut out = Vec::with_capacity( + WAL_V2_HEADER_LEN + org_id.len() + datasource.len() + frame.payload.len(), + ); + out.extend_from_slice(WAL_MAGIC); + out.push(2); + out.push(signal_tag(frame.signal)); + out.extend_from_slice(&(datasource.len() as u16).to_le_bytes()); + out.extend_from_slice(&(org_id.len() as u16).to_le_bytes()); + out.extend_from_slice(&(frame.payload.len() as u32).to_le_bytes()); + out.extend_from_slice(&(frame.row_count as u32).to_le_bytes()); + out.extend_from_slice(&checksum.to_le_bytes()); + out.extend_from_slice(org_id); + out.extend_from_slice(datasource); + out.extend_from_slice(&frame.payload); + Ok(out) +} + +struct DecodedWalFrame { + end: u64, + org_id: String, + signal: TelemetrySignal, + datasource: String, + row_count: usize, + payload: Vec, +} + +fn read_wal_frame(file: &mut File, start: u64) -> Result, String> { + let mut prefix = [0u8; 6]; + let read = file + .read(&mut prefix) + .map_err(|error| format!("read WAL header: {error}"))?; + if read == 0 { + return Ok(None); + } + if read < prefix.len() { + warn!(offset = start, "Ignoring partial WAL header"); + return Ok(None); + } + if &prefix[0..4] != WAL_MAGIC { + return Err(format!("invalid WAL magic at offset {start}")); + } + let version = prefix[4]; + let signal = signal_from_tag(prefix[5]) + .ok_or_else(|| format!("invalid WAL signal tag {} at offset {start}", prefix[5]))?; + + let (org_id_len, datasource_len, payload_len, row_count, checksum, header_len) = match version { + 1 => { + let mut rest = [0u8; WAL_V1_HEADER_LEN - 6]; + file.read_exact(&mut rest) + .map_err(|error| format!("read WAL v1 header: {error}"))?; + ( + 0usize, + u16::from_le_bytes([rest[0], rest[1]]) as usize, + u32::from_le_bytes([rest[2], rest[3], rest[4], rest[5]]) as usize, + u32::from_le_bytes([rest[6], rest[7], rest[8], rest[9]]) as usize, + u32::from_le_bytes([rest[10], rest[11], rest[12], rest[13]]), + WAL_V1_HEADER_LEN, + ) + } + 2 => { + let mut rest = [0u8; WAL_V2_HEADER_LEN - 6]; + file.read_exact(&mut rest) + .map_err(|error| format!("read WAL v2 header: {error}"))?; + ( + u16::from_le_bytes([rest[2], rest[3]]) as usize, + u16::from_le_bytes([rest[0], rest[1]]) as usize, + u32::from_le_bytes([rest[4], rest[5], rest[6], rest[7]]) as usize, + u32::from_le_bytes([rest[8], rest[9], rest[10], rest[11]]) as usize, + u32::from_le_bytes([rest[12], rest[13], rest[14], rest[15]]), + WAL_V2_HEADER_LEN, + ) + } + _ => { + return Err(format!( + "unsupported WAL version {version} at offset {start}" + )) + } + }; + + let mut org_id = vec![0u8; org_id_len]; + if org_id_len > 0 { + file.read_exact(&mut org_id) + .map_err(|error| format!("read WAL org id: {error}"))?; + } + + let mut datasource = vec![0u8; datasource_len]; + file.read_exact(&mut datasource) + .map_err(|error| format!("read WAL datasource: {error}"))?; + let mut payload = vec![0u8; payload_len]; + file.read_exact(&mut payload) + .map_err(|error| format!("read WAL payload: {error}"))?; + + let mut crc = Crc32::new(); + crc.update(&[signal_tag(signal)]); + if version >= 2 { + crc.update(&org_id); + } + crc.update(&datasource); + crc.update(&payload); + if crc.finalize() != checksum { + return Err(format!("WAL checksum mismatch at offset {start}")); + } + + let org_id = + String::from_utf8(org_id).map_err(|error| format!("WAL org id utf8 error: {error}"))?; + let datasource = String::from_utf8(datasource) + .map_err(|error| format!("WAL datasource utf8 error: {error}"))?; + let end = + start + header_len as u64 + org_id_len as u64 + datasource_len as u64 + payload_len as u64; + Ok(Some(DecodedWalFrame { + end, + org_id, + signal, + datasource, + row_count: row_count.max(1), + payload, + })) +} + +fn signal_tag(signal: TelemetrySignal) -> u8 { + match signal { + TelemetrySignal::Traces => 1, + TelemetrySignal::Logs => 2, + TelemetrySignal::Metrics => 3, + } +} + +fn signal_from_tag(tag: u8) -> Option { + match tag { + 1 => Some(TelemetrySignal::Traces), + 2 => Some(TelemetrySignal::Logs), + 3 => Some(TelemetrySignal::Metrics), + _ => None, + } +} + +fn add_org_queue_bytes(counters: &Arc>>, org_id: &str, bytes: u64) { + if org_id.is_empty() || bytes == 0 { + return; + } + let counter = counters + .entry(org_id.to_string()) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .clone(); + let current = counter.fetch_add(bytes, Ordering::AcqRel) + bytes; + gauge!("ingest_org_queue_bytes", "org_id" => org_id.to_string()).set(current as f64); +} + +fn release_org_queue_bytes( + counters: &Arc>>, + org_id: &str, + bytes: u64, +) { + if org_id.is_empty() || bytes == 0 { + return; + } + if let Some(counter) = counters.get(org_id).map(|entry| Arc::clone(entry.value())) { + let mut current = counter.load(Ordering::Relaxed); + loop { + let next = current.saturating_sub(bytes); + match counter.compare_exchange(current, next, Ordering::AcqRel, Ordering::Relaxed) { + Ok(_) => { + gauge!("ingest_org_queue_bytes", "org_id" => org_id.to_string()) + .set(next as f64); + return; + } + Err(observed) => current = observed, + } + } + } +} + +struct ExportWorker { + shard: usize, + cfg: Arc, + wal: Arc, + org_queue_bytes: Arc>>, + http: Client, + receiver: mpsc::Receiver, +} + +impl ExportWorker { + async fn run(mut self) { + let mut buffer = Vec::new(); + while let Some(frame) = self.receiver.recv().await { + buffer.push(frame); + let deadline = sleep(self.cfg.batch_max_wait); + tokio::pin!(deadline); + loop { + tokio::select! { + maybe = self.receiver.recv(), if !batch_full(&self.cfg, &buffer) => { + match maybe { + Some(frame) => buffer.push(frame), + None => break, + } + } + _ = &mut deadline => break, + } + } + + let frames = std::mem::take(&mut buffer); + if let Err(error) = self.export_and_mark(frames).await { + error!(shard = self.shard, error = %error, "Tinybird export worker failed batch"); + } + } + } + + async fn export_and_mark(&self, frames: Vec) -> Result<(), String> { + if frames.is_empty() { + return Ok(()); + } + let mut by_datasource: BTreeMap> = BTreeMap::new(); + for frame in &frames { + by_datasource + .entry(frame.datasource.clone()) + .or_default() + .push(frame); + } + + let start = Instant::now(); + let first_signal = frames[0].signal; + let first_shard = frames[0].shard; + let first_offset = frames[0].start; + for (datasource, frames) in by_datasource { + let mut body = Vec::new(); + let mut rows = 0usize; + for frame in frames { + body.extend_from_slice(&frame.payload); + if !body.ends_with(b"\n") { + body.push(b'\n'); + } + rows += frame.row_count; + } + self.post_tinybird(&datasource, body, rows).await?; + } + + let end = frames.iter().map(|frame| frame.end).max().unwrap_or(0); + self.wal.mark_exported(self.shard, end).await?; + for frame in &frames { + release_org_queue_bytes(&self.org_queue_bytes, &frame.org_id, frame.queued_bytes); + } + histogram!("ingest_export_batch_duration_seconds", "shard" => self.shard.to_string()) + .record(start.elapsed().as_secs_f64()); + histogram!( + "ingest_wal_exported_bytes", + "signal" => format!("{first_signal:?}"), + "shard" => first_shard.to_string() + ) + .record(end.saturating_sub(first_offset) as f64); + Ok(()) + } + + async fn post_tinybird( + &self, + datasource: &str, + body: Vec, + rows: usize, + ) -> Result<(), String> { + let url = format!( + "{}/v0/events?name={}", + self.cfg.endpoint.trim_end_matches('/'), + datasource + ); + let compressed = gzip(body)?; + let mut attempt = 0u32; + loop { + let started = Instant::now(); + let response = self + .http + .post(&url) + .bearer_auth(&self.cfg.token) + .header(reqwest::header::CONTENT_TYPE, "application/x-ndjson") + .header(reqwest::header::CONTENT_ENCODING, "gzip") + .body(compressed.clone()) + .send() + .await; + + match response { + Ok(response) if response.status().is_success() => { + histogram!("ingest_tinybird_export_duration_seconds", "datasource" => datasource.to_string(), "status" => "2xx") + .record(started.elapsed().as_secs_f64()); + counter!("ingest_tinybird_export_rows_total", "datasource" => datasource.to_string()) + .increment(rows as u64); + return Ok(()); + } + Ok(response) if response.status().is_client_error() => { + let status = response.status().as_u16(); + let body = response.text().await.unwrap_or_default(); + counter!("ingest_tinybird_export_dropped_total", "datasource" => datasource.to_string(), "status" => status.to_string()) + .increment(rows as u64); + warn!(datasource, status, body = %body, rows, "Dropping non-retryable Tinybird batch"); + return Ok(()); + } + Ok(response) => { + let status = response.status().as_u16(); + counter!("ingest_tinybird_export_retries_total", "datasource" => datasource.to_string(), "status" => status.to_string()) + .increment(1); + warn!(datasource, status, attempt, "Retrying Tinybird batch"); + } + Err(error) => { + counter!("ingest_tinybird_export_retries_total", "datasource" => datasource.to_string(), "status" => "transport") + .increment(1); + warn!(datasource, attempt, error = %error, "Retrying Tinybird batch after transport error"); + } + } + + attempt = attempt.saturating_add(1); + let backoff_ms = 250u64.saturating_mul(2u64.saturating_pow(attempt.min(6))); + sleep(Duration::from_millis(backoff_ms.min(30_000))).await; + } + } +} + +fn batch_full(cfg: &TinybirdConfig, frames: &[QueuedFrame]) -> bool { + let rows: usize = frames.iter().map(|frame| frame.row_count).sum(); + let bytes: usize = frames.iter().map(|frame| frame.payload.len()).sum(); + rows >= cfg.batch_max_rows || bytes >= cfg.batch_max_bytes +} + +fn gzip(body: Vec) -> Result, String> { + let mut encoder = GzEncoder::new(Vec::new(), Compression::fast()); + encoder + .write_all(&body) + .map_err(|error| format!("gzip Tinybird body: {error}"))?; + encoder + .finish() + .map_err(|error| format!("finish gzip Tinybird body: {error}")) +} + +fn encode_traces( + cfg: &TinybirdConfig, + org_id: &str, + request: &ExportTraceServiceRequest, + policy: &SamplingPolicy, +) -> Result<(Vec, AcceptStats), PipelineError> { + let mut rows = Vec::with_capacity(count_trace_rows(request)); + let mut dropped = 0usize; + let mut routing_key = hash64(org_id); + let sample_ratio = policy.clamped_ratio(); + let sample_rate = 1.0 / sample_ratio; + + for resource_spans in &request.resource_spans { + let resource = resource_spans.resource.as_ref(); + let resource_attrs = resource + .map(|resource| attr_map(&resource.attributes)) + .unwrap_or_default(); + let service_name = resource_attrs + .get("service.name") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + + for scope_spans in &resource_spans.scope_spans { + let scope = scope_spans.scope.as_ref(); + let scope_attrs = scope + .map(|scope| attr_map(&scope.attributes)) + .unwrap_or_default(); + let scope_name = scope.map(|scope| scope.name.as_str()).unwrap_or(""); + let scope_version = scope.map(|scope| scope.version.as_str()).unwrap_or(""); + + for span in &scope_spans.spans { + let trace_id = bytes_hex(&span.trace_id); + if !trace_id.is_empty() { + routing_key = hash64(&trace_id); + } + + let should_keep = should_keep_trace(org_id, &trace_id, span, policy); + if !should_keep { + dropped += 1; + continue; + } + + let mut span_attrs = attr_map(&span.attributes); + if sample_ratio < 1.0 + && !span_attrs.contains_key("SampleRate") + && !span.trace_state.contains("th:") + { + span_attrs.insert( + "SampleRate".to_string(), + json!(format_sample_rate(sample_rate)), + ); + } + + let events_timestamp: Vec = span + .events + .iter() + .map(|event| json!(format_timestamp_nano(event.time_unix_nano))) + .collect(); + let events_name: Vec = + span.events.iter().map(|event| json!(event.name)).collect(); + let events_attributes: Vec = span + .events + .iter() + .map(|event| Value::Object(attr_map(&event.attributes))) + .collect(); + let links_trace_id: Vec = span + .links + .iter() + .map(|link| json!(bytes_hex(&link.trace_id))) + .collect(); + let links_span_id: Vec = span + .links + .iter() + .map(|link| json!(bytes_hex(&link.span_id))) + .collect(); + let links_trace_state: Vec = span + .links + .iter() + .map(|link| json!(link.trace_state)) + .collect(); + let links_attributes: Vec = span + .links + .iter() + .map(|link| Value::Object(attr_map(&link.attributes))) + .collect(); + + rows.push(json_line(json!({ + "start_time": format_timestamp_nano(span.start_time_unix_nano), + "trace_id": trace_id, + "span_id": bytes_hex(&span.span_id), + "parent_span_id": bytes_hex(&span.parent_span_id), + "trace_state": span.trace_state, + "span_name": span.name, + "span_kind": span_kind(span.kind), + "service_name": service_name, + "resource_schema_url": resource_spans.schema_url, + "resource_attributes": resource_attrs, + "scope_schema_url": scope_spans.schema_url, + "scope_name": scope_name, + "scope_version": scope_version, + "scope_attributes": scope_attrs, + "duration": span.end_time_unix_nano.saturating_sub(span.start_time_unix_nano), + "status_code": status_code(span.status.as_ref().map(|status| status.code).unwrap_or_default()), + "status_message": span.status.as_ref().map(|status| status.message.as_str()).unwrap_or(""), + "span_attributes": span_attrs, + "events_timestamp": events_timestamp, + "events_name": events_name, + "events_attributes": events_attributes, + "links_trace_id": links_trace_id, + "links_span_id": links_span_id, + "links_trace_state": links_trace_state, + "links_attributes": links_attributes + }))?); + } + } + } + + let stats = AcceptStats { + rows: rows.len(), + dropped, + }; + let frames = rows_to_frames( + org_id, + routing_key, + TelemetrySignal::Traces, + cfg.datasource_traces.clone(), + rows, + ); + Ok((frames, stats)) +} + +fn encode_logs( + cfg: &TinybirdConfig, + org_id: &str, + request: &ExportLogsServiceRequest, +) -> Result<(Vec, AcceptStats), PipelineError> { + let mut rows = Vec::with_capacity(count_log_rows(request)); + let mut routing_key = hash64(org_id); + + for resource_logs in &request.resource_logs { + let resource = resource_logs.resource.as_ref(); + let resource_attrs = resource + .map(|resource| attr_map(&resource.attributes)) + .unwrap_or_default(); + let service_name = resource_attrs + .get("service.name") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + + for scope_logs in &resource_logs.scope_logs { + let scope = scope_logs.scope.as_ref(); + let scope_attrs = scope + .map(|scope| attr_map(&scope.attributes)) + .unwrap_or_default(); + let scope_name = scope.map(|scope| scope.name.as_str()).unwrap_or(""); + let scope_version = scope.map(|scope| scope.version.as_str()).unwrap_or(""); + + for log in &scope_logs.log_records { + let trace_id = bytes_hex(&log.trace_id); + if !trace_id.is_empty() { + routing_key = hash64(&trace_id); + } + rows.push(encode_log_row( + log, + &service_name, + &resource_logs.schema_url, + &resource_attrs, + &scope_logs.schema_url, + scope_name, + scope_version, + &scope_attrs, + )?); + } + } + } + + let stats = AcceptStats { + rows: rows.len(), + dropped: 0, + }; + let frames = rows_to_frames( + org_id, + routing_key, + TelemetrySignal::Logs, + cfg.datasource_logs.clone(), + rows, + ); + Ok((frames, stats)) +} + +fn encode_log_row( + log: &LogRecord, + service_name: &str, + resource_schema_url: &str, + resource_attrs: &Map, + scope_schema_url: &str, + scope_name: &str, + scope_version: &str, + scope_attrs: &Map, +) -> Result, PipelineError> { + json_line(json!({ + "timestamp": format_timestamp_nano(if log.time_unix_nano != 0 { log.time_unix_nano } else { log.observed_time_unix_nano }), + "trace_id": bytes_hex(&log.trace_id), + "span_id": bytes_hex(&log.span_id), + "flags": log.flags, + "severity_text": if log.severity_text.is_empty() { severity_number_to_text(log.severity_number) } else { log.severity_text.as_str() }, + "severity_number": log.severity_number, + "service_name": service_name, + "body": log.body.as_ref().map(any_value_string).unwrap_or_default(), + "resource_schema_url": resource_schema_url, + "resource_attributes": resource_attrs, + "scope_schema_url": scope_schema_url, + "scope_name": scope_name, + "scope_version": scope_version, + "scope_attributes": scope_attrs, + "log_attributes": attr_map(&log.attributes) + })) +} + +fn encode_metrics( + cfg: &TinybirdConfig, + org_id: &str, + request: &ExportMetricsServiceRequest, +) -> Result<(Vec, AcceptStats), PipelineError> { + let mut by_datasource: BTreeMap>> = BTreeMap::new(); + let mut row_count = 0usize; + let mut routing_key = hash64(org_id); + + for resource_metrics in &request.resource_metrics { + let resource = resource_metrics.resource.as_ref(); + let resource_attrs = resource + .map(|resource| attr_map(&resource.attributes)) + .unwrap_or_default(); + let service_name = resource_attrs + .get("service.name") + .and_then(Value::as_str) + .unwrap_or("") + .to_string(); + + for scope_metrics in &resource_metrics.scope_metrics { + let scope = scope_metrics.scope.as_ref(); + let scope_attrs = scope + .map(|scope| attr_map(&scope.attributes)) + .unwrap_or_default(); + let scope_name = scope.map(|scope| scope.name.as_str()).unwrap_or(""); + let scope_version = scope.map(|scope| scope.version.as_str()).unwrap_or(""); + + for metric in &scope_metrics.metrics { + routing_key = hash64(&metric.name); + match metric.data.as_ref() { + Some(metric::Data::Gauge(gauge)) => { + for point in &gauge.data_points { + push_metric_number_row( + &mut by_datasource, + &cfg.datasource_metrics_gauge, + point, + metric, + &service_name, + &resource_metrics.schema_url, + &resource_attrs, + &scope_metrics.schema_url, + scope_name, + scope_version, + &scope_attrs, + None, + None, + )?; + row_count += 1; + } + } + Some(metric::Data::Sum(sum)) => { + for point in &sum.data_points { + push_metric_number_row( + &mut by_datasource, + &cfg.datasource_metrics_sum, + point, + metric, + &service_name, + &resource_metrics.schema_url, + &resource_attrs, + &scope_metrics.schema_url, + scope_name, + scope_version, + &scope_attrs, + Some(sum.aggregation_temporality), + Some(sum.is_monotonic), + )?; + row_count += 1; + } + } + Some(metric::Data::Histogram(histogram_data)) => { + for point in &histogram_data.data_points { + let row = metric_common_row( + metric, + &service_name, + &resource_metrics.schema_url, + &resource_attrs, + &scope_metrics.schema_url, + scope_name, + scope_version, + &scope_attrs, + &point.attributes, + point.start_time_unix_nano, + point.time_unix_nano, + point.flags, + &point.exemplars, + ); + push_json( + &mut by_datasource, + &cfg.datasource_metrics_histogram, + extend( + row, + json!({ + "count": point.count, + "sum": point.sum.unwrap_or(0.0), + "bucket_counts": point.bucket_counts, + "explicit_bounds": point.explicit_bounds, + "min": point.min, + "max": point.max, + "aggregation_temporality": histogram_data.aggregation_temporality + }), + ), + )?; + row_count += 1; + } + } + Some(metric::Data::ExponentialHistogram(exp)) => { + for point in &exp.data_points { + let row = metric_common_row( + metric, + &service_name, + &resource_metrics.schema_url, + &resource_attrs, + &scope_metrics.schema_url, + scope_name, + scope_version, + &scope_attrs, + &point.attributes, + point.start_time_unix_nano, + point.time_unix_nano, + point.flags, + &point.exemplars, + ); + let positive = point.positive.as_ref(); + let negative = point.negative.as_ref(); + push_json( + &mut by_datasource, + &cfg.datasource_metrics_exponential_histogram, + extend( + row, + json!({ + "count": point.count, + "sum": point.sum.unwrap_or(0.0), + "scale": point.scale, + "zero_count": point.zero_count, + "positive_offset": positive.map(|b| b.offset).unwrap_or(0), + "positive_bucket_counts": positive.map(|b| b.bucket_counts.clone()).unwrap_or_default(), + "negative_offset": negative.map(|b| b.offset).unwrap_or(0), + "negative_bucket_counts": negative.map(|b| b.bucket_counts.clone()).unwrap_or_default(), + "min": point.min, + "max": point.max, + "aggregation_temporality": exp.aggregation_temporality + }), + ), + )?; + row_count += 1; + } + } + Some(metric::Data::Summary(_)) | None => { + counter!("ingest_metrics_summary_dropped_total").increment(1); + } + } + } + } + } + + let mut frames = Vec::with_capacity(by_datasource.len()); + for (datasource, rows) in by_datasource { + frames.extend(rows_to_frames( + org_id, + routing_key, + TelemetrySignal::Metrics, + datasource, + rows, + )); + } + + Ok(( + frames, + AcceptStats { + rows: row_count, + dropped: 0, + }, + )) +} + +#[allow(clippy::too_many_arguments)] +fn push_metric_number_row( + by_datasource: &mut BTreeMap>>, + datasource: &str, + point: &NumberDataPoint, + metric: &opentelemetry_proto::tonic::metrics::v1::Metric, + service_name: &str, + resource_schema_url: &str, + resource_attrs: &Map, + scope_schema_url: &str, + scope_name: &str, + scope_version: &str, + scope_attrs: &Map, + aggregation_temporality: Option, + is_monotonic: Option, +) -> Result<(), PipelineError> { + let value = match point.value { + Some(number_data_point::Value::AsDouble(value)) => value, + Some(number_data_point::Value::AsInt(value)) => value as f64, + None => 0.0, + }; + let mut extra = json!({ "value": value }); + if let Some(aggregation_temporality) = aggregation_temporality { + extra["aggregation_temporality"] = json!(aggregation_temporality); + } + if let Some(is_monotonic) = is_monotonic { + extra["is_monotonic"] = json!(is_monotonic); + } + let row = metric_common_row( + metric, + service_name, + resource_schema_url, + resource_attrs, + scope_schema_url, + scope_name, + scope_version, + scope_attrs, + &point.attributes, + point.start_time_unix_nano, + point.time_unix_nano, + point.flags, + &point.exemplars, + ); + push_json(by_datasource, datasource, extend(row, extra)) +} + +#[allow(clippy::too_many_arguments)] +fn metric_common_row( + metric: &opentelemetry_proto::tonic::metrics::v1::Metric, + service_name: &str, + resource_schema_url: &str, + resource_attrs: &Map, + scope_schema_url: &str, + scope_name: &str, + scope_version: &str, + scope_attrs: &Map, + attributes: &[KeyValue], + start_time_unix_nano: u64, + time_unix_nano: u64, + flags: u32, + exemplars: &[Exemplar], +) -> Value { + let (trace_ids, span_ids, timestamps, values, filtered_attributes) = + encode_exemplars(exemplars); + json!({ + "resource_attributes": resource_attrs, + "resource_schema_url": resource_schema_url, + "scope_name": scope_name, + "scope_version": scope_version, + "scope_attributes": scope_attrs, + "scope_schema_url": scope_schema_url, + "service_name": service_name, + "metric_name": metric.name, + "metric_description": metric.description, + "metric_unit": metric.unit, + "metric_attributes": attr_map(attributes), + "start_timestamp": format_timestamp_nano(start_time_unix_nano), + "timestamp": format_timestamp_nano(time_unix_nano), + "flags": flags, + "exemplars_trace_id": trace_ids, + "exemplars_span_id": span_ids, + "exemplars_timestamp": timestamps, + "exemplars_value": values, + "exemplars_filtered_attributes": filtered_attributes + }) +} + +fn encode_exemplars( + exemplars: &[Exemplar], +) -> (Vec, Vec, Vec, Vec, Vec) { + let mut trace_ids = Vec::with_capacity(exemplars.len()); + let mut span_ids = Vec::with_capacity(exemplars.len()); + let mut timestamps = Vec::with_capacity(exemplars.len()); + let mut values = Vec::with_capacity(exemplars.len()); + let mut filtered_attributes = Vec::with_capacity(exemplars.len()); + for exemplar in exemplars { + trace_ids.push(bytes_hex(&exemplar.trace_id)); + span_ids.push(bytes_hex(&exemplar.span_id)); + timestamps.push(format_timestamp_nano(exemplar.time_unix_nano)); + values.push(match exemplar.value { + Some(opentelemetry_proto::tonic::metrics::v1::exemplar::Value::AsDouble(value)) => { + value + } + Some(opentelemetry_proto::tonic::metrics::v1::exemplar::Value::AsInt(value)) => { + value as f64 + } + None => 0.0, + }); + filtered_attributes.push(Value::Object(attr_map(&exemplar.filtered_attributes))); + } + (trace_ids, span_ids, timestamps, values, filtered_attributes) +} + +fn push_json( + by_datasource: &mut BTreeMap>>, + datasource: &str, + value: Value, +) -> Result<(), PipelineError> { + by_datasource + .entry(datasource.to_string()) + .or_default() + .push(json_line(value)?); + Ok(()) +} + +fn extend(mut base: Value, extra: Value) -> Value { + if let (Some(base), Some(extra)) = (base.as_object_mut(), extra.as_object()) { + for (key, value) in extra { + base.insert(key.clone(), value.clone()); + } + } + base +} + +fn rows_to_frames( + org_id: &str, + routing_key: u64, + signal: TelemetrySignal, + datasource: String, + rows: Vec>, +) -> Vec { + if rows.is_empty() { + return Vec::new(); + } + let mut payload = Vec::with_capacity(rows.iter().map(Vec::len).sum::() + rows.len()); + for row in &rows { + payload.extend_from_slice(row); + payload.push(b'\n'); + } + vec![EncodedFrame { + routing_key, + org_id: org_id.to_string(), + signal, + datasource, + row_count: rows.len(), + payload, + }] +} + +fn json_line(value: T) -> Result, PipelineError> { + serde_json::to_vec(&value).map_err(|error| PipelineError::Encode(error.to_string())) +} + +fn should_keep_trace(org_id: &str, trace_id: &str, span: &Span, policy: &SamplingPolicy) -> bool { + let ratio = policy.clamped_ratio(); + if ratio >= 1.0 { + return true; + } + if policy.always_keep_error_spans + && span + .status + .as_ref() + .is_some_and(|status| status.code == status::StatusCode::Error as i32) + { + return true; + } + if let Some(ms) = policy.always_keep_slow_spans_ms { + let duration_ms = span + .end_time_unix_nano + .saturating_sub(span.start_time_unix_nano) + / 1_000_000; + if duration_ms >= ms { + return true; + } + } + let key = if trace_id.is_empty() { + format!("{org_id}:{}", bytes_hex(&span.span_id)) + } else { + format!("{org_id}:{trace_id}") + }; + let threshold = (ratio * u64::MAX as f64) as u64; + hash64(&key) <= threshold +} + +fn format_sample_rate(sample_rate: f64) -> String { + if sample_rate.fract() == 0.0 { + format!("{}", sample_rate as u64) + } else { + format!("{sample_rate:.6}") + .trim_end_matches('0') + .trim_end_matches('.') + .to_string() + } +} + +fn attr_map(attributes: &[KeyValue]) -> Map { + let mut out = Map::with_capacity(attributes.len()); + for attribute in attributes { + out.insert( + attribute.key.clone(), + json!(attribute + .value + .as_ref() + .map(any_value_string) + .unwrap_or_default()), + ); + } + out +} + +fn any_value_string(value: &AnyValue) -> String { + match value.value.as_ref() { + Some(any_value::Value::StringValue(value)) => value.clone(), + Some(any_value::Value::BoolValue(value)) => value.to_string(), + Some(any_value::Value::IntValue(value)) => value.to_string(), + Some(any_value::Value::DoubleValue(value)) => value.to_string(), + Some(any_value::Value::BytesValue(value)) => bytes_hex(value), + Some(any_value::Value::ArrayValue(value)) => { + let values: Vec = value.values.iter().map(any_value_string).collect(); + serde_json::to_string(&values).unwrap_or_default() + } + Some(any_value::Value::KvlistValue(value)) => { + let attrs = attr_map(&value.values); + serde_json::to_string(&attrs).unwrap_or_default() + } + None => String::new(), + } +} + +fn span_kind(kind: i32) -> &'static str { + match kind { + x if x == span::SpanKind::Internal as i32 => "Internal", + x if x == span::SpanKind::Server as i32 => "Server", + x if x == span::SpanKind::Client as i32 => "Client", + x if x == span::SpanKind::Producer as i32 => "Producer", + x if x == span::SpanKind::Consumer as i32 => "Consumer", + _ => "Unspecified", + } +} + +fn status_code(code: i32) -> &'static str { + match code { + x if x == status::StatusCode::Ok as i32 => "Ok", + x if x == status::StatusCode::Error as i32 => "Error", + _ => "Unset", + } +} + +fn severity_number_to_text(n: i32) -> &'static str { + match n { + 1..=4 => "TRACE", + 5..=8 => "DEBUG", + 9..=12 => "INFO", + 13..=16 => "WARN", + 17..=20 => "ERROR", + 21..=24 => "FATAL", + _ => "", + } +} + +fn bytes_hex(bytes: &[u8]) -> String { + if bytes.is_empty() || bytes.iter().all(|byte| *byte == 0) { + return String::new(); + } + const HEX: &[u8; 16] = b"0123456789abcdef"; + let mut out = String::with_capacity(bytes.len() * 2); + for byte in bytes { + out.push(HEX[(byte >> 4) as usize] as char); + out.push(HEX[(byte & 0x0f) as usize] as char); + } + out +} + +fn format_timestamp_nano(unix_nano: u64) -> String { + if unix_nano == 0 { + return "1970-01-01 00:00:00.000000000".to_string(); + } + let secs = (unix_nano / 1_000_000_000) as i64; + let nanos = (unix_nano % 1_000_000_000) as u32; + let Some(dt) = chrono::DateTime::from_timestamp(secs, nanos) else { + return "1970-01-01 00:00:00.000000000".to_string(); + }; + dt.format("%Y-%m-%d %H:%M:%S.%f").to_string() +} + +fn hash64(value: &str) -> u64 { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + value.hash(&mut hasher); + hasher.finish() +} + +fn count_trace_rows(request: &ExportTraceServiceRequest) -> usize { + request + .resource_spans + .iter() + .flat_map(|rs| &rs.scope_spans) + .map(|ss| ss.spans.len()) + .sum() +} + +fn count_log_rows(request: &ExportLogsServiceRequest) -> usize { + request + .resource_logs + .iter() + .flat_map(|rl| &rl.scope_logs) + .map(|sl| sl.log_records.len()) + .sum() +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::body::Bytes; + use axum::extract::{Query, State}; + use axum::http::header::{AUTHORIZATION, CONTENT_ENCODING}; + use axum::http::{HeaderMap, StatusCode}; + use axum::routing::post; + use axum::Router; + use flate2::read::GzDecoder; + use opentelemetry_proto::tonic::common::v1::InstrumentationScope; + use opentelemetry_proto::tonic::metrics::v1::{ + exponential_histogram_data_point, metric, AggregationTemporality, ExponentialHistogram, + ExponentialHistogramDataPoint, Gauge, Histogram, HistogramDataPoint, Metric, Sum, + }; + use opentelemetry_proto::tonic::resource::v1::Resource; + use opentelemetry_proto::tonic::trace::v1::{ResourceSpans, ScopeSpans, Span, Status}; + use std::collections::HashMap; + + #[derive(Debug)] + struct FakeTinybirdImport { + datasource: String, + authorization: String, + content_encoding: String, + body: String, + } + + fn test_cfg() -> TinybirdConfig { + TinybirdConfig { + endpoint: "http://tinybird.test".to_string(), + token: "token".to_string(), + queue_dir: std::env::temp_dir(), + queue_max_bytes: 1024 * 1024, + org_queue_max_bytes: 1024 * 1024, + queue_channel_capacity: 10, + wal_shards: 2, + batch_max_rows: 100, + batch_max_bytes: 1024 * 1024, + batch_max_wait: Duration::from_millis(10), + export_concurrency_per_shard: 1, + datasource_traces: "traces".to_string(), + datasource_logs: "logs".to_string(), + datasource_metrics_sum: "metrics_sum".to_string(), + datasource_metrics_gauge: "metrics_gauge".to_string(), + datasource_metrics_histogram: "metrics_histogram".to_string(), + datasource_metrics_exponential_histogram: "metrics_exponential_histogram".to_string(), + } + } + + fn unique_test_dir(name: &str) -> PathBuf { + std::env::temp_dir().join(format!( + "maple-ingest-{name}-{}-{}", + std::process::id(), + current_time_nanos_for_test() + )) + } + + fn current_time_nanos_for_test() -> u128 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + } + + async fn fake_tinybird_import( + State(tx): State>, + Query(query): Query>, + headers: HeaderMap, + body: Bytes, + ) -> StatusCode { + let mut decoded = String::new(); + GzDecoder::new(&body[..]) + .read_to_string(&mut decoded) + .expect("fake Tinybird should receive gzip NDJSON"); + + let _ = tx.send(FakeTinybirdImport { + datasource: query.get("name").cloned().unwrap_or_default(), + authorization: headers + .get(AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(), + content_encoding: headers + .get(CONTENT_ENCODING) + .and_then(|value| value.to_str().ok()) + .unwrap_or_default() + .to_string(), + body: decoded, + }); + + StatusCode::OK + } + + async fn wait_for_cursor(cursor_path: PathBuf) -> u64 { + tokio::time::timeout(Duration::from_secs(2), async move { + loop { + if let Ok(cursor) = std::fs::read_to_string(&cursor_path) { + if let Ok(offset) = cursor.trim().parse::() { + if offset > 0 { + return offset; + } + } + } + sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("export worker should commit a WAL cursor after Tinybird success") + } + + fn string_kv(key: &str, value: &str) -> KeyValue { + KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.to_string())), + }), + } + } + + fn bool_kv(key: &str, value: bool) -> KeyValue { + KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::BoolValue(value)), + }), + } + } + + fn frame_row(frame: &EncodedFrame) -> Value { + let payload = std::str::from_utf8(&frame.payload).unwrap(); + serde_json::from_str(payload.trim()).unwrap() + } + + fn frame_rows_by_datasource(frames: Vec) -> BTreeMap { + frames + .into_iter() + .map(|frame| (frame.datasource.clone(), frame_row(&frame))) + .collect() + } + + #[test] + fn hex_empty_for_zero_ids() { + assert_eq!(bytes_hex(&[]), ""); + assert_eq!(bytes_hex(&[0; 8]), ""); + assert_eq!(bytes_hex(&[0xab, 0xcd]), "abcd"); + } + + #[test] + fn timestamp_has_nano_precision() { + assert_eq!( + format_timestamp_nano(1_700_000_000_123_456_789), + "2023-11-14 22:13:20.123456789" + ); + } + + #[test] + fn sampling_keeps_errors_even_when_ratio_low() { + let policy = SamplingPolicy { + trace_sample_ratio: 0.000001, + always_keep_error_spans: true, + always_keep_slow_spans_ms: None, + }; + let span = Span { + trace_id: vec![1; 16], + span_id: vec![2; 8], + status: Some(Status { + code: status::StatusCode::Error as i32, + message: String::new(), + }), + ..Default::default() + }; + assert!(should_keep_trace("org_1", "trace", &span, &policy)); + } + + #[test] + fn wal_round_trips_frame() { + let frame = EncodedFrame { + routing_key: 1, + org_id: "org_1".to_string(), + signal: TelemetrySignal::Traces, + datasource: "traces".to_string(), + row_count: 1, + payload: br#"{"a":1}"#.to_vec(), + }; + let encoded = encode_wal_frame(&frame).unwrap(); + let path = std::env::temp_dir().join(format!("maple-wal-test-{}.wal", std::process::id())); + std::fs::write(&path, encoded).unwrap(); + let mut file = File::open(&path).unwrap(); + let decoded = read_wal_frame(&mut file, 0).unwrap().unwrap(); + assert_eq!(decoded.signal, TelemetrySignal::Traces); + assert_eq!(decoded.org_id, "org_1"); + assert_eq!(decoded.datasource, "traces"); + assert_eq!(decoded.payload, br#"{"a":1}"#); + assert!(decoded.end > 0); + let _ = std::fs::remove_file(path); + } + + #[test] + fn trace_encoder_matches_tinybird_row_shape() { + let request = ExportTraceServiceRequest { + resource_spans: vec![ResourceSpans { + resource: Some(Resource { + attributes: vec![ + string_kv("service.name", "checkout"), + string_kv("maple_org_id", "org_1"), + ], + dropped_attributes_count: 0, + entity_refs: Vec::new(), + }), + scope_spans: vec![ScopeSpans { + scope: Some(InstrumentationScope { + name: "maple-sdk".to_string(), + version: "1.2.3".to_string(), + attributes: vec![string_kv("scope.attr", "scope-value")], + dropped_attributes_count: 0, + }), + spans: vec![Span { + trace_id: vec![0x11; 16], + span_id: vec![0x22; 8], + parent_span_id: vec![0x33; 8], + name: "POST /checkout".to_string(), + kind: span::SpanKind::Server as i32, + start_time_unix_nano: 1_700_000_000_000_000_000, + end_time_unix_nano: 1_700_000_000_250_000_000, + attributes: vec![string_kv("http.route", "/checkout")], + status: Some(Status { + code: status::StatusCode::Ok as i32, + message: "ok".to_string(), + }), + ..Default::default() + }], + schema_url: "https://scope.schema".to_string(), + }], + schema_url: "https://resource.schema".to_string(), + }], + }; + + let (frames, stats) = + encode_traces(&test_cfg(), "org_1", &request, &SamplingPolicy::default()).unwrap(); + assert_eq!(stats.rows, 1); + assert_eq!(stats.dropped, 0); + assert_eq!(frames.len(), 1); + assert_eq!(frames[0].datasource, "traces"); + + let row = frame_row(&frames[0]); + assert_eq!(row["start_time"], "2023-11-14 22:13:20.000000000"); + assert_eq!(row["trace_id"], "11111111111111111111111111111111"); + assert_eq!(row["span_id"], "2222222222222222"); + assert_eq!(row["parent_span_id"], "3333333333333333"); + assert_eq!(row["span_name"], "POST /checkout"); + assert_eq!(row["span_kind"], "Server"); + assert_eq!(row["service_name"], "checkout"); + assert_eq!(row["duration"], 250_000_000); + assert_eq!(row["status_code"], "Ok"); + assert_eq!(row["resource_attributes"]["maple_org_id"], "org_1"); + assert_eq!(row["span_attributes"]["http.route"], "/checkout"); + assert!(row["span_attributes"].get("SampleRate").is_none()); + } + + #[test] + fn log_encoder_matches_tinybird_row_shape() { + let log = LogRecord { + time_unix_nano: 1_700_000_001_123_456_789, + observed_time_unix_nano: 1_700_000_001_123_456_789, + severity_number: 17, + severity_text: String::new(), + body: Some(AnyValue { + value: Some(any_value::Value::StringValue("payment failed".to_string())), + }), + attributes: vec![bool_kv("retryable", true)], + trace_id: vec![0xaa; 16], + span_id: vec![0xbb; 8], + flags: 1, + ..Default::default() + }; + let request = ExportLogsServiceRequest { + resource_logs: vec![opentelemetry_proto::tonic::logs::v1::ResourceLogs { + resource: Some(Resource { + attributes: vec![string_kv("service.name", "worker")], + dropped_attributes_count: 0, + entity_refs: Vec::new(), + }), + scope_logs: vec![opentelemetry_proto::tonic::logs::v1::ScopeLogs { + scope: Some(InstrumentationScope { + name: "logger".to_string(), + version: "4.5.6".to_string(), + attributes: Vec::new(), + dropped_attributes_count: 0, + }), + log_records: vec![log], + schema_url: String::new(), + }], + schema_url: String::new(), + }], + }; + + let (frames, stats) = encode_logs(&test_cfg(), "org_1", &request).unwrap(); + assert_eq!(stats.rows, 1); + assert_eq!(frames[0].datasource, "logs"); + + let row = frame_row(&frames[0]); + assert_eq!(row["timestamp"], "2023-11-14 22:13:21.123456789"); + assert_eq!(row["severity_text"], "ERROR"); + assert_eq!(row["severity_number"], 17); + assert_eq!(row["service_name"], "worker"); + assert_eq!(row["body"], "payment failed"); + assert_eq!(row["log_attributes"]["retryable"], "true"); + assert_eq!(row["trace_id"], "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"); + assert_eq!(row["span_id"], "bbbbbbbbbbbbbbbb"); + } + + #[tokio::test] + async fn pipeline_e2e_exports_gzip_ndjson_to_fake_tinybird() { + let (tx, mut rx) = mpsc::unbounded_channel(); + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let addr = listener.local_addr().unwrap(); + let app = Router::new() + .route("/v0/events", post(fake_tinybird_import)) + .with_state(tx); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + let queue_dir = unique_test_dir("fake-tinybird"); + let mut cfg = test_cfg(); + cfg.endpoint = format!("http://{addr}"); + cfg.queue_dir = queue_dir.clone(); + cfg.wal_shards = 1; + cfg.batch_max_wait = Duration::from_millis(1); + + let pipeline = TelemetryPipeline::new( + cfg, + Client::builder() + .timeout(Duration::from_secs(5)) + .build() + .unwrap(), + ) + .await + .unwrap(); + + let request = ExportLogsServiceRequest { + resource_logs: vec![opentelemetry_proto::tonic::logs::v1::ResourceLogs { + resource: Some(Resource { + attributes: vec![string_kv("service.name", "fake-e2e")], + dropped_attributes_count: 0, + entity_refs: Vec::new(), + }), + scope_logs: vec![opentelemetry_proto::tonic::logs::v1::ScopeLogs { + scope: Some(InstrumentationScope { + name: "e2e-logger".to_string(), + version: "1.0.0".to_string(), + attributes: Vec::new(), + dropped_attributes_count: 0, + }), + log_records: vec![LogRecord { + time_unix_nano: 1_700_000_002_000_000_000, + observed_time_unix_nano: 1_700_000_002_000_000_000, + severity_number: 9, + severity_text: "INFO".to_string(), + body: Some(AnyValue { + value: Some(any_value::Value::StringValue( + "hello fake tinybird".to_string(), + )), + }), + attributes: vec![string_kv("component", "pipeline-e2e")], + trace_id: vec![0xcc; 16], + span_id: vec![0xdd; 8], + ..Default::default() + }], + schema_url: String::new(), + }], + schema_url: String::new(), + }], + }; + + let stats = pipeline.accept_logs("org_e2e", &request).await.unwrap(); + assert_eq!(stats.rows, 1); + + let import = tokio::time::timeout(Duration::from_secs(2), rx.recv()) + .await + .expect("fake Tinybird should receive an import") + .expect("fake Tinybird channel should stay open"); + + assert_eq!(import.datasource, "logs"); + assert_eq!(import.authorization, "Bearer token"); + assert_eq!(import.content_encoding, "gzip"); + + let row: Value = serde_json::from_str(import.body.trim()).unwrap(); + assert_eq!(row["timestamp"], "2023-11-14 22:13:22.000000000"); + assert_eq!(row["service_name"], "fake-e2e"); + assert_eq!(row["scope_name"], "e2e-logger"); + assert_eq!(row["body"], "hello fake tinybird"); + assert_eq!(row["log_attributes"]["component"], "pipeline-e2e"); + assert_eq!(row["trace_id"], "cccccccccccccccccccccccccccccccc"); + assert_eq!(row["span_id"], "dddddddddddddddd"); + + assert!(wait_for_cursor(queue_dir.join("shard-000.cursor")).await > 0); + let _ = std::fs::remove_dir_all(queue_dir); + } + + #[test] + fn metric_encoder_matches_all_tinybird_datasource_shapes() { + let base_point = NumberDataPoint { + attributes: vec![string_kv("route", "/checkout")], + start_time_unix_nano: 1_700_000_000_000_000_000, + time_unix_nano: 1_700_000_010_000_000_000, + flags: 0, + value: Some(number_data_point::Value::AsInt(42)), + ..Default::default() + }; + let request = ExportMetricsServiceRequest { + resource_metrics: vec![opentelemetry_proto::tonic::metrics::v1::ResourceMetrics { + resource: Some(Resource { + attributes: vec![string_kv("service.name", "api")], + dropped_attributes_count: 0, + entity_refs: Vec::new(), + }), + scope_metrics: vec![opentelemetry_proto::tonic::metrics::v1::ScopeMetrics { + scope: Some(InstrumentationScope { + name: "meter".to_string(), + version: "7.8.9".to_string(), + attributes: Vec::new(), + dropped_attributes_count: 0, + }), + metrics: vec![ + Metric { + name: "requests_total".to_string(), + description: "requests".to_string(), + unit: "1".to_string(), + data: Some(metric::Data::Sum(Sum { + data_points: vec![base_point.clone()], + aggregation_temporality: AggregationTemporality::Delta as i32, + is_monotonic: true, + })), + metadata: Vec::new(), + }, + Metric { + name: "cpu_ratio".to_string(), + description: "cpu".to_string(), + unit: "1".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + value: Some(number_data_point::Value::AsDouble(0.75)), + ..base_point.clone() + }], + })), + metadata: Vec::new(), + }, + Metric { + name: "request_duration_ms".to_string(), + description: "latency".to_string(), + unit: "ms".to_string(), + data: Some(metric::Data::Histogram(Histogram { + data_points: vec![HistogramDataPoint { + attributes: vec![string_kv("route", "/checkout")], + start_time_unix_nano: base_point.start_time_unix_nano, + time_unix_nano: base_point.time_unix_nano, + count: 3, + sum: Some(123.0), + bucket_counts: vec![1, 2], + explicit_bounds: vec![100.0], + min: Some(10.0), + max: Some(90.0), + ..Default::default() + }], + aggregation_temporality: AggregationTemporality::Delta as i32, + })), + metadata: Vec::new(), + }, + Metric { + name: "payload_bytes".to_string(), + description: "payload".to_string(), + unit: "By".to_string(), + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points: vec![ExponentialHistogramDataPoint { + attributes: vec![string_kv("route", "/checkout")], + start_time_unix_nano: base_point.start_time_unix_nano, + time_unix_nano: base_point.time_unix_nano, + count: 5, + sum: Some(500.0), + scale: 2, + zero_count: 1, + positive: Some(exponential_histogram_data_point::Buckets { + offset: -1, + bucket_counts: vec![2, 3], + }), + negative: Some(exponential_histogram_data_point::Buckets { + offset: 0, + bucket_counts: vec![1], + }), + min: Some(1.0), + max: Some(250.0), + ..Default::default() + }], + aggregation_temporality: AggregationTemporality::Cumulative as i32, + })), + metadata: Vec::new(), + }, + ], + schema_url: String::new(), + }], + schema_url: String::new(), + }], + }; + + let (frames, stats) = encode_metrics(&test_cfg(), "org_1", &request).unwrap(); + assert_eq!(stats.rows, 4); + let rows = frame_rows_by_datasource(frames); + assert_eq!( + rows.keys().cloned().collect::>(), + vec![ + "metrics_exponential_histogram", + "metrics_gauge", + "metrics_histogram", + "metrics_sum", + ] + ); + + let sum = &rows["metrics_sum"]; + assert_eq!(sum["metric_name"], "requests_total"); + assert_eq!(sum["value"], 42.0); + assert_eq!(sum["aggregation_temporality"], 1); + assert_eq!(sum["is_monotonic"], true); + assert_eq!(sum["metric_attributes"]["route"], "/checkout"); + + let gauge = &rows["metrics_gauge"]; + assert_eq!(gauge["metric_name"], "cpu_ratio"); + assert_eq!(gauge["value"], 0.75); + + let histogram = &rows["metrics_histogram"]; + assert_eq!(histogram["metric_name"], "request_duration_ms"); + assert_eq!(histogram["count"], 3); + assert_eq!(histogram["sum"], 123.0); + assert_eq!(histogram["bucket_counts"], json!([1, 2])); + assert_eq!(histogram["explicit_bounds"], json!([100.0])); + assert_eq!(histogram["min"], 10.0); + assert_eq!(histogram["max"], 90.0); + + let exp = &rows["metrics_exponential_histogram"]; + assert_eq!(exp["metric_name"], "payload_bytes"); + assert_eq!(exp["count"], 5); + assert_eq!(exp["sum"], 500.0); + assert_eq!(exp["scale"], 2); + assert_eq!(exp["zero_count"], 1); + assert_eq!(exp["positive_offset"], -1); + assert_eq!(exp["positive_bucket_counts"], json!([2, 3])); + assert_eq!(exp["negative_offset"], 0); + assert_eq!(exp["negative_bucket_counts"], json!([1])); + assert_eq!(exp["aggregation_temporality"], 2); + } +} diff --git a/packages/db/drizzle/0011_org_ingest_sampling_policies.sql b/packages/db/drizzle/0011_org_ingest_sampling_policies.sql new file mode 100644 index 0000000000..5fe869ffdf --- /dev/null +++ b/packages/db/drizzle/0011_org_ingest_sampling_policies.sql @@ -0,0 +1,8 @@ +CREATE TABLE `org_ingest_sampling_policies` ( + `org_id` text PRIMARY KEY NOT NULL, + `trace_sample_ratio` real DEFAULT 1 NOT NULL, + `always_keep_error_spans` integer DEFAULT 1 NOT NULL, + `always_keep_slow_spans_ms` integer, + `created_at` integer DEFAULT (unixepoch('subsec') * 1000) NOT NULL, + `updated_at` integer DEFAULT (unixepoch('subsec') * 1000) NOT NULL +); diff --git a/packages/db/src/schema/index.ts b/packages/db/src/schema/index.ts index a909dcb15a..eb5f096d4f 100644 --- a/packages/db/src/schema/index.ts +++ b/packages/db/src/schema/index.ts @@ -6,6 +6,7 @@ export * from "./digest" export * from "./errors" export * from "./oauth-connections" export * from "./org-ingest-keys" +export * from "./org-ingest-sampling-policies" export * from "./org-openrouter-settings" export * from "./org-clickhouse-settings" export * from "./scrape-targets" diff --git a/packages/db/src/schema/org-ingest-sampling-policies.ts b/packages/db/src/schema/org-ingest-sampling-policies.ts new file mode 100644 index 0000000000..6c916f0830 --- /dev/null +++ b/packages/db/src/schema/org-ingest-sampling-policies.ts @@ -0,0 +1,15 @@ +import { integer, real, sqliteTable, text } from "drizzle-orm/sqlite-core" +import { sql } from "drizzle-orm" + +export const orgIngestSamplingPolicies = sqliteTable("org_ingest_sampling_policies", { + orgId: text("org_id").primaryKey().notNull(), + traceSampleRatio: real("trace_sample_ratio").notNull().default(1), + alwaysKeepErrorSpans: integer("always_keep_error_spans", { mode: "boolean" }).notNull().default(true), + alwaysKeepSlowSpansMs: integer("always_keep_slow_spans_ms"), + createdAt: integer("created_at") + .notNull() + .default(sql`(unixepoch('subsec') * 1000)`), + updatedAt: integer("updated_at") + .notNull() + .default(sql`(unixepoch('subsec') * 1000)`), +}) From 550139185f9d628c44f1c24f4434f41a80867f7e Mon Sep 17 00:00:00 2001 From: Makisuo Date: Wed, 13 May 2026 01:53:28 +0200 Subject: [PATCH 2/8] Run ingest tests in CI --- .github/workflows/ingest-rust-tests.yml | 43 +++++++++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 .github/workflows/ingest-rust-tests.yml diff --git a/.github/workflows/ingest-rust-tests.yml b/.github/workflows/ingest-rust-tests.yml new file mode 100644 index 0000000000..d7d8741f78 --- /dev/null +++ b/.github/workflows/ingest-rust-tests.yml @@ -0,0 +1,43 @@ +name: Ingest Rust Tests + +on: + push: + branches: [main] + paths: + - "apps/ingest/**" + - "packages/db/drizzle/**" + - "packages/db/src/schema/**" + - ".github/workflows/ingest-rust-tests.yml" + pull_request: + paths: + - "apps/ingest/**" + - "packages/db/drizzle/**" + - "packages/db/src/schema/**" + - ".github/workflows/ingest-rust-tests.yml" + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + test: + name: Cargo test + runs-on: ubuntu-latest + defaults: + run: + working-directory: apps/ingest + steps: + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 + + - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable + + - uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2 + with: + workspaces: apps/ingest + + - name: Run ingest unit and fake Tinybird e2e tests + run: cargo test From e8a886772555a5eb09a06300c75d2b732a47ad1a Mon Sep 17 00:00:00 2001 From: Makisuo Date: Wed, 13 May 2026 02:08:07 +0200 Subject: [PATCH 3/8] Add ingest load test harness --- .github/workflows/ingest-load-tests.yml | 62 +++ apps/ingest/LOAD_TESTS.md | 35 ++ apps/ingest/src/bin/load_test.rs | 525 ++++++++++++++++++++++++ 3 files changed, 622 insertions(+) create mode 100644 .github/workflows/ingest-load-tests.yml create mode 100644 apps/ingest/LOAD_TESTS.md create mode 100644 apps/ingest/src/bin/load_test.rs diff --git a/.github/workflows/ingest-load-tests.yml b/.github/workflows/ingest-load-tests.yml new file mode 100644 index 0000000000..89937317d8 --- /dev/null +++ b/.github/workflows/ingest-load-tests.yml @@ -0,0 +1,62 @@ +name: Ingest Load Tests + +on: + workflow_dispatch: + inputs: + requests: + description: Total OTLP requests to send + required: true + default: "10000" + concurrency: + description: Concurrent request workers + required: true + default: "128" + batch_logs: + description: Log records per OTLP request + required: true + default: "10" + target_rps: + description: Optional target request rate. Leave empty for unbounded. + required: false + default: "" + min_rps: + description: Optional minimum accepted request/sec threshold + required: false + default: "" + max_rss_mb: + description: Optional max ingest process RSS threshold in MiB + required: false + default: "" + +permissions: + contents: read + +jobs: + load: + name: OTLP protobuf load test + runs-on: ubuntu-latest + defaults: + run: + working-directory: apps/ingest + steps: + - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 + + - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable + + - uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2 + with: + workspaces: apps/ingest + + - name: Build ingest and load-test binaries + run: cargo build --release --bin maple-ingest --bin load_test + + - name: Run load test + env: + LOAD_TEST_INGEST_BIN: target/release/maple-ingest + LOAD_TEST_REQUESTS: ${{ inputs.requests }} + LOAD_TEST_CONCURRENCY: ${{ inputs.concurrency }} + LOAD_TEST_BATCH_LOGS: ${{ inputs.batch_logs }} + LOAD_TEST_TARGET_RPS: ${{ inputs.target_rps }} + LOAD_TEST_MIN_RPS: ${{ inputs.min_rps }} + LOAD_TEST_MAX_RSS_MB: ${{ inputs.max_rss_mb }} + run: target/release/load_test diff --git a/apps/ingest/LOAD_TESTS.md b/apps/ingest/LOAD_TESTS.md new file mode 100644 index 0000000000..ebd164922d --- /dev/null +++ b/apps/ingest/LOAD_TESTS.md @@ -0,0 +1,35 @@ +# Ingest Load Tests + +The ingest load harness starts: + +- a fake Tinybird HTTP import endpoint, +- a real `maple-ingest` process configured with the static key store, +- an OTLP protobuf log traffic generator, +- a process sampler for ingest RSS and CPU via `ps`. + +Build and run a local smoke test: + +```sh +cargo build --release --bin maple-ingest --bin load_test +LOAD_TEST_REQUESTS=10000 \ +LOAD_TEST_CONCURRENCY=128 \ +LOAD_TEST_BATCH_LOGS=10 \ +target/release/load_test +``` + +Useful knobs: + +- `LOAD_TEST_REQUESTS`: total OTLP requests to send. +- `LOAD_TEST_CONCURRENCY`: concurrent request workers. +- `LOAD_TEST_BATCH_LOGS`: log records per OTLP request. +- `LOAD_TEST_TARGET_RPS`: optional request/sec pacing target. +- `LOAD_TEST_MIN_RPS`: optional failure threshold for accepted request/sec. +- `LOAD_TEST_MAX_RSS_MB`: optional failure threshold for max ingest RSS. +- `LOAD_TEST_INGEST_BIN`: path to the ingest binary when not using the default sibling binary. +- `LOAD_TEST_QUEUE_DIR`: WAL directory override. + +The harness prints JSON with request throughput, row throughput, p50/p95/p99 +latency, export catch-up time, max RSS, max CPU, average CPU, and exported rows. + +The GitHub Actions workflow `Ingest Load Tests` is manual (`workflow_dispatch`) +so large runs do not make normal PR CI noisy or flaky. diff --git a/apps/ingest/src/bin/load_test.rs b/apps/ingest/src/bin/load_test.rs new file mode 100644 index 0000000000..c53141f854 --- /dev/null +++ b/apps/ingest/src/bin/load_test.rs @@ -0,0 +1,525 @@ +use std::collections::HashMap; +use std::io::Read; +use std::path::PathBuf; +use std::process::{Child, Command, Stdio}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use axum::body::Bytes; +use axum::extract::{Query, State}; +use axum::http::StatusCode; +use axum::routing::{get, post}; +use axum::Router; +use flate2::read::GzDecoder; +use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; +use opentelemetry_proto::tonic::common::v1::{any_value, AnyValue, InstrumentationScope, KeyValue}; +use opentelemetry_proto::tonic::logs::v1::{LogRecord, ResourceLogs, ScopeLogs}; +use opentelemetry_proto::tonic::resource::v1::Resource; +use prost::Message; +use reqwest::Client; +use serde::Serialize; +use tokio::sync::mpsc; +use tokio::time::{sleep, timeout}; + +type DynError = Box; + +const INGEST_KEY: &str = "maple_pk_loadtest"; + +#[derive(Clone, Debug)] +struct LoadConfig { + requests: u64, + concurrency: usize, + batch_logs: usize, + target_rps: Option, + ingest_port: u16, + ingest_bin: PathBuf, + max_rss_mb: Option, + min_rps: Option, + queue_dir: PathBuf, +} + +#[derive(Clone)] +struct FakeTinybirdState { + imports: Arc, + rows: Arc, + bytes: Arc, +} + +#[derive(Clone, Copy, Debug, Default)] +struct ProcessSample { + rss_kib: u64, + cpu_percent: f64, +} + +#[derive(Debug, Default)] +struct MonitorSummary { + samples: u64, + max_rss_kib: u64, + max_cpu_percent: f64, + avg_cpu_percent: f64, +} + +#[derive(Debug, Serialize)] +struct LoadSummary { + requests: u64, + successes: u64, + failures: u64, + rows_sent: u64, + rows_exported: u64, + imports: u64, + duration_seconds: f64, + export_catchup_seconds: f64, + request_rps: f64, + row_rps: f64, + p50_ms: f64, + p95_ms: f64, + p99_ms: f64, + max_rss_mb: f64, + max_cpu_percent: f64, + avg_cpu_percent: f64, +} + +#[tokio::main(flavor = "multi_thread")] +async fn main() -> Result<(), DynError> { + let cfg = LoadConfig::from_env()?; + let fake_state = FakeTinybirdState::default(); + let fake_listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)).await?; + let fake_addr = fake_listener.local_addr()?; + let fake_app = Router::new() + .route("/health", get(|| async { "OK" })) + .route("/v0/events", post(fake_tinybird_import)) + .with_state(fake_state.clone()); + tokio::spawn(async move { + if let Err(error) = axum::serve(fake_listener, fake_app).await { + eprintln!("fake Tinybird server failed: {error}"); + } + }); + + let mut ingest = spawn_ingest(&cfg, &format!("http://{fake_addr}"))?; + wait_for_ingest_health(cfg.ingest_port).await?; + + let (sample_tx, sample_rx) = mpsc::unbounded_channel(); + let monitor_pid = ingest.id(); + let monitor = tokio::spawn(async move { monitor_process(monitor_pid, sample_tx).await }); + + let payload = build_logs_payload(cfg.batch_logs)?; + let started = Instant::now(); + let (successes, failures, mut latencies) = run_load(&cfg, payload).await?; + let request_duration = started.elapsed(); + + let catchup_started = Instant::now(); + let expected_rows = successes.saturating_mul(cfg.batch_logs as u64); + wait_for_exported_rows(&fake_state, expected_rows).await?; + let export_catchup = catchup_started.elapsed(); + + let _ = ingest.kill(); + let _ = ingest.wait(); + monitor.abort(); + + let monitor_summary = summarize_samples(sample_rx); + latencies.sort_unstable(); + let summary = LoadSummary { + requests: cfg.requests, + successes, + failures, + rows_sent: expected_rows, + rows_exported: fake_state.rows.load(Ordering::Relaxed), + imports: fake_state.imports.load(Ordering::Relaxed), + duration_seconds: request_duration.as_secs_f64(), + export_catchup_seconds: export_catchup.as_secs_f64(), + request_rps: successes as f64 / request_duration.as_secs_f64().max(0.001), + row_rps: expected_rows as f64 / request_duration.as_secs_f64().max(0.001), + p50_ms: percentile_ms(&latencies, 0.50), + p95_ms: percentile_ms(&latencies, 0.95), + p99_ms: percentile_ms(&latencies, 0.99), + max_rss_mb: monitor_summary.max_rss_kib as f64 / 1024.0, + max_cpu_percent: monitor_summary.max_cpu_percent, + avg_cpu_percent: monitor_summary.avg_cpu_percent, + }; + + println!("{}", serde_json::to_string_pretty(&summary)?); + enforce_thresholds(&cfg, &summary)?; + let _ = std::fs::remove_dir_all(&cfg.queue_dir); + Ok(()) +} + +impl LoadConfig { + fn from_env() -> Result { + let ingest_bin = std::env::var("LOAD_TEST_INGEST_BIN") + .map(PathBuf::from) + .unwrap_or_else(|_| { + std::env::current_exe() + .expect("current executable path") + .with_file_name(format!("maple-ingest{}", std::env::consts::EXE_SUFFIX)) + }); + + Ok(Self { + requests: env_u64("LOAD_TEST_REQUESTS", 10_000)?, + concurrency: env_usize("LOAD_TEST_CONCURRENCY", 128)?, + batch_logs: env_usize("LOAD_TEST_BATCH_LOGS", 10)?, + target_rps: env_optional_u64("LOAD_TEST_TARGET_RPS")?, + ingest_port: env_u16("LOAD_TEST_INGEST_PORT", 3475)?, + ingest_bin, + max_rss_mb: env_optional_u64("LOAD_TEST_MAX_RSS_MB")?, + min_rps: env_optional_f64("LOAD_TEST_MIN_RPS")?, + queue_dir: std::env::var("LOAD_TEST_QUEUE_DIR") + .map(PathBuf::from) + .unwrap_or_else(|_| unique_temp_dir("maple-ingest-load-wal")), + }) + } +} + +impl Default for FakeTinybirdState { + fn default() -> Self { + Self { + imports: Arc::new(AtomicU64::new(0)), + rows: Arc::new(AtomicU64::new(0)), + bytes: Arc::new(AtomicU64::new(0)), + } + } +} + +async fn fake_tinybird_import( + State(state): State, + Query(query): Query>, + body: Bytes, +) -> StatusCode { + if !query.contains_key("name") { + return StatusCode::BAD_REQUEST; + } + + let mut decoded = String::new(); + if GzDecoder::new(&body[..]) + .read_to_string(&mut decoded) + .is_err() + { + return StatusCode::BAD_REQUEST; + } + + state.imports.fetch_add(1, Ordering::Relaxed); + state.bytes.fetch_add(body.len() as u64, Ordering::Relaxed); + state.rows.fetch_add( + decoded.lines().filter(|line| !line.is_empty()).count() as u64, + Ordering::Relaxed, + ); + StatusCode::OK +} + +fn spawn_ingest(cfg: &LoadConfig, tinybird_host: &str) -> Result { + if !cfg.ingest_bin.exists() { + return Err(format!( + "ingest binary not found at {}. Run `cargo build --release --bin maple-ingest --bin load_test` first, or set LOAD_TEST_INGEST_BIN.", + cfg.ingest_bin.display() + ) + .into()); + } + + std::fs::create_dir_all(&cfg.queue_dir)?; + let child = Command::new(&cfg.ingest_bin) + .env("INGEST_PORT", cfg.ingest_port.to_string()) + .env("INGEST_WRITE_MODE", "tinybird") + .env("TINYBIRD_HOST", tinybird_host) + .env("TINYBIRD_TOKEN", "load-test-token") + .env("INGEST_KEY_STORE_BACKEND", "static") + .env("MAPLE_ORG_ID_OVERRIDE", "org_load_test") + .env("MAPLE_INGEST_KEY_LOOKUP_HMAC_KEY", "load-test-hmac-secret") + .env("INGEST_QUEUE_DIR", &cfg.queue_dir) + .env("INGEST_WAL_SHARDS", "4") + .env("INGEST_QUEUE_CHANNEL_CAPACITY", "100000") + .env("INGEST_QUEUE_MAX_BYTES", "1073741824") + .env("INGEST_ORG_QUEUE_MAX_BYTES", "1073741824") + .env("INGEST_BATCH_MAX_ROWS", "5000") + .env("INGEST_BATCH_MAX_BYTES", "4194304") + .env("INGEST_BATCH_MAX_WAIT_MS", "10") + .env("INGEST_ORG_MAX_IN_FLIGHT", "100000") + .env("RUST_LOG", "warn") + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn()?; + Ok(child) +} + +async fn wait_for_ingest_health(port: u16) -> Result<(), DynError> { + let client = Client::new(); + let url = format!("http://127.0.0.1:{port}/health"); + timeout(Duration::from_secs(10), async { + loop { + if let Ok(response) = client.get(&url).send().await { + if response.status().is_success() { + return Ok::<(), DynError>(()); + } + } + sleep(Duration::from_millis(50)).await; + } + }) + .await? +} + +async fn run_load(cfg: &LoadConfig, payload: Vec) -> Result<(u64, u64, Vec), DynError> { + let client = Client::builder() + .pool_max_idle_per_host(cfg.concurrency) + .timeout(Duration::from_secs(30)) + .build()?; + let url = format!("http://127.0.0.1:{}/v1/logs", cfg.ingest_port); + let next_request = Arc::new(AtomicU64::new(0)); + let successes = Arc::new(AtomicU64::new(0)); + let failures = Arc::new(AtomicU64::new(0)); + let latencies = Arc::new(Mutex::new(Vec::with_capacity(cfg.requests as usize))); + let started = Instant::now(); + + let mut tasks = Vec::with_capacity(cfg.concurrency); + for _ in 0..cfg.concurrency { + let client = client.clone(); + let url = url.clone(); + let payload = payload.clone(); + let next_request = Arc::clone(&next_request); + let successes = Arc::clone(&successes); + let failures = Arc::clone(&failures); + let latencies = Arc::clone(&latencies); + let requests = cfg.requests; + let target_rps = cfg.target_rps; + tasks.push(tokio::spawn(async move { + loop { + let request_index = next_request.fetch_add(1, Ordering::Relaxed); + if request_index >= requests { + break; + } + if let Some(target_rps) = target_rps { + pace_request(started, request_index, target_rps).await; + } + + let request_started = Instant::now(); + let result = client + .post(&url) + .header( + reqwest::header::AUTHORIZATION, + format!("Bearer {INGEST_KEY}"), + ) + .header(reqwest::header::CONTENT_TYPE, "application/x-protobuf") + .body(payload.clone()) + .send() + .await; + let elapsed = request_started.elapsed().as_micros(); + latencies.lock().expect("latencies mutex").push(elapsed); + + match result { + Ok(response) if response.status().is_success() => { + successes.fetch_add(1, Ordering::Relaxed); + } + _ => { + failures.fetch_add(1, Ordering::Relaxed); + } + } + } + })); + } + + for task in tasks { + task.await?; + } + + let latencies = Arc::try_unwrap(latencies) + .map_err(|_| "latencies still shared")? + .into_inner() + .map_err(|_| "latencies mutex poisoned")?; + Ok(( + successes.load(Ordering::Relaxed), + failures.load(Ordering::Relaxed), + latencies, + )) +} + +async fn pace_request(started: Instant, request_index: u64, target_rps: u64) { + if target_rps == 0 { + return; + } + let target_elapsed = Duration::from_secs_f64(request_index as f64 / target_rps as f64); + let elapsed = started.elapsed(); + if target_elapsed > elapsed { + sleep(target_elapsed - elapsed).await; + } +} + +async fn wait_for_exported_rows( + fake_state: &FakeTinybirdState, + expected_rows: u64, +) -> Result<(), DynError> { + timeout(Duration::from_secs(30), async { + loop { + let rows = fake_state.rows.load(Ordering::Relaxed); + if rows >= expected_rows { + return Ok::<(), DynError>(()); + } + sleep(Duration::from_millis(25)).await; + } + }) + .await? +} + +async fn monitor_process(pid: u32, tx: mpsc::UnboundedSender) { + loop { + if let Some(sample) = sample_process(pid) { + let _ = tx.send(sample); + } + sleep(Duration::from_millis(500)).await; + } +} + +fn sample_process(pid: u32) -> Option { + let output = Command::new("ps") + .args(["-p", &pid.to_string(), "-o", "rss=", "-o", "%cpu="]) + .output() + .ok()?; + if !output.status.success() { + return None; + } + let stdout = String::from_utf8(output.stdout).ok()?; + let mut parts = stdout.split_whitespace(); + let rss_kib = parts.next()?.parse::().ok()?; + let cpu_percent = parts.next()?.parse::().ok()?; + Some(ProcessSample { + rss_kib, + cpu_percent, + }) +} + +fn summarize_samples(mut rx: mpsc::UnboundedReceiver) -> MonitorSummary { + let mut summary = MonitorSummary::default(); + let mut cpu_total = 0.0; + while let Ok(sample) = rx.try_recv() { + summary.samples += 1; + summary.max_rss_kib = summary.max_rss_kib.max(sample.rss_kib); + summary.max_cpu_percent = summary.max_cpu_percent.max(sample.cpu_percent); + cpu_total += sample.cpu_percent; + } + if summary.samples > 0 { + summary.avg_cpu_percent = cpu_total / summary.samples as f64; + } + summary +} + +fn build_logs_payload(batch_logs: usize) -> Result, DynError> { + let records = (0..batch_logs) + .map(|index| LogRecord { + time_unix_nano: 1_700_000_000_000_000_000 + index as u64, + observed_time_unix_nano: 1_700_000_000_000_000_000 + index as u64, + severity_number: 9, + severity_text: "INFO".to_string(), + body: Some(AnyValue { + value: Some(any_value::Value::StringValue(format!( + "load test log {index}" + ))), + }), + attributes: vec![string_kv("load_test", "true")], + ..Default::default() + }) + .collect(); + + let request = ExportLogsServiceRequest { + resource_logs: vec![ResourceLogs { + resource: Some(Resource { + attributes: vec![string_kv("service.name", "ingest-load-test")], + dropped_attributes_count: 0, + entity_refs: Vec::new(), + }), + scope_logs: vec![ScopeLogs { + scope: Some(InstrumentationScope { + name: "maple-ingest-load-test".to_string(), + version: env!("CARGO_PKG_VERSION").to_string(), + attributes: Vec::new(), + dropped_attributes_count: 0, + }), + log_records: records, + schema_url: String::new(), + }], + schema_url: String::new(), + }], + }; + Ok(request.encode_to_vec()) +} + +fn string_kv(key: &str, value: &str) -> KeyValue { + KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.to_string())), + }), + } +} + +fn percentile_ms(latencies_us: &[u128], percentile: f64) -> f64 { + if latencies_us.is_empty() { + return 0.0; + } + let index = ((latencies_us.len() - 1) as f64 * percentile).round() as usize; + latencies_us[index.min(latencies_us.len() - 1)] as f64 / 1000.0 +} + +fn enforce_thresholds(cfg: &LoadConfig, summary: &LoadSummary) -> Result<(), DynError> { + if summary.failures > 0 { + return Err(format!("load test had {} failed requests", summary.failures).into()); + } + if let Some(min_rps) = cfg.min_rps { + if summary.request_rps < min_rps { + return Err(format!( + "request RPS {:.2} was below LOAD_TEST_MIN_RPS {:.2}", + summary.request_rps, min_rps + ) + .into()); + } + } + if let Some(max_rss_mb) = cfg.max_rss_mb { + if summary.max_rss_mb > max_rss_mb as f64 { + return Err(format!( + "max RSS {:.2} MiB exceeded LOAD_TEST_MAX_RSS_MB {} MiB", + summary.max_rss_mb, max_rss_mb + ) + .into()); + } + } + Ok(()) +} + +fn env_u64(name: &str, default: u64) -> Result { + Ok(match std::env::var(name) { + Ok(value) if !value.trim().is_empty() => value.trim().parse()?, + _ => default, + }) +} + +fn env_usize(name: &str, default: usize) -> Result { + Ok(match std::env::var(name) { + Ok(value) if !value.trim().is_empty() => value.trim().parse()?, + _ => default, + }) +} + +fn env_u16(name: &str, default: u16) -> Result { + Ok(match std::env::var(name) { + Ok(value) if !value.trim().is_empty() => value.trim().parse()?, + _ => default, + }) +} + +fn env_optional_u64(name: &str) -> Result, DynError> { + Ok(match std::env::var(name) { + Ok(value) if !value.trim().is_empty() => Some(value.trim().parse()?), + _ => None, + }) +} + +fn env_optional_f64(name: &str) -> Result, DynError> { + Ok(match std::env::var(name) { + Ok(value) if !value.trim().is_empty() => Some(value.trim().parse()?), + _ => None, + }) +} + +fn unique_temp_dir(prefix: &str) -> PathBuf { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_nanos()) + .unwrap_or(0); + std::env::temp_dir().join(format!("{prefix}-{}-{nanos}", std::process::id())) +} From 7f5dde24c887450a5c2010a77d366dc7b67bf0f3 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Wed, 13 May 2026 16:42:55 +0200 Subject: [PATCH 4/8] Add ingest Criterion benchmarks --- .github/workflows/ingest-rust-tests.yml | 50 ++++- apps/ingest/Cargo.lock | 255 +++++++++++++++++++++++- apps/ingest/Cargo.toml | 7 + apps/ingest/LOAD_TESTS.md | 10 + apps/ingest/benches/ingest_bench.rs | 247 +++++++++++++++++++++++ apps/ingest/src/lib.rs | 1 + apps/ingest/src/main.rs | 3 +- 7 files changed, 568 insertions(+), 5 deletions(-) create mode 100644 apps/ingest/benches/ingest_bench.rs create mode 100644 apps/ingest/src/lib.rs diff --git a/.github/workflows/ingest-rust-tests.yml b/.github/workflows/ingest-rust-tests.yml index d7d8741f78..641c606f12 100644 --- a/.github/workflows/ingest-rust-tests.yml +++ b/.github/workflows/ingest-rust-tests.yml @@ -22,10 +22,11 @@ concurrency: permissions: contents: read + pull-requests: write jobs: test: - name: Cargo test + name: Cargo test and bench runs-on: ubuntu-latest defaults: run: @@ -40,4 +41,49 @@ jobs: workspaces: apps/ingest - name: Run ingest unit and fake Tinybird e2e tests - run: cargo test + run: | + set -o pipefail + cargo test 2>&1 | tee "$RUNNER_TEMP/ingest-cargo-test.txt" + + - name: Run Criterion ingest benchmarks + run: | + set -o pipefail + cargo bench --bench ingest_bench -- --sample-size 10 --warm-up-time 1 --measurement-time 1 2>&1 | tee "$RUNNER_TEMP/ingest-criterion-bench.txt" + + - name: Comment PR with ingest test results + if: ${{ always() && github.event_name == 'pull_request' }} + env: + GH_TOKEN: ${{ github.token }} + PR_NUMBER: ${{ github.event.pull_request.number }} + run: | + test_output="$RUNNER_TEMP/ingest-cargo-test.txt" + bench_output="$RUNNER_TEMP/ingest-criterion-bench.txt" + body="$RUNNER_TEMP/ingest-results-comment.md" + + { + echo "## Ingest Rust Results" + echo + echo "**Commit:** \`${GITHUB_SHA}\`" + echo + echo "### cargo test" + echo + echo '```text' + if [ -f "$test_output" ]; then + tail -n 80 "$test_output" + else + echo "cargo test did not produce output." + fi + echo '```' + echo + echo "### Criterion benchmarks" + echo + echo '```text' + if [ -f "$bench_output" ]; then + tail -n 120 "$bench_output" + else + echo "criterion benchmark did not run." + fi + echo '```' + } > "$body" + + gh pr comment "$PR_NUMBER" --body-file "$body" diff --git a/apps/ingest/Cargo.lock b/apps/ingest/Cargo.lock index 55b11cc143..b8b331b653 100644 --- a/apps/ingest/Cargo.lock +++ b/apps/ingest/Cargo.lock @@ -38,6 +38,18 @@ dependencies = [ "libc", ] +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + [[package]] name = "anyhow" version = "1.0.101" @@ -185,6 +197,12 @@ version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.2.56" @@ -226,6 +244,58 @@ dependencies = [ "windows-link", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "cmake" version = "0.1.57" @@ -300,6 +370,44 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "criterion" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f" +dependencies = [ + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "futures", + "is-terminal", + "itertools 0.10.5", + "num-traits", + "once_cell", + "oorandom", + "plotters", + "rayon", + "regex", + "serde", + "serde_derive", + "serde_json", + "tinytemplate", + "tokio", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1" +dependencies = [ + "cast", + "itertools 0.10.5", +] + [[package]] name = "crossbeam-channel" version = "0.5.15" @@ -309,6 +417,16 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + [[package]] name = "crossbeam-epoch" version = "0.9.18" @@ -324,6 +442,12 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.1.7" @@ -468,6 +592,20 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" +[[package]] +name = "futures" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.31" @@ -475,6 +613,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -494,6 +633,12 @@ dependencies = [ "futures-util", ] +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + [[package]] name = "futures-macro" version = "0.3.31" @@ -588,6 +733,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.14.5" @@ -897,6 +1053,26 @@ dependencies = [ "serde", ] +[[package]] +name = "is-terminal" +version = "0.4.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" +dependencies = [ + "hermit-abi", + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.14.0" @@ -1002,6 +1178,7 @@ dependencies = [ "base64", "chrono", "crc32fast", + "criterion", "dashmap", "dotenvy", "flate2", @@ -1179,6 +1356,12 @@ version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "openssl-probe" version = "0.2.1" @@ -1329,6 +1512,34 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "portable-atomic" version = "1.13.1" @@ -1394,7 +1605,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools", + "itertools 0.14.0", "proc-macro2", "quote", "syn", @@ -1542,6 +1753,26 @@ dependencies = [ "bitflags", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -1551,6 +1782,18 @@ dependencies = [ "bitflags", ] +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + [[package]] name = "regex-automata" version = "0.4.14" @@ -2063,6 +2306,16 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tinyvec" version = "1.10.0" diff --git a/apps/ingest/Cargo.toml b/apps/ingest/Cargo.toml index 74f4c3f7d7..3b89f90a5c 100644 --- a/apps/ingest/Cargo.toml +++ b/apps/ingest/Cargo.toml @@ -50,3 +50,10 @@ uuid = { version = "1", features = ["v4"] } metrics = "0.24" metrics-exporter-prometheus = "0.16" moka = { version = "0.12", features = ["future"] } + +[dev-dependencies] +criterion = { version = "0.5", features = ["async_tokio"] } + +[[bench]] +name = "ingest_bench" +harness = false diff --git a/apps/ingest/LOAD_TESTS.md b/apps/ingest/LOAD_TESTS.md index ebd164922d..301079f90d 100644 --- a/apps/ingest/LOAD_TESTS.md +++ b/apps/ingest/LOAD_TESTS.md @@ -33,3 +33,13 @@ latency, export catch-up time, max RSS, max CPU, average CPU, and exported rows. The GitHub Actions workflow `Ingest Load Tests` is manual (`workflow_dispatch`) so large runs do not make normal PR CI noisy or flaky. + +For always-on microbenchmarks, the ingest CI runs Criterion: + +```sh +cargo bench --bench ingest_bench -- --sample-size 10 --warm-up-time 1 --measurement-time 1 +``` + +Those benchmarks measure WAL-acked native accepts for representative log and +trace OTLP batches. The `Ingest Rust Tests` workflow posts the captured cargo +test and Criterion output back to pull requests. diff --git a/apps/ingest/benches/ingest_bench.rs b/apps/ingest/benches/ingest_bench.rs new file mode 100644 index 0000000000..d2d4d2d8bf --- /dev/null +++ b/apps/ingest/benches/ingest_bench.rs @@ -0,0 +1,247 @@ +use std::collections::HashMap; +use std::io::Read; +use std::path::PathBuf; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use axum::body::Bytes; +use axum::extract::{Query, State}; +use axum::http::StatusCode; +use axum::routing::post; +use axum::Router; +use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use flate2::read::GzDecoder; +use maple_ingest::telemetry::{SamplingPolicy, TelemetryPipeline, TinybirdConfig}; +use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest; +use opentelemetry_proto::tonic::common::v1::{any_value, AnyValue, InstrumentationScope, KeyValue}; +use opentelemetry_proto::tonic::logs::v1::{LogRecord, ResourceLogs, ScopeLogs}; +use opentelemetry_proto::tonic::resource::v1::Resource; +use opentelemetry_proto::tonic::{ + collector::trace::v1::ExportTraceServiceRequest, + trace::v1::{span, ResourceSpans, ScopeSpans, Span}, +}; +use reqwest::Client; +use tokio::runtime::Runtime; + +#[derive(Clone, Default)] +struct FakeTinybirdState { + rows: Arc, +} + +struct BenchFixture { + pipeline: TelemetryPipeline, + logs: ExportLogsServiceRequest, + traces: ExportTraceServiceRequest, + queue_dir: PathBuf, +} + +fn bench_ingest_accept(c: &mut Criterion) { + let runtime = Runtime::new().expect("tokio runtime"); + let fixture = runtime.block_on(BenchFixture::new()); + let mut group = c.benchmark_group("ingest_accept"); + group.sample_size(10); + group.warm_up_time(Duration::from_millis(500)); + group.measurement_time(Duration::from_secs(2)); + + group.bench_function("logs_10_rows_wal_ack", |b| { + b.to_async(&runtime).iter(|| async { + black_box( + fixture + .pipeline + .accept_logs("org_bench", black_box(&fixture.logs)) + .await + .expect("accept logs"), + ); + }); + }); + + group.bench_function("traces_10_spans_wal_ack", |b| { + b.to_async(&runtime).iter(|| async { + black_box( + fixture + .pipeline + .accept_traces( + "org_bench", + black_box(&fixture.traces), + &SamplingPolicy::default(), + ) + .await + .expect("accept traces"), + ); + }); + }); + + group.finish(); + let _ = std::fs::remove_dir_all(&fixture.queue_dir); +} + +impl BenchFixture { + async fn new() -> Self { + let fake_state = FakeTinybirdState::default(); + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .expect("fake Tinybird listener"); + let addr = listener.local_addr().expect("fake Tinybird addr"); + let app = Router::new() + .route("/v0/events", post(fake_tinybird_import)) + .with_state(fake_state); + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + let queue_dir = unique_temp_dir("maple-ingest-bench-wal"); + let pipeline = TelemetryPipeline::new( + TinybirdConfig { + endpoint: format!("http://{addr}"), + token: "bench-token".to_string(), + queue_dir: queue_dir.clone(), + queue_max_bytes: 1024 * 1024 * 1024, + org_queue_max_bytes: 1024 * 1024 * 1024, + queue_channel_capacity: 100_000, + wal_shards: 4, + batch_max_rows: 5_000, + batch_max_bytes: 4 * 1024 * 1024, + batch_max_wait: Duration::from_millis(10), + export_concurrency_per_shard: 1, + datasource_traces: "traces".to_string(), + datasource_logs: "logs".to_string(), + datasource_metrics_sum: "metrics_sum".to_string(), + datasource_metrics_gauge: "metrics_gauge".to_string(), + datasource_metrics_histogram: "metrics_histogram".to_string(), + datasource_metrics_exponential_histogram: "metrics_exponential_histogram" + .to_string(), + }, + Client::builder() + .timeout(Duration::from_secs(5)) + .build() + .expect("http client"), + ) + .await + .expect("pipeline"); + + Self { + pipeline, + logs: build_logs(10), + traces: build_traces(10), + queue_dir, + } + } +} + +async fn fake_tinybird_import( + State(state): State, + Query(query): Query>, + body: Bytes, +) -> StatusCode { + if !query.contains_key("name") { + return StatusCode::BAD_REQUEST; + } + let mut decoded = String::new(); + if GzDecoder::new(&body[..]) + .read_to_string(&mut decoded) + .is_err() + { + return StatusCode::BAD_REQUEST; + } + state.rows.fetch_add( + decoded.lines().filter(|line| !line.is_empty()).count() as u64, + Ordering::Relaxed, + ); + StatusCode::OK +} + +fn build_logs(count: usize) -> ExportLogsServiceRequest { + let records = (0..count) + .map(|index| LogRecord { + time_unix_nano: 1_700_000_000_000_000_000 + index as u64, + observed_time_unix_nano: 1_700_000_000_000_000_000 + index as u64, + severity_number: 9, + severity_text: "INFO".to_string(), + body: Some(AnyValue { + value: Some(any_value::Value::StringValue(format!( + "benchmark log {index}" + ))), + }), + attributes: vec![string_kv("benchmark", "true")], + ..Default::default() + }) + .collect(); + + ExportLogsServiceRequest { + resource_logs: vec![ResourceLogs { + resource: Some(Resource { + attributes: vec![string_kv("service.name", "ingest-bench")], + dropped_attributes_count: 0, + entity_refs: Vec::new(), + }), + scope_logs: vec![ScopeLogs { + scope: Some(InstrumentationScope { + name: "criterion".to_string(), + version: "1".to_string(), + attributes: Vec::new(), + dropped_attributes_count: 0, + }), + log_records: records, + schema_url: String::new(), + }], + schema_url: String::new(), + }], + } +} + +fn build_traces(count: usize) -> ExportTraceServiceRequest { + let spans = (0..count) + .map(|index| Span { + trace_id: vec![index as u8 + 1; 16], + span_id: vec![index as u8 + 1; 8], + name: format!("benchmark span {index}"), + kind: span::SpanKind::Server as i32, + start_time_unix_nano: 1_700_000_000_000_000_000 + index as u64, + end_time_unix_nano: 1_700_000_000_010_000_000 + index as u64, + attributes: vec![string_kv("benchmark", "true")], + ..Default::default() + }) + .collect(); + + ExportTraceServiceRequest { + resource_spans: vec![ResourceSpans { + resource: Some(Resource { + attributes: vec![string_kv("service.name", "ingest-bench")], + dropped_attributes_count: 0, + entity_refs: Vec::new(), + }), + scope_spans: vec![ScopeSpans { + scope: Some(InstrumentationScope { + name: "criterion".to_string(), + version: "1".to_string(), + attributes: Vec::new(), + dropped_attributes_count: 0, + }), + spans, + schema_url: String::new(), + }], + schema_url: String::new(), + }], + } +} + +fn string_kv(key: &str, value: &str) -> KeyValue { + KeyValue { + key: key.to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(value.to_string())), + }), + } +} + +fn unique_temp_dir(prefix: &str) -> PathBuf { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_nanos()) + .unwrap_or(0); + std::env::temp_dir().join(format!("{prefix}-{}-{nanos}", std::process::id())) +} + +criterion_group!(benches, bench_ingest_accept); +criterion_main!(benches); diff --git a/apps/ingest/src/lib.rs b/apps/ingest/src/lib.rs new file mode 100644 index 0000000000..304af1e0f4 --- /dev/null +++ b/apps/ingest/src/lib.rs @@ -0,0 +1 @@ +pub mod telemetry; diff --git a/apps/ingest/src/main.rs b/apps/ingest/src/main.rs index de1a4e5800..45af5e9ab3 100644 --- a/apps/ingest/src/main.rs +++ b/apps/ingest/src/main.rs @@ -2,7 +2,6 @@ static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; mod autumn; -mod telemetry; use std::io::{Read, Write}; use std::path::PathBuf; @@ -29,6 +28,7 @@ use flate2::read::GzDecoder; use flate2::write::GzEncoder; use flate2::Compression; use hmac::{Hmac, Mac}; +use maple_ingest::telemetry::{PipelineError, SamplingPolicy, TelemetryPipeline, TinybirdConfig}; use metrics::{counter, gauge, histogram}; use moka::future::Cache; use opentelemetry::trace::TracerProvider as _; @@ -49,7 +49,6 @@ use reqwest::Client; use serde::{Deserialize, Serialize}; use serde_json::{Map as JsonMap, Value as JsonValue}; use sha2::Sha256; -use telemetry::{PipelineError, SamplingPolicy, TelemetryPipeline, TinybirdConfig}; use tower_http::cors::{Any, CorsLayer}; use tracing::Instrument; use tracing::{debug, error, info, warn, Span}; From 7b5f180c36957729d7916e36be11e418a831d50d Mon Sep 17 00:00:00 2001 From: Makisuo Date: Wed, 13 May 2026 17:12:52 +0200 Subject: [PATCH 5/8] Compare ingest load benchmarks to main --- .../format-ingest-benchmark-comment.py | 120 ++++++++++++++++++ .github/workflows/ingest-rust-tests.yml | 111 +++++++++++----- apps/ingest/LOAD_TESTS.md | 9 +- apps/ingest/src/bin/load_test.rs | 85 +++++++++++-- 4 files changed, 278 insertions(+), 47 deletions(-) create mode 100644 .github/scripts/format-ingest-benchmark-comment.py diff --git a/.github/scripts/format-ingest-benchmark-comment.py b/.github/scripts/format-ingest-benchmark-comment.py new file mode 100644 index 0000000000..85c69ab5e3 --- /dev/null +++ b/.github/scripts/format-ingest-benchmark-comment.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +import json +import os +from pathlib import Path +from typing import Optional + + +def load_json(path: Optional[str]): + if not path: + return None + p = Path(path) + if not p.exists(): + return None + raw = p.read_text() + start = raw.find("{") + end = raw.rfind("}") + if start == -1 or end == -1 or end < start: + return None + return json.loads(raw[start : end + 1]) + + +def tail(path: Optional[str], lines: int) -> str: + if not path or not Path(path).exists(): + return "output unavailable" + content = Path(path).read_text(errors="replace").splitlines() + return "\n".join(content[-lines:]) + + +def fmt(value, suffix="", precision=2): + if value is None: + return "n/a" + if isinstance(value, int): + return f"{value}{suffix}" + return f"{value:.{precision}f}{suffix}" + + +def delta(current, base, lower_is_better=False): + if current is None or base is None: + return "n/a" + if base == 0: + return "same" if current == 0 else "baseline was 0" + if abs(base) < 0.01: + absolute = current - base + good = absolute <= 0 if lower_is_better else absolute >= 0 + marker = "better" if good else "worse" + return f"{absolute:+.2f} absolute {marker}" + pct = (current - base) / base * 100.0 + good = pct <= 0 if lower_is_better else pct >= 0 + marker = "better" if good else "worse" + sign = "+" if pct >= 0 else "" + return f"{sign}{pct:.1f}% {marker}" + + +def metric_row(label, key, base, current, suffix="", lower_is_better=False): + return ( + f"| {label} | {fmt(base.get(key) if base else None, suffix)} " + f"| {fmt(current.get(key) if current else None, suffix)} " + f"| {delta(current.get(key) if current else None, base.get(key) if base else None, lower_is_better)} |" + ) + + +def main(): + base = load_json(os.environ.get("BASE_LOAD_JSON")) + current = load_json(os.environ.get("HEAD_LOAD_JSON")) + test_output = os.environ.get("TEST_OUTPUT") + out = Path(os.environ["COMMENT_OUTPUT"]) + + rows = [ + metric_row("Requests/sec", "request_rps", base, current), + metric_row("Rows/sec", "row_rps", base, current), + metric_row("p50 latency", "p50_ms", base, current, " ms", True), + metric_row("p95 latency", "p95_ms", base, current, " ms", True), + metric_row("p99 latency", "p99_ms", base, current, " ms", True), + metric_row("Export catch-up", "export_catchup_seconds", base, current, " s", True), + metric_row("Max RSS", "max_rss_mb", base, current, " MiB", True), + metric_row("Max CPU", "max_cpu_percent", base, current, "%", True), + metric_row("Avg CPU", "avg_cpu_percent", base, current, "%", True), + metric_row("Failures", "failures", base, current, "", True), + ] + + body = [ + "## Ingest Rust Test + Benchmark Results", + "", + f"**Commit:** `{os.environ.get('GITHUB_SHA', 'unknown')}`", + "", + "### Load Benchmark Compared To `main`", + "", + "| Metric | main | PR | Delta vs main |", + "| --- | ---: | ---: | ---: |", + *rows, + "", + "Benchmark setup: the same OTLP protobuf log workload is sent to both binaries. `main` runs the collector-forwarding path against a fake collector endpoint; the PR runs native Tinybird mode against a fake Tinybird endpoint. Both runs sample the ingest process RSS and CPU via `ps`.", + "", + "
PR load benchmark JSON", + "", + "```json", + json.dumps(current, indent=2) if current else "current benchmark output unavailable", + "```", + "", + "
", + "", + "
main load benchmark JSON", + "", + "```json", + json.dumps(base, indent=2) if base else "main benchmark output unavailable", + "```", + "", + "
", + "", + "### cargo test", + "", + "```text", + tail(test_output, 80), + "```", + ] + out.write_text("\n".join(body) + "\n") + + +if __name__ == "__main__": + main() diff --git a/.github/workflows/ingest-rust-tests.yml b/.github/workflows/ingest-rust-tests.yml index 641c606f12..9646229280 100644 --- a/.github/workflows/ingest-rust-tests.yml +++ b/.github/workflows/ingest-rust-tests.yml @@ -7,12 +7,14 @@ on: - "apps/ingest/**" - "packages/db/drizzle/**" - "packages/db/src/schema/**" + - ".github/scripts/format-ingest-benchmark-comment.py" - ".github/workflows/ingest-rust-tests.yml" pull_request: paths: - "apps/ingest/**" - "packages/db/drizzle/**" - "packages/db/src/schema/**" + - ".github/scripts/format-ingest-benchmark-comment.py" - ".github/workflows/ingest-rust-tests.yml" workflow_dispatch: @@ -26,7 +28,7 @@ permissions: jobs: test: - name: Cargo test and bench + name: Cargo test and load benchmark runs-on: ubuntu-latest defaults: run: @@ -42,48 +44,93 @@ jobs: - name: Run ingest unit and fake Tinybird e2e tests run: | + set +e set -o pipefail cargo test 2>&1 | tee "$RUNNER_TEMP/ingest-cargo-test.txt" + echo "${PIPESTATUS[0]}" > "$RUNNER_TEMP/ingest-cargo-test.exit" + exit 0 - - name: Run Criterion ingest benchmarks + - name: Build PR ingest and benchmark binaries run: | + set +e + cargo build --release --bin maple-ingest --bin load_test 2>&1 | tee "$RUNNER_TEMP/ingest-head-build.txt" + echo "${PIPESTATUS[0]}" > "$RUNNER_TEMP/ingest-head-build.exit" + mkdir -p "$RUNNER_TEMP/head" + if [ -x target/release/maple-ingest ]; then + cp target/release/maple-ingest "$RUNNER_TEMP/head/maple-ingest" + fi + if [ -x target/release/load_test ]; then + cp target/release/load_test "$RUNNER_TEMP/head/load_test" + fi + exit 0 + + - name: Run PR native ingest load benchmark + run: | + set +e + set -o pipefail + LOAD_TEST_INGEST_BIN="$RUNNER_TEMP/head/maple-ingest" \ + LOAD_TEST_INGEST_MODE=tinybird \ + LOAD_TEST_REQUESTS=2000 \ + LOAD_TEST_CONCURRENCY=64 \ + LOAD_TEST_BATCH_LOGS=10 \ + LOAD_TEST_INGEST_PORT=3475 \ + LOAD_TEST_QUEUE_DIR="$RUNNER_TEMP/head-wal" \ + "$RUNNER_TEMP/head/load_test" 2>&1 | tee "$RUNNER_TEMP/ingest-head-load.json" + echo "${PIPESTATUS[0]}" > "$RUNNER_TEMP/ingest-head-load.exit" + exit 0 + + - name: Build main ingest binary + if: ${{ github.event_name == 'pull_request' }} + run: | + set +e + git fetch origin "${{ github.base_ref }}" --depth=1 + git worktree add "$RUNNER_TEMP/base" FETCH_HEAD + ( + cd "$RUNNER_TEMP/base/apps/ingest" + cargo build --release --bin maple-ingest + ) 2>&1 | tee "$RUNNER_TEMP/ingest-base-build.txt" + echo "${PIPESTATUS[0]}" > "$RUNNER_TEMP/ingest-base-build.exit" + exit 0 + + - name: Run main collector-forwarding load benchmark + if: ${{ github.event_name == 'pull_request' }} + run: | + set +e set -o pipefail - cargo bench --bench ingest_bench -- --sample-size 10 --warm-up-time 1 --measurement-time 1 2>&1 | tee "$RUNNER_TEMP/ingest-criterion-bench.txt" + LOAD_TEST_INGEST_BIN="$RUNNER_TEMP/base/apps/ingest/target/release/maple-ingest" \ + LOAD_TEST_INGEST_MODE=forward \ + LOAD_TEST_REQUESTS=2000 \ + LOAD_TEST_CONCURRENCY=64 \ + LOAD_TEST_BATCH_LOGS=10 \ + LOAD_TEST_INGEST_PORT=3476 \ + LOAD_TEST_QUEUE_DIR="$RUNNER_TEMP/base-wal" \ + "$RUNNER_TEMP/head/load_test" 2>&1 | tee "$RUNNER_TEMP/ingest-base-load.json" + echo "${PIPESTATUS[0]}" > "$RUNNER_TEMP/ingest-base-load.exit" + exit 0 - name: Comment PR with ingest test results if: ${{ always() && github.event_name == 'pull_request' }} env: GH_TOKEN: ${{ github.token }} PR_NUMBER: ${{ github.event.pull_request.number }} + TEST_OUTPUT: ${{ runner.temp }}/ingest-cargo-test.txt + HEAD_LOAD_JSON: ${{ runner.temp }}/ingest-head-load.json + BASE_LOAD_JSON: ${{ runner.temp }}/ingest-base-load.json + COMMENT_OUTPUT: ${{ runner.temp }}/ingest-results-comment.md run: | - test_output="$RUNNER_TEMP/ingest-cargo-test.txt" - bench_output="$RUNNER_TEMP/ingest-criterion-bench.txt" - body="$RUNNER_TEMP/ingest-results-comment.md" + python "$GITHUB_WORKSPACE/.github/scripts/format-ingest-benchmark-comment.py" + gh pr comment "$PR_NUMBER" --body-file "$COMMENT_OUTPUT" - { - echo "## Ingest Rust Results" - echo - echo "**Commit:** \`${GITHUB_SHA}\`" - echo - echo "### cargo test" - echo - echo '```text' - if [ -f "$test_output" ]; then - tail -n 80 "$test_output" - else - echo "cargo test did not produce output." - fi - echo '```' - echo - echo "### Criterion benchmarks" - echo - echo '```text' - if [ -f "$bench_output" ]; then - tail -n 120 "$bench_output" - else - echo "criterion benchmark did not run." + - name: Fail if ingest checks failed + if: ${{ always() }} + run: | + failed=0 + for file in "$RUNNER_TEMP"/ingest-*.exit; do + [ -f "$file" ] || continue + code=$(cat "$file") + if [ "$code" != "0" ]; then + echo "::error file=$file::step failed with exit code $code" + failed=1 fi - echo '```' - } > "$body" - - gh pr comment "$PR_NUMBER" --body-file "$body" + done + exit "$failed" diff --git a/apps/ingest/LOAD_TESTS.md b/apps/ingest/LOAD_TESTS.md index 301079f90d..29ac7e01d5 100644 --- a/apps/ingest/LOAD_TESTS.md +++ b/apps/ingest/LOAD_TESTS.md @@ -34,12 +34,15 @@ latency, export catch-up time, max RSS, max CPU, average CPU, and exported rows. The GitHub Actions workflow `Ingest Load Tests` is manual (`workflow_dispatch`) so large runs do not make normal PR CI noisy or flaky. -For always-on microbenchmarks, the ingest CI runs Criterion: +For local microbenchmarks, the ingest crate also has Criterion benches: ```sh cargo bench --bench ingest_bench -- --sample-size 10 --warm-up-time 1 --measurement-time 1 ``` Those benchmarks measure WAL-acked native accepts for representative log and -trace OTLP batches. The `Ingest Rust Tests` workflow posts the captured cargo -test and Criterion output back to pull requests. +trace OTLP batches. + +On pull requests, the `Ingest Rust Tests` workflow runs a comparative load +benchmark against `main` and posts request throughput, row throughput, latency, +export catch-up, RSS, CPU, and failure deltas back to the PR. diff --git a/apps/ingest/src/bin/load_test.rs b/apps/ingest/src/bin/load_test.rs index c53141f854..587b4cbaa9 100644 --- a/apps/ingest/src/bin/load_test.rs +++ b/apps/ingest/src/bin/load_test.rs @@ -8,7 +8,8 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use axum::body::Bytes; use axum::extract::{Query, State}; -use axum::http::StatusCode; +use axum::http::header::CONTENT_ENCODING; +use axum::http::{HeaderMap, StatusCode}; use axum::routing::{get, post}; use axum::Router; use flate2::read::GzDecoder; @@ -28,6 +29,7 @@ const INGEST_KEY: &str = "maple_pk_loadtest"; #[derive(Clone, Debug)] struct LoadConfig { + ingest_mode: IngestMode, requests: u64, concurrency: usize, batch_logs: usize, @@ -39,6 +41,12 @@ struct LoadConfig { queue_dir: PathBuf, } +#[derive(Clone, Copy, Debug)] +enum IngestMode { + Tinybird, + Forward, +} + #[derive(Clone)] struct FakeTinybirdState { imports: Arc, @@ -62,6 +70,7 @@ struct MonitorSummary { #[derive(Debug, Serialize)] struct LoadSummary { + ingest_mode: &'static str, requests: u64, successes: u64, failures: u64, @@ -89,6 +98,7 @@ async fn main() -> Result<(), DynError> { let fake_app = Router::new() .route("/health", get(|| async { "OK" })) .route("/v0/events", post(fake_tinybird_import)) + .route("/v1/logs", post(fake_collector_logs)) .with_state(fake_state.clone()); tokio::spawn(async move { if let Err(error) = axum::serve(fake_listener, fake_app).await { @@ -120,6 +130,7 @@ async fn main() -> Result<(), DynError> { let monitor_summary = summarize_samples(sample_rx); latencies.sort_unstable(); let summary = LoadSummary { + ingest_mode: cfg.ingest_mode.as_str(), requests: cfg.requests, successes, failures, @@ -155,6 +166,7 @@ impl LoadConfig { }); Ok(Self { + ingest_mode: IngestMode::from_env()?, requests: env_u64("LOAD_TEST_REQUESTS", 10_000)?, concurrency: env_usize("LOAD_TEST_CONCURRENCY", 128)?, batch_logs: env_usize("LOAD_TEST_BATCH_LOGS", 10)?, @@ -170,6 +182,27 @@ impl LoadConfig { } } +impl IngestMode { + fn from_env() -> Result { + let raw = std::env::var("LOAD_TEST_INGEST_MODE") + .unwrap_or_else(|_| "tinybird".to_string()) + .trim() + .to_ascii_lowercase(); + match raw.as_str() { + "tinybird" | "native" => Ok(Self::Tinybird), + "forward" | "collector" => Ok(Self::Forward), + _ => Err("LOAD_TEST_INGEST_MODE must be tinybird or forward".into()), + } + } + + fn as_str(self) -> &'static str { + match self { + Self::Tinybird => "tinybird", + Self::Forward => "forward", + } + } +} + impl Default for FakeTinybirdState { fn default() -> Self { Self { @@ -183,19 +216,16 @@ impl Default for FakeTinybirdState { async fn fake_tinybird_import( State(state): State, Query(query): Query>, + headers: HeaderMap, body: Bytes, ) -> StatusCode { if !query.contains_key("name") { return StatusCode::BAD_REQUEST; } - let mut decoded = String::new(); - if GzDecoder::new(&body[..]) - .read_to_string(&mut decoded) - .is_err() - { + let Some(decoded) = decode_body(&headers, &body) else { return StatusCode::BAD_REQUEST; - } + }; state.imports.fetch_add(1, Ordering::Relaxed); state.bytes.fetch_add(body.len() as u64, Ordering::Relaxed); @@ -206,6 +236,36 @@ async fn fake_tinybird_import( StatusCode::OK } +async fn fake_collector_logs(State(state): State, body: Bytes) -> StatusCode { + let Ok(request) = ExportLogsServiceRequest::decode(&body[..]) else { + return StatusCode::BAD_REQUEST; + }; + let rows = request + .resource_logs + .iter() + .flat_map(|resource_logs| &resource_logs.scope_logs) + .map(|scope_logs| scope_logs.log_records.len() as u64) + .sum::(); + state.imports.fetch_add(1, Ordering::Relaxed); + state.bytes.fetch_add(body.len() as u64, Ordering::Relaxed); + state.rows.fetch_add(rows, Ordering::Relaxed); + StatusCode::OK +} + +fn decode_body(headers: &HeaderMap, body: &[u8]) -> Option { + let content_encoding = headers + .get(CONTENT_ENCODING) + .and_then(|value| value.to_str().ok()) + .map(|value| value.trim().to_ascii_lowercase()); + if content_encoding.as_deref() == Some("gzip") { + let mut decoded = String::new(); + GzDecoder::new(body).read_to_string(&mut decoded).ok()?; + Some(decoded) + } else { + String::from_utf8(body.to_vec()).ok() + } +} + fn spawn_ingest(cfg: &LoadConfig, tinybird_host: &str) -> Result { if !cfg.ingest_bin.exists() { return Err(format!( @@ -216,9 +276,11 @@ fn spawn_ingest(cfg: &LoadConfig, tinybird_host: &str) -> Result Result Result<(), DynError> { From 8180e4f902321242d64313847b8f73bcfc27152f Mon Sep 17 00:00:00 2001 From: Makisuo Date: Thu, 14 May 2026 10:42:51 +0200 Subject: [PATCH 6/8] Fix ingest CI benchmark: median-of-3, sticky comment, no fake regressions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The PR comment was comparing tinybird-mode-PR against forward-mode-main and labelling everything "-83% worse" — a meaningless cross-mode comparison. CPU samples were also stuck at 0% because the base run finished before the 500ms ps poller landed any sample. Changes: - Drop duplicate Rust (ingest) job from ci.yml; ingest-rust-tests.yml already covers cargo check/test/build via path filter - Run the load harness 3x on PR and 3x on base; report median + spread - Probe whether the base branch supports the head's LOAD_TEST_INGEST_MODE before running a baseline; skip baseline + render absolute-only table when it doesn't (so the introducing PR no longer shows fake deltas) - Sticky PR comment via gh api PATCH using a stable HTML marker; no more comment spam on every push - Strip ANSI escapes from embedded cargo test output (CARGO_TERM_COLOR=never) - Drop misleading Max/Avg CPU rows from the comment table; ps -o %cpu is lifetime-average on Linux and useless at sub-second scales - Add LOAD_TEST_REPORT_PATH to write JSON to a file, removing the fragile "find first { to last }" stdout parser - Wire the existing Criterion microbench into CI and embed bencher output in the comment Co-Authored-By: Claude Opus 4.7 (1M context) --- .../format-ingest-benchmark-comment.py | 267 +++++++++++++----- .github/workflows/ci.yml | 23 +- .github/workflows/ingest-rust-tests.yml | 163 +++++++++-- apps/ingest/LOAD_TESTS.md | 23 +- apps/ingest/src/bin/load_test.rs | 9 +- 5 files changed, 361 insertions(+), 124 deletions(-) diff --git a/.github/scripts/format-ingest-benchmark-comment.py b/.github/scripts/format-ingest-benchmark-comment.py index 85c69ab5e3..9ed3e1aaf7 100644 --- a/.github/scripts/format-ingest-benchmark-comment.py +++ b/.github/scripts/format-ingest-benchmark-comment.py @@ -1,118 +1,247 @@ #!/usr/bin/env python3 +"""Render the ingest load benchmark comment posted on PRs. + +Inputs (env): +- HEAD_LOAD_JSON comma-separated paths to per-iteration LoadSummary JSON files + for the PR. Multiple paths => median + min/max. +- BASE_LOAD_JSON same shape, for the base branch. Empty/missing => no + baseline available (rendered as absolute-only table). +- TEST_OUTPUT path to a cargo test log (tail is embedded). +- CRITERION_OUTPUT optional path to `cargo bench --output-format=bencher` output. +- COMMENT_OUTPUT path to write the rendered markdown. +- GITHUB_SHA optional commit SHA shown in the header. + +The first line of the output is a stable HTML marker so the workflow can find +and update the existing PR comment instead of appending new ones. +""" import json import os +import re from pathlib import Path +from statistics import median from typing import Optional +MARKER = "" +ANSI_RE = re.compile(r"\x1b\[[0-9;?]*[A-Za-z]") -def load_json(path: Optional[str]): - if not path: - return None - p = Path(path) - if not p.exists(): - return None - raw = p.read_text() - start = raw.find("{") - end = raw.rfind("}") - if start == -1 or end == -1 or end < start: + +def parse_paths(raw: Optional[str]) -> list[Path]: + if not raw: + return [] + return [Path(piece.strip()) for piece in raw.split(",") if piece.strip()] + + +def load_runs(raw: Optional[str]) -> list[dict]: + out = [] + for path in parse_paths(raw): + if not path.exists(): + continue + try: + out.append(json.loads(path.read_text())) + except json.JSONDecodeError: + continue + return out + + +def numeric(runs: list[dict], key: str) -> list[float]: + values: list[float] = [] + for run in runs: + value = run.get(key) + if isinstance(value, (int, float)): + values.append(float(value)) + return values + + +def aggregate(runs: list[dict], key: str) -> Optional[float]: + values = numeric(runs, key) + if not values: return None - return json.loads(raw[start : end + 1]) + return median(values) -def tail(path: Optional[str], lines: int) -> str: - if not path or not Path(path).exists(): - return "output unavailable" - content = Path(path).read_text(errors="replace").splitlines() - return "\n".join(content[-lines:]) +def value_range(runs: list[dict], key: str) -> Optional[tuple[float, float]]: + values = numeric(runs, key) + if not values: + return None + return min(values), max(values) -def fmt(value, suffix="", precision=2): +def fmt(value: Optional[float], suffix: str = "", precision: int = 2) -> str: if value is None: return "n/a" - if isinstance(value, int): - return f"{value}{suffix}" + if float(value).is_integer() and not suffix.endswith("ms") and suffix not in (" MiB", "%"): + return f"{int(value)}{suffix}" return f"{value:.{precision}f}{suffix}" -def delta(current, base, lower_is_better=False): +def fmt_range(span: Optional[tuple[float, float]], suffix: str, precision: int) -> str: + if span is None: + return "" + lo, hi = span + return f" (min {lo:.{precision}f}, max {hi:.{precision}f}{suffix})" + + +def delta(current: Optional[float], base: Optional[float], lower_is_better: bool) -> str: if current is None or base is None: return "n/a" if base == 0: - return "same" if current == 0 else "baseline was 0" + if current == 0: + return "same" + return "baseline 0" if abs(base) < 0.01: absolute = current - base good = absolute <= 0 if lower_is_better else absolute >= 0 - marker = "better" if good else "worse" - return f"{absolute:+.2f} absolute {marker}" + return f"{absolute:+.2f} {'better' if good else 'worse'}" pct = (current - base) / base * 100.0 good = pct <= 0 if lower_is_better else pct >= 0 - marker = "better" if good else "worse" sign = "+" if pct >= 0 else "" - return f"{sign}{pct:.1f}% {marker}" + return f"{sign}{pct:.1f}% {'better' if good else 'worse'}" -def metric_row(label, key, base, current, suffix="", lower_is_better=False): - return ( - f"| {label} | {fmt(base.get(key) if base else None, suffix)} " - f"| {fmt(current.get(key) if current else None, suffix)} " - f"| {delta(current.get(key) if current else None, base.get(key) if base else None, lower_is_better)} |" - ) +def tail(path: Optional[str], lines: int) -> str: + if not path or not Path(path).exists(): + return "output unavailable" + content = Path(path).read_text(errors="replace") + content = ANSI_RE.sub("", content).rstrip("\n") + return "\n".join(content.splitlines()[-lines:]) -def main(): - base = load_json(os.environ.get("BASE_LOAD_JSON")) - current = load_json(os.environ.get("HEAD_LOAD_JSON")) - test_output = os.environ.get("TEST_OUTPUT") - out = Path(os.environ["COMMENT_OUTPUT"]) +METRICS = [ + ("Requests/sec", "request_rps", "", 2, False), + ("Rows/sec", "row_rps", "", 2, False), + ("p50 latency", "p50_ms", " ms", 2, True), + ("p95 latency", "p95_ms", " ms", 2, True), + ("p99 latency", "p99_ms", " ms", 2, True), + ("Export catch-up", "export_catchup_seconds", " s", 3, True), + ("Max RSS", "max_rss_mb", " MiB", 2, True), + ("Failures", "failures", "", 0, True), +] + +def comparison_table(base: list[dict], head: list[dict]) -> list[str]: rows = [ - metric_row("Requests/sec", "request_rps", base, current), - metric_row("Rows/sec", "row_rps", base, current), - metric_row("p50 latency", "p50_ms", base, current, " ms", True), - metric_row("p95 latency", "p95_ms", base, current, " ms", True), - metric_row("p99 latency", "p99_ms", base, current, " ms", True), - metric_row("Export catch-up", "export_catchup_seconds", base, current, " s", True), - metric_row("Max RSS", "max_rss_mb", base, current, " MiB", True), - metric_row("Max CPU", "max_cpu_percent", base, current, "%", True), - metric_row("Avg CPU", "avg_cpu_percent", base, current, "%", True), - metric_row("Failures", "failures", base, current, "", True), + "| Metric | main (median) | PR (median) | Delta |", + "| --- | ---: | ---: | ---: |", ] + for label, key, suffix, precision, lower in METRICS: + b = aggregate(base, key) + h = aggregate(head, key) + rows.append( + f"| {label} " + f"| {fmt(b, suffix, precision)} " + f"| {fmt(h, suffix, precision)} " + f"| {delta(h, b, lower)} |" + ) + return rows + + +def absolute_table(head: list[dict]) -> list[str]: + rows = [ + "| Metric | PR (median) | Spread |", + "| --- | ---: | ---: |", + ] + for label, key, suffix, precision, _lower in METRICS: + h = aggregate(head, key) + spread = value_range(head, key) + rows.append( + f"| {label} " + f"| {fmt(h, suffix, precision)} " + f"| {fmt_range(spread, suffix, precision).strip() or '—'} |" + ) + return rows + + +def render_runs_json(runs: list[dict]) -> str: + if not runs: + return "no runs captured" + return json.dumps(runs, indent=2) + + +def main() -> None: + head_runs = load_runs(os.environ.get("HEAD_LOAD_JSON")) + base_runs = load_runs(os.environ.get("BASE_LOAD_JSON")) + test_output = os.environ.get("TEST_OUTPUT") + criterion_output = os.environ.get("CRITERION_OUTPUT") + out = Path(os.environ["COMMENT_OUTPUT"]) - body = [ - "## Ingest Rust Test + Benchmark Results", - "", - f"**Commit:** `{os.environ.get('GITHUB_SHA', 'unknown')}`", - "", - "### Load Benchmark Compared To `main`", - "", - "| Metric | main | PR | Delta vs main |", - "| --- | ---: | ---: | ---: |", - *rows, - "", - "Benchmark setup: the same OTLP protobuf log workload is sent to both binaries. `main` runs the collector-forwarding path against a fake collector endpoint; the PR runs native Tinybird mode against a fake Tinybird endpoint. Both runs sample the ingest process RSS and CPU via `ps`.", - "", - "
PR load benchmark JSON", - "", - "```json", - json.dumps(current, indent=2) if current else "current benchmark output unavailable", - "```", - "", - "
", - "", - "
main load benchmark JSON", + head_mode = head_runs[0].get("ingest_mode", "tinybird") if head_runs else "tinybird" + base_mode = base_runs[0].get("ingest_mode") if base_runs else None + same_mode = bool(base_runs) and base_mode == head_mode + + body: list[str] = [MARKER, "## Ingest Rust Test + Benchmark Results", ""] + body.append(f"**Commit:** `{os.environ.get('GITHUB_SHA', 'unknown')}`") + body.append("") + + if same_mode: + body += [ + f"### Load Benchmark — `{head_mode}` mode, median of {len(head_runs)} run(s) vs main", + "", + *comparison_table(base_runs, head_runs), + "", + "Same code path on both sides (same `LOAD_TEST_INGEST_MODE`), so the delta column is meaningful. Numbers come from `ubuntu-latest`, which is noisy — treat single-digit-percent deltas as noise.", + "", + ] + else: + head_label = f"`{head_mode}`" + if base_runs and not same_mode: + note = ( + f"Baseline ran in `{base_mode}` mode while PR ran in {head_label} mode — " + "cross-mode comparison is misleading, so only absolute PR numbers are shown." + ) + else: + note = ( + f"No baseline benchmark on main yet (this PR introduces {head_label} mode). " + "Once it lands, future PRs will get a real PR-vs-main delta automatically." + ) + body += [ + f"### Load Benchmark — `{head_mode}` mode, {len(head_runs)} run(s)", + "", + *absolute_table(head_runs), + "", + note, + "", + ] + + body += [ + "
PR load benchmark JSON (per-iteration)", "", "```json", - json.dumps(base, indent=2) if base else "main benchmark output unavailable", + render_runs_json(head_runs), "```", "", "
", "", + ] + if base_runs: + body += [ + "
main load benchmark JSON (per-iteration)", + "", + "```json", + render_runs_json(base_runs), + "```", + "", + "
", + "", + ] + + if criterion_output and Path(criterion_output).exists(): + body += [ + "### WAL-acked microbench (`cargo bench --bench ingest_bench`)", + "", + "```text", + tail(criterion_output, 60), + "```", + "", + ] + + body += [ "### cargo test", "", "```text", tail(test_output, 80), "```", ] + out.write_text("\n".join(body) + "\n") diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6f4559d821..02409ebbb8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,23 +26,6 @@ jobs: - run: bun turbo typecheck test build --filter='!@maple/ingest' - rust: - name: Rust (ingest) - runs-on: ubuntu-latest - defaults: - run: - working-directory: apps/ingest - steps: - - uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4 - - - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable - - - uses: Swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2 - with: - workspaces: apps/ingest - - - run: cargo check - - - run: cargo test - - - run: cargo build --release + # Note: the `apps/ingest` Rust crate is built and tested by + # `.github/workflows/ingest-rust-tests.yml`, which also runs the load + # benchmark and posts results to the PR. diff --git a/.github/workflows/ingest-rust-tests.yml b/.github/workflows/ingest-rust-tests.yml index 9646229280..da800938f2 100644 --- a/.github/workflows/ingest-rust-tests.yml +++ b/.github/workflows/ingest-rust-tests.yml @@ -26,6 +26,14 @@ permissions: contents: read pull-requests: write +env: + CARGO_TERM_COLOR: never + LOAD_TEST_REQUESTS: "2000" + LOAD_TEST_CONCURRENCY: "64" + LOAD_TEST_BATCH_LOGS: "10" + LOAD_BENCH_ITERATIONS: "3" + HEAD_INGEST_MODE: tinybird + jobs: test: name: Cargo test and load benchmark @@ -64,22 +72,44 @@ jobs: fi exit 0 - - name: Run PR native ingest load benchmark + - name: Run PR ingest load benchmark (median of N) + run: | + set +e + set -o pipefail + rm -f "$RUNNER_TEMP"/ingest-head-load-*.json + iterations="${LOAD_BENCH_ITERATIONS:-3}" + overall=0 + for i in $(seq 1 "$iterations"); do + report="$RUNNER_TEMP/ingest-head-load-$i.json" + queue_dir="$RUNNER_TEMP/head-wal-$i" + LOAD_TEST_INGEST_BIN="$RUNNER_TEMP/head/maple-ingest" \ + LOAD_TEST_INGEST_MODE="$HEAD_INGEST_MODE" \ + LOAD_TEST_REQUESTS="$LOAD_TEST_REQUESTS" \ + LOAD_TEST_CONCURRENCY="$LOAD_TEST_CONCURRENCY" \ + LOAD_TEST_BATCH_LOGS="$LOAD_TEST_BATCH_LOGS" \ + LOAD_TEST_INGEST_PORT="$((3475 + i))" \ + LOAD_TEST_QUEUE_DIR="$queue_dir" \ + LOAD_TEST_REPORT_PATH="$report" \ + "$RUNNER_TEMP/head/load_test" 2>&1 \ + | tee "$RUNNER_TEMP/ingest-head-load-$i.log" + code="${PIPESTATUS[0]}" + if [ "$code" != "0" ]; then + overall="$code" + fi + done + echo "$overall" > "$RUNNER_TEMP/ingest-head-load.exit" + exit 0 + + - name: Run Criterion WAL-ack microbench run: | set +e set -o pipefail - LOAD_TEST_INGEST_BIN="$RUNNER_TEMP/head/maple-ingest" \ - LOAD_TEST_INGEST_MODE=tinybird \ - LOAD_TEST_REQUESTS=2000 \ - LOAD_TEST_CONCURRENCY=64 \ - LOAD_TEST_BATCH_LOGS=10 \ - LOAD_TEST_INGEST_PORT=3475 \ - LOAD_TEST_QUEUE_DIR="$RUNNER_TEMP/head-wal" \ - "$RUNNER_TEMP/head/load_test" 2>&1 | tee "$RUNNER_TEMP/ingest-head-load.json" - echo "${PIPESTATUS[0]}" > "$RUNNER_TEMP/ingest-head-load.exit" + cargo bench --bench ingest_bench -- --output-format=bencher 2>&1 \ + | tee "$RUNNER_TEMP/ingest-criterion.txt" + echo "${PIPESTATUS[0]}" > "$RUNNER_TEMP/ingest-criterion.exit" exit 0 - - name: Build main ingest binary + - name: Build base ingest binary if: ${{ github.event_name == 'pull_request' }} run: | set +e @@ -87,39 +117,116 @@ jobs: git worktree add "$RUNNER_TEMP/base" FETCH_HEAD ( cd "$RUNNER_TEMP/base/apps/ingest" - cargo build --release --bin maple-ingest + cargo build --release --bin maple-ingest --bin load_test ) 2>&1 | tee "$RUNNER_TEMP/ingest-base-build.txt" echo "${PIPESTATUS[0]}" > "$RUNNER_TEMP/ingest-base-build.exit" + mkdir -p "$RUNNER_TEMP/base-bin" + if [ -x "$RUNNER_TEMP/base/apps/ingest/target/release/maple-ingest" ]; then + cp "$RUNNER_TEMP/base/apps/ingest/target/release/maple-ingest" "$RUNNER_TEMP/base-bin/maple-ingest" + fi + if [ -x "$RUNNER_TEMP/base/apps/ingest/target/release/load_test" ]; then + cp "$RUNNER_TEMP/base/apps/ingest/target/release/load_test" "$RUNNER_TEMP/base-bin/load_test" + fi exit 0 - - name: Run main collector-forwarding load benchmark + - name: Probe whether base supports head ingest mode + if: ${{ github.event_name == 'pull_request' }} + run: | + set +e + echo "0" > "$RUNNER_TEMP/base-supports-mode" + if [ ! -x "$RUNNER_TEMP/base-bin/maple-ingest" ] || [ ! -x "$RUNNER_TEMP/head/load_test" ]; then + echo "base binaries missing — skipping probe" + exit 0 + fi + probe_dir="$RUNNER_TEMP/base-probe-wal" + rm -rf "$probe_dir" + LOAD_TEST_INGEST_BIN="$RUNNER_TEMP/base-bin/maple-ingest" \ + LOAD_TEST_INGEST_MODE="$HEAD_INGEST_MODE" \ + LOAD_TEST_REQUESTS=10 \ + LOAD_TEST_CONCURRENCY=4 \ + LOAD_TEST_BATCH_LOGS=1 \ + LOAD_TEST_INGEST_PORT=3490 \ + LOAD_TEST_QUEUE_DIR="$probe_dir" \ + LOAD_TEST_REPORT_PATH="$RUNNER_TEMP/base-probe.json" \ + timeout 60 "$RUNNER_TEMP/head/load_test" >/dev/null 2>&1 + if [ -s "$RUNNER_TEMP/base-probe.json" ]; then + echo "1" > "$RUNNER_TEMP/base-supports-mode" + fi + rm -rf "$probe_dir" + exit 0 + + - name: Run base ingest load benchmark (median of N) if: ${{ github.event_name == 'pull_request' }} run: | set +e set -o pipefail - LOAD_TEST_INGEST_BIN="$RUNNER_TEMP/base/apps/ingest/target/release/maple-ingest" \ - LOAD_TEST_INGEST_MODE=forward \ - LOAD_TEST_REQUESTS=2000 \ - LOAD_TEST_CONCURRENCY=64 \ - LOAD_TEST_BATCH_LOGS=10 \ - LOAD_TEST_INGEST_PORT=3476 \ - LOAD_TEST_QUEUE_DIR="$RUNNER_TEMP/base-wal" \ - "$RUNNER_TEMP/head/load_test" 2>&1 | tee "$RUNNER_TEMP/ingest-base-load.json" - echo "${PIPESTATUS[0]}" > "$RUNNER_TEMP/ingest-base-load.exit" + rm -f "$RUNNER_TEMP"/ingest-base-load-*.json + if [ ! -f "$RUNNER_TEMP/base-supports-mode" ] || [ "$(cat "$RUNNER_TEMP/base-supports-mode")" != "1" ]; then + echo "base does not support $HEAD_INGEST_MODE mode — skipping baseline benchmark" + echo "0" > "$RUNNER_TEMP/ingest-base-load.exit" + exit 0 + fi + iterations="${LOAD_BENCH_ITERATIONS:-3}" + overall=0 + for i in $(seq 1 "$iterations"); do + report="$RUNNER_TEMP/ingest-base-load-$i.json" + queue_dir="$RUNNER_TEMP/base-wal-$i" + LOAD_TEST_INGEST_BIN="$RUNNER_TEMP/base-bin/maple-ingest" \ + LOAD_TEST_INGEST_MODE="$HEAD_INGEST_MODE" \ + LOAD_TEST_REQUESTS="$LOAD_TEST_REQUESTS" \ + LOAD_TEST_CONCURRENCY="$LOAD_TEST_CONCURRENCY" \ + LOAD_TEST_BATCH_LOGS="$LOAD_TEST_BATCH_LOGS" \ + LOAD_TEST_INGEST_PORT="$((3500 + i))" \ + LOAD_TEST_QUEUE_DIR="$queue_dir" \ + LOAD_TEST_REPORT_PATH="$report" \ + "$RUNNER_TEMP/head/load_test" 2>&1 \ + | tee "$RUNNER_TEMP/ingest-base-load-$i.log" + code="${PIPESTATUS[0]}" + if [ "$code" != "0" ]; then + overall="$code" + fi + done + echo "$overall" > "$RUNNER_TEMP/ingest-base-load.exit" exit 0 - - name: Comment PR with ingest test results + - name: Render PR comment body + if: ${{ always() && github.event_name == 'pull_request' }} + env: + TEST_OUTPUT: ${{ runner.temp }}/ingest-cargo-test.txt + CRITERION_OUTPUT: ${{ runner.temp }}/ingest-criterion.txt + COMMENT_OUTPUT: ${{ runner.temp }}/ingest-results-comment.md + run: | + set -e + head_json=$(ls "$RUNNER_TEMP"/ingest-head-load-*.json 2>/dev/null | paste -sd, -) + base_json=$(ls "$RUNNER_TEMP"/ingest-base-load-*.json 2>/dev/null | paste -sd, -) + HEAD_LOAD_JSON="$head_json" \ + BASE_LOAD_JSON="$base_json" \ + python "$GITHUB_WORKSPACE/.github/scripts/format-ingest-benchmark-comment.py" + + - name: Post or update PR comment (sticky) if: ${{ always() && github.event_name == 'pull_request' }} env: GH_TOKEN: ${{ github.token }} PR_NUMBER: ${{ github.event.pull_request.number }} - TEST_OUTPUT: ${{ runner.temp }}/ingest-cargo-test.txt - HEAD_LOAD_JSON: ${{ runner.temp }}/ingest-head-load.json - BASE_LOAD_JSON: ${{ runner.temp }}/ingest-base-load.json COMMENT_OUTPUT: ${{ runner.temp }}/ingest-results-comment.md + REPO: ${{ github.repository }} run: | - python "$GITHUB_WORKSPACE/.github/scripts/format-ingest-benchmark-comment.py" - gh pr comment "$PR_NUMBER" --body-file "$COMMENT_OUTPUT" + set -e + if [ ! -s "$COMMENT_OUTPUT" ]; then + echo "no comment body rendered, skipping" + exit 0 + fi + existing=$(gh api "repos/$REPO/issues/$PR_NUMBER/comments" --paginate \ + --jq '[.[] | select(.body | startswith(""))][0].id') + if [ -n "$existing" ] && [ "$existing" != "null" ]; then + payload="$RUNNER_TEMP/ingest-comment-payload.json" + jq -Rs '{body: .}' < "$COMMENT_OUTPUT" > "$payload" + gh api -X PATCH "repos/$REPO/issues/comments/$existing" \ + --input "$payload" >/dev/null + echo "updated existing comment $existing" + else + gh pr comment "$PR_NUMBER" --body-file "$COMMENT_OUTPUT" + fi - name: Fail if ingest checks failed if: ${{ always() }} diff --git a/apps/ingest/LOAD_TESTS.md b/apps/ingest/LOAD_TESTS.md index 29ac7e01d5..fcf954251e 100644 --- a/apps/ingest/LOAD_TESTS.md +++ b/apps/ingest/LOAD_TESTS.md @@ -26,10 +26,15 @@ Useful knobs: - `LOAD_TEST_MIN_RPS`: optional failure threshold for accepted request/sec. - `LOAD_TEST_MAX_RSS_MB`: optional failure threshold for max ingest RSS. - `LOAD_TEST_INGEST_BIN`: path to the ingest binary when not using the default sibling binary. +- `LOAD_TEST_INGEST_MODE`: `tinybird` (default) or `forward`. - `LOAD_TEST_QUEUE_DIR`: WAL directory override. +- `LOAD_TEST_REPORT_PATH`: when set, the `LoadSummary` JSON is also written here + (in addition to stdout). CI uses this to avoid parsing JSON out of mixed stdout. The harness prints JSON with request throughput, row throughput, p50/p95/p99 -latency, export catch-up time, max RSS, max CPU, average CPU, and exported rows. +latency, export catch-up time, max RSS, and exported rows. CPU samples come +from `ps` and are unreliable for short runs on Linux — they are kept in the +JSON for reference but the CI comment intentionally omits them. The GitHub Actions workflow `Ingest Load Tests` is manual (`workflow_dispatch`) so large runs do not make normal PR CI noisy or flaky. @@ -41,8 +46,14 @@ cargo bench --bench ingest_bench -- --sample-size 10 --warm-up-time 1 --measurem ``` Those benchmarks measure WAL-acked native accepts for representative log and -trace OTLP batches. - -On pull requests, the `Ingest Rust Tests` workflow runs a comparative load -benchmark against `main` and posts request throughput, row throughput, latency, -export catch-up, RSS, CPU, and failure deltas back to the PR. +trace OTLP batches. CI runs them on every PR via `--output-format=bencher` +and embeds the result in the PR comment. + +On pull requests, the `Ingest Rust Tests` workflow runs the load harness +**three times** on the PR and (when the same `LOAD_TEST_INGEST_MODE` is +supported on the base branch) three times on the base branch, then posts +median throughput, row throughput, latency, export catch-up, RSS, and failure +deltas back to the PR as a single sticky comment that updates on each push. +If the base branch does not support the head's ingest mode (e.g. this PR +introduces a new mode), the cross-mode comparison is skipped and only +absolute PR numbers are shown. diff --git a/apps/ingest/src/bin/load_test.rs b/apps/ingest/src/bin/load_test.rs index 587b4cbaa9..cf22d29ea4 100644 --- a/apps/ingest/src/bin/load_test.rs +++ b/apps/ingest/src/bin/load_test.rs @@ -149,7 +149,14 @@ async fn main() -> Result<(), DynError> { avg_cpu_percent: monitor_summary.avg_cpu_percent, }; - println!("{}", serde_json::to_string_pretty(&summary)?); + let pretty = serde_json::to_string_pretty(&summary)?; + println!("{pretty}"); + if let Ok(report_path) = std::env::var("LOAD_TEST_REPORT_PATH") { + let report_path = report_path.trim(); + if !report_path.is_empty() { + std::fs::write(report_path, &pretty)?; + } + } enforce_thresholds(&cfg, &summary)?; let _ = std::fs::remove_dir_all(&cfg.queue_dir); Ok(()) From 4fb206b4f5cc27e9500a97d3550cebab9294661d Mon Sep 17 00:00:00 2001 From: Makisuo Date: Thu, 14 May 2026 14:15:06 +0200 Subject: [PATCH 7/8] Add schema-parity + e2e tests for Tinybird ingest output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously the ingest's Rust encoder and the Tinybird datasource schemas (packages/domain/src/tinybird/datasources.ts) were only loosely coupled: they asserted invariants independently, and a rename or jsonPath edit on either side would silently land rows in the wrong column (or drop them) with no test failure. Pin both sides against the same canonical JSON-key contract: - apps/ingest/src/telemetry.rs::tests::schema_contract holds the expected top-level JSON keys per ingested datasource (logs, traces, metrics_sum, metrics_gauge, metrics_histogram, metrics_exponential_histogram). New tests assert the encoder emits *exactly* that set — no missing, no extras. - packages/domain/src/tinybird/datasources.contract.test.ts introspects each datasource's schema via `getColumnJsonPath`, derives the top-level jsonPath roots, and asserts the same canonical set. If either side drifts, the failing side's test prints the diff and points at the paired file to update. Also fills coverage gaps that left non-logs paths under-tested: - pipeline_e2e_exports_traces_to_fake_tinybird: full pipeline through the fake Tinybird import endpoint for traces (logs already covered). - pipeline_e2e_exports_metrics_to_fake_tinybird: same for each of the four metric datasources, asserting the JSON keys per datasource. - custom_datasource_names_propagate_to_frames: confirms the INGEST_TINYBIRD_DATASOURCE_* overrides flow through to the frame's `datasource` field (which becomes the Tinybird `?name=` query param). - timestamps_match_clickhouse_datetime64_nine_format: enforces the `YYYY-MM-DD HH:MM:SS.fffffffff` shape ClickHouse `DateTime64(9)` expects across logs.timestamp, traces.start_time, and metric timestamps. - logs_severity_text_falls_back_to_mapped_number: pins the severity-number -> text fallback the encoder applies when the SDK only sent a numeric severity. - logs_use_observed_time_when_time_unix_nano_is_zero: pins the OTLP-spec observed_time fallback so the Timestamp column is never epoch. - metrics_summary_data_points_are_dropped: pins the intentional Summary drop, so adding a Summary datasource won't go unnoticed. Total: 10 new tests in the Rust ingest crate (`cargo test` 18 -> 28 in the telemetry mod) + 7 new tests in the TS domain package, all passing. Co-Authored-By: Claude Opus 4.7 (1M context) --- apps/ingest/src/telemetry.rs | 656 ++++++++++++++++++ .../src/tinybird/datasources.contract.test.ts | 191 +++++ 2 files changed, 847 insertions(+) create mode 100644 packages/domain/src/tinybird/datasources.contract.test.ts diff --git a/apps/ingest/src/telemetry.rs b/apps/ingest/src/telemetry.rs index af2b4b7cb2..6dd087cdc7 100644 --- a/apps/ingest/src/telemetry.rs +++ b/apps/ingest/src/telemetry.rs @@ -2120,4 +2120,660 @@ mod tests { assert_eq!(exp["negative_bucket_counts"], json!([1])); assert_eq!(exp["aggregation_temporality"], 2); } + + // ----------------------------------------------------------------------- + // Schema-parity contract. + // + // The lists below are the JSON top-level keys each ingest datasource must + // populate. They MUST stay in lockstep with the `jsonPath` declarations in + // packages/domain/src/tinybird/datasources.ts — a TS-side test + // (`datasources.contract.test.ts`) pins the same lists from the other + // direction, so changing either side without updating both will fail CI. + // + // Keys come from the *roots* of jsonPath strings, deduplicated: + // "$.foo" -> "foo" + // "$.foo[:]" -> "foo" + // "$.foo.bar" -> "foo" (only the top level) + // ResourceAttributes uses `$.resource_attributes.maple_org_id` for OrgId, + // which is already covered by the `resource_attributes` map. + // ----------------------------------------------------------------------- + mod schema_contract { + pub const LOGS: &[&str] = &[ + "timestamp", + "trace_id", + "span_id", + "flags", + "severity_text", + "severity_number", + "service_name", + "body", + "resource_schema_url", + "resource_attributes", + "scope_schema_url", + "scope_name", + "scope_version", + "scope_attributes", + "log_attributes", + ]; + + pub const TRACES: &[&str] = &[ + "start_time", + "trace_id", + "span_id", + "parent_span_id", + "trace_state", + "span_name", + "span_kind", + "service_name", + "resource_schema_url", + "resource_attributes", + "scope_schema_url", + "scope_name", + "scope_version", + "scope_attributes", + "duration", + "status_code", + "status_message", + "span_attributes", + "events_timestamp", + "events_name", + "events_attributes", + "links_trace_id", + "links_span_id", + "links_trace_state", + "links_attributes", + ]; + + const METRIC_COMMON: &[&str] = &[ + "resource_attributes", + "resource_schema_url", + "scope_name", + "scope_version", + "scope_attributes", + "scope_schema_url", + "service_name", + "metric_name", + "metric_description", + "metric_unit", + "metric_attributes", + "start_timestamp", + "timestamp", + "flags", + "exemplars_trace_id", + "exemplars_span_id", + "exemplars_timestamp", + "exemplars_value", + "exemplars_filtered_attributes", + ]; + + fn with(extra: &[&'static str]) -> Vec<&'static str> { + let mut v: Vec<&'static str> = METRIC_COMMON.to_vec(); + v.extend(extra.iter().copied()); + v + } + + pub fn metrics_sum() -> Vec<&'static str> { + with(&["value", "aggregation_temporality", "is_monotonic"]) + } + pub fn metrics_gauge() -> Vec<&'static str> { + with(&["value"]) + } + pub fn metrics_histogram() -> Vec<&'static str> { + with(&[ + "count", + "sum", + "bucket_counts", + "explicit_bounds", + "min", + "max", + "aggregation_temporality", + ]) + } + pub fn metrics_exponential_histogram() -> Vec<&'static str> { + with(&[ + "count", + "sum", + "scale", + "zero_count", + "positive_offset", + "positive_bucket_counts", + "negative_offset", + "negative_bucket_counts", + "min", + "max", + "aggregation_temporality", + ]) + } + } + + fn assert_row_keys_match(row: &Value, expected: &[&str], datasource: &str) { + use std::collections::BTreeSet; + let obj = row + .as_object() + .unwrap_or_else(|| panic!("{datasource} row must be a JSON object")); + let actual: BTreeSet<&str> = obj.keys().map(String::as_str).collect(); + let expected: BTreeSet<&str> = expected.iter().copied().collect(); + let missing: Vec<&&str> = expected.difference(&actual).collect(); + let extra: Vec<&&str> = actual.difference(&expected).collect(); + assert!( + missing.is_empty() && extra.is_empty(), + "datasource '{datasource}' JSON key mismatch.\n missing (declared in datasources.ts but not emitted): {missing:?}\n extra (emitted but not declared in datasources.ts): {extra:?}\n If you intentionally changed the schema, update packages/domain/src/tinybird/datasources.ts AND the schema_contract module here AND packages/domain/src/tinybird/datasources.contract.test.ts." + ); + } + + fn populated_log() -> LogRecord { + LogRecord { + time_unix_nano: 1_700_000_001_123_456_789, + observed_time_unix_nano: 1_700_000_001_123_456_789, + severity_number: 17, + severity_text: "ERROR".to_string(), + body: Some(AnyValue { + value: Some(any_value::Value::StringValue("payment failed".into())), + }), + attributes: vec![string_kv("component", "billing")], + trace_id: vec![0xaa; 16], + span_id: vec![0xbb; 8], + flags: 1, + ..Default::default() + } + } + + fn populated_log_request() -> ExportLogsServiceRequest { + ExportLogsServiceRequest { + resource_logs: vec![opentelemetry_proto::tonic::logs::v1::ResourceLogs { + resource: Some(Resource { + attributes: vec![ + string_kv("service.name", "billing"), + string_kv("maple_org_id", "org_contract"), + ], + dropped_attributes_count: 0, + entity_refs: Vec::new(), + }), + scope_logs: vec![opentelemetry_proto::tonic::logs::v1::ScopeLogs { + scope: Some(InstrumentationScope { + name: "billing-logger".to_string(), + version: "2.0.1".to_string(), + attributes: vec![string_kv("scope.key", "scope-value")], + dropped_attributes_count: 0, + }), + log_records: vec![populated_log()], + schema_url: "https://scope.schema/logs".to_string(), + }], + schema_url: "https://resource.schema/logs".to_string(), + }], + } + } + + fn populated_trace_request() -> ExportTraceServiceRequest { + ExportTraceServiceRequest { + resource_spans: vec![ResourceSpans { + resource: Some(Resource { + attributes: vec![ + string_kv("service.name", "checkout"), + string_kv("maple_org_id", "org_contract"), + ], + dropped_attributes_count: 0, + entity_refs: Vec::new(), + }), + scope_spans: vec![ScopeSpans { + scope: Some(InstrumentationScope { + name: "checkout-tracer".to_string(), + version: "3.4.5".to_string(), + attributes: vec![string_kv("scope.key", "scope-value")], + dropped_attributes_count: 0, + }), + spans: vec![Span { + trace_id: vec![0x11; 16], + span_id: vec![0x22; 8], + parent_span_id: vec![0x33; 8], + trace_state: "vendor=foo".to_string(), + name: "POST /checkout".to_string(), + kind: span::SpanKind::Server as i32, + start_time_unix_nano: 1_700_000_000_000_000_000, + end_time_unix_nano: 1_700_000_000_250_000_000, + attributes: vec![string_kv("http.route", "/checkout")], + status: Some(Status { + code: status::StatusCode::Ok as i32, + message: "ok".to_string(), + }), + ..Default::default() + }], + schema_url: "https://scope.schema/traces".to_string(), + }], + schema_url: "https://resource.schema/traces".to_string(), + }], + } + } + + fn one_of_each_metric_request() -> ExportMetricsServiceRequest { + let base = NumberDataPoint { + attributes: vec![string_kv("route", "/checkout")], + start_time_unix_nano: 1_700_000_000_000_000_000, + time_unix_nano: 1_700_000_010_000_000_000, + flags: 0, + value: Some(number_data_point::Value::AsInt(42)), + ..Default::default() + }; + ExportMetricsServiceRequest { + resource_metrics: vec![opentelemetry_proto::tonic::metrics::v1::ResourceMetrics { + resource: Some(Resource { + attributes: vec![ + string_kv("service.name", "api"), + string_kv("maple_org_id", "org_contract"), + ], + dropped_attributes_count: 0, + entity_refs: Vec::new(), + }), + scope_metrics: vec![opentelemetry_proto::tonic::metrics::v1::ScopeMetrics { + scope: Some(InstrumentationScope { + name: "meter".to_string(), + version: "7.8.9".to_string(), + attributes: Vec::new(), + dropped_attributes_count: 0, + }), + metrics: vec![ + Metric { + name: "requests_total".to_string(), + description: "requests".to_string(), + unit: "1".to_string(), + data: Some(metric::Data::Sum(Sum { + data_points: vec![base.clone()], + aggregation_temporality: AggregationTemporality::Delta as i32, + is_monotonic: true, + })), + metadata: Vec::new(), + }, + Metric { + name: "cpu_ratio".to_string(), + description: "cpu".to_string(), + unit: "1".to_string(), + data: Some(metric::Data::Gauge(Gauge { + data_points: vec![NumberDataPoint { + value: Some(number_data_point::Value::AsDouble(0.75)), + ..base.clone() + }], + })), + metadata: Vec::new(), + }, + Metric { + name: "request_duration_ms".to_string(), + description: "latency".to_string(), + unit: "ms".to_string(), + data: Some(metric::Data::Histogram(Histogram { + data_points: vec![HistogramDataPoint { + attributes: vec![string_kv("route", "/checkout")], + start_time_unix_nano: base.start_time_unix_nano, + time_unix_nano: base.time_unix_nano, + count: 3, + sum: Some(123.0), + bucket_counts: vec![1, 2], + explicit_bounds: vec![100.0], + min: Some(10.0), + max: Some(90.0), + ..Default::default() + }], + aggregation_temporality: AggregationTemporality::Delta as i32, + })), + metadata: Vec::new(), + }, + Metric { + name: "payload_bytes".to_string(), + description: "payload".to_string(), + unit: "By".to_string(), + data: Some(metric::Data::ExponentialHistogram(ExponentialHistogram { + data_points: vec![ExponentialHistogramDataPoint { + attributes: vec![string_kv("route", "/checkout")], + start_time_unix_nano: base.start_time_unix_nano, + time_unix_nano: base.time_unix_nano, + count: 5, + sum: Some(500.0), + scale: 2, + zero_count: 1, + positive: Some(exponential_histogram_data_point::Buckets { + offset: -1, + bucket_counts: vec![2, 3], + }), + negative: Some(exponential_histogram_data_point::Buckets { + offset: 0, + bucket_counts: vec![1], + }), + min: Some(1.0), + max: Some(250.0), + ..Default::default() + }], + aggregation_temporality: AggregationTemporality::Cumulative as i32, + })), + metadata: Vec::new(), + }, + ], + schema_url: String::new(), + }], + schema_url: String::new(), + }], + } + } + + #[test] + fn logs_emit_exactly_the_jsonpaths_declared_in_datasources_ts() { + let (frames, _) = encode_logs(&test_cfg(), "org_contract", &populated_log_request()).unwrap(); + let row = frame_row(&frames[0]); + assert_row_keys_match(&row, schema_contract::LOGS, "logs"); + // Spot-check the resource_attributes carry the org id, since OrgId's + // jsonPath in TS is `$.resource_attributes.maple_org_id`. + assert_eq!(row["resource_attributes"]["maple_org_id"], "org_contract"); + } + + #[test] + fn traces_emit_exactly_the_jsonpaths_declared_in_datasources_ts() { + let (frames, _) = + encode_traces(&test_cfg(), "org_contract", &populated_trace_request(), &SamplingPolicy::default()) + .unwrap(); + let row = frame_row(&frames[0]); + assert_row_keys_match(&row, schema_contract::TRACES, "traces"); + assert_eq!(row["resource_attributes"]["maple_org_id"], "org_contract"); + } + + #[test] + fn metrics_emit_exactly_the_jsonpaths_declared_in_datasources_ts() { + let (frames, _) = encode_metrics(&test_cfg(), "org_contract", &one_of_each_metric_request()).unwrap(); + let rows = frame_rows_by_datasource(frames); + assert_row_keys_match(&rows["metrics_sum"], &schema_contract::metrics_sum(), "metrics_sum"); + assert_row_keys_match(&rows["metrics_gauge"], &schema_contract::metrics_gauge(), "metrics_gauge"); + assert_row_keys_match( + &rows["metrics_histogram"], + &schema_contract::metrics_histogram(), + "metrics_histogram", + ); + assert_row_keys_match( + &rows["metrics_exponential_histogram"], + &schema_contract::metrics_exponential_histogram(), + "metrics_exponential_histogram", + ); + } + + #[test] + fn timestamps_match_clickhouse_datetime64_nine_format() { + // ClickHouse `DateTime64(9)` accepts `YYYY-MM-DD HH:MM:SS.fffffffff` + // (nanosecond precision). Every timestamp-bearing column must match. + let pattern = |s: &str| { + s.len() == 29 + && &s[4..5] == "-" + && &s[7..8] == "-" + && &s[10..11] == " " + && &s[13..14] == ":" + && &s[16..17] == ":" + && &s[19..20] == "." + && s[..4].chars().all(|c| c.is_ascii_digit()) + && s[5..7].chars().all(|c| c.is_ascii_digit()) + && s[8..10].chars().all(|c| c.is_ascii_digit()) + && s[11..13].chars().all(|c| c.is_ascii_digit()) + && s[14..16].chars().all(|c| c.is_ascii_digit()) + && s[17..19].chars().all(|c| c.is_ascii_digit()) + && s[20..].chars().all(|c| c.is_ascii_digit()) + }; + + let (log_frames, _) = + encode_logs(&test_cfg(), "org_contract", &populated_log_request()).unwrap(); + let log_row = frame_row(&log_frames[0]); + let ts = log_row["timestamp"].as_str().unwrap(); + assert!(pattern(ts), "logs.timestamp not DateTime64(9): {ts:?}"); + + let (trace_frames, _) = + encode_traces(&test_cfg(), "org_contract", &populated_trace_request(), &SamplingPolicy::default()) + .unwrap(); + let trace_row = frame_row(&trace_frames[0]); + let ts = trace_row["start_time"].as_str().unwrap(); + assert!(pattern(ts), "traces.start_time not DateTime64(9): {ts:?}"); + + let (metric_frames, _) = + encode_metrics(&test_cfg(), "org_contract", &one_of_each_metric_request()).unwrap(); + for frame in &metric_frames { + let row = frame_row(frame); + let start_ts = row["start_timestamp"].as_str().unwrap(); + let ts = row["timestamp"].as_str().unwrap(); + assert!( + pattern(start_ts), + "{}.start_timestamp not DateTime64(9): {start_ts:?}", + frame.datasource + ); + assert!( + pattern(ts), + "{}.timestamp not DateTime64(9): {ts:?}", + frame.datasource + ); + } + } + + #[test] + fn custom_datasource_names_propagate_to_frames() { + // Operators can rebind each datasource via `INGEST_TINYBIRD_DATASOURCE_*` + // env vars. The frames emitted by the encoders must carry the + // configured name (which becomes the `name` query parameter on the + // Tinybird import call), not the hardcoded default. + let mut cfg = test_cfg(); + cfg.datasource_logs = "tenant_logs_v2".into(); + cfg.datasource_traces = "tenant_traces_v2".into(); + cfg.datasource_metrics_sum = "tenant_sum_v2".into(); + cfg.datasource_metrics_gauge = "tenant_gauge_v2".into(); + cfg.datasource_metrics_histogram = "tenant_hist_v2".into(); + cfg.datasource_metrics_exponential_histogram = "tenant_exp_v2".into(); + + let (log_frames, _) = encode_logs(&cfg, "org_contract", &populated_log_request()).unwrap(); + assert_eq!(log_frames[0].datasource, "tenant_logs_v2"); + + let (trace_frames, _) = + encode_traces(&cfg, "org_contract", &populated_trace_request(), &SamplingPolicy::default()) + .unwrap(); + assert_eq!(trace_frames[0].datasource, "tenant_traces_v2"); + + let (metric_frames, _) = + encode_metrics(&cfg, "org_contract", &one_of_each_metric_request()).unwrap(); + let datasources: std::collections::BTreeSet<_> = + metric_frames.iter().map(|f| f.datasource.as_str()).collect(); + assert!(datasources.contains("tenant_sum_v2")); + assert!(datasources.contains("tenant_gauge_v2")); + assert!(datasources.contains("tenant_hist_v2")); + assert!(datasources.contains("tenant_exp_v2")); + } + + #[test] + fn logs_severity_text_falls_back_to_mapped_number() { + // Tinybird's `SeverityText` column expects a non-empty string when the + // SDK left only a numeric severity; the encoder maps the OTLP severity + // number to the canonical text label. + let mut request = populated_log_request(); + request.resource_logs[0].scope_logs[0].log_records[0].severity_text.clear(); + request.resource_logs[0].scope_logs[0].log_records[0].severity_number = 9; // INFO + let (frames, _) = encode_logs(&test_cfg(), "org_contract", &request).unwrap(); + let row = frame_row(&frames[0]); + assert_eq!(row["severity_text"], "INFO"); + } + + #[test] + fn logs_use_observed_time_when_time_unix_nano_is_zero() { + // Per OTLP spec the receiver should fall back to observed_time when the + // emitter didn't set time_unix_nano. The encoder honors this so that + // the `Timestamp` column is never the epoch. + let mut request = populated_log_request(); + request.resource_logs[0].scope_logs[0].log_records[0].time_unix_nano = 0; + request.resource_logs[0].scope_logs[0].log_records[0].observed_time_unix_nano = + 1_700_000_100_000_000_000; + let (frames, _) = encode_logs(&test_cfg(), "org_contract", &request).unwrap(); + let row = frame_row(&frames[0]); + assert_eq!(row["timestamp"], "2023-11-14 22:15:00.000000000"); + } + + #[test] + fn metrics_summary_data_points_are_dropped() { + // The encoder does not support Summary metrics (no Tinybird datasource); + // dropping silently is intentional but worth pinning so an accidental + // schema addition is noticed. + use opentelemetry_proto::tonic::metrics::v1::{summary_data_point, Summary, SummaryDataPoint}; + let request = ExportMetricsServiceRequest { + resource_metrics: vec![opentelemetry_proto::tonic::metrics::v1::ResourceMetrics { + resource: Some(Resource { + attributes: vec![string_kv("service.name", "api")], + dropped_attributes_count: 0, + entity_refs: Vec::new(), + }), + scope_metrics: vec![opentelemetry_proto::tonic::metrics::v1::ScopeMetrics { + scope: Some(InstrumentationScope { + name: "meter".to_string(), + version: "1.0.0".to_string(), + attributes: Vec::new(), + dropped_attributes_count: 0, + }), + metrics: vec![Metric { + name: "summary_metric".to_string(), + description: "".into(), + unit: "".into(), + data: Some(metric::Data::Summary(Summary { + data_points: vec![SummaryDataPoint { + attributes: vec![], + start_time_unix_nano: 0, + time_unix_nano: 1_700_000_000_000_000_000, + count: 1, + sum: 1.0, + quantile_values: vec![summary_data_point::ValueAtQuantile { + quantile: 0.5, + value: 1.0, + }], + flags: 0, + }], + })), + metadata: Vec::new(), + }], + schema_url: String::new(), + }], + schema_url: String::new(), + }], + }; + let (frames, stats) = encode_metrics(&test_cfg(), "org_contract", &request).unwrap(); + assert!(frames.is_empty(), "summary metrics should not produce frames"); + assert_eq!(stats.rows, 0); + } + + #[tokio::test] + async fn pipeline_e2e_exports_traces_to_fake_tinybird() { + let (tx, mut rx) = mpsc::unbounded_channel(); + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let addr = listener.local_addr().unwrap(); + let app = Router::new() + .route("/v0/events", post(fake_tinybird_import)) + .with_state(tx); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + let queue_dir = unique_test_dir("fake-tinybird-traces"); + let mut cfg = test_cfg(); + cfg.endpoint = format!("http://{addr}"); + cfg.queue_dir = queue_dir.clone(); + cfg.wal_shards = 1; + cfg.batch_max_wait = Duration::from_millis(1); + + let pipeline = TelemetryPipeline::new( + cfg, + Client::builder() + .timeout(Duration::from_secs(5)) + .build() + .unwrap(), + ) + .await + .unwrap(); + + let request = populated_trace_request(); + let stats = pipeline + .accept_traces("org_contract", &request, &SamplingPolicy::default()) + .await + .unwrap(); + assert_eq!(stats.rows, 1); + + let import = tokio::time::timeout(Duration::from_secs(2), rx.recv()) + .await + .expect("fake Tinybird should receive a traces import") + .expect("fake Tinybird channel should stay open"); + + assert_eq!(import.datasource, "traces"); + assert_eq!(import.authorization, "Bearer token"); + assert_eq!(import.content_encoding, "gzip"); + + let row: Value = serde_json::from_str(import.body.trim()).unwrap(); + assert_row_keys_match(&row, schema_contract::TRACES, "traces"); + assert_eq!(row["span_kind"], "Server"); + assert_eq!(row["status_code"], "Ok"); + + assert!(wait_for_cursor(queue_dir.join("shard-000.cursor")).await > 0); + let _ = std::fs::remove_dir_all(queue_dir); + } + + #[tokio::test] + async fn pipeline_e2e_exports_metrics_to_fake_tinybird() { + let (tx, mut rx) = mpsc::unbounded_channel(); + let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)) + .await + .unwrap(); + let addr = listener.local_addr().unwrap(); + let app = Router::new() + .route("/v0/events", post(fake_tinybird_import)) + .with_state(tx); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + let queue_dir = unique_test_dir("fake-tinybird-metrics"); + let mut cfg = test_cfg(); + cfg.endpoint = format!("http://{addr}"); + cfg.queue_dir = queue_dir.clone(); + cfg.wal_shards = 1; + cfg.batch_max_wait = Duration::from_millis(1); + + let pipeline = TelemetryPipeline::new( + cfg, + Client::builder() + .timeout(Duration::from_secs(5)) + .build() + .unwrap(), + ) + .await + .unwrap(); + + let stats = pipeline + .accept_metrics("org_contract", &one_of_each_metric_request()) + .await + .unwrap(); + assert_eq!(stats.rows, 4); + + // Collect 4 imports — one per metric datasource. + let mut by_datasource: BTreeMap = BTreeMap::new(); + for _ in 0..4 { + let import = tokio::time::timeout(Duration::from_secs(2), rx.recv()) + .await + .expect("fake Tinybird should receive each metric datasource import") + .expect("fake Tinybird channel should stay open"); + by_datasource.insert(import.datasource.clone(), import); + } + + for (datasource, expected_keys) in [ + ("metrics_sum", schema_contract::metrics_sum()), + ("metrics_gauge", schema_contract::metrics_gauge()), + ("metrics_histogram", schema_contract::metrics_histogram()), + ( + "metrics_exponential_histogram", + schema_contract::metrics_exponential_histogram(), + ), + ] { + let import = by_datasource + .remove(datasource) + .unwrap_or_else(|| panic!("missing import for {datasource}")); + assert_eq!(import.authorization, "Bearer token"); + assert_eq!(import.content_encoding, "gzip"); + let row: Value = serde_json::from_str(import.body.trim()).unwrap(); + assert_row_keys_match(&row, &expected_keys, datasource); + } + + let _ = std::fs::remove_dir_all(queue_dir); + } } diff --git a/packages/domain/src/tinybird/datasources.contract.test.ts b/packages/domain/src/tinybird/datasources.contract.test.ts new file mode 100644 index 0000000000..ed2fbd7d52 --- /dev/null +++ b/packages/domain/src/tinybird/datasources.contract.test.ts @@ -0,0 +1,191 @@ +import { getColumnJsonPath, type DatasourceDefinition } from "@tinybirdco/sdk" +import { describe, expect, it } from "vitest" +import { + logs, + metricsExponentialHistogram, + metricsGauge, + metricsHistogram, + metricsSum, + traces, +} from "./datasources" + +/** + * Schema-parity contract for the Rust ingest service. + * + * The Rust encoder in `apps/ingest/src/telemetry.rs` writes one JSON object + * per row to Tinybird. Each top-level JSON key is bound to a ClickHouse column + * via the `jsonPath` declared on that column in `./datasources.ts`. If the two + * sides drift, rows silently land in the wrong columns (or are dropped). + * + * This file pins the TS side. The Rust side is pinned by `mod schema_contract` + * inside `apps/ingest/src/telemetry.rs::tests`. The two lists below MUST stay + * identical — if you intentionally change a column's `jsonPath`, update both. + * + * Conversion rule: a jsonPath like `$.foo`, `$.foo[:]`, or `$.foo.bar` all + * collapse to the top-level key `foo`. The encoder only writes top-level keys + * (nested values are written by writing the parent map). + */ +const EXPECTED_TOPLEVEL_KEYS = { + logs: new Set([ + "timestamp", + "trace_id", + "span_id", + "flags", + "severity_text", + "severity_number", + "service_name", + "body", + "resource_schema_url", + "resource_attributes", + "scope_schema_url", + "scope_name", + "scope_version", + "scope_attributes", + "log_attributes", + ]), + traces: new Set([ + "start_time", + "trace_id", + "span_id", + "parent_span_id", + "trace_state", + "span_name", + "span_kind", + "service_name", + "resource_schema_url", + "resource_attributes", + "scope_schema_url", + "scope_name", + "scope_version", + "scope_attributes", + "duration", + "status_code", + "status_message", + "span_attributes", + "events_timestamp", + "events_name", + "events_attributes", + "links_trace_id", + "links_span_id", + "links_trace_state", + "links_attributes", + ]), + metrics_sum: new Set([ + ...metricCommonKeys(), + "value", + "aggregation_temporality", + "is_monotonic", + ]), + metrics_gauge: new Set([...metricCommonKeys(), "value"]), + metrics_histogram: new Set([ + ...metricCommonKeys(), + "count", + "sum", + "bucket_counts", + "explicit_bounds", + "min", + "max", + "aggregation_temporality", + ]), + metrics_exponential_histogram: new Set([ + ...metricCommonKeys(), + "count", + "sum", + "scale", + "zero_count", + "positive_offset", + "positive_bucket_counts", + "negative_offset", + "negative_bucket_counts", + "min", + "max", + "aggregation_temporality", + ]), +} + +function metricCommonKeys(): string[] { + return [ + "resource_attributes", + "resource_schema_url", + "scope_name", + "scope_version", + "scope_attributes", + "scope_schema_url", + "service_name", + "metric_name", + "metric_description", + "metric_unit", + "metric_attributes", + "start_timestamp", + "timestamp", + "flags", + "exemplars_trace_id", + "exemplars_span_id", + "exemplars_timestamp", + "exemplars_value", + "exemplars_filtered_attributes", + ] +} + +function topLevelKey(jsonPath: string): string | null { + // `$.foo` / `$.foo[:]` / `$.foo.bar.baz` -> `foo` + if (!jsonPath.startsWith("$.")) return null + const tail = jsonPath.slice(2) + const head = tail.split(/[.\[]/, 1)[0] + return head || null +} + +function emittedTopLevelKeys( + datasource: DatasourceDefinition, +): Set { + const keys = new Set() + for (const column of Object.values(datasource.options.schema)) { + const path = getColumnJsonPath(column) + if (!path) continue // defaultExpr-only columns (SampleRate, IsEntryPoint) are computed in CH, not ingested + const top = topLevelKey(path) + if (top) keys.add(top) + } + return keys +} + +function diff(a: Set, b: Set): { missing: string[]; extra: string[] } { + const missing = [...a].filter((k) => !b.has(k)).sort() + const extra = [...b].filter((k) => !a.has(k)).sort() + return { missing, extra } +} + +const driftHint = + "If you intentionally changed the schema, update BOTH this file AND `mod schema_contract` in apps/ingest/src/telemetry.rs::tests." + +describe("Tinybird datasource ↔ Rust ingest encoder JSON contract", () => { + const datasources: Array<{ + name: keyof typeof EXPECTED_TOPLEVEL_KEYS + datasource: DatasourceDefinition + }> = [ + { name: "logs", datasource: logs }, + { name: "traces", datasource: traces }, + { name: "metrics_sum", datasource: metricsSum }, + { name: "metrics_gauge", datasource: metricsGauge }, + { name: "metrics_histogram", datasource: metricsHistogram }, + { name: "metrics_exponential_histogram", datasource: metricsExponentialHistogram }, + ] + + for (const { name, datasource } of datasources) { + it(`${name} jsonPath roots match the Rust encoder contract`, () => { + const expected = EXPECTED_TOPLEVEL_KEYS[name] + const actual = emittedTopLevelKeys(datasource) + const { missing, extra } = diff(expected, actual) + expect( + { missing, extra }, + `Drift between datasources.ts (${name}) and the Rust contract.\n${driftHint}`, + ).toEqual({ missing: [], extra: [] }) + }) + } + + it("every ingest datasource has at least one jsonPath declared", () => { + for (const { name, datasource } of datasources) { + const actual = emittedTopLevelKeys(datasource) + expect(actual.size, `${name} should have jsonPath-bound columns`).toBeGreaterThan(0) + } + }) +}) From 27be26765addd81a30b5ad0b37febb1e4aa81dd9 Mon Sep 17 00:00:00 2001 From: Makisuo Date: Thu, 14 May 2026 14:26:16 +0200 Subject: [PATCH 8/8] ci: don't build load_test on the base branch The base may not have a load_test binary at all (this PR introduces it). Building --bin load_test on the worktree checkout of main fails with "no bin target named load_test", which marks the step as failed even though the probe + benchmark steps only need the head's load_test binary pointed at the base's maple-ingest. Co-Authored-By: Claude Opus 4.7 (1M context) --- .github/workflows/ingest-rust-tests.yml | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ingest-rust-tests.yml b/.github/workflows/ingest-rust-tests.yml index da800938f2..9afab77244 100644 --- a/.github/workflows/ingest-rust-tests.yml +++ b/.github/workflows/ingest-rust-tests.yml @@ -115,18 +115,19 @@ jobs: set +e git fetch origin "${{ github.base_ref }}" --depth=1 git worktree add "$RUNNER_TEMP/base" FETCH_HEAD + # Only build maple-ingest; load_test may not exist on the base + # branch (it's added by the PR that introduces the harness). + # The probe + benchmark steps always use the head's load_test + # binary pointed at the base's maple-ingest. ( cd "$RUNNER_TEMP/base/apps/ingest" - cargo build --release --bin maple-ingest --bin load_test + cargo build --release --bin maple-ingest ) 2>&1 | tee "$RUNNER_TEMP/ingest-base-build.txt" echo "${PIPESTATUS[0]}" > "$RUNNER_TEMP/ingest-base-build.exit" mkdir -p "$RUNNER_TEMP/base-bin" if [ -x "$RUNNER_TEMP/base/apps/ingest/target/release/maple-ingest" ]; then cp "$RUNNER_TEMP/base/apps/ingest/target/release/maple-ingest" "$RUNNER_TEMP/base-bin/maple-ingest" fi - if [ -x "$RUNNER_TEMP/base/apps/ingest/target/release/load_test" ]; then - cp "$RUNNER_TEMP/base/apps/ingest/target/release/load_test" "$RUNNER_TEMP/base-bin/load_test" - fi exit 0 - name: Probe whether base supports head ingest mode