From 40fcdd1d4da660afe34313d8a02d52f08272ef06 Mon Sep 17 00:00:00 2001 From: zz_y Date: Mon, 10 Aug 2026 06:57:36 -0600 Subject: [PATCH] feat(otap): real plugin integration + Arrow IPC network transport Stacked on feat/schema-dictionary-record-codec, which implemented the Schema/Dictionary/Record codec (SeriesDictionary/SeriesDictionaryDecoder) but only proved it end-to-end via direct calls in one process, in-memory -- not through the actual OTAP plugin lifecycle, and not across a real serialize/transmit/deserialize hop. This PR closes both gaps: otap::wire (new): Arrow-IPC serializes a SketchStreamBatch's four RecordBatches into a length-prefixed frame, plus async send_stream_batch/ recv_stream_batch over a TcpStream. Each sub-batch is its own self-contained IPC stream (schema + one record batch + EOS); recv_stream_batch distinguishes a clean EOF between frames from a truncated one mid-frame. AsapSketchesPlugin::start_from_envelopes (new): the receiver-role counterpart to the existing producer-role start(). Consumes Stream instead of Stream, decodes via a persistent SeriesDictionaryDecoder, and routes reconstructed envelopes through Precompute::observe_envelope (merge, never expand to samples) -- reusing the same ticker/control-task/graceful-drain machinery as the producer role via a new shared spawn_lifecycle helper. A receiver configured with transmit_sketch=false naturally re-emits query-mode (quantile) estimates instead of sketch bytes through its own emit channel, so a chain of AsapSketchesPlugins can compose without any new machinery. examples/sketch_producer_node.rs + sketch_receiver_node.rs (new): two separate binaries -- real AsapSketchesPlugin producer and receiver roles, connected over a real TCP socket via otap::wire, not the in-process mpsc channel sketch_pipeline_demo.rs uses. The producer feeds a real OTAP-shaped input stream (records::flatten + decode_batch, not a direct observe() call) and lets the plugin's actual Wakeup-style Tokio ticker close windows on its own wall-clock schedule. Verified running both together: producer emits 5 windows (window 0 carries SCHEMA+DICTIONARY+LABELS, windows 1-4 carry RECORD only), receiver receives and decodes all 5 over the socket, merges them, and prints a correct p99 gauge. Co-Authored-By: Claude Sonnet 5 --- asap-precompute-rs/Cargo.lock | 176 +++++++- asap-precompute-rs/Cargo.toml | 27 +- .../examples/sketch_producer_node.rs | 214 ++++++++++ .../examples/sketch_receiver_node.rs | 133 ++++++ asap-precompute-rs/src/otap/lifecycle.rs | 232 ++++++++++- asap-precompute-rs/src/otap/mod.rs | 1 + asap-precompute-rs/src/otap/wire.rs | 382 ++++++++++++++++++ 7 files changed, 1154 insertions(+), 11 deletions(-) create mode 100644 asap-precompute-rs/examples/sketch_producer_node.rs create mode 100644 asap-precompute-rs/examples/sketch_receiver_node.rs create mode 100644 asap-precompute-rs/src/otap/wire.rs diff --git a/asap-precompute-rs/Cargo.lock b/asap-precompute-rs/Cargo.lock index f579db4..db673fb 100644 --- a/asap-precompute-rs/Cargo.lock +++ b/asap-precompute-rs/Cargo.lock @@ -64,6 +64,26 @@ dependencies = [ "num", ] +[[package]] +name = "arrow-cast" +version = "53.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6365f8527d4f87b133eeb862f9b8093c009d41a210b8f101f91aa2392f61daac" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "atoi", + "base64", + "chrono", + "half", + "lexical-core", + "num", + "ryu", +] + [[package]] name = "arrow-data" version = "53.4.1" @@ -76,17 +96,46 @@ dependencies = [ "num", ] +[[package]] +name = "arrow-ipc" +version = "53.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3527365b24372f9c948f16e53738eb098720eea2093ae73c7af04ac5e30a39b" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-schema", + "flatbuffers", +] + [[package]] name = "arrow-schema" version = "53.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35b0f9c0c3582dd55db0f136d3b44bfa0189df07adcf7dc7f2f2e74db0f52eb8" +[[package]] +name = "arrow-select" +version = "53.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92fc337f01635218493c23da81a364daf38c694b05fc20569c3193c11c561984" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "num", +] + [[package]] name = "asap-precompute-rs" version = "0.1.0" dependencies = [ "arrow-array", + "arrow-ipc", "arrow-schema", "asap_sketchlib", "futures", @@ -96,6 +145,7 @@ dependencies = [ "serde_json", "thiserror", "tokio", + "tokio-stream", ] [[package]] @@ -114,12 +164,27 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + [[package]] name = "autocfg" version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + [[package]] name = "bit-set" version = "0.8.0" @@ -135,6 +200,12 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.13.0" @@ -241,6 +312,16 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "flatbuffers" +version = "24.12.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f1baf0dbf96932ec9a3038d57900329c015b0bfb7b63d904f3bc27e2b02a096" +dependencies = [ + "bitflags 1.3.2", + "rustc_version", +] + [[package]] name = "fnv" version = "1.0.7" @@ -437,6 +518,63 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lexical-core" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d8d125a277f807e55a77304455eb7b1cb52f2b18c143b60e766c120bd64a594" +dependencies = [ + "lexical-parse-float", + "lexical-parse-integer", + "lexical-util", + "lexical-write-float", + "lexical-write-integer", +] + +[[package]] +name = "lexical-parse-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52a9f232fbd6f550bc0137dcb5f99ab674071ac2d690ac69704593cb4abbea56" +dependencies = [ + "lexical-parse-integer", + "lexical-util", +] + +[[package]] +name = "lexical-parse-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a7a039f8fb9c19c996cd7b2fcce303c1b2874fe1aca544edc85c4a5f8489b34" +dependencies = [ + "lexical-util", +] + +[[package]] +name = "lexical-util" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2604dd126bb14f13fb5d1bd6a66155079cb9fa655b37f875b3a742c705dbed17" + +[[package]] +name = "lexical-write-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50c438c87c013188d415fbabbb1dceb44249ab81664efbd31b14ae55dabb6361" +dependencies = [ + "lexical-util", + "lexical-write-integer", +] + +[[package]] +name = "lexical-write-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "409851a618475d2d5796377cad353802345cba92c867d9fbcde9cf4eac4e14df" +dependencies = [ + "lexical-util", +] + [[package]] name = "libc" version = "0.2.186" @@ -622,7 +760,7 @@ checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" dependencies = [ "bit-set", "bit-vec", - "bitflags", + "bitflags 2.13.0", "num-traits", "rand", "rand_chacha", @@ -727,7 +865,7 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags", + "bitflags 2.13.0", ] [[package]] @@ -755,13 +893,22 @@ dependencies = [ "serde", ] +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "rustix" version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags", + "bitflags 2.13.0", "errno", "libc", "linux-raw-sys", @@ -786,12 +933,24 @@ dependencies = [ "wait-timeout", ] +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + [[package]] name = "scopeguard" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "semver" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" + [[package]] name = "serde" version = "1.0.228" @@ -963,6 +1122,17 @@ dependencies = [ "syn", ] +[[package]] +name = "tokio-stream" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "twox-hash" version = "2.1.2" diff --git a/asap-precompute-rs/Cargo.toml b/asap-precompute-rs/Cargo.toml index f93de8a..8d10bd3 100644 --- a/asap-precompute-rs/Cargo.toml +++ b/asap-precompute-rs/Cargo.toml @@ -17,13 +17,17 @@ thiserror = "1" # non-default to keep their build cheap. arrow-array = { version = "53", optional = true } arrow-schema = { version = "53", optional = true } +# Only pulled in for otap::wire's Arrow IPC serialization of +# SketchStreamBatch across a real transport (see docs/data_model.md's +# "crosses a node or network boundary" framing). +arrow-ipc = { version = "53", optional = true } # Tokio drives the plugin lifecycle (Stream consumer # task + Wakeup-driven flush ticker + control-channel poll task + -# graceful drain). Only the `rt`, `sync`, `time`, and `macros` -# sub-features are needed; the full set is pulled in via the `otap` -# feature gate so default-feature consumers don't pay for it. -tokio = { version = "1", features = ["rt", "rt-multi-thread", "sync", "time", "macros"], optional = true } +# graceful drain), plus `otap::wire`'s TCP transport (`net`, +# `io-util`). The full set is pulled in via the `otap` feature gate so +# default-feature consumers don't pay for it. +tokio = { version = "1", features = ["rt", "rt-multi-thread", "sync", "time", "macros", "net", "io-util"], optional = true } futures = { version = "0.3", optional = true } [dev-dependencies] @@ -32,14 +36,27 @@ serde_json = "1" # The full tokio runtime is required for #[tokio::test]-driven # lifecycle tests under `tests/otap_*.rs`. tokio = { version = "1", features = ["full"] } +# Only used to adapt an `mpsc::UnboundedReceiver` into a `Stream` for +# chaining two `AsapSketchesPlugin`s (producer's EmitReceiver -> +# receiver's start_from_envelopes) in tests and the network-transport +# examples. The library itself is Stream-source-agnostic. +tokio-stream = "0.1" [features] default = [] # Enables the OTAP codec + plugin lifecycle at `crate::otap` # (decode_batch / encode_batch over `arrow::RecordBatch`, plus the # `AsapSketchesPlugin` Tokio runtime). -otap = ["dep:arrow-array", "dep:arrow-schema", "dep:tokio", "dep:futures"] +otap = ["dep:arrow-array", "dep:arrow-schema", "dep:arrow-ipc", "dep:tokio", "dep:futures"] [[example]] name = "sketch_pipeline_demo" required-features = ["otap"] + +[[example]] +name = "sketch_producer_node" +required-features = ["otap"] + +[[example]] +name = "sketch_receiver_node" +required-features = ["otap"] diff --git a/asap-precompute-rs/examples/sketch_producer_node.rs b/asap-precompute-rs/examples/sketch_producer_node.rs new file mode 100644 index 0000000..a00b6ad --- /dev/null +++ b/asap-precompute-rs/examples/sketch_producer_node.rs @@ -0,0 +1,214 @@ +//! Network-transport half of the `docs/data_model.md` demo: a +//! **sketch creation processor** running as a real `AsapSketchesPlugin` +//! (not called directly like `sketch_pipeline_demo.rs` does), pushing +//! each emitted [`SketchStreamBatch`] over a real TCP socket — Arrow +//! IPC-serialized via `otap::wire` — to `sketch_receiver_node` +//! (`examples/sketch_receiver_node.rs`), which must already be +//! listening. +//! +//! Run (in one terminal, first): +//! ```text +//! cargo run --example sketch_receiver_node --features otap +//! ``` +//! Then (in a second terminal): +//! ```text +//! cargo run --example sketch_producer_node --features otap +//! ``` +//! +//! Unlike `sketch_pipeline_demo.rs`'s in-process `mpsc` channel, what +//! crosses the wire here is genuinely serialized bytes over a socket +//! — the actual "crosses a node or network boundary" hop +//! `docs/data_model.md` opens with. See `otap::wire`'s module doc for +//! the exact frame layout. +//! +//! Unlike `sketch_pipeline_demo.rs` (which force-closes windows with +//! explicit `drain()` calls), this binary feeds observations through +//! a real OTAP-shaped input stream (`records::flatten` + +//! `decode_batch`, `AsapSketchesPlugin::start`'s input task) and lets +//! the plugin's real `Wakeup`-style Tokio ticker close windows on its +//! own wall-clock schedule — a `window_size` short enough (300ms) to +//! see several windows roll in one run. + +use std::time::Duration; + +use arrow_array::{BinaryArray, Float64Array, RecordBatch, StringArray, UInt32Array, UInt64Array}; +use arrow_schema::{DataType, Field, Schema}; +use asap_precompute_rs::otap::config::PluginConfig; +use asap_precompute_rs::otap::records::{ + OtapMetricRecords, ATTR_BATCH_BYTES, ATTR_BATCH_INT, ATTR_BATCH_KEY, ATTR_BATCH_PARENT_ID, + ATTR_BATCH_STR, +}; +use asap_precompute_rs::otap::wire::send_stream_batch; +use asap_precompute_rs::otap::{ + AsapSketchesPlugin, StartOptions, COLUMN_METRIC, COLUMN_TIME_UNIX_NANO, COLUMN_VALUE, +}; +use tokio::net::TcpStream; +use tokio::sync::mpsc; +use tokio_stream::wrappers::UnboundedReceiverStream; + +/// Must match `sketch_receiver_node`'s listen address. +const RECEIVER_ADDR: &str = "127.0.0.1:47821"; +const AGG_ID: u64 = 1; +const NUM_WINDOWS: usize = 4; +const WINDOW_SIZE: Duration = Duration::from_millis(300); +/// Longer than `WINDOW_SIZE` so each window's data has already landed +/// before the real ticker rotates it out — real wall-clock pacing, not +/// a guaranteed lockstep boundary, so this is a margin, not a promise. +const PACING: Duration = Duration::from_millis(400); + +#[tokio::main] +async fn main() { + let mut socket = connect_with_retry(RECEIVER_ADDR).await; + println!("[producer] connected to {RECEIVER_ADDR}"); + + let plugin_cfg = PluginConfig { + sketch_type: "ddsketch".into(), + window_size: WINDOW_SIZE, + output_metric_name: "http_request_duration_ms".into(), + agg_id: AGG_ID, + sketch_params: [("relative_accuracy".to_string(), 0.01)] + .into_iter() + .collect(), + ..Default::default() + }; + let plugin = AsapSketchesPlugin::from_plugin_config(&plugin_cfg).expect("producer config"); + + // Bridge a paced feed of synthetic OTAP-shaped input records into + // the Stream AsapSketchesPlugin::start + // wants -- this is the plugin's real ingest path (records::flatten + // + decode_batch), not a direct Precompute::observe() call. + let (input_tx, input_rx) = mpsc::unbounded_channel::(); + let feed_task = tokio::spawn(async move { + for window_idx in 0..NUM_WINDOWS { + // One series (path=/api), latency drifting upward window + // to window so the receiver's printed p99 visibly moves. + for i in 0..200u64 { + let base = 10.0 + (window_idx as f64) * 8.0; + let latency = base + (i % 25) as f64; + let records = + build_scalar_records("http_request_duration_ms", latency, now_ms(), "/api"); + if input_tx.send(records).is_err() { + return; // plugin shut down early. + } + } + // Let the ticker rotate this window out before the next + // window's data starts arriving. + tokio::time::sleep(PACING).await; + } + // Dropping input_tx here ends the plugin's input stream. + }); + + let (handle, mut emit_rx) = plugin.start( + UnboundedReceiverStream::new(input_rx), + None, + StartOptions::default(), + ); + + // Forward every emitted SketchStreamBatch over the socket as it + // arrives, concurrently with the feed task still running. + let forward_task = tokio::spawn(async move { + let mut window_idx = 0; + while let Some(batch) = emit_rx.recv().await { + println!( + "[producer] window {window_idx}: schema={} dictionary={} labels={} record={} row(s) -- sending over the wire", + batch.schema.num_rows(), + batch.dictionary.num_rows(), + batch.labels.num_rows(), + batch.record.num_rows(), + ); + send_stream_batch(&mut socket, &batch) + .await + .expect("send over socket"); + window_idx += 1; + } + socket + }); + + feed_task.await.expect("feed task"); + // Give the last window's data one more full period to roll + // naturally, then shut down -- the final drain flushes any + // residue that hadn't hit a tick boundary yet. + tokio::time::sleep(PACING).await; + handle.shutdown().await.expect("producer shutdown"); + + let socket = forward_task.await.expect("forward task"); + drop(socket); // close the connection -> receiver sees a clean EOF. + println!("[producer] done, connection closed"); +} + +/// Builds a one-row [`OtapMetricRecords`] for a raw scalar +/// observation with a single `path` label -- the OTAP-Metrics-shaped +/// input `AsapSketchesPlugin::start`'s input task consumes from a +/// real upstream OTAP source (Telegraf / Vector / another OTAP +/// collector). +fn build_scalar_records( + metric: &str, + value: f64, + timestamp_ms: u64, + path: &str, +) -> OtapMetricRecords { + let metrics_schema = std::sync::Arc::new(Schema::new(vec![ + Field::new(COLUMN_TIME_UNIX_NANO, DataType::UInt64, false), + Field::new(COLUMN_METRIC, DataType::Utf8, false), + Field::new(COLUMN_VALUE, DataType::Float64, false), + Field::new(ATTR_BATCH_PARENT_ID, DataType::UInt32, false), + ])); + let metrics = RecordBatch::try_new( + metrics_schema, + vec![ + std::sync::Arc::new(UInt64Array::from(vec![timestamp_ms * 1_000_000])), + std::sync::Arc::new(StringArray::from(vec![metric])), + std::sync::Arc::new(Float64Array::from(vec![value])), + std::sync::Arc::new(UInt32Array::from(vec![0_u32])), + ], + ) + .expect("metrics batch"); + let attributes_schema = std::sync::Arc::new(Schema::new(vec![ + Field::new(ATTR_BATCH_PARENT_ID, DataType::UInt32, false), + Field::new(ATTR_BATCH_KEY, DataType::Utf8, false), + Field::new(ATTR_BATCH_BYTES, DataType::Binary, true), + Field::new(ATTR_BATCH_STR, DataType::Utf8, true), + Field::new(ATTR_BATCH_INT, DataType::UInt64, true), + ])); + let attributes = RecordBatch::try_new( + attributes_schema, + vec![ + std::sync::Arc::new(UInt32Array::from(vec![0_u32])), + std::sync::Arc::new(StringArray::from(vec!["path"])), + std::sync::Arc::new(BinaryArray::from_opt_vec(vec![None as Option<&[u8]>])), + std::sync::Arc::new(StringArray::from(vec![Some(path)])), + std::sync::Arc::new(UInt64Array::from(vec![None as Option])), + ], + ) + .expect("attributes batch"); + OtapMetricRecords { + metrics, + attributes, + } +} + +/// Retries the connection a few times -- `sketch_receiver_node` may +/// not have bound its listener yet if both binaries are started at +/// nearly the same moment. +async fn connect_with_retry(addr: &str) -> TcpStream { + for attempt in 0..20 { + match TcpStream::connect(addr).await { + Ok(s) => return s, + Err(e) => { + if attempt == 0 { + println!("[producer] waiting for {addr} to accept connections ({e})..."); + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + } + } + panic!("could not connect to {addr} after retries -- is sketch_receiver_node running?"); +} + +fn now_ms() -> u64 { + use std::time::{SystemTime, UNIX_EPOCH}; + SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock before Unix epoch") + .as_millis() as u64 +} diff --git a/asap-precompute-rs/examples/sketch_receiver_node.rs b/asap-precompute-rs/examples/sketch_receiver_node.rs new file mode 100644 index 0000000..a92713f --- /dev/null +++ b/asap-precompute-rs/examples/sketch_receiver_node.rs @@ -0,0 +1,133 @@ +//! Network-transport half of the `docs/data_model.md` demo: a +//! **receive processor** running as a real `AsapSketchesPlugin` in +//! its receiver role (`start_from_envelopes` — see +//! `src/otap/lifecycle.rs`'s module doc, "The other role"), reading +//! [`SketchStreamBatch`]es off a real TCP socket (Arrow IPC-decoded +//! via `otap::wire`) sent by `sketch_producer_node` +//! (`examples/sketch_producer_node.rs`). +//! +//! Merges every reconstructed envelope via `Precompute::observe_envelope`, +//! then — because this plugin's own config sets `transmit_sketch = +//! false`, `quantiles = [0.99]` — its own ticker/drain naturally +//! produces p99 *estimate* envelopes instead of re-emitting sketch +//! bytes. Those come back out through this plugin's own emit channel +//! as another `SketchStreamBatch`, which this binary decodes and +//! prints as Prometheus text (the "Prometheus backend" stage — see +//! `sketch_pipeline_demo.rs`'s module doc for why printing stands in +//! for a real `/metrics` HTTP handler). +//! +//! Run this first (see `sketch_producer_node.rs`'s module doc): +//! ```text +//! cargo run --example sketch_receiver_node --features otap +//! ``` + +use std::time::Duration; + +use asap_precompute_rs::envelope::SketchEnvelope; +use asap_precompute_rs::otap::config::PluginConfig; +use asap_precompute_rs::otap::wire::recv_stream_batch; +use asap_precompute_rs::otap::{ + AsapSketchesPlugin, SeriesDictionaryDecoder, SketchStreamBatch, StartOptions, +}; +use tokio::net::TcpListener; +use tokio::sync::mpsc; +use tokio_stream::wrappers::UnboundedReceiverStream; + +/// Must match `sketch_producer_node`'s connect address. +const LISTEN_ADDR: &str = "127.0.0.1:47821"; +const AGG_ID: u64 = 1; + +#[tokio::main] +async fn main() { + let listener = TcpListener::bind(LISTEN_ADDR).await.expect("bind"); + println!("[receiver] listening on {LISTEN_ADDR}, waiting for sketch_producer_node..."); + let (mut socket, peer) = listener.accept().await.expect("accept"); + println!("[receiver] accepted connection from {peer}"); + + // Bridge "read framed batches off the socket" into the + // Stream start_from_envelopes wants. + let (batch_tx, batch_rx) = mpsc::unbounded_channel::(); + let socket_task = tokio::spawn(async move { + loop { + match recv_stream_batch(&mut socket).await { + Ok(Some(batch)) => { + println!( + "[receiver] received over the wire: schema={} dictionary={} labels={} record={} row(s)", + batch.schema.num_rows(), + batch.dictionary.num_rows(), + batch.labels.num_rows(), + batch.record.num_rows(), + ); + if batch_tx.send(batch).is_err() { + return; // plugin already gone. + } + } + Ok(None) => { + println!("[receiver] producer closed the connection"); + return; + } + Err(e) => { + eprintln!("[receiver] wire error: {e}"); + return; + } + } + } + }); + + let plugin_cfg = PluginConfig { + sketch_type: "ddsketch".into(), + window_size: Duration::from_secs(60), // driven by drain() on shutdown, same as the producer. + output_metric_name: "http_request_duration_ms_p99".into(), + agg_id: AGG_ID, + transmit_sketch: false, // query mode: emits quantile estimates, not sketch bytes. + quantiles: vec![0.99], + ..Default::default() + }; + let plugin = AsapSketchesPlugin::from_plugin_config(&plugin_cfg).expect("receiver config"); + let (handle, mut emit_rx) = plugin.start_from_envelopes( + UnboundedReceiverStream::new(batch_rx), + None, + StartOptions::default(), + ); + + // Wait for the producer to finish sending (socket EOF) before + // asking for the final drain -- otherwise shutdown could race a + // still-in-flight batch. + socket_task.await.expect("socket task"); + handle.shutdown().await.expect("receiver shutdown"); + + let mut decoder = SeriesDictionaryDecoder::new(); + while let Ok(Some(batch)) = + tokio::time::timeout(Duration::from_millis(500), emit_rx.recv()).await + { + for estimate in decoder.decode(&batch).expect("decode estimate batch") { + print!("{}", format_prometheus_gauge(&estimate)); + } + } +} + +/// Formats one estimate-mode [`SketchEnvelope`] (`payload` empty, +/// `value` set — see `docs/data_model.md`'s `RECORD.value`) as a +/// Prometheus text-exposition gauge sample. Stands in for a real +/// `/metrics` HTTP handler. +fn format_prometheus_gauge(env: &SketchEnvelope) -> String { + let mut labels: Vec = env + .labels + .iter() + .map(|kv| format!("{}=\"{}\"", kv.key, kv.value)) + .collect(); + labels.sort(); + let label_str = if labels.is_empty() { + String::new() + } else { + format!("{{{}}}", labels.join(",")) + }; + format!( + "# HELP {name} sketch-derived quantile estimate\n\ + # TYPE {name} gauge\n\ + {name}{label_str} {value} {ts}\n", + name = env.metric_name, + value = env.value, + ts = env.window_end_ms, + ) +} diff --git a/asap-precompute-rs/src/otap/lifecycle.rs b/asap-precompute-rs/src/otap/lifecycle.rs index 85165ca..9cb7ae3 100644 --- a/asap-precompute-rs/src/otap/lifecycle.rs +++ b/asap-precompute-rs/src/otap/lifecycle.rs @@ -45,6 +45,23 @@ //! boundary silently drops in-flight observations. The drain runs on //! the same emit channel as the flush ticker, so the consumer sees //! exactly one final batch carrying the residue. +//! +//! # The other role: [`AsapSketchesPlugin::start_from_envelopes`] +//! +//! Everything above is the *producer* role — raw observations in, +//! `SketchStreamBatch`es out. A node receiving from another +//! `asap_sketches` node instead needs the *receiver* role: +//! [`AsapSketchesPlugin::start_from_envelopes`] swaps the input task +//! for one that consumes `Stream`, decodes +//! each via a [`super::dictionary::SeriesDictionaryDecoder`], and +//! routes the reconstructed envelopes through +//! `Precompute::observe_envelope` (merge, never expand to samples). +//! It reuses the exact same flush ticker / control-channel task / +//! graceful-drain machinery as the producer role — a receiver is +//! still free to have its own `Precompute` config (e.g. +//! `transmit_sketch = false` to *query* the merged sketch every window +//! instead of re-emitting it), so its own emit channel carries +//! whatever that config produces. use std::sync::Arc; use std::time::Duration; @@ -58,7 +75,7 @@ use crate::envelope::SketchEnvelope; use crate::precompute::{Precompute, PrecomputeError, PrecomputeImpl, StatsSnapshot}; use super::config::{resolve, ConfigError, PluginConfig}; -use super::dictionary::{SeriesDictionary, SketchStreamBatch}; +use super::dictionary::{SeriesDictionary, SeriesDictionaryDecoder, SketchStreamBatch}; use super::records::{flatten, OtapMetricRecords, OtapRecordsError}; use super::{decode_batch, OtapDecodeError, OtapEncodeError}; @@ -181,7 +198,8 @@ impl AsapSketchesPlugin { } /// Launch the plugin's three lifecycle tasks against an OTAP - /// input stream. + /// input stream — the **producer** role (raw observations in, + /// `SketchStreamBatch`es out). /// /// `input` is the host-supplied stream of `OtapMetricRecords` /// — the OTAP shell wraps the runtime's `Stream` to @@ -201,6 +219,64 @@ impl AsapSketchesPlugin { ) -> (PluginHandle, EmitReceiver) where S: futures::Stream + Send + Unpin + 'static, + { + let precompute = self.inner.clone(); + self.spawn_lifecycle( + move |cancel| spawn_input_task(precompute, input, cancel), + control, + opts, + ) + } + + /// Launch the plugin's three lifecycle tasks against a stream of + /// pre-aggregated envelopes — the **receiver** role: another + /// `asap_sketches` node's `SketchStreamBatch` output in, this + /// plugin's own `SketchStreamBatch` output out (which, depending + /// on this plugin's own `Precompute` config, might carry + /// re-emitted sketch state, or — with `transmit_sketch = false` — + /// query-mode estimates of the merged sketch). + /// + /// Decodes each batch through a fresh + /// [`SeriesDictionaryDecoder`] retained for the life of this + /// plugin instance, and routes every reconstructed envelope + /// through `Precompute::observe_envelope` (merge, never expand to + /// samples — see the module doc's "The other role" section). + /// Otherwise identical to [`Self::start`]: same ticker / control / + /// graceful-drain machinery, same [`PluginHandle`] / + /// [`EmitReceiver`] return shape. + pub fn start_from_envelopes( + self, + input: S, + control: Option>, + opts: StartOptions, + ) -> (PluginHandle, EmitReceiver) + where + S: futures::Stream + Send + Unpin + 'static, + { + let precompute = self.inner.clone(); + let decoder = Arc::new(Mutex::new(SeriesDictionaryDecoder::new())); + self.spawn_lifecycle( + move |cancel| spawn_envelope_input_task(precompute, decoder, input, cancel), + control, + opts, + ) + } + + /// Shared tail of [`Self::start`] / [`Self::start_from_envelopes`]: + /// wires up the emit channel, shutdown signal, ticker task, + /// optional control task, and the graceful-drain supervisor — + /// everything except *which* input task to spawn, which the two + /// public entry points supply as `spawn_input` (given the shared + /// [`Cancellation`] token so the input task honors shutdown the + /// same way the others do). + fn spawn_lifecycle( + self, + spawn_input: F, + control: Option>, + opts: StartOptions, + ) -> (PluginHandle, EmitReceiver) + where + F: FnOnce(Cancellation) -> JoinHandle<()>, { let (emit_tx, emit_rx) = mpsc::unbounded_channel(); let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); @@ -212,7 +288,7 @@ impl AsapSketchesPlugin { let dictionary = self.dictionary.clone(); let opts = Arc::new(opts); - let input_task = spawn_input_task(precompute.clone(), input, cancellation.clone()); + let input_task = spawn_input(cancellation.clone()); let ticker_task = spawn_ticker_task( precompute.clone(), dictionary.clone(), @@ -363,6 +439,62 @@ fn ingest_one_batch( Ok(()) } +fn spawn_envelope_input_task( + precompute: Arc, + decoder: Arc>, + mut input: S, + cancel: Cancellation, +) -> JoinHandle<()> +where + S: futures::Stream + Send + Unpin + 'static, +{ + use futures::StreamExt; + tokio::spawn(async move { + loop { + tokio::select! { + _ = cancel.cancelled() => return, + next = input.next() => match next { + None => return, + Some(batch) => { + if let Err(_e) = ingest_one_stream_batch(&*precompute, &decoder, &batch).await { + // Same "drop the bad batch, keep the + // plugin alive" policy as ingest_one_batch. + // A decode error here means this stream's + // continuity contract was violated (see + // `OtapDecodeError::UnknownSeriesId` / + // `UnknownAggId`) — Phase D routes this + // onto OTAP's effect-handler error channel. + } + } + }, + } + } + }) +} + +async fn ingest_one_stream_batch( + precompute: &dyn Precompute, + decoder: &Mutex, + batch: &SketchStreamBatch, +) -> Result<(), PluginError> { + let envelopes = { + let mut d = decoder.lock().await; + d.decode(batch)? + }; + for env in &envelopes { + // Merge only — the runtime never expands envelope bytes back + // into scalar samples (the bandwidth invariant). + match precompute.observe_envelope(env) { + Ok(()) => {} + Err(PrecomputeError::LateData) | Err(PrecomputeError::SeriesCapExceeded) => { + continue; + } + Err(e) => return Err(PluginError::Precompute(e)), + } + } + Ok(()) +} + fn spawn_ticker_task( precompute: Arc, dictionary: Arc>, @@ -532,4 +664,98 @@ mod tests { // And that the SketchType enum landed correctly via update_config. let _ = SketchType::KLLSketch; } + + #[tokio::test] + async fn receiver_role_smoke_test_drop_aborts_supervisor() { + // Mirrors handle_drop_aborts_supervisor for the receiver role: + // start_from_envelopes must compile and run against an empty + // SketchStreamBatch stream without deadlocking or panicking. + let cfg = PluginConfig { + sketch_type: "ddsketch".into(), + window_size: Duration::from_millis(50), + ..Default::default() + }; + let plugin = AsapSketchesPlugin::from_plugin_config(&cfg).expect("config"); + let input = futures::stream::empty::(); + let (handle, _rx) = plugin.start_from_envelopes(input, None, StartOptions::default()); + drop(handle); + } + + #[tokio::test] + async fn receiver_role_merges_producer_role_output_end_to_end() { + // Full producer -> receiver chain, both AsapSketchesPlugin, + // connected by an in-process channel (the network-transport + // version of this same chain lives in + // examples/sketch_producer_node.rs / + // examples/sketch_receiver_node.rs). + use crate::observation::{KeyValue, Observation, ObservationValue}; + use tokio_stream::wrappers::UnboundedReceiverStream; + + let producer_cfg = PluginConfig { + sketch_type: "ddsketch".into(), + window_size: Duration::from_secs(60), // drained explicitly below. + output_metric_name: "latency_ms".into(), + agg_id: 1, + ..Default::default() + }; + let producer = + AsapSketchesPlugin::from_plugin_config(&producer_cfg).expect("producer config"); + for v in [1.0_f64, 2.0, 3.0, 4.0, 5.0] { + let obs = Observation::new( + 1_000, + "latency_ms", + vec![], + vec![KeyValue::new("host", "h1")], + ObservationValue::float(v), + ); + producer.precompute().observe(&obs).expect("observe"); + } + let (producer_handle, producer_rx) = producer.start( + futures::stream::pending::(), + None, + StartOptions::default(), + ); + + let receiver_cfg = PluginConfig { + sketch_type: "ddsketch".into(), + window_size: Duration::from_secs(60), + output_metric_name: "latency_ms_p99".into(), + agg_id: 1, + transmit_sketch: false, + quantiles: vec![0.99], + ..Default::default() + }; + let receiver = + AsapSketchesPlugin::from_plugin_config(&receiver_cfg).expect("receiver config"); + let (receiver_handle, mut receiver_rx) = receiver.start_from_envelopes( + UnboundedReceiverStream::new(producer_rx), + None, + StartOptions::default(), + ); + + // Shut the producer down first: its final drain pushes the + // one window it accumulated onto producer_rx, which the + // receiver's envelope-input task picks up and merges. + producer_handle.shutdown().await.expect("producer shutdown"); + // Now shut the receiver down: its final drain (transmit_sketch + // = false) turns the merged sketch into a p99 estimate batch. + receiver_handle.shutdown().await.expect("receiver shutdown"); + + let mut decoder = SeriesDictionaryDecoder::new(); + let mut saw_estimate = false; + while let Ok(Some(batch)) = + tokio::time::timeout(Duration::from_secs(2), receiver_rx.recv()).await + { + for env in decoder.decode(&batch).expect("decode") { + assert_eq!(env.metric_name, "latency_ms_p99"); + assert!( + env.payload.is_empty(), + "estimate mode carries no sketch bytes" + ); + assert!(env.value > 0.0, "p99 of {{1..5}} must be positive"); + saw_estimate = true; + } + } + assert!(saw_estimate, "receiver never emitted a p99 estimate"); + } } diff --git a/asap-precompute-rs/src/otap/mod.rs b/asap-precompute-rs/src/otap/mod.rs index fb96837..31e7f88 100644 --- a/asap-precompute-rs/src/otap/mod.rs +++ b/asap-precompute-rs/src/otap/mod.rs @@ -131,6 +131,7 @@ mod schema; pub mod config; pub mod lifecycle; pub mod records; +pub mod wire; pub use decode::{decode_batch, OtapDecodeError}; pub use dictionary::{ diff --git a/asap-precompute-rs/src/otap/wire.rs b/asap-precompute-rs/src/otap/wire.rs new file mode 100644 index 0000000..a307dd6 --- /dev/null +++ b/asap-precompute-rs/src/otap/wire.rs @@ -0,0 +1,382 @@ +//! Arrow-IPC serialization + a minimal length-prefixed framing for +//! carrying a [`SketchStreamBatch`] across a real transport (a TCP +//! socket here) — the actual "crosses a node or network boundary" +//! hop `docs/data_model.md` opens with, rather than the in-process +//! `mpsc` channel [`crate::otap::lifecycle`]'s tests and +//! `examples/sketch_pipeline_demo.rs` use. +//! +//! # Wire shape +//! +//! One [`SketchStreamBatch`] is framed as: +//! +//! ```text +//! [u32 total_len] +//! [u32 schema_len] [schema_len bytes: Arrow IPC stream] +//! [u32 dictionary_len] [dictionary_len bytes: Arrow IPC stream] +//! [u32 labels_len] [labels_len bytes: Arrow IPC stream] +//! [u32 record_len] [record_len bytes: Arrow IPC stream] +//! ``` +//! +//! Each sub-batch is its own self-contained Arrow IPC *stream* +//! (schema message + one record-batch message + EOS) via +//! [`arrow_ipc::writer::StreamWriter`] — not a shared/continuous +//! Arrow IPC stream across the whole session. That's a deliberate +//! simplification: a real continuous-stream transport would let the +//! four sub-streams themselves carry the Schema/Dictionary economics +//! at the Arrow IPC layer too (per `docs/data_model.md`'s closing +//! "Open design question"), but framing each `SketchStreamBatch` as +//! four independent one-shot streams keeps this module's job to +//! exactly "get the same four `RecordBatch`es to the other side +//! intact," leaving `SeriesDictionary`/`SeriesDictionaryDecoder` (not +//! this module) responsible for the actual dedup. +//! +//! All four sub-batch lengths (and the leading `total_len`) are +//! big-endian `u32`s. A zero-row batch still serializes to a valid +//! (small) Arrow IPC stream — a schema message plus an empty record +//! batch — so a wire-level frame always carries exactly four +//! sub-streams even when, say, `schema`/`dictionary`/`labels` are +//! empty because the series was already known. + +use std::io; + +use arrow_array::RecordBatch; +use thiserror::Error; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpStream; + +use super::dictionary::SketchStreamBatch; + +/// Failure modes for [`encode_stream_batch`] / [`decode_stream_batch`] +/// / [`send_stream_batch`] / [`recv_stream_batch`]. +#[derive(Debug, Error)] +pub enum WireError { + /// Arrow IPC encode/decode failed (malformed batch, schema + /// mismatch inside a sub-stream, etc.). + #[error("otap wire: arrow ipc error: {0}")] + Arrow(#[from] arrow_schema::ArrowError), + + /// A sub-stream decoded to zero record batches (`StreamReader` + /// yielded nothing) — a well-formed IPC stream always carries + /// exactly one, even for zero rows. + #[error("otap wire: sub-batch {which:?} decoded no record batches")] + EmptyRecordBatch { + /// Which of the four sub-batches was empty. + which: &'static str, + }, + + /// The frame's length prefix didn't match the bytes actually + /// available — a truncated or corrupt frame. + #[error("otap wire: frame truncated: expected {expected} bytes, got {actual}")] + Truncated { + /// Bytes the length prefix promised. + expected: usize, + /// Bytes actually present. + actual: usize, + }, + + /// Network I/O failed while sending/receiving a frame. + #[error("otap wire: io error: {0}")] + Io(#[from] io::Error), +} + +/// Serializes one [`RecordBatch`] as a self-contained Arrow IPC +/// stream (schema + one record batch + EOS). +fn write_ipc_bytes(batch: &RecordBatch) -> Result, WireError> { + let mut buf: Vec = Vec::new(); + { + let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut buf, &batch.schema())?; + writer.write(batch)?; + writer.finish()?; + } + Ok(buf) +} + +/// Deserializes one [`RecordBatch`] from bytes written by +/// [`write_ipc_bytes`]. `which` names the sub-batch for error +/// messages only. +fn read_ipc_bytes(bytes: &[u8], which: &'static str) -> Result { + let mut reader = arrow_ipc::reader::StreamReader::try_new(bytes, None)?; + match reader.next() { + Some(batch) => Ok(batch?), + None => Err(WireError::EmptyRecordBatch { which }), + } +} + +/// Serializes a whole [`SketchStreamBatch`] into one length-prefixed +/// frame — see the module doc for the exact byte layout. Does *not* +/// include the leading `total_len` prefix; that's added by +/// [`send_stream_batch`] (or by a caller framing its own transport, +/// e.g. writing this to a file). +pub fn encode_stream_batch(batch: &SketchStreamBatch) -> Result, WireError> { + let parts = [ + write_ipc_bytes(&batch.schema)?, + write_ipc_bytes(&batch.dictionary)?, + write_ipc_bytes(&batch.labels)?, + write_ipc_bytes(&batch.record)?, + ]; + let mut out = Vec::with_capacity(parts.iter().map(|p| p.len() + 4).sum()); + for part in &parts { + out.extend_from_slice(&(part.len() as u32).to_be_bytes()); + out.extend_from_slice(part); + } + Ok(out) +} + +/// Inverse of [`encode_stream_batch`]: reconstructs a +/// [`SketchStreamBatch`] from a frame's body bytes (i.e. everything +/// after the leading `total_len` prefix, if any). +pub fn decode_stream_batch(bytes: &[u8]) -> Result { + let mut cursor = bytes; + let names = ["schema", "dictionary", "labels", "record"]; + let mut parts: Vec = Vec::with_capacity(4); + for which in names { + let mut len_buf = [0u8; 4]; + std::io::Read::read_exact(&mut cursor, &mut len_buf).map_err(|_| WireError::Truncated { + expected: 4, + actual: cursor.len(), + })?; + let len = u32::from_be_bytes(len_buf) as usize; + if cursor.len() < len { + return Err(WireError::Truncated { + expected: len, + actual: cursor.len(), + }); + } + let (part, rest) = cursor.split_at(len); + parts.push(read_ipc_bytes(part, which)?); + cursor = rest; + } + let mut parts = parts.into_iter(); + Ok(SketchStreamBatch { + schema: parts.next().expect("schema part"), + dictionary: parts.next().expect("dictionary part"), + labels: parts.next().expect("labels part"), + record: parts.next().expect("record part"), + }) +} + +/// Sends one [`SketchStreamBatch`] over `stream`, framed with a +/// leading big-endian `u32` total length. The receiving side reads it +/// back with [`recv_stream_batch`]. +pub async fn send_stream_batch( + stream: &mut TcpStream, + batch: &SketchStreamBatch, +) -> Result<(), WireError> { + let body = encode_stream_batch(batch)?; + stream.write_all(&(body.len() as u32).to_be_bytes()).await?; + stream.write_all(&body).await?; + Ok(()) +} + +/// Reads one [`SketchStreamBatch`] previously written by +/// [`send_stream_batch`]. Returns `Ok(None)` on a clean EOF at a +/// frame boundary (the sender closed the connection after its last +/// batch) rather than an error. +pub async fn recv_stream_batch( + stream: &mut TcpStream, +) -> Result, WireError> { + let mut len_buf = [0u8; 4]; + if !read_exact_or_eof(stream, &mut len_buf).await? { + return Ok(None); + } + let len = u32::from_be_bytes(len_buf) as usize; + let mut body = vec![0u8; len]; + stream.read_exact(&mut body).await?; + Ok(Some(decode_stream_batch(&body)?)) +} + +/// Like [`tokio::io::AsyncReadExt::read_exact`], but distinguishes "EOF +/// before any byte of this frame" (returns `Ok(false)` — a clean +/// stream close between frames) from "EOF partway through a frame's +/// length prefix" (still surfaced as an error by the subsequent +/// `read_exact` inside [`recv_stream_batch`], since that's a +/// truncated frame, not a clean close). +async fn read_exact_or_eof(stream: &mut TcpStream, buf: &mut [u8]) -> Result { + let mut filled = 0; + while filled < buf.len() { + let n = stream.read(&mut buf[filled..]).await?; + if n == 0 { + return if filled == 0 { + Ok(false) + } else { + Err(WireError::Truncated { + expected: buf.len(), + actual: filled, + }) + }; + } + filled += n; + } + Ok(true) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::envelope::{Encoding, SketchEnvelope, SketchType}; + use crate::observation::KeyValue; + use crate::otap::dictionary::{SeriesDictionary, SeriesDictionaryDecoder}; + + fn envelope() -> SketchEnvelope { + SketchEnvelope { + schema_version: 1, + sketch_type: SketchType::DDSketch, + agg_id: 7, + resource_labels: Vec::new(), + labels: vec![KeyValue::new("path", "/api")], + window_start_ms: 1_000, + window_end_ms: 11_000, + encoding: Encoding::ProtoFull, + payload: vec![1, 2, 3, 4, 5], + hash_spec: None, + metric_name: "http_request_duration_ms".into(), + count: 42, + aggregation_temporality: 1, + value: 0.0, + } + } + + #[test] + fn encode_decode_round_trips_a_populated_batch() { + let mut dict = SeriesDictionary::new(); + let env = envelope(); + let batch = dict + .encode(std::slice::from_ref(&env), None) + .expect("encode"); + + let bytes = encode_stream_batch(&batch).expect("wire encode"); + let decoded_batch = decode_stream_batch(&bytes).expect("wire decode"); + + assert_eq!(decoded_batch.schema.num_rows(), batch.schema.num_rows()); + assert_eq!( + decoded_batch.dictionary.num_rows(), + batch.dictionary.num_rows() + ); + assert_eq!(decoded_batch.labels.num_rows(), batch.labels.num_rows()); + assert_eq!(decoded_batch.record.num_rows(), batch.record.num_rows()); + + // The whole point: joining the IPC-round-tripped batch back + // through SeriesDictionaryDecoder reconstructs the original + // envelope exactly. + let mut decoder = SeriesDictionaryDecoder::new(); + let out = decoder.decode(&decoded_batch).expect("dictionary decode"); + assert_eq!(out.len(), 1); + assert_eq!(out[0].payload, env.payload); + assert_eq!(out[0].labels, env.labels); + assert_eq!(out[0].metric_name, env.metric_name); + assert_eq!(out[0].sketch_type, env.sketch_type); + } + + #[test] + fn encode_decode_round_trips_an_empty_batch() { + let mut dict = SeriesDictionary::new(); + let batch = dict.encode(&[], None).expect("encode empty"); + let bytes = encode_stream_batch(&batch).expect("wire encode"); + let decoded = decode_stream_batch(&bytes).expect("wire decode"); + assert!(decoded.schema.num_rows() == 0); + assert!(decoded.dictionary.num_rows() == 0); + assert!(decoded.labels.num_rows() == 0); + assert!(decoded.record.num_rows() == 0); + } + + #[test] + fn encode_decode_round_trips_repeat_window_dictionary_free_batch() { + // The batch that matters most: window 2+ for an + // already-known series, where schema/dictionary/labels are + // genuinely empty and only `record` carries a row. + let mut dict = SeriesDictionary::new(); + let env1 = envelope(); + let _ = dict + .encode(std::slice::from_ref(&env1), None) + .expect("window 1"); + let env2 = SketchEnvelope { + window_start_ms: 11_000, + window_end_ms: 21_000, + payload: vec![9, 9, 9], + ..envelope() + }; + let batch2 = dict + .encode(std::slice::from_ref(&env2), None) + .expect("window 2"); + assert_eq!(batch2.schema.num_rows(), 0); + assert_eq!(batch2.dictionary.num_rows(), 0); + assert_eq!(batch2.labels.num_rows(), 0); + assert_eq!(batch2.record.num_rows(), 1); + + let bytes = encode_stream_batch(&batch2).expect("wire encode"); + let decoded_batch = decode_stream_batch(&bytes).expect("wire decode"); + + let mut decoder = SeriesDictionaryDecoder::new(); + // Must ingest window 1 first — this decoded batch alone has + // no DICTIONARY entry to resolve series_id 0 against. + let bytes1 = encode_stream_batch(&dict_only_window1(&env1)).expect("encode w1"); + let decoded1 = decode_stream_batch(&bytes1).expect("decode w1"); + decoder.decode(&decoded1).expect("decode w1 into decoder"); + + let out = decoder.decode(&decoded_batch).expect("decode w2"); + assert_eq!(out.len(), 1); + assert_eq!(out[0].payload, env2.payload); + assert_eq!(out[0].window_start_ms, 11_000); + assert_eq!(out[0].window_end_ms, 21_000); + } + + /// Helper for the test above: re-derives window 1's batch from a + /// fresh dictionary so it can be fed to a fresh decoder + /// independently of the outer test's `dict` state. + fn dict_only_window1(env: &SketchEnvelope) -> SketchStreamBatch { + let mut dict = SeriesDictionary::new(); + dict.encode(std::slice::from_ref(env), None) + .expect("encode") + } + + #[tokio::test] + async fn send_recv_round_trips_over_a_real_tcp_loopback_socket() { + use tokio::net::{TcpListener, TcpStream}; + + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("local_addr"); + + let mut dict = SeriesDictionary::new(); + let env = envelope(); + let batch = dict + .encode(std::slice::from_ref(&env), None) + .expect("encode"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accept"); + recv_stream_batch(&mut socket) + .await + .expect("recv") + .expect("Some(batch)") + }); + + let mut client = TcpStream::connect(addr).await.expect("connect"); + send_stream_batch(&mut client, &batch).await.expect("send"); + drop(client); // signal EOF after the one frame. + + let received = server.await.expect("server task"); + let mut decoder = SeriesDictionaryDecoder::new(); + let out = decoder.decode(&received).expect("decode"); + assert_eq!(out.len(), 1); + assert_eq!(out[0].payload, env.payload); + } + + #[tokio::test] + async fn recv_returns_none_on_clean_eof_between_frames() { + use tokio::net::{TcpListener, TcpStream}; + + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("local_addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accept"); + recv_stream_batch(&mut socket).await + }); + + let client = TcpStream::connect(addr).await.expect("connect"); + drop(client); // close immediately, no frames sent. + + let result = server.await.expect("server task").expect("no error"); + assert!(result.is_none()); + } +}