datacollector integration

Author: GnaneshGnani

asap-query-engine/Cargo.toml

@@ -56,6 +56,7 @@ lazy_static = "1.4"
 zstd = "0.13"
 reqwest = { version = "0.11", features = ["json"] }
 tracing-appender = "0.2"
+sketchlib-rust = { git = "https://github.com/ProjectASAP/sketchlib-rust", branch = "protobuf-xxh3-shared" }
 
 [dev-dependencies]
 ctor = "0.2"

asap-query-engine/src/drivers/ingest/otel.rs

@@ -5,14 +5,18 @@
 //! handoff to precompute engine as TODO.
 
 use std::collections::HashMap;
+use std::io::Read;
 
 use axum::{body::Bytes, extract::State, routing::post, Json, Router};
+use flate2::read::GzDecoder;
 use opentelemetry_proto::tonic::collector::metrics::v1::{
     metrics_service_server::MetricsService, ExportMetricsServiceRequest,
     ExportMetricsServiceResponse,
 };
+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 sketchlib_rust::proto::sketchlib::{sketch_envelope, SketchEnvelope};
 use tonic::{Request, Response, Status};
 use tracing::{debug, error, info};
 
@@ -80,25 +84,49 @@ impl MetricsService for MetricsServiceImpl {
         &self,
         request: Request<ExportMetricsServiceRequest>,
     ) -> Result<Response<ExportMetricsServiceResponse>, Status> {
+        debug!("OTLP received request via gRPC");
         let req = request.into_inner();
-        process_otlp_request(&req);
+        process_otlp_request(&req, "gRPC");
+        debug!("OTLP sending response via gRPC");
         Ok(Response::new(ExportMetricsServiceResponse {
             partial_success: None,
         }))
     }
 }
 
 async fn handle_otlp_http(
+    headers: axum::http::HeaderMap,
     State(_state): State<()>,
     body: Bytes,
 ) -> Result<Json<serde_json::Value>, (axum::http::StatusCode, String)> {
+    debug!("OTLP received request via HTTP, body_bytes={}", body.len());
+    let body = if let Some(enc) = headers.get(axum::http::header::CONTENT_ENCODING) {
+        let enc = enc.to_str().unwrap_or("").trim().to_ascii_lowercase();
+        if enc == "gzip" {
+            let mut decoder = GzDecoder::new(body.as_ref());
+            let mut out = Vec::new();
+            decoder.read_to_end(&mut out).map_err(|e| {
+                (
+                    axum::http::StatusCode::BAD_REQUEST,
+                    format!("Gzip decode error: {}", e),
+                )
+            })?;
+            Bytes::from(out)
+        } else {
+            body
+        }
+    } else {
+        body
+    };
+
     let req = ExportMetricsServiceRequest::decode(body.as_ref()).map_err(|e| {
         (
             axum::http::StatusCode::BAD_REQUEST,
             format!("Protobuf decode error: {}", e),
         )
     })?;
-    process_otlp_request(&req);
+    process_otlp_request(&req, "HTTP");
+    debug!("OTLP sending response via HTTP");
     Ok(Json(serde_json::json!({"rejected": 0})))
 }
 
@@ -111,15 +139,101 @@ pub struct MetricPoint {
     pub value: f64,
 }
 
