diff --git a/asap-query-engine/src/drivers/ingest/otel.rs b/asap-query-engine/src/drivers/ingest/otel.rs index 18d764bd..8f0aed3b 100644 --- a/asap-query-engine/src/drivers/ingest/otel.rs +++ b/asap-query-engine/src/drivers/ingest/otel.rs @@ -1,13 +1,33 @@ //! OTLP ingest driver. //! -//! Accepts OTLP metrics via gRPC (4317) and HTTP (4318, POST /v1/metrics), -//! parses ExportMetricsServiceRequest, logs counts at DEBUG, and leaves -//! handoff to precompute engine as TODO. +//! Accepts OTLP metrics via gRPC (4317) and HTTP (4318, POST /v1/metrics). +//! Parses `ExportMetricsServiceRequest`, and — when wired to a precompute +//! engine via [`OtlpReceiver::with_ingest_state`] — routes both raw metric +//! points and pre-built sketches through the precompute engine's worker +//! pool. The precompute engine then performs window-aligned aggregation +//! per `StreamingConfig` and writes results to `SimpleMapStore`. +//! +//! Architectural flow: +//! ```text +//! DataCollector OTel collector +//! → OTLP gRPC/HTTP (this receiver) +//! → precompute engine ingest router +//! → workers (per (agg_id, group_key) panes) +//! → StoreOutputSink → SimpleMapStore +//! → query engine +//! ``` +//! +//! Labels from the OTLP wire format are preserved all the way into +//! `KeyByLabelValues` via the standard `series_key` → grouping-label +//! extraction used by the Prometheus/VictoriaMetrics ingest paths. use std::collections::HashMap; use std::io::Read; -use std::sync::Arc; +use crate::data_model::AggregateCore; +use crate::precompute_engine::series_router::WorkerMessage; +use crate::precompute_engine::IngestState; +use crate::precompute_operators::sketch_envelope_accumulator::SketchEnvelopeAccumulator; use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; use axum::{body::Bytes, extract::State, routing::post, Json, Router}; use flate2::read::GzDecoder; @@ -18,13 +38,11 @@ use opentelemetry_proto::tonic::collector::metrics::v1::{ use opentelemetry_proto::tonic::common::v1::any_value::Value as AnyValueVariant; use opentelemetry_proto::tonic::metrics::v1::number_data_point::Value as NumberValue; use prost::Message; +use std::sync::Arc; +use std::time::Instant; use tonic::{Request, Response, Status}; use tracing::{debug, error, info, warn}; -use crate::data_model::{KeyByLabelValues, PrecomputedOutput, StreamingConfig}; -use crate::precompute_operators::SketchEnvelopeAccumulator; -use crate::stores::Store; - /// Configuration for the OTLP receiver. #[derive(Debug, Clone)] pub struct OtlpReceiverConfig { @@ -32,30 +50,39 @@ pub struct OtlpReceiverConfig { pub http_port: u16, } -/// Shared state passed to both gRPC and HTTP handlers. +/// Shared state accessible by both gRPC and HTTP handlers. #[derive(Clone)] struct OtlpSharedState { - store: Arc, - streaming_config: Arc, + /// Handle into the precompute engine's worker pool. When `Some`, + /// OTLP metrics and sketches are routed through the engine; when + /// `None` the receiver accepts data but only logs it (no storage). + ingest_state: Option>, } /// OTLP receiver that accepts metrics via gRPC and HTTP. pub struct OtlpReceiver { config: OtlpReceiverConfig, - store: Arc, - streaming_config: Arc, + ingest_state: Option>, } impl OtlpReceiver { - pub fn new( - config: OtlpReceiverConfig, - store: Arc, - streaming_config: Arc, - ) -> Self { + /// Construct a receiver without a backend. Metrics are parsed and + /// logged but not stored — useful for smoke-testing the OTLP pipe. + pub fn new(config: OtlpReceiverConfig) -> Self { Self { config, - store, - streaming_config, + ingest_state: None, + } + } + + /// Construct a receiver wired to a precompute engine's ingest state. + /// Incoming metrics and sketches are routed through the engine's + /// worker pool, where they are merged into per-`(agg_id, group_key)` + /// panes and eventually emitted to the store. + pub fn with_ingest_state(config: OtlpReceiverConfig, ingest_state: Arc) -> Self { + Self { + config, + ingest_state: Some(ingest_state), } } @@ -64,8 +91,7 @@ impl OtlpReceiver { let http_addr = std::net::SocketAddr::from(([0, 0, 0, 0], self.config.http_port)); let shared = Arc::new(OtlpSharedState { - store: self.store.clone(), - streaming_config: self.streaming_config.clone(), + ingest_state: self.ingest_state.clone(), }); let grpc_svc = MetricsServiceImpl { @@ -78,7 +104,7 @@ impl OtlpReceiver { let app = Router::new() .route("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/v1/metrics", post(handle_otlp_http)) - .with_state(shared.clone()); + .with_state(shared); let grpc_listener = tokio::net::TcpListener::bind(grpc_addr).await?; let http_listener = tokio::net::TcpListener::bind(http_addr).await?; @@ -117,7 +143,10 @@ impl MetricsService for MetricsServiceImpl { ) -> Result, Status> { debug!("OTLP received request via gRPC"); let req = request.into_inner(); - process_otlp_request(&req, "gRPC", &self.shared); + process_otlp_request(&req, "gRPC"); + if let Some(state) = &self.shared.ingest_state { + route_otlp_to_precompute(&req, state).await; + } debug!("OTLP sending response via gRPC"); Ok(Response::new(ExportMetricsServiceResponse { partial_success: None, @@ -156,7 +185,10 @@ async fn handle_otlp_http( format!("Protobuf decode error: {}", e), ) })?; - process_otlp_request(&req, "HTTP", &shared); + process_otlp_request(&req, "HTTP"); + if let Some(state) = &shared.ingest_state { + route_otlp_to_precompute(&req, state).await; + } debug!("OTLP sending response via HTTP"); Ok(Json(serde_json::json!({"rejected": 0}))) } @@ -170,15 +202,17 @@ pub struct MetricPoint { pub value: f64, } -/// A parsed sketch data point: metric name, attribute key, raw payload bytes, -/// timestamp in nanoseconds, and label map from the data point. +/// A parsed sketch payload extracted from OTLP attributes. Carries the +/// metric name, attribute name (identifies the sketch kind on the wire), +/// labels (preserved from the OTLP DataPoint attributes + resource/scope), +/// wire timestamp, and the opaque `SketchEnvelope` protobuf bytes. #[derive(Debug)] pub struct SketchPoint { - pub metric_name: String, + pub name: String, pub attr_name: String, - pub payload: Vec, - pub timestamp_nanos: u64, pub labels: HashMap, + pub timestamp_nanos: u64, + pub payload: Vec, } type OtlpParseResult = (Vec, Vec); @@ -242,31 +276,7 @@ fn log_sketch_envelope_type(attr_name: &str, payload: &[u8], metric_name: &str) } } -/// Build a reverse lookup from metric name -> aggregation_id using the -/// streaming config. Called once per request (configs are small). -fn build_metric_to_aggregation_id(streaming_config: &StreamingConfig) -> HashMap { - let mut map = HashMap::new(); - for (agg_id, config) in streaming_config.get_all_aggregation_configs() { - map.entry(config.metric.clone()).or_insert(*agg_id); - } - map -} - -/// Stable hash of a metric name to use as aggregation_id when the metric is -/// not found in the streaming config. Uses the same FNV-1a 64-bit hash that -/// other parts of the pipeline use for deterministic IDs. -fn hash_metric_name(name: &str) -> u64 { - use std::hash::{Hash, Hasher}; - let mut hasher = std::collections::hash_map::DefaultHasher::new(); - name.hash(&mut hasher); - hasher.finish() -} - -fn process_otlp_request( - request: &ExportMetricsServiceRequest, - transport: &str, - shared: &OtlpSharedState, -) { +fn process_otlp_request(request: &ExportMetricsServiceRequest, transport: &str) { let resource_count = request.resource_metrics.len(); let total_points = otlp_to_record_count(request); if resource_count > 0 || total_points > 0 { @@ -276,93 +286,18 @@ fn process_otlp_request( ); } - let (points, sketch_points) = otlp_to_metric_points_and_sketches(request); + let (points, sketch_payloads) = otlp_to_metric_points_and_sketches(request); - for sp in &sketch_points { - log_sketch_envelope_type(&sp.attr_name, &sp.payload, &sp.metric_name); + for sketch in &sketch_payloads { + log_sketch_envelope_type(&sketch.attr_name, &sketch.payload, &sketch.name); } - if !sketch_points.is_empty() { + if !sketch_payloads.is_empty() { debug!( "OTLP Sketch Payload Flow: received {} sketch payload(s), decoded successfully", - sketch_points.len() + sketch_payloads.len() ); } - // --- Store sketch data into SimpleMapStore --- - if !sketch_points.is_empty() { - let metric_to_agg = build_metric_to_aggregation_id(&shared.streaming_config); - let mut batch: Vec<(PrecomputedOutput, Box)> = - Vec::with_capacity(sketch_points.len()); - - for sp in &sketch_points { - // Resolve aggregation_id: prefer streaming config, fall back to hash. - let aggregation_id = metric_to_agg - .get(&sp.metric_name) - .copied() - .unwrap_or_else(|| { - let h = hash_metric_name(&sp.metric_name); - warn!( - "OTLP ingest: metric '{}' not found in streaming config, using hash {} as aggregation_id", - sp.metric_name, h - ); - h - }); - - // Convert nanoseconds to milliseconds for store timestamps. - let ts_ms = sp.timestamp_nanos / 1_000_000; - - // Build KeyByLabelValues from the label map (sorted for determinism). - let mut sorted_labels: Vec<(&String, &String)> = sp.labels.iter().collect(); - sorted_labels.sort_by_key(|(k, _)| *k); - let label_values: Vec = sorted_labels - .iter() - .map(|(k, v)| format!("{}={}", k, v)) - .collect(); - let key = if label_values.is_empty() { - None - } else { - Some(KeyByLabelValues::new_with_labels(label_values)) - }; - - let output = PrecomputedOutput::new(ts_ms, ts_ms, key, aggregation_id); - - match SketchEnvelopeAccumulator::from_proto_bytes(sp.payload.clone()) { - Ok(accumulator) => { - debug!( - "OTLP ingest: storing sketch metric='{}' agg_id={} ts_ms={} type={}", - sp.metric_name, aggregation_id, ts_ms, accumulator.sketch_type - ); - batch.push((output, Box::new(accumulator))); - } - Err(e) => { - error!( - "OTLP ingest: failed to decode SketchEnvelope for metric '{}': {}", - sp.metric_name, e - ); - } - } - } - - if !batch.is_empty() { - let batch_len = batch.len(); - match shared.store.insert_precomputed_output_batch(batch) { - Ok(_) => { - info!( - "OTLP ingest: stored {} sketch precompute(s) (transport={})", - batch_len, transport - ); - } - Err(e) => { - error!( - "OTLP ingest: failed to insert sketch batch into store: {}", - e - ); - } - } - } - } - - // --- Log raw metric points (non-sketch) --- let mut by_series: HashMap = HashMap::new(); for point in &points { let key = format_series_key(&point.name, &point.labels); @@ -385,6 +320,176 @@ fn process_otlp_request( } } +/// Route parsed OTLP data through the precompute engine's worker pool. +/// +/// Both raw metric points and pre-built sketch payloads are dispatched via +/// `WorkerMessage::GroupSamples` / `WorkerMessage::AccumulatorInput`, with +/// `(agg_id, group_key)` derived from the wire labels using the same logic +/// as the Prometheus/VictoriaMetrics paths. This preserves full label +/// semantics and lets the precompute engine perform config-driven window +/// aggregation over both streams. +/// +/// Metrics whose name does not match any aggregation in the streaming +/// config are dropped with a debug log — the precompute engine only +/// maintains state for configured metrics. +async fn route_otlp_to_precompute( + request: &ExportMetricsServiceRequest, + ingest_state: &Arc, +) { + let ingest_received_at = Instant::now(); + let (points, sketch_payloads) = otlp_to_metric_points_and_sketches(request); + + // Build (agg_id, group_key) → Vec<(series_key, ts_ms, value)> for raw points. + type GroupKey = (u64, String); + type SampleTuple = (String, i64, f64); + let mut by_group: HashMap> = HashMap::new(); + let mut raw_matched = 0usize; + let mut raw_unmatched = 0usize; + + for point in &points { + let series_key = format_series_key(&point.name, &point.labels); + let ts_ms = (point.timestamp_nanos / 1_000_000) as i64; + let mut matched = false; + for config in &ingest_state.agg_configs { + if config.metric != point.name + && config.spatial_filter_normalized != point.name + && config.spatial_filter != point.name + { + continue; + } + let group_key = IngestState::extract_group_key_for(&series_key, config); + by_group + .entry((config.aggregation_id, group_key)) + .or_default() + .push((series_key.clone(), ts_ms, point.value)); + matched = true; + } + if matched { + raw_matched += 1; + } else { + raw_unmatched += 1; + } + } + + let raw_messages: Vec = by_group + .into_iter() + .map( + |((agg_id, group_key), samples)| WorkerMessage::GroupSamples { + agg_id, + group_key, + samples, + ingest_received_at, + }, + ) + .collect(); + + if !raw_messages.is_empty() { + if let Err(e) = ingest_state + .router + .route_group_batch(raw_messages, ingest_received_at) + .await + { + warn!("OTLP raw-sample routing error: {}", e); + } + } + + // Build AccumulatorInput messages for pre-built sketches. + // Each sketch payload carries a metric name that must match an + // aggregation config; labels come from the point's attributes. The + // payload is decoded enough to identify the sketch type for logging; + // the accumulator the worker receives is a conservative placeholder + // until per-variant `SketchEnvelope → concrete accumulator` decoders + // are wired up (see TODO below). + let mut sketch_messages: Vec = Vec::new(); + let mut sketch_matched = 0usize; + let mut sketch_unmatched = 0usize; + for point in &sketch_payloads { + let series_key = format_series_key(&point.name, &point.labels); + let ts_ms = (point.timestamp_nanos / 1_000_000) as i64; + let sketch_type = identify_sketch_type(&point.payload); + let mut matched = false; + for config in &ingest_state.agg_configs { + if config.metric != point.name + && config.spatial_filter_normalized != point.name + && config.spatial_filter != point.name + { + continue; + } + let group_key = IngestState::extract_group_key_for(&series_key, config); + // Wrap the raw SketchEnvelope bytes in a SketchEnvelopeAccumulator + // so the precompute engine receives the opaque sketch as-is. This + // preserves all sketch state end-to-end; per-variant decoding + // (e.g. CountMin → CountMinSketchAccumulator) can layer on top + // later without changing the routing contract. + let accumulator: Box = + match SketchEnvelopeAccumulator::from_proto_bytes(point.payload.clone()) { + Ok(acc) => Box::new(acc), + Err(e) => { + warn!( + "OTLP sketch decode failed for metric='{}' attr='{}': {}", + point.name, point.attr_name, e + ); + continue; + } + }; + sketch_messages.push(WorkerMessage::AccumulatorInput { + agg_id: config.aggregation_id, + group_key, + timestamp_ms: ts_ms, + accumulator, + ingest_received_at, + }); + matched = true; + } + if matched { + sketch_matched += 1; + debug!( + "OTLP sketch routed to precompute engine: metric='{}' attr='{}' type={} bytes={}", + point.name, + point.attr_name, + sketch_type, + point.payload.len() + ); + } else { + sketch_unmatched += 1; + } + } + + if !sketch_messages.is_empty() { + if let Err(e) = ingest_state + .router + .route_group_batch(sketch_messages, ingest_received_at) + .await + { + warn!("OTLP sketch routing error: {}", e); + } + } + + if raw_unmatched > 0 || sketch_unmatched > 0 { + debug!( + "OTLP ingest: {} raw samples + {} sketches dropped (no matching aggregation config); \ + {} raw + {} sketches routed to precompute engine", + raw_unmatched, sketch_unmatched, raw_matched, sketch_matched + ); + } +} + +/// Identify the concrete sketch type inside a `SketchEnvelope` payload, +/// returning a human-readable name for logging. Returns `"Unknown"` if the +/// payload does not decode or the `sketch_state` variant is unset. +fn identify_sketch_type(payload: &[u8]) -> &'static str { + match SketchEnvelope::decode(payload) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::Kll(_)) => "KLL", + Some(sketch_envelope::SketchState::CountMin(_)) => "CountMin", + Some(sketch_envelope::SketchState::CountSketch(_)) => "CountSketch", + Some(_) => "Other", + None => "Unknown", + }, + Err(_) => "Unknown", + } +} + /// Count total data points by traversing resource_metrics -> scope_metrics -> metrics. /// Reuses the same conversion traversal as asap-otel-ingest (Gauge, Sum, Histogram, etc.). fn otlp_to_record_count(request: &ExportMetricsServiceRequest) -> usize { @@ -438,7 +543,7 @@ fn otlp_to_record_count(request: &ExportMetricsServiceRequest) -> usize { /// separately for Sketch Payload Flow processing. fn otlp_to_metric_points_and_sketches(request: &ExportMetricsServiceRequest) -> OtlpParseResult { let mut points = Vec::new(); - let mut sketch_points = Vec::new(); + let mut sketch_payloads = Vec::new(); for resource_metrics in &request.resource_metrics { let resource_attrs = resource_metrics .resource @@ -471,13 +576,12 @@ fn otlp_to_metric_points_and_sketches(request: &ExportMetricsServiceRequest) -> if let Some((attr_name, payload)) = get_sketch_payload_from_attrs(&dp.attributes) { - let labels = merge_point_attributes(&base_labels, &dp.attributes); - sketch_points.push(SketchPoint { - metric_name: metric.name.clone(), + sketch_payloads.push(SketchPoint { + name: metric.name.clone(), attr_name, - payload, + labels: merge_point_attributes(&base_labels, &dp.attributes), timestamp_nanos: dp.time_unix_nano, - labels, + payload, }); continue; } @@ -496,13 +600,12 @@ fn otlp_to_metric_points_and_sketches(request: &ExportMetricsServiceRequest) -> if let Some((attr_name, payload)) = get_sketch_payload_from_attrs(&dp.attributes) { - let labels = merge_point_attributes(&base_labels, &dp.attributes); - sketch_points.push(SketchPoint { - metric_name: metric.name.clone(), + sketch_payloads.push(SketchPoint { + name: metric.name.clone(), attr_name, - payload, + labels: merge_point_attributes(&base_labels, &dp.attributes), timestamp_nanos: dp.time_unix_nano, - labels, + payload, }); continue; } @@ -521,13 +624,12 @@ fn otlp_to_metric_points_and_sketches(request: &ExportMetricsServiceRequest) -> if let Some((attr_name, payload)) = get_sketch_payload_from_attrs(&dp.attributes) { - let labels = merge_point_attributes(&base_labels, &dp.attributes); - sketch_points.push(SketchPoint { - metric_name: metric.name.clone(), + sketch_payloads.push(SketchPoint { + name: metric.name.clone(), attr_name, - payload, + labels: merge_point_attributes(&base_labels, &dp.attributes), timestamp_nanos: dp.time_unix_nano, - labels, + payload, }); continue; } @@ -553,13 +655,12 @@ fn otlp_to_metric_points_and_sketches(request: &ExportMetricsServiceRequest) -> if let Some((attr_name, payload)) = get_sketch_payload_from_attrs(&dp.attributes) { - let labels = merge_point_attributes(&base_labels, &dp.attributes); - sketch_points.push(SketchPoint { - metric_name: metric.name.clone(), + sketch_payloads.push(SketchPoint { + name: metric.name.clone(), attr_name, - payload, + labels: merge_point_attributes(&base_labels, &dp.attributes), timestamp_nanos: dp.time_unix_nano, - labels, + payload, }); continue; } @@ -585,13 +686,12 @@ fn otlp_to_metric_points_and_sketches(request: &ExportMetricsServiceRequest) -> if let Some((attr_name, payload)) = get_sketch_payload_from_attrs(&dp.attributes) { - let labels = merge_point_attributes(&base_labels, &dp.attributes); - sketch_points.push(SketchPoint { - metric_name: metric.name.clone(), + sketch_payloads.push(SketchPoint { + name: metric.name.clone(), attr_name, - payload, + labels: merge_point_attributes(&base_labels, &dp.attributes), timestamp_nanos: dp.time_unix_nano, - labels, + payload, }); continue; } @@ -615,7 +715,7 @@ fn otlp_to_metric_points_and_sketches(request: &ExportMetricsServiceRequest) -> } } } - (points, sketch_points) + (points, sketch_payloads) } fn merge_point_attributes( diff --git a/asap-query-engine/src/main.rs b/asap-query-engine/src/main.rs index f39dfafa..533b0541 100644 --- a/asap-query-engine/src/main.rs +++ b/asap-query-engine/src/main.rs @@ -293,31 +293,15 @@ async fn main() -> Result<()> { } }; - // Setup OTLP receiver - let otel_handle = if args.enable_otel_ingest { - let otel_config = OtlpReceiverConfig { - grpc_port: args.otel_grpc_port, - http_port: args.otel_http_port, - }; - let receiver = OtlpReceiver::new(otel_config, store.clone(), streaming_config.clone()); - info!( - "Starting OTLP receiver (gRPC port {}, HTTP port {})", - args.otel_grpc_port, args.otel_http_port - ); - Some(tokio::spawn(async move { - if let Err(e) = receiver.run().await { - error!("OTLP receiver error: {}", e); - } - })) - } else { - None - }; - // Setup precompute engine (replaces standalone Prometheus remote write server) - // Automatically enable when using precompute streaming engine + // Automatically enable when using precompute streaming engine. + // + // NOTE: precompute is constructed BEFORE the OTLP receiver so the receiver + // can obtain an `Arc` handle and push OTLP metrics / sketches + // into the same worker pool (and not just write directly to the store). let enable_precompute = args.enable_prometheus_remote_write || args.streaming_engine == StreamingEngine::Precompute; - let precompute_handle = if enable_precompute { + let (precompute_handle, precompute_ingest_state) = if enable_precompute { let precompute_config = PrecomputeEngineConfig { num_workers: args.precompute_num_workers, ingest_port: args.prometheus_remote_write_port, @@ -333,6 +317,7 @@ async fn main() -> Result<()> { let engine = PrecomputeEngine::new(precompute_config, streaming_config.clone(), output_sink); let worker_diagnostics = engine.diagnostics(); + let ingest_state = engine.ingest_state(); info!( "Starting precompute engine on port {}", args.prometheus_remote_write_port @@ -344,17 +329,51 @@ async fn main() -> Result<()> { spawn_memory_diagnostics(diag_store, Some(worker_diagnostics)).await; }); - Some(tokio::spawn(async move { + let handle = tokio::spawn(async move { if let Err(e) = engine.run().await { error!("Precompute engine error: {}", e); } - })) + }); + (Some(handle), Some(ingest_state)) } else { // Even without precompute, log store diagnostics let diag_store = store.clone(); tokio::spawn(async move { spawn_memory_diagnostics(diag_store, None).await; }); + (None, None) + }; + + // Setup OTLP receiver (after precompute engine so it can share the ingest state) + let otel_handle = if args.enable_otel_ingest { + let otel_config = OtlpReceiverConfig { + grpc_port: args.otel_grpc_port, + http_port: args.otel_http_port, + }; + let receiver = match precompute_ingest_state.clone() { + Some(ingest_state) => { + info!( + "Starting OTLP receiver wired to precompute engine \ + (gRPC port {}, HTTP port {})", + args.otel_grpc_port, args.otel_http_port + ); + OtlpReceiver::with_ingest_state(otel_config, ingest_state) + } + None => { + info!( + "Starting OTLP receiver in log-only mode \ + (precompute engine not enabled; gRPC port {}, HTTP port {})", + args.otel_grpc_port, args.otel_http_port + ); + OtlpReceiver::new(otel_config) + } + }; + Some(tokio::spawn(async move { + if let Err(e) = receiver.run().await { + error!("OTLP receiver error: {}", e); + } + })) + } else { None }; diff --git a/asap-query-engine/src/precompute_engine/engine.rs b/asap-query-engine/src/precompute_engine/engine.rs index 8fd45b8d..b7f48985 100644 --- a/asap-query-engine/src/precompute_engine/engine.rs +++ b/asap-query-engine/src/precompute_engine/engine.rs @@ -24,11 +24,17 @@ pub struct PrecomputeWorkerDiagnostics { /// The top-level precompute engine orchestrator. /// /// Creates worker threads, the series router, and the Axum ingest server. +/// The ingest state (router + agg configs) is built eagerly in `new()` so +/// that other ingest sources (e.g. OTLP) can hold a handle and push data +/// into the same worker pool. pub struct PrecomputeEngine { config: PrecomputeEngineConfig, - streaming_config: Arc, output_sink: Arc, diagnostics: Arc, + ingest_state: Arc, + agg_configs_map: HashMap>, + /// Worker receivers, one per worker. Taken by `run()` when spawning workers. + receivers: Vec>, } impl PrecomputeEngine { @@ -47,26 +53,10 @@ impl PrecomputeEngine { worker_group_counts, worker_watermarks, }); - Self { - config, - streaming_config, - output_sink, - diagnostics, - } - } - - /// Get a handle to worker diagnostics, readable even after `run()` starts. - pub fn diagnostics(&self) -> Arc { - self.diagnostics.clone() - } - /// Start the precompute engine. This spawns worker tasks and the HTTP - /// ingest server, then blocks until shutdown. - pub async fn run(self) -> Result<(), Box> { - let num_workers = self.config.num_workers; - let channel_size = self.config.channel_buffer_size; - - // Build MPSC channels for each worker + // Build MPSC channels for each worker up front. + let num_workers = config.num_workers; + let channel_size = config.channel_buffer_size; let mut senders = Vec::with_capacity(num_workers); let mut receivers = Vec::with_capacity(num_workers); for _ in 0..num_workers { @@ -75,29 +65,66 @@ impl PrecomputeEngine { receivers.push(rx); } - // Build the router + // Build the router that owns the senders; it will be shared via IngestState. let router = SeriesRouter::new(senders); - // Build aggregation config map from streaming config, wrapping each config - // in Arc so all workers share one copy per aggregation (no N×M deep clones). - let agg_configs: HashMap> = self - .streaming_config + // Resolve all aggregation configs from the streaming config. Wrap each + // in Arc so workers can share one copy per aggregation. + let agg_configs_map: HashMap> = streaming_config .get_all_aggregation_configs() .iter() .map(|(&id, cfg)| (id, Arc::new(cfg.clone()))) .collect(); + let agg_configs_vec: Vec> = + agg_configs_map.values().cloned().collect(); + + // Ingest state holds the router and all configs; this is the single + // handle any ingest source uses to push work onto the worker pool. + let ingest_state = Arc::new(IngestState { + router, + samples_ingested: std::sync::atomic::AtomicU64::new(0), + agg_configs: agg_configs_vec, + pass_raw_samples: config.pass_raw_samples, + }); + + Self { + config, + output_sink, + diagnostics, + ingest_state, + agg_configs_map, + receivers, + } + } + + /// Get a handle to worker diagnostics, readable even after `run()` starts. + pub fn diagnostics(&self) -> Arc { + self.diagnostics.clone() + } - // Build a Vec> for the ingest handler - let agg_configs_vec: Vec> = agg_configs.values().cloned().collect(); + /// Get a clonable handle to the shared ingest state. Other ingest sources + /// (OTLP, Kafka, etc.) call this before `run()` to push into the same + /// worker pool as the built-in Prometheus/VM HTTP server. + pub fn ingest_state(&self) -> Arc { + self.ingest_state.clone() + } + + /// Start the precompute engine. This spawns worker tasks and the HTTP + /// ingest server, then blocks until shutdown. + pub async fn run(mut self) -> Result<(), Box> { + let num_workers = self.config.num_workers; + + // Take ownership of receivers (they can only be used once). + let receivers = std::mem::take(&mut self.receivers); - // Spawn workers + // Spawn workers. let mut worker_handles = Vec::with_capacity(num_workers); for (id, rx) in receivers.into_iter().enumerate() { let worker = Worker::new( id, rx, self.output_sink.clone(), - agg_configs.clone(), + self.agg_configs_map.clone(), WorkerRuntimeConfig { max_buffer_per_series: self.config.max_buffer_per_series, allowed_lateness_ms: self.config.allowed_lateness_ms, @@ -120,15 +147,9 @@ impl PrecomputeEngine { num_workers, self.config.ingest_port ); - // Build the ingest state - let ingest_state = Arc::new(IngestState { - router, - samples_ingested: std::sync::atomic::AtomicU64::new(0), - agg_configs: agg_configs_vec, - pass_raw_samples: self.config.pass_raw_samples, - }); + let ingest_state = self.ingest_state.clone(); - // Start flush timer + // Start flush timer. let flush_state = ingest_state.clone(); let flush_interval_ms = self.config.flush_interval_ms; tokio::spawn(async move { @@ -143,7 +164,7 @@ impl PrecomputeEngine { } }); - // Start the Axum HTTP server for ingest (Prometheus + VictoriaMetrics) + // Start the Axum HTTP server for ingest (Prometheus + VictoriaMetrics). let app = Router::new() .route("/api/v1/write", post(handle_prometheus_ingest)) .route("/api/v1/import", post(handle_victoriametrics_ingest)) @@ -155,7 +176,7 @@ impl PrecomputeEngine { let listener = TcpListener::bind(&addr).await?; axum::serve(listener, app).await?; - // Wait for workers to finish (this only happens on shutdown) + // Wait for workers to finish (this only happens on shutdown). for handle in worker_handles { let _ = handle.await; } diff --git a/asap-query-engine/src/precompute_engine/ingest_handler.rs b/asap-query-engine/src/precompute_engine/ingest_handler.rs index 03b9ac51..daa3275f 100644 --- a/asap-query-engine/src/precompute_engine/ingest_handler.rs +++ b/asap-query-engine/src/precompute_engine/ingest_handler.rs @@ -10,13 +10,27 @@ use std::time::Instant; use tracing::warn; /// Shared state for the ingest HTTP handler. -pub(crate) struct IngestState { - pub(crate) router: SeriesRouter, - pub(crate) samples_ingested: std::sync::atomic::AtomicU64, +/// +/// Holds the worker router plus the aggregation configs needed for group-key +/// extraction. A single instance is shared by the Prometheus/VictoriaMetrics +/// HTTP ingest server and any other ingest source that wants to route into +/// the same worker pool (e.g. the OTLP receiver). +pub struct IngestState { + pub router: SeriesRouter, + pub samples_ingested: std::sync::atomic::AtomicU64, /// Aggregation configs for group-key extraction. - pub(crate) agg_configs: Vec>, + pub agg_configs: Vec>, /// When true, skip group-key extraction and pass raw samples through. - pub(crate) pass_raw_samples: bool, + pub pass_raw_samples: bool, +} + +impl IngestState { + /// Extract the group key for a series key against a given aggregation + /// config. Re-exports the module-private helper so that out-of-module + /// ingest sources (e.g. OTLP) can reuse it. + pub fn extract_group_key_for(series_key: &str, config: &AggregationConfig) -> String { + extract_group_key(series_key, config) + } } /// Extract the group key (grouping label values joined by semicolons) diff --git a/asap-query-engine/src/precompute_engine/mod.rs b/asap-query-engine/src/precompute_engine/mod.rs index 0edda2fb..275bc93e 100644 --- a/asap-query-engine/src/precompute_engine/mod.rs +++ b/asap-query-engine/src/precompute_engine/mod.rs @@ -1,7 +1,7 @@ pub mod accumulator_factory; pub mod config; mod engine; -mod ingest_handler; +pub mod ingest_handler; pub mod output_sink; pub mod series_buffer; pub mod series_router; @@ -9,3 +9,4 @@ pub mod window_manager; pub mod worker; pub use engine::{PrecomputeEngine, PrecomputeWorkerDiagnostics}; +pub use ingest_handler::IngestState; diff --git a/asap-query-engine/src/precompute_engine/series_router.rs b/asap-query-engine/src/precompute_engine/series_router.rs index 45ca1b2a..39936c16 100644 --- a/asap-query-engine/src/precompute_engine/series_router.rs +++ b/asap-query-engine/src/precompute_engine/series_router.rs @@ -1,11 +1,12 @@ +use crate::data_model::AggregateCore; use futures::future::try_join_all; use std::collections::HashMap; +use std::fmt; use std::time::Instant; use tokio::sync::mpsc; use xxhash_rust::xxh64::xxh64; /// A message sent from the router to a worker. -#[derive(Debug)] pub enum WorkerMessage { /// A batch of samples for the same series, routed by series key. /// Used in `pass_raw_samples` mode where no aggregation is needed. @@ -28,12 +29,75 @@ pub enum WorkerMessage { samples: Vec<(String, i64, f64)>, ingest_received_at: Instant, }, + /// A pre-built accumulator destined for a specific (agg_id, group_key) + /// pane. The worker merges it into that pane's existing accumulator + /// (or inserts it if the pane is empty) via `AggregateCore::merge_with`. + /// + /// Produced by ingest sources that deliver pre-aggregated sketches — + /// e.g. the OTLP receiver when DataCollector emits KLL / CountMin / + /// CountSketch payloads on a `SketchEnvelope`. Lets the precompute + /// engine perform further window-aligned aggregation on sketches the + /// same way it does on raw samples. + AccumulatorInput { + agg_id: u64, + /// Grouping label values joined by semicolons, matching the + /// format produced by `IngestState::extract_group_key_for`. + group_key: String, + /// Wall-clock timestamp the sketch refers to (millis since epoch). + /// Used to place the sketch into the correct pane. + timestamp_ms: i64, + /// The incoming accumulator to be merged into the pane. + accumulator: Box, + ingest_received_at: Instant, + }, /// Signal the worker to flush/check idle windows. Flush, /// Graceful shutdown. Shutdown, } +impl fmt::Debug for WorkerMessage { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::RawSamples { + series_key, + samples, + .. + } => f + .debug_struct("RawSamples") + .field("series_key", series_key) + .field("sample_count", &samples.len()) + .finish(), + Self::GroupSamples { + agg_id, + group_key, + samples, + .. + } => f + .debug_struct("GroupSamples") + .field("agg_id", agg_id) + .field("group_key", group_key) + .field("sample_count", &samples.len()) + .finish(), + Self::AccumulatorInput { + agg_id, + group_key, + timestamp_ms, + accumulator, + .. + } => f + .debug_struct("AccumulatorInput") + .field("agg_id", agg_id) + .field("group_key", group_key) + .field("timestamp_ms", timestamp_ms) + .field("accumulator_type", &accumulator.type_name()) + .finish(), + Self::Flush => f.write_str("Flush"), + Self::Shutdown => f.write_str("Shutdown"), + } + } +} + /// Routes incoming samples to one of N workers based on a consistent hash. pub struct SeriesRouter { senders: Vec>, @@ -65,6 +129,9 @@ impl SeriesRouter { WorkerMessage::GroupSamples { agg_id, group_key, .. } => self.worker_for_group(*agg_id, group_key), + WorkerMessage::AccumulatorInput { + agg_id, group_key, .. + } => self.worker_for_group(*agg_id, group_key), WorkerMessage::RawSamples { series_key, .. } => self.worker_for(series_key), _ => 0, }; diff --git a/asap-query-engine/src/precompute_engine/worker.rs b/asap-query-engine/src/precompute_engine/worker.rs index 3e86980f..29af67da 100644 --- a/asap-query-engine/src/precompute_engine/worker.rs +++ b/asap-query-engine/src/precompute_engine/worker.rs @@ -23,8 +23,15 @@ use tracing::{debug, debug_span, info, warn}; struct GroupState { config: Arc, window_manager: WindowManager, - /// Active panes keyed by pane_start_ms. + /// Active panes for raw-sample accumulation, keyed by pane_start_ms. active_panes: BTreeMap>, + /// Active panes for pre-built accumulator inputs (e.g. OTLP-delivered + /// sketches), keyed by pane_start_ms. Each entry is the running merge + /// of every accumulator that landed in that pane's time range. Kept + /// separate from `active_panes` because sketches come in as opaque + /// `Box` objects and do not share the updater + /// machinery used for incremental sample updates. + sketch_panes: BTreeMap>, /// Per-group watermark: tracks the maximum timestamp seen across all /// series in this group on this worker. previous_watermark_ms: i64, @@ -157,6 +164,38 @@ impl Worker { "e2e: ingest->worker complete (raw)" ); } + WorkerMessage::AccumulatorInput { + agg_id, + group_key, + timestamp_ms, + accumulator, + ingest_received_at, + } => { + let _span = debug_span!( + "worker_process_accumulator", + worker_id = self.id, + agg_id, + group = %group_key, + timestamp_ms, + accumulator_type = accumulator.type_name(), + ) + .entered(); + if let Err(e) = self.process_accumulator_input( + agg_id, + &group_key, + timestamp_ms, + accumulator, + ) { + warn!( + "Worker {} accumulator input error for ({}, {}): {}", + self.id, agg_id, group_key, e + ); + } + debug!( + e2e_latency_us = ingest_received_at.elapsed().as_micros() as u64, + "e2e: ingest->worker complete (accumulator)" + ); + } WorkerMessage::Flush => { if let Err(e) = self.flush_all() { warn!("Worker {} flush error: {}", self.id, e); @@ -193,6 +232,7 @@ impl Worker { window_manager: WindowManager::new(config.window_size, config.slide_interval), config: Arc::clone(config), active_panes: BTreeMap::new(), + sketch_panes: BTreeMap::new(), previous_watermark_ms: i64::MIN, }; self.group_states.insert(key.clone(), gs); @@ -337,6 +377,150 @@ impl Worker { Ok(()) } + /// Process a pre-built accumulator (e.g. an OTLP-delivered sketch) for a + /// specific (agg_id, group_key) pane. + /// + /// The incoming accumulator is merged into `sketch_panes[pane_start]` via + /// `AggregateCore::merge_with`. If the pane is empty the accumulator is + /// installed as-is. If the pane's window has already closed, the late + /// data policy decides whether to drop it or forward it to the sink as + /// a standalone output so no data is silently lost. + /// + /// Unlike `process_group_samples`, this path does not touch + /// `active_panes` — sketches live in their own pane map and get merged + /// at window close (see `merge_sketch_panes_for_window`). + pub fn process_accumulator_input( + &mut self, + agg_id: u64, + group_key: &str, + timestamp_ms: i64, + incoming: Box, + ) -> Result<(), Box> { + let worker_id = self.id; + let allowed_lateness_ms = self.allowed_lateness_ms; + let late_data_policy = self.late_data_policy; + + if self.get_or_create_group_state(agg_id, group_key).is_none() { + warn!( + "Worker {} skipping accumulator input for unknown agg_id={}, group_key={}", + self.id, agg_id, group_key + ); + return Ok(()); + } + let state = self + .group_states + .get_mut(&(agg_id, group_key.to_string())) + .unwrap(); + + let previous_wm = state.previous_watermark_ms; + let current_wm = if timestamp_ms > previous_wm { + timestamp_ms + } else { + previous_wm + }; + + let mut emit_batch: Vec<(PrecomputedOutput, Box)> = Vec::new(); + + // Late-arrival check against the existing watermark. + let too_late = previous_wm != i64::MIN && timestamp_ms < previous_wm - allowed_lateness_ms; + let pane_start = state.window_manager.pane_start_for(timestamp_ms); + let pane_closed = !state.sketch_panes.contains_key(&pane_start) + && current_wm >= pane_start + state.window_manager.window_size_ms(); + + if too_late || pane_closed { + match late_data_policy { + LateDataPolicy::Drop => { + debug!( + "Worker {} dropping late accumulator input for group ({}, {}): ts={} watermark={}", + worker_id, agg_id, group_key, timestamp_ms, previous_wm + ); + } + LateDataPolicy::ForwardToStore => { + let window_start = pane_start; + let window_end = pane_start + state.window_manager.window_size_ms(); + let key = build_group_key_label_values(group_key); + let output = PrecomputedOutput::new( + window_start as u64, + window_end as u64, + Some(key), + agg_id, + ); + emit_batch.push((output, incoming)); + debug!( + "Forwarding late accumulator input to store for evicted pane [{}, {})", + pane_start, + pane_start + state.window_manager.slide_interval_ms() + ); + self.output_sink.emit_batch(emit_batch)?; + } + } + return Ok(()); + } + + // Merge into the sketch pane covering this timestamp. + match state.sketch_panes.remove(&pane_start) { + Some(existing) => { + let merged = existing + .merge_with(incoming.as_ref()) + .map_err(|e| format!("merge_with failed for pane {pane_start}: {e}"))?; + state.sketch_panes.insert(pane_start, merged); + } + None => { + state.sketch_panes.insert(pane_start, incoming); + } + } + + // Check for closed windows and emit merged outputs. + let closed = state.window_manager.closed_windows(previous_wm, current_wm); + for window_start in &closed { + let (_, window_end) = state.window_manager.window_bounds(*window_start); + let pane_starts = state.window_manager.panes_for_window(*window_start); + + // Emit from the raw-sample pane map (in case both sources are + // populated for the same group; rare but supported). + if let Some(accumulator) = merge_panes_for_window(&mut state.active_panes, &pane_starts) + { + let key = build_group_key_label_values(group_key); + let output = PrecomputedOutput::new( + *window_start as u64, + window_end as u64, + Some(key), + agg_id, + ); + emit_batch.push((output, accumulator)); + } + + // Emit from the sketch pane map. + if let Some(accumulator) = + merge_sketch_panes_for_window(&mut state.sketch_panes, &pane_starts) + { + let key = build_group_key_label_values(group_key); + let output = PrecomputedOutput::new( + *window_start as u64, + window_end as u64, + Some(key), + agg_id, + ); + emit_batch.push((output, accumulator)); + } + } + + state.previous_watermark_ms = current_wm; + + if !emit_batch.is_empty() { + debug!( + "Worker {} emitting {} sketch outputs for group ({}, {})", + worker_id, + emit_batch.len(), + agg_id, + group_key + ); + self.output_sink.emit_batch(emit_batch)?; + } + + Ok(()) + } + /// Raw fast-path: emit each sample as a standalone `SumAccumulator`. pub fn process_samples_raw( &self, @@ -428,6 +612,19 @@ impl Worker { ); emit_batch.push((output, accumulator)); } + + if let Some(accumulator) = + merge_sketch_panes_for_window(&mut state.sketch_panes, &pane_starts) + { + let key = build_group_key_label_values(group_key); + let output = PrecomputedOutput::new( + *window_start as u64, + window_end as u64, + Some(key), + *agg_id, + ); + emit_batch.push((output, accumulator)); + } } // Update group watermark to reflect the advancement. @@ -637,6 +834,38 @@ fn merge_panes_for_window( merged } +/// Merge pre-built accumulator panes for a window. +/// +/// Equivalent to `merge_panes_for_window` but operating on the sketch pane +/// map (`Box` directly). The oldest pane is destructively +/// taken (it will never be needed by a later window); subsequent panes are +/// cloned so that still-open overlapping windows can still read them. +fn merge_sketch_panes_for_window( + sketch_panes: &mut BTreeMap>, + pane_starts: &[i64], +) -> Option> { + let mut merged: Option> = None; + + for (i, &ps) in pane_starts.iter().enumerate() { + let pane_acc: Option> = if i == 0 { + // Oldest pane: destructive take + evict + sketch_panes.remove(&ps) + } else { + // Shared pane: non-destructive clone + sketch_panes.get(&ps).map(|acc| acc.clone_boxed_core()) + }; + + if let Some(acc) = pane_acc { + merged = Some(match merged { + None => acc, + Some(existing) => existing.merge_with(acc.as_ref()).unwrap_or(existing), + }); + } + } + + merged +} + #[cfg(test)] mod tests { use super::*;