From 4b4425e4d3af5277945b68c6c42d96fad72b2b71 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Tue, 14 Apr 2026 15:16:54 -0400 Subject: [PATCH 1/2] feat: MessagePack encoding parity for modified-OTLP sketch hot path (PR I) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds MessagePack as a parallel encoding option alongside the protobuf-encoded sketchlib `*State` path that PR B and PR C wired up for all five modified-OTLP `Metric.data` sketch variants. The new encoding is a cross-language contract between DataCollector (Go, sketchlib-go) and ASAPQuery-backend (Rust, sketch-core) — both sides serialize the same cross-language sketch-core wire struct via MessagePack. This gives sketchlib-go a practical alternative to prost-encoded sketchlib proto without touching existing paths. ## Proto Adds two new enum variants to each of the five per-sketch encoding enums in the vendored `opentelemetry.proto.metrics.v1`: * `*_ENCODING_MSGPACK = 3` * `*_ENCODING_MSGPACK_DELTA = 4` Backward-compatible: existing producers emitting tags 0-2 see no change. The vendored proto is the consumer side — DataCollector's upstream proto can adopt matching enum names in a parallel PR without breaking either side, since protobuf enums are just integer tags. ## Dispatcher (`drivers/ingest/otel.rs`) `decode_modified_otlp_sketch_bytes` now matches on the full `encoding` i32 value rather than only accepting `PROTO = 1`: * `PROTO = 1` → per-sketch `from_sketchlib_proto_bytes` (existing) * `MSGPACK = 3` → per-sketch `from_msgpack_bytes` (new) * `PROTO_DELTA = 2`, `MSGPACK_DELTA = 4` → deferred, caller falls through to §5.2 fallback ## Per-sketch decoders Each of the five accumulators gains a `from_msgpack_bytes(buffer)` constructor that wraps the sketch-core `deserialize_msgpack` already present on the underlying type: * `CountMinSketchAccumulator::from_msgpack_bytes` * `CountSketchAccumulator::from_msgpack_bytes` * `DatasketchesKLLAccumulator::from_msgpack_bytes` — and this is notable: unlike the `_ENCODING_PROTO` path (which does lossy statistical reconstruction via `update()` replay because sketchlib's `KllState` proto elides level structure the Rust backend types keep private), the msgpack path is a **bit-identical round-trip** because sketch-core's `KllSketch` serializes its full internal state to msgpack. This gives Go producers a way to send KLL sketches losslessly to the Rust backend. * `DDSketchAccumulator::from_msgpack_bytes` * `HllSketchAccumulator::from_msgpack_bytes` ## Tests Unit tests — 6 new, 2 per decoder (round-trip + garbage rejection) on `CountSketchAccumulator`, `DDSketchAccumulator`, and `HllSketchAccumulator`. CountMin + KLL msgpack paths already had coverage through `deserialize_from_bytes_arroyo` tests. E2e test — new `e2e_count_min_sketch_msgpack_modified_otlp_path` in `tests/e2e_modified_otlp_sketch_path.rs`. Builds a real `CountMinSketch` in sketch-core, serializes via msgpack, POSTs through the OTLP HTTP receiver with `encoding = COUNT_MIN_SKETCH_ENCODING_MSGPACK`, closes a window, and verifies the stored accumulator preserves per-key counts. Covers the dispatcher path end-to-end for the msgpack branch; the other four variants share the same branch so per-variant smoke tests aren't needed for dispatcher coverage (per-variant msgpack round-trips are covered by the unit tests). ## Drive-by: fix two clippy warnings inherited from main `cargo clippy --all-targets -- -D warnings` on rust 1.91 flags two pre-existing issues in the persistence PR #4 code path: * `cache.rs` unnecessary `u64 as u64` casts * `part.rs` test `&[snap.clone()]` → `std::slice::from_ref(&snap)` Fixed in this PR since the CI clippy gate would otherwise block PR I. ## Validation - cargo check --all-targets: clean - cargo clippy --all-targets -- -D warnings: clean - cargo fmt --check: clean - cargo test -p sketch-core --lib: 54 passed - cargo test -p query_engine_rust --lib: 480 passed (includes 6 new msgpack unit tests) - cargo test -p query_engine_rust --test e2e_modified_otlp_sketch_path: 6 passed (5 existing PROTO variants + 1 new MSGPACK variant) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../proto/metrics/v1/metrics.proto | 23 +++ asap-query-engine/src/drivers/ingest/otel.rs | 82 +++++---- .../count_min_sketch_accumulator.rs | 12 ++ .../count_sketch_accumulator.rs | 32 ++++ .../datasketches_kll_accumulator.rs | 15 ++ .../dd_sketch_accumulator.rs | 30 ++++ .../hll_sketch_accumulator.rs | 36 ++++ .../simple_map_store/persistence/cache.rs | 2 +- .../simple_map_store/persistence/part.rs | 3 +- .../tests/e2e_modified_otlp_sketch_path.rs | 158 ++++++++++++++++++ 10 files changed, 357 insertions(+), 36 deletions(-) diff --git a/asap-common/dependencies/rs/asap_otel_proto/proto/opentelemetry/proto/metrics/v1/metrics.proto b/asap-common/dependencies/rs/asap_otel_proto/proto/opentelemetry/proto/metrics/v1/metrics.proto index 3a909e47..bdb08d01 100644 --- a/asap-common/dependencies/rs/asap_otel_proto/proto/opentelemetry/proto/metrics/v1/metrics.proto +++ b/asap-common/dependencies/rs/asap_otel_proto/proto/opentelemetry/proto/metrics/v1/metrics.proto @@ -384,6 +384,11 @@ enum DDSketchEncoding { // github.com/DataDog/sketches-go/ddsketch/pb/sketchpb.DDSketch message. DDSKETCH_ENCODING_PROTO = 1; DDSKETCH_ENCODING_PROTO_DELTA = 2; + // DDSKETCH_ENCODING_MSGPACK indicates the sketch bytes are the MessagePack + // serialization of the cross-language sketch-core `DdSketch` struct + // (rmp_serde encoding on Rust, matching struct on Go via sketchlib-go). + DDSKETCH_ENCODING_MSGPACK = 3; + DDSKETCH_ENCODING_MSGPACK_DELTA = 4; } // KLLSketch represents the type of a metric that encodes measurements using the KLL quantile sketch. @@ -415,6 +420,10 @@ message KLLSketchDataPoint { enum KLLSketchEncoding { KLL_SKETCH_ENCODING_UNSPECIFIED = 0; KLL_SKETCH_ENCODING_PROTO = 1; + // KLL_SKETCH_ENCODING_MSGPACK indicates the sketch bytes are the MessagePack + // serialization of the cross-language sketch-core `KllSketch` struct. + KLL_SKETCH_ENCODING_MSGPACK = 3; + KLL_SKETCH_ENCODING_MSGPACK_DELTA = 4; } // CountSketch represents the type of a metric that encodes frequency estimations using CountSketch. @@ -446,6 +455,10 @@ enum CountSketchEncoding { COUNT_SKETCH_ENCODING_UNSPECIFIED = 0; COUNT_SKETCH_ENCODING_PROTO = 1; COUNT_SKETCH_ENCODING_DELTA = 2; + // COUNT_SKETCH_ENCODING_MSGPACK indicates the sketch bytes are the MessagePack + // serialization of the cross-language sketch-core `CountSketch` struct. + COUNT_SKETCH_ENCODING_MSGPACK = 3; + COUNT_SKETCH_ENCODING_MSGPACK_DELTA = 4; } // CountMinSketch represents the type of a metric encoding frequency estimations using CountMinSketch. @@ -477,6 +490,11 @@ enum CountMinSketchEncoding { COUNT_MIN_SKETCH_ENCODING_UNSPECIFIED = 0; COUNT_MIN_SKETCH_ENCODING_PROTO = 1; COUNT_MIN_SKETCH_ENCODING_DELTA = 2; + // COUNT_MIN_SKETCH_ENCODING_MSGPACK indicates the sketch bytes are the + // MessagePack serialization of the cross-language sketch-core + // `CountMinSketch` wire struct. + COUNT_MIN_SKETCH_ENCODING_MSGPACK = 3; + COUNT_MIN_SKETCH_ENCODING_MSGPACK_DELTA = 4; } // HLLSketch represents the type of a metric that encodes cardinality estimations using HyperLogLog. @@ -508,6 +526,11 @@ enum HLLSketchEncoding { HLL_SKETCH_ENCODING_UNSPECIFIED = 0; HLL_SKETCH_ENCODING_PROTO = 1; HLL_SKETCH_ENCODING_DELTA = 2; + // HLL_SKETCH_ENCODING_MSGPACK indicates the sketch bytes are the + // MessagePack serialization of the cross-language sketch-core + // `HllSketch` struct (registers + variant + precision + HIP state). + HLL_SKETCH_ENCODING_MSGPACK = 3; + HLL_SKETCH_ENCODING_MSGPACK_DELTA = 4; } // AggregationTemporality defines how a metric aggregator reports aggregated diff --git a/asap-query-engine/src/drivers/ingest/otel.rs b/asap-query-engine/src/drivers/ingest/otel.rs index cf8a42c7..ec4b8e80 100644 --- a/asap-query-engine/src/drivers/ingest/otel.rs +++ b/asap-query-engine/src/drivers/ingest/otel.rs @@ -701,43 +701,57 @@ fn decode_modified_otlp_sketch_bytes( DatasketchesKLLAccumulator, HllSketchAccumulator, }; - // The encoding value is the raw i32 from the proto enum. We only - // accept ENCODING_PROTO (= 1) for now; ENCODING_PROTO_DELTA (= 2) - // and ENCODING_MSGPACK / ENCODING_MSGPACK_DELTA (PR I) are routed - // to a "not yet implemented" Err so the caller falls through to - // the §5.2 fallback. + // The encoding value is the raw i32 from the per-sketch encoding + // enum. All five sketch variants share the same wire tag layout for + // tags 1–4, so we can match on one set of constants here: + // + // 1 — ENCODING_PROTO (PR B / PR C: sketchlib proto-encoded) + // 2 — ENCODING_PROTO_DELTA (delta transmission; deferred — falls + // through to §5.2 fallback) + // 3 — ENCODING_MSGPACK (PR I: cross-language sketch-core + // msgpack wire format — this dispatcher) + // 4 — ENCODING_MSGPACK_DELTA (deferred same as PROTO_DELTA) const ENCODING_PROTO: i32 = 1; - - if encoding != ENCODING_PROTO { - return Err(format!( + const ENCODING_MSGPACK: i32 = 3; + + match encoding { + ENCODING_PROTO => match kind { + SketchKind::CountMin => Ok(Box::new( + CountMinSketchAccumulator::from_sketchlib_proto_bytes(bytes)?, + )), + SketchKind::CountSketch => Ok(Box::new( + CountSketchAccumulator::from_sketchlib_proto_bytes(bytes)?, + )), + SketchKind::Kll => Ok(Box::new( + DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(bytes)?, + )), + SketchKind::DdSketch => Ok(Box::new(DDSketchAccumulator::from_sketchlib_proto_bytes( + bytes, + )?)), + SketchKind::Hll => Ok(Box::new(HllSketchAccumulator::from_sketchlib_proto_bytes( + bytes, + )?)), + }, + ENCODING_MSGPACK => match kind { + SketchKind::CountMin => Ok(Box::new(CountMinSketchAccumulator::from_msgpack_bytes( + bytes, + )?)), + SketchKind::CountSketch => { + Ok(Box::new(CountSketchAccumulator::from_msgpack_bytes(bytes)?)) + } + SketchKind::Kll => Ok(Box::new(DatasketchesKLLAccumulator::from_msgpack_bytes( + bytes, + )?)), + SketchKind::DdSketch => Ok(Box::new(DDSketchAccumulator::from_msgpack_bytes(bytes)?)), + SketchKind::Hll => Ok(Box::new(HllSketchAccumulator::from_msgpack_bytes(bytes)?)), + }, + _ => Err(format!( "modified-OTLP sketch encoding {encoding} not yet supported \ - (only ENCODING_PROTO = 1 is wired today; deltas tracked in PR C, \ - msgpack tracked in PR I)" + (PROTO = 1 and MSGPACK = 3 are wired; PROTO_DELTA = 2 and \ + MSGPACK_DELTA = 4 are deferred — caller falls through to §5.2 \ + fallback)" ) - .into()); - } - - match kind { - SketchKind::CountMin => { - let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(bytes)?; - Ok(Box::new(acc)) - } - SketchKind::CountSketch => { - let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(bytes)?; - Ok(Box::new(acc)) - } - SketchKind::Kll => { - let acc = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(bytes)?; - Ok(Box::new(acc)) - } - SketchKind::DdSketch => { - let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(bytes)?; - Ok(Box::new(acc)) - } - SketchKind::Hll => { - let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(bytes)?; - Ok(Box::new(acc)) - } + .into()), } } diff --git a/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs index 0a8774c3..b0fe7374 100644 --- a/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs @@ -68,6 +68,18 @@ impl CountMinSketchAccumulator { }) } + /// Decode from the modified OTLP wire format's + /// `CountMinSketchDataPoint.sketch` bytes when + /// `encoding = COUNT_MIN_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `CountMinSketch` wire struct (same format the legacy Arroyo path + /// uses — this method is the modified-OTLP entrypoint for PR I). + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: CountMinSketch::deserialize_msgpack(buffer)?, + }) + } + /// Decode from the modified OTLP wire format's /// `CountMinSketchDataPoint.sketch` bytes — i.e. the protobuf-encoded /// `asap_sketchlib::proto::sketchlib::CountMinState` message used by diff --git a/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs index f05936f2..7e5d0896 100644 --- a/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_sketch_accumulator.rs @@ -37,6 +37,18 @@ impl CountSketchAccumulator { } } + /// Decode from the modified OTLP wire format's + /// `CountSketchDataPoint.sketch` bytes when + /// `encoding = COUNT_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `CountSketch` struct — PR I parity entrypoint. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: CountSketch::deserialize_msgpack(buffer) + .map_err(|e| format!("deserialize CountSketch msgpack: {e}"))?, + }) + } + /// Decode from the modified OTLP wire format's /// `CountSketchDataPoint.sketch` bytes — the protobuf-encoded /// `asap_sketchlib::proto::sketchlib::CountSketchState` message @@ -301,4 +313,24 @@ mod tests { let result = cs.merge_with(&cms); assert!(result.is_err()); } + + #[test] + fn test_from_msgpack_bytes_round_trip() { + let original = CountSketch::from_legacy_matrix( + vec![vec![1.0, -2.0, 3.0], vec![-4.0, 5.0, -6.0]], + 2, + 3, + ); + let bytes = original.serialize_msgpack(); + let acc = CountSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.row_num, 2); + assert_eq!(acc.inner.col_num, 3); + assert_eq!(acc.inner.sketch(), original.sketch()); + } + + #[test] + fn test_from_msgpack_bytes_rejects_garbage() { + let result = CountSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); + assert!(result.is_err()); + } } diff --git a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs index 317403d2..da1f3f90 100644 --- a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs @@ -70,6 +70,21 @@ impl DatasketchesKLLAccumulator { }) } + /// Decode from the modified OTLP wire format's + /// `KLLSketchDataPoint.sketch` bytes when + /// `encoding = KLL_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `KllSketch` struct — PR I parity entrypoint. Unlike the + /// `_ENCODING_PROTO` path (which does lossy statistical + /// reconstruction via `update()` replay), the msgpack path is a + /// bit-identical round-trip because sketch-core's `KllSketch` + /// serializes its full internal state to msgpack. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: KllSketch::deserialize_msgpack(buffer)?, + }) + } + /// Decode from the modified OTLP wire format's /// `KLLSketchDataPoint.sketch` bytes — the protobuf-encoded /// `asap_sketchlib::proto::sketchlib::KllState` message that diff --git a/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs index de8f6e24..48f0dc36 100644 --- a/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/dd_sketch_accumulator.rs @@ -29,6 +29,18 @@ impl DDSketchAccumulator { } } + /// Decode from the modified OTLP wire format's + /// `DDSketchDataPoint.sketch` bytes when + /// `encoding = DDSKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `DdSketch` struct — PR I parity entrypoint. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: DdSketch::deserialize_msgpack(buffer) + .map_err(|e| format!("deserialize DdSketch msgpack: {e}"))?, + }) + } + /// Decode from the modified OTLP wire format's /// `DDSketchDataPoint.sketch` bytes — the protobuf-encoded /// `asap_sketchlib::proto::sketchlib::DDSketchState` message that @@ -206,4 +218,22 @@ mod tests { let cs = CountSketchAccumulator::new(2, 3); assert!(dd.merge_with(&cs).is_err()); } + + #[test] + fn test_from_msgpack_bytes_round_trip() { + let original = DdSketch::from_raw(0.01, vec![5, 10, 15, 20], -2, 50, 150.0, 0.25, 8.0); + let bytes = original.serialize_msgpack(); + let acc = DDSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.alpha, 0.01); + assert_eq!(acc.inner.store_counts, vec![5, 10, 15, 20]); + assert_eq!(acc.inner.store_offset, -2); + assert_eq!(acc.inner.count, 50); + assert_eq!(acc.inner.sum, 150.0); + } + + #[test] + fn test_from_msgpack_bytes_rejects_garbage() { + let result = DDSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); + assert!(result.is_err()); + } } diff --git a/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs index 3647759b..f1eff42f 100644 --- a/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/hll_sketch_accumulator.rs @@ -29,6 +29,18 @@ impl HllSketchAccumulator { } } + /// Decode from the modified OTLP wire format's + /// `HLLSketchDataPoint.sketch` bytes when + /// `encoding = HLL_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `HllSketch` struct — PR I parity entrypoint. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: HllSketch::deserialize_msgpack(buffer) + .map_err(|e| format!("deserialize HllSketch msgpack: {e}"))?, + }) + } + /// Decode from the modified OTLP wire format's /// `HLLSketchDataPoint.sketch` bytes — the protobuf-encoded /// `asap_sketchlib::proto::sketchlib::HyperLogLogState` message @@ -258,4 +270,28 @@ mod tests { let cs = CountSketchAccumulator::new(2, 3); assert!(hll.merge_with(&cs).is_err()); } + + #[test] + fn test_from_msgpack_bytes_round_trip() { + let original = HllSketch::from_raw( + HllVariant::Hip, + 3, + vec![0, 1, 2, 3, 4, 5, 6, 7], + 1.5, + 2.5, + 42.0, + ); + let bytes = original.serialize_msgpack(); + let acc = HllSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.variant, HllVariant::Hip); + assert_eq!(acc.inner.precision, 3); + assert_eq!(acc.inner.registers, vec![0, 1, 2, 3, 4, 5, 6, 7]); + assert_eq!(acc.inner.hip_kxq0, 1.5); + } + + #[test] + fn test_from_msgpack_bytes_rejects_garbage() { + let result = HllSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); + assert!(result.is_err()); + } } diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/cache.rs b/asap-query-engine/src/stores/simple_map_store/persistence/cache.rs index c00ecbb8..5d1e4723 100644 --- a/asap-query-engine/src/stores/simple_map_store/persistence/cache.rs +++ b/asap-query-engine/src/stores/simple_map_store/persistence/cache.rs @@ -36,7 +36,7 @@ impl PartCache { .weigher(|_k: &PartId, v: &LoadedPart| -> u32 { // Weight = sum of mmap'd bytes. Moka's weigher // returns u32, so clamp large parts. - let len = v.meta.data_len as u64 + v.meta.index_len as u64; + let len = v.meta.data_len + v.meta.index_len; len.min(u32::MAX as u64) as u32 }) .max_capacity(byte_budget) diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/part.rs b/asap-query-engine/src/stores/simple_map_store/persistence/part.rs index 768d78fe..dc297309 100644 --- a/asap-query-engine/src/stores/simple_map_store/persistence/part.rs +++ b/asap-query-engine/src/stores/simple_map_store/persistence/part.rs @@ -631,7 +631,8 @@ mod tests { let tmp = TempDir::new().unwrap(); let part_dir = tmp.path().join("0000000000000001"); let snap = make_snapshot(); - let report = PartWriter::write_part(&part_dir, 1, &[snap.clone()]).expect("write_part"); + let report = + PartWriter::write_part(&part_dir, 1, std::slice::from_ref(&snap)).expect("write_part"); assert_eq!(report.part_id, 1); assert_eq!(report.num_entries, 2); diff --git a/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs b/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs index 09f76d58..6996a233 100644 --- a/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs +++ b/asap-query-engine/tests/e2e_modified_otlp_sketch_path.rs @@ -1081,3 +1081,161 @@ async fn e2e_hll_sketch_modified_otlp_path() { assert_eq!(hll_acc.inner.precision, precision); assert_eq!(hll_acc.inner.registers, registers); } + +// ─── MessagePack encoding path (PR I) ──────────────────────────────────── +// +// Smoke test for `encoding = COUNT_MIN_SKETCH_ENCODING_MSGPACK`. The +// dispatcher should recognize the new `MSGPACK = 3` tag and route to +// `CountMinSketchAccumulator::from_msgpack_bytes`, which deserializes +// the cross-language sketch-core msgpack wire format. The other four +// sketch variants go through the same dispatcher branch, so one smoke +// test is sufficient for dispatcher coverage — per-variant msgpack +// round-trips are validated in the accumulator unit tests. + +fn build_count_min_msgpack_export_request( + metric_name: &str, + service_label: &str, + time_unix_nano: u64, + sketch_bytes: Vec, +) -> ExportMetricsServiceRequest { + let dp = CountMinSketchDataPoint { + attributes: vec![KeyValue { + key: "service".to_string(), + value: Some(AnyValue { + value: Some(any_value::Value::StringValue(service_label.to_string())), + }), + }], + start_time_unix_nano: 0, + time_unix_nano, + sample_count: 0, + sketch: sketch_bytes, + encoding: CountMinSketchEncoding::Msgpack as i32, + rows: 0, + cols: 0, + flags: 0, + series_id: 0, + }; + ExportMetricsServiceRequest { + resource_metrics: vec![ResourceMetrics { + resource: None, + scope_metrics: vec![ScopeMetrics { + scope: None, + metrics: vec![Metric { + name: metric_name.to_string(), + description: String::new(), + unit: String::new(), + metadata: Vec::new(), + data: Some(Data::Countminsketch(CountMinSketch { + data_points: vec![dp], + aggregation_temporality: 0, + })), + }], + schema_url: String::new(), + }], + schema_url: String::new(), + }], + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn e2e_count_min_sketch_msgpack_modified_otlp_path() { + let agg_id = 47u64; + let metric_name = "http_requests_msgpack"; + let service_label = "auth"; + let window_secs = 1u64; + let rows = 2u32; + let cols = 4u32; + + let precompute_port = 19550u16; + let otlp_grpc_port = 19551u16; + let otlp_http_port = 19552u16; + + let cms_config = make_count_min_agg_config( + agg_id, + metric_name, + window_secs, + vec!["service"], + rows as usize, + cols as usize, + ); + let mut agg_map = HashMap::new(); + agg_map.insert(agg_id, cms_config); + let streaming_config = Arc::new(StreamingConfig::new(agg_map)); + + let sink = Arc::new(CapturingOutputSink::new()); + let engine = PrecomputeEngine::new( + engine_config(precompute_port), + streaming_config, + sink.clone(), + ); + let ingest_state = engine.ingest_state(); + + tokio::spawn(async move { + let _ = engine.run().await; + }); + + let otlp_receiver = OtlpReceiver::with_ingest_state( + OtlpReceiverConfig { + grpc_port: otlp_grpc_port, + http_port: otlp_http_port, + }, + ingest_state, + ); + tokio::spawn(async move { + let _ = otlp_receiver.run().await; + }); + + tokio::time::sleep(tokio::time::Duration::from_millis(400)).await; + + // Build a known sketch in sketch-core and serialize with msgpack — this + // is what the Go producer (sketchlib-go) will emit once PR I's matching + // Go-side work lands. + let mut cms = sketch_core::count_min::CountMinSketch::new(rows as usize, cols as usize); + cms.update("user_a", 1.0); + cms.update("user_b", 1.0); + cms.update("user_a", 1.0); + let sketch_bytes = cms.serialize_msgpack(); + + let client = reqwest::Client::new(); + let req = build_count_min_msgpack_export_request( + metric_name, + service_label, + 100_000_000, + sketch_bytes, + ); + post_otlp_http(&client, otlp_http_port, req).await; + + // Watermark advance using an empty msgpack sketch. + let empty = sketch_core::count_min::CountMinSketch::new(rows as usize, cols as usize); + let watermark_req = build_count_min_msgpack_export_request( + metric_name, + service_label, + 2_000_000_000, + empty.serialize_msgpack(), + ); + post_otlp_http(&client, otlp_http_port, watermark_req).await; + + tokio::time::sleep(tokio::time::Duration::from_millis(800)).await; + + let captured = sink.drain(); + assert!(!captured.is_empty(), "expected at least one output"); + + let (window0_output, window0_acc_box) = captured + .iter() + .find(|(out, _)| out.start_timestamp == 0) + .expect("no captured output for window 0"); + + assert_eq!(window0_output.aggregation_id, agg_id); + + let cms_acc = window0_acc_box + .as_any() + .downcast_ref::() + .expect("captured accumulator should be CountMinSketchAccumulator"); + + // user_a was updated twice → query_key("user_a") should estimate ≥ 2. + assert!( + cms_acc.inner.query_key("user_a") >= 2.0, + "CountMinSketch msgpack round-trip should preserve user_a count (got {})", + cms_acc.inner.query_key("user_a") + ); +} From 7744cd7a59c63a797507dc799f945506f1e7a948 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Tue, 14 Apr 2026 16:14:01 -0400 Subject: [PATCH 2/2] test(persistence): de-flake hard_cap_back_pressure via explicit counter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The existing `hard_cap_back_pressure_blocks_inserts_until_flusher_drains` test asserts that at least one insert takes ≥3 ms wall-clock time when sealed memory hits the hard cap — the intuition being that a blocked insert must wait ≥1 flush tick for the condvar to fire. On CI (fast runners, small test data) the flush tick completes well under 1 ms, so the slow path runs and fully drains before the 3 ms threshold ever trips. Result: the assertion is flaky even though back-pressure is wired up correctly. Replace wall-clock timing with a dedicated monotonic counter. ## Changes * `FlusherShared` gains `back_pressure_wait_count: AtomicU64` * `FlusherHandle::wait_for_memory_under` increments the counter exactly once per call that observes `mem_counter >= cap` (i.e. every call that enters the blocking wait loop — the fast-path early return does not bump the counter) * `FlusherHandle::back_pressure_wait_count()` getter exposes it * `SimpleMapStorePerKey::back_pressure_wait_count()` proxies through to the flusher, returning 0 when persistence is disabled * The test now asserts `store.back_pressure_wait_count() > 0` after the 200-insert loop — a timing-independent signal ## Why a counter, not a timing threshold Tests that assert "this path is slower than X ms" assume the path is slow for a reason other than what it's actually measuring. Here, `wait_for_memory_under` blocks until the flusher drains, and the drain is fast, so the total wait can be sub-millisecond even though the mechanism is firing exactly as designed. Timing thresholds trade signal for noise on fast hardware. A counter sidesteps that tradeoff: it is 0 if the mechanism never fires and positive if it fires even once — exactly the semantic the test cares about. ## Validation Ran 5× back-to-back: all 5 passed (previously observed flake on CI). Full query_engine_rust --lib suite: 480 passed. Clippy + fmt clean. Drive-by fix: inherited from main (the persistence PR #4 merge), not introduced by this PR. Landing it here because the PR is already touching enough persistence-adjacent surface (clippy fixes) that CI needs to be unblocked on the full test gate, and a real fix is cleaner than a retry loop or `#[ignore]`. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/stores/simple_map_store/per_key.rs | 15 +++++ .../simple_map_store/persistence/flusher.rs | 24 +++++++ .../tests/persistence_integration_tests.rs | 65 ++++++++----------- 3 files changed, 66 insertions(+), 38 deletions(-) diff --git a/asap-query-engine/src/stores/simple_map_store/per_key.rs b/asap-query-engine/src/stores/simple_map_store/per_key.rs index 05ad9a0f..c344f42c 100644 --- a/asap-query-engine/src/stores/simple_map_store/per_key.rs +++ b/asap-query-engine/src/stores/simple_map_store/per_key.rs @@ -347,6 +347,21 @@ impl SimpleMapStorePerKey { }) } + /// Total number of times the insert hot path has blocked on + /// flusher back-pressure (i.e. observed sealed memory at or above + /// `hard_cap_bytes` and entered `wait_for_memory_under`'s wait + /// loop). Returns 0 when persistence is disabled. Monotonic. + /// + /// Used as a timing-independent signal in back-pressure tests — + /// wall-clock thresholds are flaky on fast CI runners where a + /// flush tick completes in under a millisecond. + pub fn back_pressure_wait_count(&self) -> u64 { + match self.persistence.as_ref() { + Some(state) => state.flusher.back_pressure_wait_count(), + None => 0, + } + } + /// Collect diagnostic info about store contents. pub fn diagnostic_info(&self) -> super::StoreDiagnostics { use super::{AggregationDiagnostic, StoreDiagnostics}; diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs b/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs index 10f5fb11..eb1f23d5 100644 --- a/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs +++ b/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs @@ -45,6 +45,13 @@ pub(crate) struct FlusherShared { /// the flusher when it finishes a tick. pub pressure_cv: Condvar, pub pressure_mutex: Mutex<()>, + /// Monotonic counter of how many times `wait_for_memory_under` + /// has entered its blocking wait loop because `mem_counter >= cap`. + /// Incremented exactly once per back-pressure-triggered call; used + /// as a timing-independent signal in tests (wall-clock thresholds + /// are flaky on fast CI runners where a flush tick completes in + /// <1 ms). + pub back_pressure_wait_count: AtomicU64, } impl FlusherHandle { @@ -75,6 +82,7 @@ impl FlusherHandle { shutdown: AtomicBool::new(false), pressure_cv: Condvar::new(), pressure_mutex: Mutex::new(()), + back_pressure_wait_count: AtomicU64::new(0), }); let shared_for_thread = Arc::clone(&shared); @@ -129,6 +137,14 @@ impl FlusherHandle { if mem_counter.load(Ordering::Relaxed) < cap { return true; } + // Slow path: the cap was crossed and we are about to block. + // Bump a counter so tests can assert that back-pressure fired + // without relying on wall-clock thresholds (the actual wait is + // ≤1 flush tick, which is sub-millisecond on fast CI runners + // and makes wall-clock-based signals flaky). + self.inner + .back_pressure_wait_count + .fetch_add(1, Ordering::Relaxed); // Kick the flusher once so it tries to drain right now instead // of waiting for its own interval tick. self.inner.pressure_cv.notify_all(); @@ -163,6 +179,14 @@ impl FlusherHandle { Arc::clone(&self.inner.manifest) } + /// Total number of times the insert hot path has entered + /// `wait_for_memory_under`'s blocking wait loop (i.e. observed + /// sealed memory at or above `hard_cap_bytes`). Monotonic counter, + /// used by back-pressure tests as a timing-independent signal. + pub fn back_pressure_wait_count(&self) -> u64 { + self.inner.back_pressure_wait_count.load(Ordering::Relaxed) + } + pub fn disk_path(&self) -> &std::path::Path { &self.inner.cfg.disk_path } diff --git a/asap-query-engine/src/tests/persistence_integration_tests.rs b/asap-query-engine/src/tests/persistence_integration_tests.rs index f6e14e9d..12f42b8e 100644 --- a/asap-query-engine/src/tests/persistence_integration_tests.rs +++ b/asap-query-engine/src/tests/persistence_integration_tests.rs @@ -202,17 +202,21 @@ fn construct_and_drop_shuts_flusher_cleanly() { #[test] fn hard_cap_back_pressure_blocks_inserts_until_flusher_drains() { - // Construct a store with a very small hard cap and a flusher that - // is deliberately slow (long flush interval). The first few - // inserts will push mem_bytes_sealed past the cap; subsequent - // inserts must block in `wait_for_memory_under` until the flusher - // evicts a sealed epoch. + // Construct a store with a very small hard cap and a flusher + // whose tick interval is long enough that at least one insert + // will catch the cap before a tick has time to drain it. The + // first few inserts push mem_bytes_sealed past the cap; any + // subsequent insert that arrives while mem_bytes_sealed is still + // at or above the cap enters `wait_for_memory_under`'s blocking + // wait loop, which increments a dedicated counter. // - // Test strategy: measure wall-clock time of an insert that we - // *know* will hit the cap. If back-pressure is wired up, it - // must be longer than the flusher's tick interval (because it - // waits at least one tick for the condvar notify). If it's not - // wired up, the insert returns in a handful of microseconds. + // We assert on that counter rather than on wall-clock timing. + // A previous version of this test measured the slowest insert + // and required it to be ≥3 ms, but on fast CI runners a flush + // tick completes in well under a millisecond, so the slow path + // would run and fully drain before any 3 ms threshold tripped — + // leaving the assertion flaky even though back-pressure was + // firing correctly. The counter is a timing-independent signal. let dir = TempDir::new().unwrap(); let cfg = make_streaming_config(1); @@ -224,8 +228,9 @@ fn hard_cap_back_pressure_blocks_inserts_until_flusher_drains() { hard_cap_bytes: 640, hot_window_ms: Some(0), delete_older_than_ms: None, - // Flusher tick is 200ms — long enough that a blocking insert - // is clearly distinguishable from a non-blocking one. + // Flusher tick is 200ms — long enough that at least one + // insert inside the 200-iteration loop can race ahead of + // the flusher and observe the cap. flush_interval: Duration::from_millis(200), disk_path: dir.path().to_path_buf(), part_cache_bytes: 0, @@ -234,41 +239,25 @@ fn hard_cap_back_pressure_blocks_inserts_until_flusher_drains() { .expect("with_persistence"); // Push well past the cap — 200 items × ~16 bytes each, vs. a - // 640-byte cap — so the insert path is forced to block on the - // flusher's condvar at least once. num_aggregates_to_retain=2 - // means every 2 distinct windows triggers a seal, so sealed - // epochs accumulate fast. - // - // A normal insert on the hot path returns in <1 ms; anything - // ≥3 ms unambiguously indicates a condvar wait fired. We also - // capture the median as a baseline to make the assertion - // message informative when it fails. - let mut elapsed_all: Vec = Vec::with_capacity(200); + // 640-byte cap. num_aggregates_to_retain=2 means every 2 distinct + // windows triggers a seal, so sealed epochs accumulate fast. for i in 0..200u64 { let start = i * 1_000; let end = start + 1_000; let batch = vec![sum_entry(1, start, end, i as f64)]; - let t = Instant::now(); store .insert_precomputed_output_batch(batch) .expect("insert"); - elapsed_all.push(t.elapsed()); } - let mut sorted = elapsed_all.clone(); - sorted.sort(); - let median = sorted[sorted.len() / 2]; - let max = *sorted.last().unwrap(); - let slow_count = sorted - .iter() - .filter(|d| **d >= Duration::from_millis(3)) - .count(); + + let back_pressure_waits = store.back_pressure_wait_count(); assert!( - slow_count > 0, - "expected ≥1 insert to block on back-pressure; all 200 returned fast \ - (median={:?}, max={:?}, slow≥3ms count={})", - median, - max, - slow_count + back_pressure_waits > 0, + "expected ≥1 insert to enter the back-pressure wait loop; \ + wait_count={} (cap={}, 200 inserts × ~16B vs. 640B cap \ + should force at least one observed over-cap)", + back_pressure_waits, + 640, ); // Sanity: after the loop, wait a bit and confirm the flusher