-fn process_otlp_request(request: &ExportMetricsServiceRequest) {
+fn format_series_key(name: &str, labels: &HashMap<String, String>) -> String {
+    let mut pairs: Vec<_> = labels.iter().collect();
+    pairs.sort_by_key(|(k, _)| *k);
+    let labels_str = pairs
+        .iter()
+        .map(|(k, v)| format!("{}={}", k, v))
+        .collect::<Vec<_>>()
+        .join(",");
+    format!("{}{{{}}}", name, labels_str)
+}
+
+fn get_sketch_payload_from_attrs(
+    attrs: &[opentelemetry_proto::tonic::common::v1::KeyValue],
+) -> Option<(String, Vec<u8>)> {
+    for kv in attrs {
+        match kv.key.as_str() {
+            "kll.sketch_payload" | "cms.sketch_payload" | "countsketch.sketch_payload" => {
+                if let Some(value) = &kv.value {
+                    if let Some(AnyValueVariant::BytesValue(bytes)) = &value.value {
+                        return Some((kv.key.clone(), bytes.clone()));
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+    None
+}
+
+fn log_sketch_envelope_type(attr_name: &str, payload: &[u8], metric_name: &str) {
+    match SketchEnvelope::decode(payload) {
+        Ok(env) => {
+            let sketch_type = 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",
+            };
+            debug!(
+                "OTLP Sketches: metric='{}' attr='{}' payload_bytes={} sketch_type={}",
+                metric_name,
+                attr_name,
+                payload.len(),
+                sketch_type
+            );
+        }
+        Err(e) => {
+            debug!(
+                "OTLP Sketches: metric='{}' attr='{}' payload_bytes={} decode_error='{}'",
+                metric_name,
+                attr_name,
+                payload.len(),
+                e
+            );
+        }
+    }
+}
+
+fn process_otlp_request(request: &ExportMetricsServiceRequest, transport: &str) {
     let resource_count = request.resource_metrics.len();
     let total_points = otlp_to_record_count(request);
-    debug!(
-        "OTLP ingest: received {} resource metrics, {} total data points",
-        resource_count, total_points
-    );
+    if resource_count > 0 || total_points > 0 {
+        debug!(
+            "OTLP ingest: received {} resource metrics, {} total data points (transport={})",
+            resource_count, total_points, transport
+        );
+    }
+
+    let (points, sketch_payloads) = otlp_to_metric_points_and_sketches(request);
+
+    for (metric_name, attr_name, payload) in &sketch_payloads {
+        log_sketch_envelope_type(attr_name, payload, metric_name);
+    }
+    if !sketch_payloads.is_empty() {
+        debug!(
+            "OTLP Sketch Payload Flow: received {} sketch payload(s), decoded successfully",
+            sketch_payloads.len()
+        );
+    }
 
-    let points = otlp_to_metric_points(request);
+    let mut by_series: HashMap<String, usize> = HashMap::new();
+    for point in &points {
+        let key = format_series_key(&point.name, &point.labels);
+        *by_series.entry(key).or_insert(0) += 1;
+    }
+    if !by_series.is_empty() {
+        debug!(
+            "OTLP Raw Metrics Flow: received {} raw metric series",
+            by_series.len()
+        );
+        for (series, count) in by_series {
+            debug!("OTLP Raw Metrics Flow: series {} count={}", series, count);
+        }
+    }
     if let Some(first) = points.first() {
         debug!(
             "OTLP parse example: {} {:?} @{}ns = {}",
@@ -178,8 +292,13 @@ fn otlp_to_record_count(request: &ExportMetricsServiceRequest) -> usize {
 }
 
 /// Parse OTLP request and convert to metric data points (name, labels, timestamp, value).
-fn otlp_to_metric_points(request: &ExportMetricsServiceRequest) -> Vec<MetricPoint> {
+/// Data points with sketch payloads in attributes are excluded from points and returned
+/// separately for Sketch Payload Flow processing.
+fn otlp_to_metric_points_and_sketches(
+    request: &ExportMetricsServiceRequest,
+) -> (Vec<MetricPoint>, Vec<(String, String, Vec<u8>)>) {
     let mut points = Vec::new();
+    let mut sketch_payloads = Vec::new();
     for resource_metrics in &request.resource_metrics {
         let resource_attrs = resource_metrics
             .resource
@@ -209,6 +328,12 @@ fn otlp_to_metric_points(request: &ExportMetricsServiceRequest) -> Vec<MetricPoi
                 match &metric.data {
                     Some(Data::Gauge(g)) => {
                         for dp in &g.data_points {
+                            if let Some((attr_name, payload)) =
+                                get_sketch_payload_from_attrs(&dp.attributes)
+                            {
+                                sketch_payloads.push((metric.name.clone(), attr_name, payload));
+                                continue;
+                            }
                             let labels = merge_point_attributes(&base_labels, &dp.attributes);
                             let value = number_value_to_f64(&dp.value);
                             points.push(MetricPoint {
@@ -221,6 +346,12 @@ fn otlp_to_metric_points(request: &ExportMetricsServiceRequest) -> Vec<MetricPoi
                     }
                     Some(Data::Sum(s)) => {
                         for dp in &s.data_points {
+                            if let Some((attr_name, payload)) =
+                                get_sketch_payload_from_attrs(&dp.attributes)
+                            {
+                                sketch_payloads.push((metric.name.clone(), attr_name, payload));
+                                continue;
+                            }
                             let labels = merge_point_attributes(&base_labels, &dp.attributes);
                             let value = number_value_to_f64(&dp.value);
                             points.push(MetricPoint {
@@ -233,6 +364,12 @@ fn otlp_to_metric_points(request: &ExportMetricsServiceRequest) -> Vec<MetricPoi
                     }
                     Some(Data::Histogram(hist)) => {
                         for dp in &hist.data_points {
+                            if let Some((attr_name, payload)) =
+                                get_sketch_payload_from_attrs(&dp.attributes)
+                            {
+                                sketch_payloads.push((metric.name.clone(), attr_name, payload));
+                                continue;
+                            }
                             let labels = merge_point_attributes(&base_labels, &dp.attributes);
                             if let Some(sum) = dp.sum {
                                 points.push(MetricPoint {
@@ -252,6 +389,12 @@ fn otlp_to_metric_points(request: &ExportMetricsServiceRequest) -> Vec<MetricPoi
                     }
                     Some(Data::ExponentialHistogram(eh)) => {
                         for dp in &eh.data_points {
+                            if let Some((attr_name, payload)) =
+                                get_sketch_payload_from_attrs(&dp.attributes)
+                            {
+                                sketch_payloads.push((metric.name.clone(), attr_name, payload));
+                                continue;
+                            }
                             let labels = merge_point_attributes(&base_labels, &dp.attributes);
                             if let Some(sum) = dp.sum {
                                 points.push(MetricPoint {
@@ -271,6 +414,12 @@ fn otlp_to_metric_points(request: &ExportMetricsServiceRequest) -> Vec<MetricPoi
                     }
                     Some(Data::Summary(sm)) => {
                         for dp in &sm.data_points {
+                            if let Some((attr_name, payload)) =
+                                get_sketch_payload_from_attrs(&dp.attributes)
+                            {
+                                sketch_payloads.push((metric.name.clone(), attr_name, payload));
+                                continue;
+                            }
                             let labels = merge_point_attributes(&base_labels, &dp.attributes);
                             points.push(MetricPoint {
                                 name: format!("{}_sum", metric.name),
@@ -291,7 +440,7 @@ fn otlp_to_metric_points(request: &ExportMetricsServiceRequest) -> Vec<MetricPoi
             }
         }
     }
-    points
+    (points, sketch_payloads)
 }
 
 fn merge_point_attributes(