From c4f317f2b0df2701c4987273d33f1a2e3b158760 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Wed, 6 May 2026 00:49:06 -0400 Subject: [PATCH 1/3] =?UTF-8?q?feat(cold-store):=20GorillaS3ColdStore=20?= =?UTF-8?q?=E2=80=94=20list+read=20chunks=20from=20S3=20via=20asap-gorilla?= =?UTF-8?q?=20(Phase=203)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3 of the Gorilla-S3-cold-engine — adds a `ColdStore` adapter that fetches per-hour `index.json` catalogs from an S3-compatible bucket, prunes them by time range, and decodes the selected `GORILLA1` chunks via the freshly-merged `asap-gorilla` crate (ASAPCollector PR #281). Changes ------- * `cold_store/mod.rs` — additive trait extension. New `ChunkRef` descriptor + `list_chunks` / `read_chunk` methods carry default impls returning `ColdStoreError::Unsupported`, so `LocalFsColdStore` and the existing `s3_adapter::ColdFallback` chain compile unchanged. Two new `ColdStoreError` variants (`Backend(String)` for transport failures, `Unsupported(&'static str)` for the default-impl errors). * `cold_store/gorilla_s3.rs` — new module. `GorillaS3ColdStore` + `GorillaS3Config` (with `from_env`) + `ObjectStore` trait + a `rust-s3`-backed production impl (`S3ObjectStore`, with MinIO path-style support) and an in-memory mock used by the tests. LRU cache (default 256 chunks) keyed on the chunk object key, holding pre-decoded `RawSample` lists so repeated reads skip the Gorilla decode pass entirely. * `asap-query-engine/Cargo.toml` — new deps: `asap-gorilla` (path), `rust-s3 = 0.37` (default-features off, `tokio-rustls-tls` to share the rustls backend reqwest already pulls in), `lru = 0.12`. S3 client choice ---------------- Picked `rust-s3` over `aws-sdk-s3` for the lighter dep tree (no full AWS SDK fanout) and first-class MinIO support (`with_path_style()` is the supported config rather than a workaround). The trait-based `ObjectStore` indirection means the choice is replaceable without touching the `ColdStore` impl. Tests (10, all green) --------------------- * `list_chunks_via_indexfile_prunes_by_time` — hand-crafted index with three chunks; window overlaps only the middle one. * `read_chunk_decodes_via_asap_gorilla` — round-trip through a real `GorillaEncoder` block; samples + labels match. * `cache_hit_skips_s3_fetch` — second `read_chunk` on the same key must not trigger another S3 GET (mock counter assertion). * `lru_eviction_under_pressure` — `cache_capacity = 2`, fill three, re-read the first — must trigger a fresh GET. * `index_json_corrupted_returns_error` — clean `Malformed`, no panic. * `s3_unavailable_returns_error` — clean `Backend(...)`, no panic. * `missing_index_is_empty_not_error` — producer hasn't flushed yet; empty result, not error. * `scan_filters_to_requested_range` — chunk overlaps but inner samples filter to zero. * `list_chunks_spans_two_hour_buckets` — request crosses an hour boundary; both hour indexes are fetched. * `from_env_requires_bucket` — config validation surfaces missing env. Co-Authored-By: Claude Opus 4.7 (1M context) --- Cargo.lock | 643 ++++++++++- asap-query-engine/Cargo.toml | 11 + .../query/fallback/cold_store/gorilla_s3.rs | 1000 +++++++++++++++++ .../drivers/query/fallback/cold_store/mod.rs | 85 ++ 4 files changed, 1721 insertions(+), 18 deletions(-) create mode 100644 asap-query-engine/src/drivers/query/fallback/cold_store/gorilla_s3.rs diff --git a/Cargo.lock b/Cargo.lock index b8172af3..4f977a85 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -385,6 +385,15 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "asap-gorilla" +version = "0.1.0" +dependencies = [ + "serde", + "serde_json", + "thiserror 1.0.69", +] + [[package]] name = "asap-precompute-rs" version = "0.1.0" @@ -418,7 +427,7 @@ dependencies = [ "pretty_assertions", "promql-parser", "promql_utilities", - "reqwest", + "reqwest 0.11.27", "serde", "serde_json", "serde_yaml", @@ -527,6 +536,22 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" +[[package]] +name = "attohttpc" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16e2cdb6d5ed835199484bb92bb8b3edd526effe995c61732580439c1a67e2e9" +dependencies = [ + "base64 0.22.1", + "http 1.4.0", + "log", + "rustls 0.23.40", + "serde", + "serde_json", + "url", + "webpki-roots 1.0.7", +] + [[package]] name = "atty" version = "0.2.14" @@ -544,6 +569,54 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +[[package]] +name = "aws-creds" +version = "0.39.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca3b85155d265df828f84e53886ed9e427aed979dd8a39f5b8b2162c77e142d7" +dependencies = [ + "attohttpc", + "home", + "log", + "quick-xml", + "rust-ini", + "serde", + "thiserror 2.0.18", + "time", + "url", +] + +[[package]] +name = "aws-lc-rs" +version = "1.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec6fb3fe69024a75fa7e1bfb48aa6cf59706a101658ea01bfd33b2b248a038f" +dependencies = [ + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.40.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f50037ee5e1e41e7b8f9d161680a725bd1626cb6f8c7e901f91f942850852fe7" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", +] + +[[package]] +name = "aws-region" +version = "0.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "838b36c8dc927b6db1b6c6b8f5d05865f2213550b9e83bf92fa99ed6525472c0" +dependencies = [ + "thiserror 2.0.18", +] + [[package]] name = "axum" version = "0.7.9" @@ -762,6 +835,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "cfgrammar" version = "0.13.10" @@ -883,6 +962,15 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + [[package]] name = "colorchoice" version = "1.0.5" @@ -1534,6 +1622,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" dependencies = [ "powerfmt", + "serde_core", ] [[package]] @@ -1564,6 +1653,21 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "dlv-list" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "442039f5147480ba31067cb00ada1adae6892028e40e45fc5de7b7df6dcc1b5f" +dependencies = [ + "const-random", +] + +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + [[package]] name = "either" version = "1.15.0" @@ -1692,6 +1796,12 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + [[package]] name = "futures" version = "0.3.32" @@ -1806,8 +1916,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" dependencies = [ "cfg-if", + "js-sys", "libc", "wasi", + "wasm-bindgen", ] [[package]] @@ -1817,9 +1929,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi 5.3.0", "wasip2", + "wasm-bindgen", ] [[package]] @@ -1913,6 +2027,8 @@ version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ + "allocator-api2", + "equivalent", "foldhash", ] @@ -1967,6 +2083,24 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest", +] + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "http" version = "0.2.12" @@ -2095,9 +2229,25 @@ dependencies = [ "futures-util", "http 0.2.12", "hyper 0.14.32", - "rustls", + "rustls 0.21.12", "tokio", - "tokio-rustls", + "tokio-rustls 0.24.1", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http 1.4.0", + "hyper 1.9.0", + "hyper-util", + "rustls 0.23.40", + "tokio", + "tokio-rustls 0.26.4", + "tower-service", + "webpki-roots 1.0.7", ] [[package]] @@ -2119,13 +2269,16 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ + "base64 0.22.1", "bytes", "futures-channel", "futures-util", "http 1.4.0", "http-body 1.0.1", "hyper 1.9.0", + "ipnet", "libc", + "percent-encoding", "pin-project-lite", "socket2 0.6.3", "tokio", @@ -2154,7 +2307,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core", + "windows-core 0.62.2", ] [[package]] @@ -2631,6 +2784,21 @@ dependencies = [ "vob", ] +[[package]] +name = "lru" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" +dependencies = [ + "hashbrown 0.15.5", +] + +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + [[package]] name = "lz4_flex" version = "0.11.6" @@ -2666,6 +2834,17 @@ version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" +[[package]] +name = "maybe-async" +version = "0.2.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cf92c10c7e361d6b99666ec1c6f9805b0bea2c3bd8c78dc6fe98ac5bd78db11" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "md-5" version = "0.10.6" @@ -2676,6 +2855,12 @@ dependencies = [ "digest", ] +[[package]] +name = "md5" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae960838283323069879657ca3de837e9f7bbb4c7bf6ea7f1b290d5e9476d2e0" + [[package]] name = "memchr" version = "2.8.0" @@ -2741,6 +2926,15 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" +[[package]] +name = "ntapi" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3b335231dfd352ffb0f8017f3b6027a4917f7df785ea2143d8af2adc66980ae" +dependencies = [ + "winapi", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -2871,6 +3065,25 @@ dependencies = [ "libc", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags 2.11.1", +] + +[[package]] +name = "objc2-io-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33fafba39597d6dc1fb709123dfa8289d39406734be322956a69f0931c73bb15" +dependencies = [ + "libc", + "objc2-core-foundation", +] + [[package]] name = "object" version = "0.37.3" @@ -2928,6 +3141,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "ordered-multimap" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49203cdcae0030493bad186b28da2fa25645fa276a51b6fec8010d281e02ef79" +dependencies = [ + "dlv-list", + "hashbrown 0.14.5", +] + [[package]] name = "packedvec" version = "1.2.5" @@ -2958,7 +3181,7 @@ dependencies = [ "libc", "redox_syscall 0.5.18", "smallvec", - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -3420,6 +3643,7 @@ dependencies = [ "anyhow", "arc-swap", "arrow", + "asap-gorilla", "asap-precompute-rs", "asap_otel_proto", "asap_planner", @@ -3442,6 +3666,7 @@ dependencies = [ "futures", "hex", "lazy_static", + "lru", "memmap2", "moka", "prometheus", @@ -3450,9 +3675,10 @@ dependencies = [ "prost", "rdkafka", "regex", - "reqwest", + "reqwest 0.11.27", "rmp-serde", "rusqlite", + "rust-s3", "serde", "serde_json", "serde_yaml", @@ -3474,6 +3700,71 @@ dependencies = [ "zstd", ] +[[package]] +name = "quick-xml" +version = "0.38.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b66c2058c55a409d601666cffe35f04333cf1013010882cec174a7467cd4e21c" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "quinn" +version = "0.11.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash", + "rustls 0.23.40", + "socket2 0.6.3", + "thiserror 2.0.18", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" +dependencies = [ + "bytes", + "getrandom 0.3.4", + "lru-slab", + "rand 0.9.4", + "ring", + "rustc-hash", + "rustls 0.23.40", + "rustls-pki-types", + "slab", + "thiserror 2.0.18", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2 0.6.3", + "tracing", + "windows-sys 0.52.0", +] + [[package]] name = "quote" version = "1.0.45" @@ -3686,7 +3977,7 @@ dependencies = [ "http 0.2.12", "http-body 0.4.6", "hyper 0.14.32", - "hyper-rustls", + "hyper-rustls 0.24.2", "ipnet", "js-sys", "log", @@ -3694,7 +3985,7 @@ dependencies = [ "once_cell", "percent-encoding", "pin-project-lite", - "rustls", + "rustls 0.21.12", "rustls-pemfile", "serde", "serde_json", @@ -3702,16 +3993,57 @@ dependencies = [ "sync_wrapper 0.1.2", "system-configuration", "tokio", - "tokio-rustls", + "tokio-rustls 0.24.1", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", "web-sys", - "webpki-roots", + "webpki-roots 0.25.4", "winreg", ] +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "futures-util", + "http 1.4.0", + "http-body 1.0.1", + "http-body-util", + "hyper 1.9.0", + "hyper-rustls 0.27.9", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls 0.23.40", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper 1.0.2", + "tokio", + "tokio-rustls 0.26.4", + "tokio-util", + "tower 0.5.3", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", + "webpki-roots 1.0.7", +] + [[package]] name = "ring" version = "0.17.14" @@ -3759,6 +4091,56 @@ dependencies = [ "smallvec", ] +[[package]] +name = "rust-ini" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "796e8d2b6696392a43bea58116b667fb4c29727dc5abd27d6acf338bb4f688c7" +dependencies = [ + "cfg-if", + "ordered-multimap", +] + +[[package]] +name = "rust-s3" +version = "0.37.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aeedb13abdaa7e48d391de05b0569b37fa0a7a64a668dff6ffb2141ad0c2527e" +dependencies = [ + "async-trait", + "aws-creds", + "aws-region", + "base64 0.22.1", + "bytes", + "cfg-if", + "futures-util", + "hex", + "hmac", + "http 1.4.0", + "log", + "maybe-async", + "md5", + "percent-encoding", + "quick-xml", + "reqwest 0.12.28", + "serde", + "serde_derive", + "serde_json", + "sha2", + "sysinfo", + "thiserror 2.0.18", + "time", + "tokio", + "tokio-stream", + "url", +] + +[[package]] +name = "rustc-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" + [[package]] name = "rustc_version" version = "0.4.1" @@ -3789,10 +4171,26 @@ checksum = "3f56a14d1f48b391359b22f731fd4bd7e43c97f3c50eee276f3aa09c94784d3e" dependencies = [ "log", "ring", - "rustls-webpki", + "rustls-webpki 0.101.7", "sct", ] +[[package]] +name = "rustls" +version = "0.23.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" +dependencies = [ + "aws-lc-rs", + "log", + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki 0.103.13", + "subtle", + "zeroize", +] + [[package]] name = "rustls-pemfile" version = "1.0.4" @@ -3802,6 +4200,16 @@ dependencies = [ "base64 0.21.7", ] +[[package]] +name = "rustls-pki-types" +version = "1.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +dependencies = [ + "web-time", + "zeroize", +] + [[package]] name = "rustls-webpki" version = "0.101.7" @@ -3812,6 +4220,18 @@ dependencies = [ "untrusted", ] +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "aws-lc-rs", + "ring", + "rustls-pki-types", + "untrusted", +] + [[package]] name = "rustversion" version = "1.0.22" @@ -4237,6 +4657,9 @@ name = "sync_wrapper" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] [[package]] name = "synstructure" @@ -4249,6 +4672,20 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "sysinfo" +version = "0.37.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16607d5caffd1c07ce073528f9ed972d88db15dd44023fa57142963be3feb11f" +dependencies = [ + "libc", + "memchr", + "ntapi", + "objc2-core-foundation", + "objc2-io-kit", + "windows", +] + [[package]] name = "system-configuration" version = "0.5.1" @@ -4420,6 +4857,21 @@ dependencies = [ "serde_json", ] +[[package]] +name = "tinyvec" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + [[package]] name = "tokio" version = "1.52.2" @@ -4454,7 +4906,17 @@ version = "0.24.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" dependencies = [ - "rustls", + "rustls 0.21.12", + "tokio", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls 0.23.40", "tokio", ] @@ -4603,6 +5065,24 @@ dependencies = [ "tracing", ] +[[package]] +name = "tower-http" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a28f0d049ccfaa566e14e9663d304d8577427b368cb4710a20528690287a738b" +dependencies = [ + "bitflags 2.11.1", + "bytes", + "futures-util", + "http 1.4.0", + "http-body 1.0.1", + "pin-project-lite", + "tower 0.5.3", + "tower-layer", + "tower-service", + "url", +] + [[package]] name = "tower-layer" version = "0.3.3" @@ -4969,6 +5449,19 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "wasmparser" version = "0.244.0" @@ -4991,12 +5484,31 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "webpki-roots" version = "0.25.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5f20c57d8d7db6d3b86154206ae5d8fba62dd39573114de97c2cb0578251f8e1" +[[package]] +name = "webpki-roots" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52f5ee44c96cf55f1b349600768e3ece3a8f26010c05265ab73f945bb1a2eb9d" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "winapi" version = "0.3.9" @@ -5028,6 +5540,41 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows" +version = "0.61.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9babd3a767a4c1aef6900409f85f5d53ce2544ccdfaa86dad48c91782c6d6893" +dependencies = [ + "windows-collections", + "windows-core 0.61.2", + "windows-future", + "windows-link 0.1.3", + "windows-numerics", +] + +[[package]] +name = "windows-collections" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3beeceb5e5cfd9eb1d76b381630e82c4241ccd0d27f1a39ed41b2760b255c5e8" +dependencies = [ + "windows-core 0.61.2", +] + +[[package]] +name = "windows-core" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0fdd3ddb90610c7638aa2b3a3ab2904fb9e5cdbecc643ddb3647212781c4ae3" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link 0.1.3", + "windows-result 0.3.4", + "windows-strings 0.4.2", +] + [[package]] name = "windows-core" version = "0.62.2" @@ -5036,9 +5583,20 @@ checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" dependencies = [ "windows-implement", "windows-interface", - "windows-link", - "windows-result", - "windows-strings", + "windows-link 0.2.1", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + +[[package]] +name = "windows-future" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc6a41e98427b19fe4b73c550f060b59fa592d7d686537eebf9385621bfbad8e" +dependencies = [ + "windows-core 0.61.2", + "windows-link 0.1.3", + "windows-threading", ] [[package]] @@ -5063,19 +5621,53 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "windows-link" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-numerics" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9150af68066c4c5c07ddc0ce30421554771e528bde427614c61038bc2c92c2b1" +dependencies = [ + "windows-core 0.61.2", + "windows-link 0.1.3", +] + +[[package]] +name = "windows-result" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56f42bd332cc6c8eac5af113fc0c1fd6a8fd2aa08a0119358686e5160d0586c6" +dependencies = [ + "windows-link 0.1.3", +] + [[package]] name = "windows-result" version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" dependencies = [ - "windows-link", + "windows-link 0.2.1", +] + +[[package]] +name = "windows-strings" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56e6c93f3a0c3b36176cb1327a4958a0353d5d166c2a35cb268ace15e91d3b57" +dependencies = [ + "windows-link 0.1.3", ] [[package]] @@ -5084,7 +5676,7 @@ version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" dependencies = [ - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -5111,7 +5703,7 @@ version = "0.61.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" dependencies = [ - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -5145,6 +5737,15 @@ dependencies = [ "windows_x86_64_msvc 0.52.6", ] +[[package]] +name = "windows-threading" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b66463ad2e0ea3bbf808b7f1d371311c80e115c0b71d60efc142cafbcfb057a6" +dependencies = [ + "windows-link 0.1.3", +] + [[package]] name = "windows_aarch64_gnullvm" version = "0.48.5" @@ -5448,6 +6049,12 @@ dependencies = [ "synstructure", ] +[[package]] +name = "zeroize" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" + [[package]] name = "zerotrie" version = "0.2.4" diff --git a/asap-query-engine/Cargo.toml b/asap-query-engine/Cargo.toml index cfe6e59f..08d53de7 100644 --- a/asap-query-engine/Cargo.toml +++ b/asap-query-engine/Cargo.toml @@ -86,6 +86,17 @@ asap-precompute-rs = { path = "../../ASAPCollector/asap-precompute-rs" } moka = { version = "0.12", features = ["sync"] } memmap2 = "0.9" crc32fast = "1.4" +# Phase 3 — Gorilla-S3 cold engine. `asap-gorilla` (path-dep, mirrors +# the `asap-precompute-rs` sibling-checkout pattern) provides the +# `GORILLA1` block decoder + per-hour `index.json` catalog. `rust-s3` +# (`s3` package on crates.io) is a single-crate S3-compatible client +# with first-class MinIO support; we hold it behind +# `default-features = false` + `tokio-rustls-tls` so it shares the +# rustls backend already pulled in by `reqwest`. `lru` powers the +# Phase 3 `ChunkCache` keyed on chunk object key. +asap-gorilla = { path = "../../ASAPCollector/asap-gorilla" } +s3 = { version = "0.37", package = "rust-s3", default-features = false, features = ["tokio-rustls-tls"] } +lru = "0.12" [[bin]] name = "precompute_engine" diff --git a/asap-query-engine/src/drivers/query/fallback/cold_store/gorilla_s3.rs b/asap-query-engine/src/drivers/query/fallback/cold_store/gorilla_s3.rs new file mode 100644 index 00000000..02368662 --- /dev/null +++ b/asap-query-engine/src/drivers/query/fallback/cold_store/gorilla_s3.rs @@ -0,0 +1,1000 @@ +//! Gorilla-on-S3 [`ColdStore`] adapter — Phase 3 of the +//! Gorilla-S3-cold-engine. +//! +//! Lists per-hour `index.json` catalogs out of an S3-compatible +//! bucket, prunes them by time range, then fetches + decodes the +//! selected `GORILLA1` chunks via the freshly-merged +//! [`asap_gorilla`] crate (`ASAPCollector` PR #281). +//! +//! Sits alongside [`super::LocalFsColdStore`] — both impls satisfy +//! the same [`super::ColdStore`] trait, so the existing +//! `s3_adapter::ColdFallback` query path can swap between them +//! without code change. The Phase 3 trait extension +//! ([`super::ColdStore::list_chunks`] / [`super::ColdStore::read_chunk`]) +//! lets the upcoming Phase 4 `GorillaQueryEngine` pull chunks one +//! at a time without materialising every sample. +//! +//! # Object key layout +//! +//! `GorillaS3ColdStore` is **agnostic** about the on-S3 chunk-key +//! shape. Two layouts are known to coexist (see PR #281): +//! +//! * design.md canonical: +//! `//YYYY/MM/DD/HH/part-NNNNNN.gor` +//! * Telegraf-side `gorilla_s3` output: +//! `/block---.gorilla` (random suffix) +//! +//! The per-hour `index.json` is the source of truth for what keys +//! exist; we treat [`asap_gorilla::IndexEntry::key`] as opaque and +//! do not try to parse it. The `prefix_template` config field +//! controls only where the **index** files live, not the chunks. +//! +//! # S3 client +//! +//! Backed by the `rust-s3` crate (`s3 = "0.37"`) — single-crate +//! dep, MinIO-friendly out of the box (no AWS-specific signing +//! quirks, supports custom endpoint URLs + path-style addressing). +//! Hidden behind the [`ObjectStore`] trait below so tests use an +//! in-memory mock and do not need a live MinIO. + +use std::num::NonZeroUsize; +use std::sync::Arc; + +#[cfg(test)] +use std::collections::HashMap; + +use async_trait::async_trait; +use chrono::{DateTime, Datelike, Timelike, Utc}; +use lru::LruCache; +use thiserror::Error; +use tokio::sync::Mutex; +use tracing::debug; + +use asap_gorilla::{GorillaDecoder, IndexFile}; + +use super::{ChunkRef, ColdStore, ColdStoreError, RawSample}; + +// ───────────────────────────────────────────────────────────────────── +// Public config +// ───────────────────────────────────────────────────────────────────── + +/// Tunable configuration for [`GorillaS3ColdStore`]. +/// +/// Use [`GorillaS3Config::from_env`] to pull values from environment +/// variables in deployment, or build manually for tests. +#[derive(Debug, Clone)] +pub struct GorillaS3Config { + /// `None` for AWS S3 (the SDK uses the standard regional + /// endpoint), `Some("http://minio:9000")` for MinIO / a custom + /// S3-compatible endpoint. + pub endpoint: Option, + /// Bucket name to list / read from. + pub bucket: String, + /// Tenant identifier prepended to every index-file prefix. + /// Empty string is allowed for single-tenant deployments. + pub tenant: String, + /// Prefix template for per-hour `index.json` files. Supports + /// the placeholders `{tenant}`, `{metric}`, `{year}`, `{month}`, + /// `{day}`, `{hour}` (zero-padded). Default: + /// `"{tenant}/{metric}/{year}/{month}/{day}/{hour}/"`. + pub prefix_template: String, + /// AWS-region the bucket lives in (e.g. `"us-east-1"`). For + /// MinIO any non-empty placeholder works. + pub region: String, + /// Optional static credential override. Both fields must be set + /// together; if either is `None` the underlying SDK falls back + /// to its environment / IMDS resolution. + pub access_key_id: Option, + /// See [`Self::access_key_id`]. + pub secret_access_key: Option, + /// LRU capacity (in number of decoded chunks). Default `256`. + pub cache_capacity: usize, + /// `false` switches the SDK to plain HTTP — required for local + /// MinIO / docker-compose smoke tests. Default `true`. + pub use_ssl: bool, +} + +impl Default for GorillaS3Config { + fn default() -> Self { + Self { + endpoint: None, + bucket: String::new(), + tenant: String::new(), + prefix_template: "{tenant}/{metric}/{year}/{month}/{day}/{hour}/".to_string(), + region: "us-east-1".to_string(), + access_key_id: None, + secret_access_key: None, + cache_capacity: 256, + use_ssl: true, + } + } +} + +/// Errors raised by [`GorillaS3Config::from_env`]. +#[derive(Debug, Error)] +pub enum GorillaS3ConfigError { + /// A required environment variable was missing. + #[error("missing required env var: {0}")] + MissingEnv(&'static str), + /// `ASAP_GORILLA_S3_CACHE_CAPACITY` could not be parsed as a + /// positive `usize`. + #[error("invalid env var {var}: {value} ({source})")] + InvalidEnv { + /// Variable name. + var: &'static str, + /// Raw value the user supplied. + value: String, + /// Underlying parse error. + source: std::num::ParseIntError, + }, +} + +impl GorillaS3Config { + /// Read a config from process environment variables. Required: + /// + /// * `ASAP_GORILLA_S3_BUCKET` + /// * `ASAP_GORILLA_S3_REGION` + /// + /// Optional (with defaults shown above): + /// + /// * `ASAP_GORILLA_S3_ENDPOINT` + /// * `ASAP_GORILLA_S3_TENANT` + /// * `ASAP_GORILLA_S3_PREFIX_TEMPLATE` + /// * `ASAP_GORILLA_S3_ACCESS_KEY_ID` / `..._SECRET_ACCESS_KEY` + /// * `ASAP_GORILLA_S3_CACHE_CAPACITY` + /// * `ASAP_GORILLA_S3_USE_SSL` (`"true"` / `"false"`, + /// case-insensitive) + pub fn from_env() -> Result { + let bucket = std::env::var("ASAP_GORILLA_S3_BUCKET") + .map_err(|_| GorillaS3ConfigError::MissingEnv("ASAP_GORILLA_S3_BUCKET"))?; + let region = std::env::var("ASAP_GORILLA_S3_REGION") + .map_err(|_| GorillaS3ConfigError::MissingEnv("ASAP_GORILLA_S3_REGION"))?; + let endpoint = std::env::var("ASAP_GORILLA_S3_ENDPOINT").ok(); + let tenant = std::env::var("ASAP_GORILLA_S3_TENANT").unwrap_or_default(); + let prefix_template = std::env::var("ASAP_GORILLA_S3_PREFIX_TEMPLATE") + .unwrap_or_else(|_| "{tenant}/{metric}/{year}/{month}/{day}/{hour}/".to_string()); + let access_key_id = std::env::var("ASAP_GORILLA_S3_ACCESS_KEY_ID").ok(); + let secret_access_key = std::env::var("ASAP_GORILLA_S3_SECRET_ACCESS_KEY").ok(); + let cache_capacity = match std::env::var("ASAP_GORILLA_S3_CACHE_CAPACITY") { + Ok(s) => s + .parse::() + .map_err(|e| GorillaS3ConfigError::InvalidEnv { + var: "ASAP_GORILLA_S3_CACHE_CAPACITY", + value: s, + source: e, + })?, + Err(_) => 256, + }; + let use_ssl = std::env::var("ASAP_GORILLA_S3_USE_SSL") + .map(|s| !matches!(s.trim().to_ascii_lowercase().as_str(), "false" | "0" | "no")) + .unwrap_or(true); + Ok(Self { + endpoint, + bucket, + tenant, + prefix_template, + region, + access_key_id, + secret_access_key, + cache_capacity, + use_ssl, + }) + } +} + +// ───────────────────────────────────────────────────────────────────── +// ObjectStore — internal trait so tests don't need real S3 +// ───────────────────────────────────────────────────────────────────── + +/// Minimal async object-fetch interface. +/// +/// Sized + `Send + Sync` so [`GorillaS3ColdStore`] can hold one +/// behind an `Arc` regardless of how it's backed. +/// Production callers use [`S3ObjectStore`] (rust-s3); tests use the +/// in-memory mock at the bottom of this file. +#[async_trait] +pub trait ObjectStore: Send + Sync { + /// Fetch the full object body for `key`. + /// + /// Returns [`ColdStoreError::Backend`] for transport errors and + /// [`ColdStoreError::Backend`] (with a `not found` substring) + /// for missing keys; callers distinguish via + /// [`ObjectStore::object_missing`] if they need to. + async fn get_object(&self, key: &str) -> Result, ColdStoreError>; + + /// True iff `err` was raised because the requested key did not + /// exist (vs. a transport / permission failure). Used by the + /// list path to treat a missing `index.json` as "no chunks for + /// this hour" rather than a hard error. + fn object_missing(&self, err: &ColdStoreError) -> bool { + matches!(err, ColdStoreError::Backend(msg) if msg.contains("not found")) + } +} + +// ───────────────────────────────────────────────────────────────────── +// rust-s3 backed production ObjectStore +// ───────────────────────────────────────────────────────────────────── + +mod rust_s3_backend { + use super::*; + use s3::creds::Credentials; + use s3::region::Region as S3Region; + use s3::Bucket; + + /// `rust-s3`-backed [`ObjectStore`]. Default production choice. + pub struct S3ObjectStore { + bucket: Box, + } + + impl S3ObjectStore { + /// Build from a [`GorillaS3Config`]. Sets + /// `path_style = true` whenever a custom endpoint is + /// configured (MinIO mandates path-style addressing). + pub fn new(cfg: &GorillaS3Config) -> Result { + let region = match &cfg.endpoint { + Some(ep) => { + let endpoint = if ep.starts_with("http://") || ep.starts_with("https://") { + ep.clone() + } else if cfg.use_ssl { + format!("https://{}", ep) + } else { + format!("http://{}", ep) + }; + S3Region::Custom { + region: cfg.region.clone(), + endpoint, + } + } + None => cfg + .region + .parse::() + .map_err(|e| ColdStoreError::Backend(format!("region parse: {e}")))?, + }; + let creds = match (&cfg.access_key_id, &cfg.secret_access_key) { + (Some(ak), Some(sk)) => { + Credentials::new(Some(ak), Some(sk), None, None, None).map_err(|e| { + ColdStoreError::Backend(format!("credentials: {e}")) + })? + } + _ => Credentials::default().map_err(|e| { + ColdStoreError::Backend(format!("default credentials: {e}")) + })?, + }; + let bucket = Bucket::new(&cfg.bucket, region, creds) + .map_err(|e| ColdStoreError::Backend(format!("bucket: {e}")))?; + // MinIO + most S3-compatibles require path-style addressing + // when a custom endpoint is in play. AWS S3 supports both, + // so leaving it on for the AWS path is safe but slightly + // less efficient — only flip when an endpoint is set. + let bucket = if cfg.endpoint.is_some() { + bucket.with_path_style() + } else { + bucket + }; + Ok(Self { bucket }) + } + } + + #[async_trait] + impl ObjectStore for S3ObjectStore { + async fn get_object(&self, key: &str) -> Result, ColdStoreError> { + let resp = self + .bucket + .get_object(key) + .await + .map_err(|e| ColdStoreError::Backend(format!("s3 get {key}: {e}")))?; + if resp.status_code() == 404 { + return Err(ColdStoreError::Backend(format!( + "s3 get {key}: not found" + ))); + } + if !(200..300).contains(&resp.status_code()) { + return Err(ColdStoreError::Backend(format!( + "s3 get {key}: status {}", + resp.status_code() + ))); + } + Ok(resp.to_vec()) + } + } +} + +pub use rust_s3_backend::S3ObjectStore; + +// ───────────────────────────────────────────────────────────────────── +// GorillaS3ColdStore +// ───────────────────────────────────────────────────────────────────── + +/// LRU cache keyed by chunk object key. Stored values are +/// pre-decoded `RawSample` lists so repeated reads of the same +/// chunk skip the Gorilla decode pass entirely. +type ChunkCache = Mutex>>>; + +/// `ColdStore` adapter that reads `GORILLA1`-format chunks out of +/// an S3-compatible bucket. See module docs for layout + S3 client +/// notes. +pub struct GorillaS3ColdStore { + object_store: Arc, + config: GorillaS3Config, + cache: ChunkCache, +} + +impl GorillaS3ColdStore { + /// Build with an explicit object-store backend. The production + /// constructor [`Self::with_default_backend`] wires up + /// `S3ObjectStore` from `cfg`; tests inject the in-memory mock. + pub fn new(object_store: Arc, config: GorillaS3Config) -> Self { + let cap = NonZeroUsize::new(config.cache_capacity.max(1)) + .unwrap_or(NonZeroUsize::new(1).unwrap()); + Self { + object_store, + config, + cache: Mutex::new(LruCache::new(cap)), + } + } + + /// Build from a [`GorillaS3Config`] using the default + /// `rust-s3`-backed [`ObjectStore`]. + pub fn with_default_backend(config: GorillaS3Config) -> Result { + let backend = Arc::new(S3ObjectStore::new(&config)?); + Ok(Self::new(backend, config)) + } + + /// Borrow the active config — useful for diagnostics. + pub fn config(&self) -> &GorillaS3Config { + &self.config + } + + /// Render the configured `prefix_template` for one + /// `(metric, hour)` bucket and append `index.json`. + fn index_key(&self, metric: &str, ts_ms: i64) -> String { + let dt: DateTime = DateTime::::from_timestamp_millis(ts_ms) + .unwrap_or_else(|| DateTime::::from_timestamp(0, 0).unwrap()); + let prefix = self + .config + .prefix_template + .replace("{tenant}", &self.config.tenant) + .replace("{metric}", metric) + .replace("{year}", &format!("{:04}", dt.year())) + .replace("{month}", &format!("{:02}", dt.month())) + .replace("{day}", &format!("{:02}", dt.day())) + .replace("{hour}", &format!("{:02}", dt.hour())); + let mut key = prefix; + if !key.ends_with('/') { + key.push('/'); + } + key.push_str("index.json"); + key + } + + /// Iterate the wall-clock-hour starts (in ms) covered by + /// `[start_ms, end_ms)`. Always emits at least one bucket. + fn hour_starts(start_ms: i64, end_ms: i64) -> Vec { + const HOUR_MS: i64 = 3_600_000; + if end_ms <= start_ms { + let h = (start_ms / HOUR_MS) * HOUR_MS; + return vec![h]; + } + let first = (start_ms / HOUR_MS) * HOUR_MS; + let last = ((end_ms - 1) / HOUR_MS) * HOUR_MS; + let mut out = Vec::new(); + let mut cur = first; + while cur <= last { + out.push(cur); + cur += HOUR_MS; + } + out + } + + /// Fetch + parse one hour's `index.json`. Missing index = empty + /// catalog (the producer may not have flushed yet); transport + /// failure surfaces as `ColdStoreError::Backend`. + async fn fetch_index(&self, metric: &str, hour_ms: i64) -> Result { + let key = self.index_key(metric, hour_ms); + match self.object_store.get_object(&key).await { + Ok(bytes) => IndexFile::read(bytes.as_slice()).map_err(|e| { + ColdStoreError::Malformed(format!("index.json at {key}: {e}")) + }), + Err(e) if self.object_store.object_missing(&e) => { + debug!(key = %key, "gorilla-s3: index.json missing for hour bucket; skipping"); + Ok(IndexFile::new(0)) + } + Err(e) => Err(e), + } + } +} + +#[async_trait] +impl ColdStore for GorillaS3ColdStore { + async fn scan( + &self, + metric: &str, + start_ms: i64, + end_ms: i64, + ) -> Result, ColdStoreError> { + let chunks = self.list_chunks(metric, start_ms, end_ms).await?; + let mut out = Vec::new(); + for chunk in chunks { + let samples = self.read_chunk(&chunk).await?; + for s in samples { + if s.ts_ms >= start_ms && s.ts_ms < end_ms { + out.push(s); + } + } + } + Ok(out) + } + + async fn list_chunks( + &self, + metric: &str, + start_ms: i64, + end_ms: i64, + ) -> Result, ColdStoreError> { + // Convert the request window to the nanosecond unit the + // index file uses (`IndexEntry.time_range` is `(ns, ns)`, + // mirroring the Go encoder's `time.Time.UnixNano()` source). + let start_ns = (start_ms as i128).saturating_mul(1_000_000) as u64; + // `end_ms` is exclusive on the ms side; the index iter + // overlap test is inclusive so subtract 1 ns to keep the + // semantics aligned. If `end_ms == start_ms` we still want + // to scan the bucket containing `start_ms`. + let end_ns = if end_ms <= start_ms { + start_ns + } else { + ((end_ms as i128).saturating_mul(1_000_000) - 1).max(0) as u64 + }; + + let mut out = Vec::new(); + for hour_ms in Self::hour_starts(start_ms, end_ms) { + let idx = self.fetch_index(metric, hour_ms).await?; + for entry in idx.prune_by_time((start_ns, end_ns)) { + let (entry_start_ms, entry_end_ms) = ( + (entry.time_range.0 / 1_000_000) as i64, + (entry.time_range.1 / 1_000_000) as i64, + ); + out.push(ChunkRef { + key: entry.key.clone(), + metric: metric.to_string(), + time_range_ms: (entry_start_ms, entry_end_ms), + label_hash: entry.label_hash, + sample_count: entry.sample_count, + size_bytes: entry.size_bytes, + }); + } + } + Ok(out) + } + + async fn read_chunk(&self, chunk: &ChunkRef) -> Result, ColdStoreError> { + // Cache hit fast path. + { + let mut guard = self.cache.lock().await; + if let Some(cached) = guard.get(&chunk.key).cloned() { + return Ok((*cached).clone()); + } + } + + let bytes = self.object_store.get_object(&chunk.key).await?; + let samples = decode_block(&bytes) + .map_err(|e| ColdStoreError::Malformed(format!("decode {}: {e}", chunk.key)))?; + + let arc = Arc::new(samples.clone()); + { + let mut guard = self.cache.lock().await; + guard.put(chunk.key.clone(), arc); + } + Ok(samples) + } +} + +// ───────────────────────────────────────────────────────────────────── +// Decoding helper — converts a GORILLA1 block to RawSample units +// ───────────────────────────────────────────────────────────────────── + +/// Decode a single in-memory `GORILLA1` block into [`RawSample`]s. +/// +/// Walks every series in the block; multi-series blocks are +/// flattened into one `Vec`. Timestamps are converted from the +/// on-wire nanoseconds (Go `time.Time.UnixNano()` source) to the +/// [`RawSample::ts_ms`] millisecond unit. +fn decode_block(bytes: &[u8]) -> Result, asap_gorilla::DecodeError> { + let mut decoder = GorillaDecoder::from_reader(bytes)?; + let mut out: Vec = Vec::new(); + while let Some(header) = decoder.header().cloned() { + let labels: std::collections::BTreeMap = + header.labels.iter().cloned().collect(); + for sample in decoder.samples() { + let (ts_ns, value) = sample?; + out.push(RawSample { + ts_ms: (ts_ns / 1_000_000) as i64, + labels: labels.clone(), + value, + }); + } + if !decoder.next_series()? { + break; + } + } + Ok(out) +} + +// ───────────────────────────────────────────────────────────────────── +// In-memory ObjectStore mock — pub(crate) so tests in sibling files +// can exercise the same fixture without a live MinIO. +// ───────────────────────────────────────────────────────────────────── + +/// In-memory [`ObjectStore`] used by `gorilla_s3` tests. +/// +/// Holds a `HashMap>` plus a per-key fetch counter so +/// cache-hit assertions are first-class. Optionally fails every +/// `get_object` call for the network-error test. +#[cfg(test)] +#[derive(Default)] +pub(crate) struct InMemoryObjectStore { + inner: Mutex, +} + +#[cfg(test)] +#[derive(Default)] +struct InMemoryState { + objects: HashMap>, + fetch_counts: HashMap, + fail_all: Option, +} + +#[cfg(test)] +impl InMemoryObjectStore { + pub(crate) fn new() -> Self { + Self::default() + } + + pub(crate) async fn put(&self, key: impl Into, body: Vec) { + let mut g = self.inner.lock().await; + g.objects.insert(key.into(), body); + } + + pub(crate) async fn get_count(&self, key: &str) -> usize { + let g = self.inner.lock().await; + g.fetch_counts.get(key).copied().unwrap_or(0) + } + + /// Make every subsequent `get_object` fail with a backend error + /// containing `msg`. Used by the network-error test. + pub(crate) async fn fail_all(&self, msg: impl Into) { + let mut g = self.inner.lock().await; + g.fail_all = Some(msg.into()); + } +} + +#[cfg(test)] +#[async_trait] +impl ObjectStore for InMemoryObjectStore { + async fn get_object(&self, key: &str) -> Result, ColdStoreError> { + let mut g = self.inner.lock().await; + if let Some(msg) = g.fail_all.clone() { + return Err(ColdStoreError::Backend(msg)); + } + *g.fetch_counts.entry(key.to_string()).or_insert(0) += 1; + match g.objects.get(key) { + Some(b) => Ok(b.clone()), + None => Err(ColdStoreError::Backend(format!("get {key}: not found"))), + } + } +} + +// Static `Send` assertion — `GorillaS3ColdStore` must be storable +// behind an `Arc` in the existing s3_adapter chain. +const _: fn() = || { + fn _assert_send() {} + _assert_send::(); +}; + +// ───────────────────────────────────────────────────────────────────── +// Tests +// ───────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + + use asap_gorilla::{GorillaEncoder, IndexEntry, IndexFile}; + use chrono::TimeZone; + + /// Build a minimal index.json fixture. + fn make_index(entries: Vec) -> Vec { + let mut idx = IndexFile::new(0); + idx.entries = entries; + let mut buf = Vec::new(); + idx.write(&mut buf).unwrap(); + buf + } + + /// Encode a single-series Gorilla block from `(ts_ms, value)` pairs. + fn make_block(metric: &str, labels: &[(&str, &str)], samples: &[(i64, f64)]) -> Vec { + let mut enc = GorillaEncoder::new( + metric.to_string(), + labels + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect(), + ); + for (ts_ms, v) in samples { + // ts_ms → ts_ns + enc.append((*ts_ms as u64) * 1_000_000, *v); + } + enc.finalize().unwrap() + } + + fn ms(year: i32, month: u32, day: u32, hour: u32, min: u32, sec: u32) -> i64 { + Utc.with_ymd_and_hms(year, month, day, hour, min, sec) + .unwrap() + .timestamp_millis() + } + + fn cfg() -> GorillaS3Config { + GorillaS3Config { + endpoint: Some("http://mock".to_string()), + bucket: "test-bucket".to_string(), + tenant: "tenant1".to_string(), + prefix_template: "{tenant}/{metric}/{year}/{month}/{day}/{hour}/".to_string(), + region: "us-east-1".to_string(), + access_key_id: None, + secret_access_key: None, + cache_capacity: 4, + use_ssl: false, + } + } + + /// Layout: hour bucket H, three chunks A/B/C in time order, the + /// requested window only overlaps B → list returns B alone. + #[tokio::test] + async fn list_chunks_via_indexfile_prunes_by_time() { + let store = InMemoryObjectStore::new(); + + let h0 = ms(2026, 5, 6, 12, 0, 0); + let metric = "node_cpu_seconds_total"; + + let key_a = "tenant1/node_cpu_seconds_total/2026/05/06/12/part-A.gor".to_string(); + let key_b = "tenant1/node_cpu_seconds_total/2026/05/06/12/part-B.gor".to_string(); + let key_c = "tenant1/node_cpu_seconds_total/2026/05/06/12/part-C.gor".to_string(); + + let entries = vec![ + IndexEntry { + key: key_a.clone(), + time_range: ((h0) as u64 * 1_000_000, (h0 + 999) as u64 * 1_000_000), + sample_count: 10, + label_hash: 0xAAAA, + size_bytes: 100, + }, + IndexEntry { + key: key_b.clone(), + time_range: ((h0 + 5_000) as u64 * 1_000_000, (h0 + 6_000) as u64 * 1_000_000), + sample_count: 11, + label_hash: 0xBBBB, + size_bytes: 110, + }, + IndexEntry { + key: key_c.clone(), + time_range: ((h0 + 10_000) as u64 * 1_000_000, (h0 + 11_000) as u64 * 1_000_000), + sample_count: 12, + label_hash: 0xCCCC, + size_bytes: 120, + }, + ]; + + store + .put( + "tenant1/node_cpu_seconds_total/2026/05/06/12/index.json", + make_index(entries), + ) + .await; + + let cs = GorillaS3ColdStore::new(Arc::new(store), cfg()); + let chunks = cs + .list_chunks(metric, h0 + 5_500, h0 + 5_800) + .await + .unwrap(); + + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].key, key_b); + assert_eq!(chunks[0].sample_count, 11); + assert_eq!(chunks[0].label_hash, 0xBBBB); + assert_eq!(chunks[0].metric, metric); + } + + #[tokio::test] + async fn read_chunk_decodes_via_asap_gorilla() { + let store = InMemoryObjectStore::new(); + + let h0 = ms(2026, 5, 6, 12, 0, 0); + let metric = "node_cpu_seconds_total"; + let labels = &[("instance", "i-1"), ("mode", "user")]; + + let block = make_block( + metric, + labels, + &[(h0 + 1_000, 0.5), (h0 + 2_000, 0.7), (h0 + 3_000, 0.7)], + ); + let chunk_key = "tenant1/node_cpu_seconds_total/2026/05/06/12/part-001.gor".to_string(); + + store.put(chunk_key.clone(), block.clone()).await; + store + .put( + "tenant1/node_cpu_seconds_total/2026/05/06/12/index.json", + make_index(vec![IndexEntry { + key: chunk_key.clone(), + time_range: ( + (h0 + 1_000) as u64 * 1_000_000, + (h0 + 3_000) as u64 * 1_000_000, + ), + sample_count: 3, + label_hash: 0x1234, + size_bytes: block.len() as u32, + }]), + ) + .await; + + let cs = GorillaS3ColdStore::new(Arc::new(store), cfg()); + let chunks = cs.list_chunks(metric, h0, h0 + 60_000).await.unwrap(); + assert_eq!(chunks.len(), 1); + + let samples = cs.read_chunk(&chunks[0]).await.unwrap(); + assert_eq!(samples.len(), 3); + assert_eq!(samples[0].ts_ms, h0 + 1_000); + assert_eq!(samples[0].value, 0.5); + assert_eq!(samples[1].ts_ms, h0 + 2_000); + assert_eq!(samples[1].value, 0.7); + assert_eq!(samples[2].ts_ms, h0 + 3_000); + assert_eq!(samples[2].value, 0.7); + assert_eq!(samples[0].labels.get("instance").map(String::as_str), Some("i-1")); + assert_eq!(samples[0].labels.get("mode").map(String::as_str), Some("user")); + } + + #[tokio::test] + async fn cache_hit_skips_s3_fetch() { + let store = Arc::new(InMemoryObjectStore::new()); + + let h0 = ms(2026, 5, 6, 12, 0, 0); + let metric = "m"; + let block = make_block(metric, &[], &[(h0 + 1_000, 1.0), (h0 + 2_000, 2.0)]); + let chunk_key = "tenant1/m/2026/05/06/12/part-X.gor".to_string(); + store.put(chunk_key.clone(), block.clone()).await; + store + .put( + "tenant1/m/2026/05/06/12/index.json", + make_index(vec![IndexEntry { + key: chunk_key.clone(), + time_range: ( + (h0 + 1_000) as u64 * 1_000_000, + (h0 + 2_000) as u64 * 1_000_000, + ), + sample_count: 2, + label_hash: 0, + size_bytes: block.len() as u32, + }]), + ) + .await; + + let cs = GorillaS3ColdStore::new(store.clone(), cfg()); + let chunks = cs.list_chunks(metric, h0, h0 + 60_000).await.unwrap(); + assert_eq!(chunks.len(), 1); + + let _ = cs.read_chunk(&chunks[0]).await.unwrap(); + let count_after_first = store.get_count(&chunk_key).await; + let _ = cs.read_chunk(&chunks[0]).await.unwrap(); + let count_after_second = store.get_count(&chunk_key).await; + + assert_eq!(count_after_first, 1); + assert_eq!( + count_after_second, 1, + "second read_chunk must hit cache and skip S3 GET" + ); + } + + #[tokio::test] + async fn lru_eviction_under_pressure() { + // cache_capacity=2, fill with three chunks then re-read the + // first → that triggers an S3 GET because the LRU evicted it. + let store = Arc::new(InMemoryObjectStore::new()); + let h0 = ms(2026, 5, 6, 12, 0, 0); + + let mut chunk_refs: Vec = Vec::new(); + let mut entries: Vec = Vec::new(); + for i in 0..3i64 { + let block = make_block( + "m", + &[("i", &i.to_string())], + &[(h0 + i * 1_000, i as f64), (h0 + i * 1_000 + 100, i as f64 + 0.5)], + ); + let key = format!("tenant1/m/2026/05/06/12/part-{i}.gor"); + store.put(key.clone(), block.clone()).await; + entries.push(IndexEntry { + key: key.clone(), + time_range: ( + ((h0 + i * 1_000) as u64) * 1_000_000, + ((h0 + i * 1_000 + 100) as u64) * 1_000_000, + ), + sample_count: 2, + label_hash: i as u64, + size_bytes: block.len() as u32, + }); + chunk_refs.push(ChunkRef { + key, + metric: "m".to_string(), + time_range_ms: (h0 + i * 1_000, h0 + i * 1_000 + 100), + label_hash: i as u64, + sample_count: 2, + size_bytes: block.len() as u32, + }); + } + store + .put("tenant1/m/2026/05/06/12/index.json", make_index(entries)) + .await; + + let mut config = cfg(); + config.cache_capacity = 2; + let cs = GorillaS3ColdStore::new(store.clone(), config); + + cs.read_chunk(&chunk_refs[0]).await.unwrap(); + cs.read_chunk(&chunk_refs[1]).await.unwrap(); + cs.read_chunk(&chunk_refs[2]).await.unwrap(); // evicts chunk_refs[0] + + let before = store.get_count(&chunk_refs[0].key).await; + cs.read_chunk(&chunk_refs[0]).await.unwrap(); + let after = store.get_count(&chunk_refs[0].key).await; + assert_eq!( + after, + before + 1, + "evicted chunk must trigger a fresh S3 GET" + ); + } + + #[tokio::test] + async fn index_json_corrupted_returns_error() { + let store = InMemoryObjectStore::new(); + let h0 = ms(2026, 5, 6, 12, 0, 0); + store + .put( + "tenant1/m/2026/05/06/12/index.json", + b"this is not json {{{".to_vec(), + ) + .await; + + let cs = GorillaS3ColdStore::new(Arc::new(store), cfg()); + let res = cs.list_chunks("m", h0, h0 + 60_000).await; + match res { + Err(ColdStoreError::Malformed(msg)) => { + assert!(msg.contains("index.json"), "msg should name the key: {msg}") + } + other => panic!("expected Malformed, got {other:?}"), + } + } + + #[tokio::test] + async fn s3_unavailable_returns_error() { + let store = Arc::new(InMemoryObjectStore::new()); + store.fail_all("simulated network outage").await; + let cs = GorillaS3ColdStore::new(store, cfg()); + let h0 = ms(2026, 5, 6, 12, 0, 0); + let res = cs.list_chunks("m", h0, h0 + 60_000).await; + match res { + Err(ColdStoreError::Backend(msg)) => assert!(msg.contains("simulated network outage")), + other => panic!("expected Backend, got {other:?}"), + } + } + + #[tokio::test] + async fn missing_index_is_empty_not_error() { + let store = InMemoryObjectStore::new(); + let cs = GorillaS3ColdStore::new(Arc::new(store), cfg()); + let h0 = ms(2026, 5, 6, 12, 0, 0); + let chunks = cs.list_chunks("never_written", h0, h0 + 60_000).await.unwrap(); + assert!(chunks.is_empty()); + let samples = cs.scan("never_written", h0, h0 + 60_000).await.unwrap(); + assert!(samples.is_empty()); + } + + #[tokio::test] + async fn scan_filters_to_requested_range() { + // Chunk has samples at h0+1_000 and h0+10_000; request only + // [h0+5_000, h0+9_000) — chunk overlaps the request, but the + // matching sample is *outside* the inner filter, so scan + // returns 0 samples (read_chunk would still load + cache the + // chunk). + let store = Arc::new(InMemoryObjectStore::new()); + let h0 = ms(2026, 5, 6, 12, 0, 0); + let block = make_block( + "m", + &[], + &[(h0 + 1_000, 1.0), (h0 + 10_000, 2.0)], + ); + let key = "tenant1/m/2026/05/06/12/part-Z.gor".to_string(); + store.put(key.clone(), block.clone()).await; + store + .put( + "tenant1/m/2026/05/06/12/index.json", + make_index(vec![IndexEntry { + key: key.clone(), + time_range: ( + (h0 + 1_000) as u64 * 1_000_000, + (h0 + 10_000) as u64 * 1_000_000, + ), + sample_count: 2, + label_hash: 0, + size_bytes: block.len() as u32, + }]), + ) + .await; + + let cs = GorillaS3ColdStore::new(store, cfg()); + let samples = cs.scan("m", h0 + 5_000, h0 + 9_000).await.unwrap(); + assert!(samples.is_empty(), "no sample inside [5_000, 9_000) ms"); + + let samples = cs.scan("m", h0, h0 + 60_000).await.unwrap(); + assert_eq!(samples.len(), 2); + } + + #[tokio::test] + async fn list_chunks_spans_two_hour_buckets() { + let store = InMemoryObjectStore::new(); + let h12 = ms(2026, 5, 6, 12, 0, 0); + let h13 = ms(2026, 5, 6, 13, 0, 0); + let key12 = "tenant1/m/2026/05/06/12/part-1.gor".to_string(); + let key13 = "tenant1/m/2026/05/06/13/part-1.gor".to_string(); + store + .put( + "tenant1/m/2026/05/06/12/index.json", + make_index(vec![IndexEntry { + key: key12.clone(), + time_range: ( + (h12 + 3_500_000) as u64 * 1_000_000, + (h12 + 3_590_000) as u64 * 1_000_000, + ), + sample_count: 1, + label_hash: 0, + size_bytes: 50, + }]), + ) + .await; + store + .put( + "tenant1/m/2026/05/06/13/index.json", + make_index(vec![IndexEntry { + key: key13.clone(), + time_range: ( + (h13 + 1_000) as u64 * 1_000_000, + (h13 + 30_000) as u64 * 1_000_000, + ), + sample_count: 1, + label_hash: 0, + size_bytes: 50, + }]), + ) + .await; + let cs = GorillaS3ColdStore::new(Arc::new(store), cfg()); + let chunks = cs + .list_chunks("m", h12 + 3_500_000, h13 + 30_000) + .await + .unwrap(); + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].key, key12); + assert_eq!(chunks[1].key, key13); + } + + #[test] + fn from_env_requires_bucket() { + // Don't pollute global env in a unit test; just exercise the + // missing-var path. + let prev_bucket = std::env::var("ASAP_GORILLA_S3_BUCKET").ok(); + std::env::remove_var("ASAP_GORILLA_S3_BUCKET"); + let res = GorillaS3Config::from_env(); + if let Some(v) = prev_bucket { + std::env::set_var("ASAP_GORILLA_S3_BUCKET", v); + } + match res { + Err(GorillaS3ConfigError::MissingEnv("ASAP_GORILLA_S3_BUCKET")) => {} + other => panic!("expected MissingEnv(BUCKET), got {other:?}"), + } + } +} diff --git a/asap-query-engine/src/drivers/query/fallback/cold_store/mod.rs b/asap-query-engine/src/drivers/query/fallback/cold_store/mod.rs index c9a8a679..77bd8e6c 100644 --- a/asap-query-engine/src/drivers/query/fallback/cold_store/mod.rs +++ b/asap-query-engine/src/drivers/query/fallback/cold_store/mod.rs @@ -34,9 +34,11 @@ use std::collections::BTreeMap; use thiserror::Error; pub mod format; +pub mod gorilla_s3; pub mod local_fs; pub use format::{part_path_prefix, RawSample}; +pub use gorilla_s3::{GorillaS3ColdStore, GorillaS3Config, GorillaS3ConfigError}; pub use local_fs::LocalFsColdStore; /// Error surface for cold-store scans. @@ -46,6 +48,51 @@ pub enum ColdStoreError { Io(#[from] std::io::Error), #[error("malformed raw record: {0}")] Malformed(String), + /// Backend-storage error (e.g. an S3 GET failed) that is not + /// itself a `std::io::Error`. Phase 3 introduced this variant for + /// the Gorilla-S3 cold store; the local-FS path keeps using + /// [`ColdStoreError::Io`]. + #[error("backend error: {0}")] + Backend(String), + /// A trait method that this `ColdStore` impl does not support. + /// Returned by the default `list_chunks` / `read_chunk` impls on + /// JSONL-only stores; Gorilla-S3 / future chunk-native stores + /// override. + #[error("unsupported cold-store operation: {0}")] + Unsupported(&'static str), +} + +/// Descriptor for a single immutable cold-store chunk. +/// +/// Returned by [`ColdStore::list_chunks`] for chunk-native backends +/// (Phase 3+ Gorilla-S3). Carries enough metadata for callers to +/// prune by time / label without reading the chunk body. +#[derive(Debug, Clone, PartialEq)] +pub struct ChunkRef { + /// Opaque object key (e.g. an S3 key). The Telegraf-side + /// `gorilla_s3` output uses + /// `/block---.gorilla`; the + /// design.md-style layout is `//YYYY/MM/DD/HH/ + /// part-NNNNNN.gor`. Either is fine — the index file is the + /// source of truth for what keys exist. + pub key: String, + /// Metric name the chunk was fetched against. Recovered from + /// the caller's `list_chunks` request rather than the on-wire + /// chunk metadata, since not all backends require chunks to be + /// metric-pure. + pub metric: String, + /// `(start_unix_ms, end_unix_ms)` covered by the chunk — + /// converted from the on-wire nanosecond range so it can be + /// directly compared with [`ColdStore::scan`]'s + /// `[start_ms, end_ms)` window. + pub time_range_ms: (i64, i64), + /// 64-bit canonical-label-set hash — for prune-by-label-equality + /// without fetching the chunk. + pub label_hash: u64, + /// Number of samples in the chunk. + pub sample_count: u32, + /// On-wire size of the chunk object in bytes. + pub size_bytes: u32, } /// Read-only view over a cold raw-sample store. @@ -62,6 +109,17 @@ pub enum ColdStoreError { /// Label matching is **not** pushed down here — callers filter /// samples client-side. This keeps the trait small and makes the /// local-FS / S3 impls trivially swappable. +/// +/// # Phase-3 trait extension +/// +/// The `list_chunks` / `read_chunk` pair is additive (default impls +/// return [`ColdStoreError::Unsupported`]) so the existing JSONL +/// `LocalFsColdStore` keeps compiling unchanged. Chunk-native +/// backends (Gorilla-S3) override both so the upcoming +/// `GorillaQueryEngine` can iterate chunks one at a time without +/// materialising every sample up front. See +/// [`docs/design-gorilla-s3-cold-engine.md` §7.2](#) for the +/// rationale. #[async_trait] pub trait ColdStore: Send + Sync { /// Return all samples for `metric` whose timestamp lies in @@ -72,6 +130,33 @@ pub trait ColdStore: Send + Sync { start_ms: i64, end_ms: i64, ) -> Result, ColdStoreError>; + + /// List chunk descriptors covering `[start_ms, end_ms)` without + /// decoding any bodies. Default impl returns + /// [`ColdStoreError::Unsupported`] — only chunk-native backends + /// (e.g. [`GorillaS3ColdStore`]) override. + async fn list_chunks( + &self, + _metric: &str, + _start_ms: i64, + _end_ms: i64, + ) -> Result, ColdStoreError> { + Err(ColdStoreError::Unsupported("list_chunks")) + } + + /// Decode a single chunk into an owned `Vec`. + /// + /// Returning `Vec` rather than a streaming iterator keeps the + /// trait object-safe and matches the existing `scan` contract; + /// chunks are bounded-size in practice (Phase 1 emits one series + /// per ~1 hour). The decoded samples can also be cached cheaply + /// by the impl. Default returns [`ColdStoreError::Unsupported`]. + async fn read_chunk( + &self, + _chunk: &ChunkRef, + ) -> Result, ColdStoreError> { + Err(ColdStoreError::Unsupported("read_chunk")) + } } /// Convenience alias: a label set as stored in a [`RawSample`]. From 36ca4af44c82408746cb442e4572f6ca45edd18a Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Wed, 6 May 2026 01:07:21 -0400 Subject: [PATCH 2/3] =?UTF-8?q?feat(engine):=20GorillaQueryEngine=20?= =?UTF-8?q?=E2=80=94=20exact=20PromQL=20over=20Gorilla-S3=20chunks=20(Phas?= =?UTF-8?q?e=204)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 4 of the Gorilla-S3-cold-engine — adds a sibling engine to `SimpleEngine` that executes PromQL exactly against the Phase-3 `GorillaS3ColdStore`. Returns answers carrying an `AccuracyEnvelope { kind: Exact, ε: 0, δ: 0 }` and a `data_source: gorilla_archive` info marker so callers / dashboards can distinguish cold-tier exact answers from warm-tier sketch answers. Architecture ------------ `asap-query-engine/src/engines/gorilla_engine/` (new): * `mod.rs` — `GorillaQueryEngine` (holds `Arc` so tests can inject mocks; production constructor `with_gorilla_s3` keeps the design.md type signature). Public surface: `execute(query)` + `execute_at(query, now_ms)`. Wraps results via `wrap_result` (exact accuracy envelope + window). `GorillaEngineConfig { max_buffered_samples, query_timeout_secs }` with sensible defaults (10M / 30s). `EngineError` enum covers Plan / ColdStore / TooManySamples / Timeout. `ExecutionOutcome` carries `(value, samples_scanned, chunks_fetched)` + an `info_lines` builder pinning the on-wire info strings. * `query_planner.rs` — minimal PromQL → `QueryPlan` translator (metric, half-open `[start_ms, end_ms)`, `QueryStatistic`). Supports `sum/count/avg/min/max_over_time`, `rate`, `increase`, `quantile_over_time(φ, m[range])`, and `topk(k, )` (PromQL grammar requires the inner be a vector so `topk(k, sum_over_time(m[range]))` is the legal spelling). Tests + caller can use `plan_query_at` to pin `now_ms`. * `exact_executor.rs` — `ExactExecutor` dispatches per `QueryStatistic`: - **Streaming additive** (`Sum/Count/Avg/Min/Max/Rate/Increase`): `list_chunks` → for each chunk `read_chunk` → fold into a bounded `AdditiveAccumulator` → drop the decoded chunk before fetching the next. Memory is O(1) per query. Rate / Increase track first/last `(ts, value)` and divide by `range_seconds` at finalisation. - **Buffered** (`Quantile / TopK`): `collect_buffered_samples` materialises every in-range sample up to `max_buffered_samples` (errors with `TooManySamples` otherwise). Quantile sorts + nearest-rank index; TopK sorts descending and returns sum of the top-k values. * `tests.rs` — 20 unit tests via an in-process `MockColdStore` satisfying `ColdStore` (no S3 / disk dependency). Pinned NOW via `execute_at` so time math is deterministic. Covers every `QueryStatistic` happy path, the empty-data sentinel, the outside-range filter, the buffered-budget guard, the result wrapping (accuracy envelope + `data_source: gorilla_archive` marker), and the timeout path. Test additions (15 spec'd + 5 planner) -------------------------------------- * `execute_sum_over_time_streaming` * `execute_count_over_time` * `execute_avg_over_time` * `execute_min_over_time` / `execute_max_over_time` * `execute_rate_basic` * `execute_increase_basic` * `execute_quantile_buffered_basic` * `execute_quantile_too_many_samples_errors` * `execute_topk_basic` * `execute_empty_chunks_returns_zero_or_nan` * `execute_chunks_partially_outside_range_filtered` * `result_carries_exact_accuracy_envelope` * `result_includes_data_source_gorilla_archive` * `engine_respects_config_timeout` * `query_planner::tests::{plans_sum_over_time, plans_quantile_over_time, plans_topk, rejects_binary_expression, streaming_classification}` Touches only the new directory + `engines/mod.rs` (re-exports). Zero changes to `simple_engine.rs`, the cold store, or the sketch warm-tier. No new Cargo deps — `asap-gorilla` / `tokio` were already pulled in by PR #84. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../engines/gorilla_engine/exact_executor.rs | 316 +++++++++ .../src/engines/gorilla_engine/mod.rs | 266 ++++++++ .../engines/gorilla_engine/query_planner.rs | 307 +++++++++ .../src/engines/gorilla_engine/tests.rs | 615 ++++++++++++++++++ asap-query-engine/src/engines/mod.rs | 4 + 5 files changed, 1508 insertions(+) create mode 100644 asap-query-engine/src/engines/gorilla_engine/exact_executor.rs create mode 100644 asap-query-engine/src/engines/gorilla_engine/mod.rs create mode 100644 asap-query-engine/src/engines/gorilla_engine/query_planner.rs create mode 100644 asap-query-engine/src/engines/gorilla_engine/tests.rs diff --git a/asap-query-engine/src/engines/gorilla_engine/exact_executor.rs b/asap-query-engine/src/engines/gorilla_engine/exact_executor.rs new file mode 100644 index 00000000..6398b0dd --- /dev/null +++ b/asap-query-engine/src/engines/gorilla_engine/exact_executor.rs @@ -0,0 +1,316 @@ +//! Per-statistic executors for the Phase-4 Gorilla engine. +//! +//! Two strategies, picked by [`super::query_planner::QueryStatistic::is_streaming_additive`]: +//! +//! * **Streaming-additive** — `Sum / Count / Avg / Min / Max / +//! Rate / Increase`. Walk chunks one at a time, fold each +//! sample into a tiny accumulator, drop the decoded chunk +//! before fetching the next one. Memory is O(1) per query. +//! * **Buffered** — `Quantile / TopK / Cardinality`. Materialise +//! every in-range sample, then sort or otherwise post-process. +//! Bounded by [`super::GorillaEngineConfig::max_buffered_samples`]; +//! over-budget queries fail fast with +//! [`super::EngineError::TooManySamples`] rather than OOM. + +use std::sync::Arc; + +use tracing::debug; + +use crate::drivers::query::fallback::cold_store::{ColdStore, RawSample}; + +use super::query_planner::{QueryPlan, QueryStatistic}; +use super::{EngineError, ExecutionOutcome, GorillaEngineConfig}; + +/// Streaming-additive operation tag — what the per-sample fold +/// does. Pulled out so [`ExactExecutor::execute_streaming_additive`] +/// is a single function regardless of which stat is being computed. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdditiveOp { + Sum, + Count, + /// `(sum, count)` — the engine divides at the end. + Avg, + Min, + Max, + /// `last - first` over the time-ordered samples. + Increase, + /// `(last - first) / range_seconds`. + Rate, +} + +/// Per-statistic executor. Holds an `Arc` so the +/// engine + executor share the same cold-tier handle without +/// re-implementing trait dispatch. +pub struct ExactExecutor { + cold_store: Arc, + config: GorillaEngineConfig, +} + +impl ExactExecutor { + pub fn new(cold_store: Arc, config: GorillaEngineConfig) -> Self { + Self { cold_store, config } + } + + /// Top-level dispatch — picks streaming vs buffered based on + /// the plan's statistic. + pub async fn execute_plan(&self, plan: &QueryPlan) -> Result { + match &plan.statistic { + QueryStatistic::SumOverTime => { + self.execute_streaming_additive(plan, AdditiveOp::Sum).await + } + QueryStatistic::CountOverTime => { + self.execute_streaming_additive(plan, AdditiveOp::Count) + .await + } + QueryStatistic::AvgOverTime => { + self.execute_streaming_additive(plan, AdditiveOp::Avg).await + } + QueryStatistic::MinOverTime => { + self.execute_streaming_additive(plan, AdditiveOp::Min).await + } + QueryStatistic::MaxOverTime => { + self.execute_streaming_additive(plan, AdditiveOp::Max).await + } + QueryStatistic::Rate => self.execute_streaming_additive(plan, AdditiveOp::Rate).await, + QueryStatistic::Increase => { + self.execute_streaming_additive(plan, AdditiveOp::Increase) + .await + } + QueryStatistic::QuantileOverTime { phi } => self.execute_quantile(plan, *phi).await, + QueryStatistic::TopK { k } => self.execute_topk(plan, *k).await, + } + } + + /// Streaming additive path. Reads chunks one at a time, applies + /// the per-sample fold, drops the decoded chunk before fetching + /// the next one. Memory is O(1) per query, regardless of how + /// many samples the time range covers. + pub async fn execute_streaming_additive( + &self, + plan: &QueryPlan, + op: AdditiveOp, + ) -> Result { + let (start_ms, end_ms) = plan.time_range_ms; + let chunks = self + .cold_store + .list_chunks(&plan.metric, start_ms, end_ms) + .await?; + let chunks_fetched = chunks.len(); + debug!( + metric = plan.metric.as_str(), + chunks = chunks_fetched, + op = ?op, + "gorilla-engine: streaming additive over chunks" + ); + + let mut acc = AdditiveAccumulator::new(op); + let mut samples_scanned: usize = 0; + for chunk in chunks { + let samples = self.cold_store.read_chunk(&chunk).await?; + for s in samples { + if s.ts_ms >= start_ms && s.ts_ms < end_ms { + acc.observe(&s); + samples_scanned += 1; + } + } + } + + let value = acc.finalize(plan, op); + Ok(ExecutionOutcome { + value, + samples_scanned, + chunks_fetched, + }) + } + + /// Buffered quantile path. Materialises every in-range sample + /// up to [`GorillaEngineConfig::max_buffered_samples`], sorts + /// the value column, and picks the φ-rank using a + /// nearest-rank rule (matches Prometheus's + /// `quantile_over_time` semantics for the linear-interp-free + /// midpoint case — the float index rounds to nearest). + pub async fn execute_quantile( + &self, + plan: &QueryPlan, + phi: f64, + ) -> Result { + let (samples, chunks_fetched) = self.collect_buffered_samples(plan).await?; + if samples.is_empty() { + return Ok(ExecutionOutcome { + value: f64::NAN, + samples_scanned: 0, + chunks_fetched, + }); + } + + let mut values: Vec = samples.iter().map(|s| s.value).collect(); + values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + + let n = values.len() as f64; + let phi = phi.clamp(0.0, 1.0); + let raw_idx = ((n - 1.0) * phi).round() as i64; + let idx = raw_idx.clamp(0, values.len() as i64 - 1) as usize; + Ok(ExecutionOutcome { + value: values[idx], + samples_scanned: values.len(), + chunks_fetched, + }) + } + + /// Buffered top-k path. Materialises every in-range sample, + /// sorts the value column descending, and returns the SUM of + /// the top-`k` values. The Phase-5 capability router will + /// extend this to per-group top-k once spatial grouping + /// lands; the MVP scalar return value matches the existing + /// `ExecutionOutcome` shape. + pub async fn execute_topk( + &self, + plan: &QueryPlan, + k: usize, + ) -> Result { + let (samples, chunks_fetched) = self.collect_buffered_samples(plan).await?; + if samples.is_empty() { + return Ok(ExecutionOutcome { + value: f64::NAN, + samples_scanned: 0, + chunks_fetched, + }); + } + + let mut values: Vec = samples.iter().map(|s| s.value).collect(); + values.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal)); + let take = k.min(values.len()); + let topk_sum: f64 = values.iter().take(take).sum(); + Ok(ExecutionOutcome { + value: topk_sum, + samples_scanned: values.len(), + chunks_fetched, + }) + } + + /// Shared helper for the buffered paths: walk every chunk, + /// keep in-range samples, enforce the + /// [`GorillaEngineConfig::max_buffered_samples`] ceiling. + async fn collect_buffered_samples( + &self, + plan: &QueryPlan, + ) -> Result<(Vec, usize), EngineError> { + let (start_ms, end_ms) = plan.time_range_ms; + let chunks = self + .cold_store + .list_chunks(&plan.metric, start_ms, end_ms) + .await?; + let chunks_fetched = chunks.len(); + let mut buffer: Vec = Vec::new(); + let limit = self.config.max_buffered_samples; + for chunk in chunks { + let samples = self.cold_store.read_chunk(&chunk).await?; + for s in samples { + if s.ts_ms >= start_ms && s.ts_ms < end_ms { + if buffer.len() >= limit { + // Surface the over-budget sample count + // (limit + 1) so callers can pin + // `count > limit` in tests; we don't + // bother walking the rest of the chunks + // just to tighten the number. + return Err(EngineError::TooManySamples { + count: buffer.len() + 1, + limit, + }); + } + buffer.push(s); + } + } + } + Ok((buffer, chunks_fetched)) + } +} + +/// Per-sample fold for the streaming-additive path. Fields are +/// kept in raw f64 (sum / min / max) + i64 (count) so the +/// finaliser can pick the right arithmetic per op. +/// +/// `op` is taken at `new` time rather than carried as a struct +/// field — the finaliser receives it as an argument, keeping the +/// struct itself op-agnostic and shrinking the per-fold footprint. +#[derive(Debug, Clone, Copy)] +struct AdditiveAccumulator { + sum: f64, + count: i64, + min: f64, + max: f64, + /// Earliest observed `(ts_ms, value)` — used by Rate / Increase. + first: Option<(i64, f64)>, + /// Latest observed `(ts_ms, value)` — used by Rate / Increase. + last: Option<(i64, f64)>, +} + +impl AdditiveAccumulator { + fn new(_op: AdditiveOp) -> Self { + Self { + sum: 0.0, + count: 0, + min: f64::INFINITY, + max: f64::NEG_INFINITY, + first: None, + last: None, + } + } + + fn observe(&mut self, s: &RawSample) { + self.sum += s.value; + self.count += 1; + if s.value < self.min { + self.min = s.value; + } + if s.value > self.max { + self.max = s.value; + } + match self.first { + None => self.first = Some((s.ts_ms, s.value)), + Some((ts, _)) if s.ts_ms < ts => self.first = Some((s.ts_ms, s.value)), + _ => {} + } + match self.last { + None => self.last = Some((s.ts_ms, s.value)), + Some((ts, _)) if s.ts_ms > ts => self.last = Some((s.ts_ms, s.value)), + _ => {} + } + } + + /// Convert the running accumulator into a final scalar. + /// Returns `NaN` for the empty-time-range case so downstream + /// formatting stays consistent. + fn finalize(&self, plan: &QueryPlan, op: AdditiveOp) -> f64 { + if self.count == 0 { + return match op { + AdditiveOp::Count => 0.0, + _ => f64::NAN, + }; + } + match op { + AdditiveOp::Sum => self.sum, + AdditiveOp::Count => self.count as f64, + AdditiveOp::Avg => self.sum / self.count as f64, + AdditiveOp::Min => self.min, + AdditiveOp::Max => self.max, + AdditiveOp::Increase => match (self.first, self.last) { + (Some((_, fv)), Some((_, lv))) => lv - fv, + _ => f64::NAN, + }, + AdditiveOp::Rate => match (self.first, self.last) { + (Some((_, fv)), Some((_, lv))) => { + let (start_ms, end_ms) = plan.time_range_ms; + let range_secs = ((end_ms - start_ms).max(1)) as f64 / 1000.0; + if range_secs <= 0.0 { + f64::NAN + } else { + (lv - fv) / range_secs + } + } + _ => f64::NAN, + }, + } + } +} + diff --git a/asap-query-engine/src/engines/gorilla_engine/mod.rs b/asap-query-engine/src/engines/gorilla_engine/mod.rs new file mode 100644 index 00000000..bc2c3ce2 --- /dev/null +++ b/asap-query-engine/src/engines/gorilla_engine/mod.rs @@ -0,0 +1,266 @@ +//! Phase 4: `GorillaQueryEngine` — exact PromQL execution over the +//! Gorilla-S3 cold tier. +//! +//! This engine is a SIBLING of [`crate::engines::simple_engine::SimpleEngine`]. +//! Both consume the same PromQL surface, but where `SimpleEngine` +//! answers from warm-tier sketches (approximate, ε/δ-bounded), the +//! `GorillaQueryEngine` answers exactly from per-hour Gorilla +//! chunks landed on S3 / MinIO via the Phase-3 +//! [`crate::drivers::query::fallback::cold_store::gorilla_s3::GorillaS3ColdStore`]. +//! +//! Result wrapping pins three things: +//! +//! 1. an [`crate::stores::sketch_db::AccuracyEnvelope`] with +//! `kind = Exact`, ε = 0, δ = 0, +//! 2. a `data_source: gorilla_archive` info line, +//! 3. cheap diagnostics (`samples_scanned`, `chunks_fetched`). +//! +//! See `docs/design-gorilla-s3-cold-engine.md` §6. +//! +//! ## Two execution strategies +//! +//! Per-statistic dispatch in [`exact_executor`]: +//! +//! * **Streaming-additive** — `Sum`, `Count`, `Min`, `Max`, `Rate`, +//! `Increase` (and `Avg` derived as Sum/Count). One chunk at a +//! time, fold into a small accumulator, drop the decoded samples +//! before fetching the next chunk. Memory cost is O(1) per group. +//! * **Buffered** — `Quantile`, `TopK`, `Cardinality`. Materialise +//! every in-range sample, then sort / count. Bounded by +//! [`GorillaEngineConfig::max_buffered_samples`]; over-budget +//! queries fail fast with [`EngineError::TooManySamples`]. + +pub mod exact_executor; +pub mod query_planner; + +#[cfg(test)] +mod tests; + +use std::sync::Arc; +use std::time::Duration; + +use thiserror::Error; +use tokio::time::error::Elapsed; +use tracing::debug; + +use crate::data_model::KeyByLabelValues; +use crate::drivers::query::fallback::cold_store::{ColdStore, ColdStoreError}; +use crate::engines::query_result::{InstantVectorElement, QueryResult}; +use crate::stores::sketch_db::accuracy::{AccuracyEnvelope, AccuracyProfile}; + +pub use exact_executor::{AdditiveOp, ExactExecutor}; +pub use query_planner::{plan_query, plan_query_at, QueryPlan, QueryStatistic}; + +/// Marker line that every `GorillaQueryEngine` answer carries on +/// its `infos` array. Pinned so dashboards / Phase-5 capability +/// routers can byte-compare without parsing. +pub const DATA_SOURCE_GORILLA_ARCHIVE: &str = "data_source: gorilla_archive"; + +/// Tunable runtime knobs for the Gorilla query engine. +/// +/// Call sites typically construct via `Default::default()`; tests +/// override `max_buffered_samples` to exercise the bounded-buffer +/// guard. +#[derive(Debug, Clone)] +pub struct GorillaEngineConfig { + /// Hard cap on the number of samples a buffered-aggregate + /// query (quantile / topk / cardinality) is allowed to + /// materialise in memory. Default `10_000_000` + /// (~160 MB at 16 B per `(ts, value)` pair). + pub max_buffered_samples: usize, + /// Wall-clock query timeout, in seconds. Default `30`. + pub query_timeout_secs: u64, +} + +impl Default for GorillaEngineConfig { + fn default() -> Self { + Self { + max_buffered_samples: 10_000_000, + query_timeout_secs: 30, + } + } +} + +/// Error surface returned by [`GorillaQueryEngine::execute`]. +#[derive(Debug, Error)] +pub enum EngineError { + /// PromQL string failed to parse, or used a construct outside + /// the engine's supported surface (see [`query_planner`]). + #[error("query planning failed: {0}")] + Plan(String), + /// Cold-store fetch / decode failed. + #[error("cold-store error: {0}")] + ColdStore(#[from] ColdStoreError), + /// Buffered-aggregate budget exceeded — query asked for more + /// samples than [`GorillaEngineConfig::max_buffered_samples`] + /// will allow. The user should narrow the time range or + /// lower the cardinality. + #[error( + "buffered-aggregate budget exceeded: {count} samples > limit {limit}; \ + narrow the time range or lower the metric cardinality" + )] + TooManySamples { + /// Samples the engine attempted to materialise. + count: usize, + /// Configured ceiling. + limit: usize, + }, + /// Wall-clock timeout fired before the query finished. + #[error("query timed out after {0:?}")] + Timeout(Duration), +} + +impl From for EngineError { + fn from(_: Elapsed) -> Self { + Self::Timeout(Duration::from_secs(0)) + } +} + +/// Phase-4 cold-tier exact engine. +/// +/// Holds an `Arc` rather than a concrete +/// `Arc` so tests can inject in-memory mocks +/// and so future cold backends (local-FS chunks, multi-region +/// fan-out) drop in without changing the engine surface. The +/// production constructor [`GorillaQueryEngine::with_gorilla_s3`] +/// keeps the design.md type signature working at the call site. +pub struct GorillaQueryEngine { + cold_store: Arc, + config: GorillaEngineConfig, +} + +impl GorillaQueryEngine { + /// Build with an arbitrary cold-store implementation. Used by + /// tests + the Phase-5 capability router (which may swap the + /// concrete impl based on routing decisions). + pub fn new(cold_store: Arc, config: GorillaEngineConfig) -> Self { + Self { + cold_store, + config, + } + } + + /// Convenience constructor for the production + /// [`crate::drivers::query::fallback::cold_store::gorilla_s3::GorillaS3ColdStore`] + /// path. Mirrors the design.md type signature. + pub fn with_gorilla_s3( + cold_store: Arc, + config: GorillaEngineConfig, + ) -> Self { + Self::new(cold_store as Arc, config) + } + + /// Read-only access to the configured limits — useful for + /// diagnostics + the Phase-5 router's cost estimator. + pub fn config(&self) -> &GorillaEngineConfig { + &self.config + } + + /// Execute a parsed PromQL query against the cold tier. + /// + /// The query string is parsed via [`query_planner::plan_query`], + /// the resulting plan dispatches to either the streaming + /// additive or the buffered execution path, and the answer is + /// wrapped with the exact-accuracy envelope + the + /// `data_source: gorilla_archive` annotation. + pub async fn execute(&self, query: &str) -> Result { + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::SystemTime::UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0); + self.execute_at(query, now_ms).await + } + + /// Like [`Self::execute`], with a caller-supplied `now_ms` + /// pinning the right edge of the request window. Used by + /// tests + by the (future) Phase-5 router that wants to back- + /// date a query against historical chunks. + pub async fn execute_at( + &self, + query: &str, + now_ms: i64, + ) -> Result { + let timeout = Duration::from_secs(self.config.query_timeout_secs.max(1)); + tokio::time::timeout(timeout, self.execute_inner(query, now_ms)) + .await + .map_err(|_| EngineError::Timeout(timeout))? + } + + async fn execute_inner( + &self, + query: &str, + now_ms: i64, + ) -> Result { + let plan = query_planner::plan_query_at(query, now_ms).map_err(EngineError::Plan)?; + debug!( + metric = plan.metric.as_str(), + stat = ?plan.statistic, + start_ms = plan.time_range_ms.0, + end_ms = plan.time_range_ms.1, + "gorilla-engine: executing plan" + ); + + let executor = ExactExecutor::new(self.cold_store.clone(), self.config.clone()); + let outcome = executor.execute_plan(&plan).await?; + + Ok(wrap_result(&plan, outcome)) + } +} + +/// Wrap a finished `(scalar value, sample / chunk counts)` into a +/// `QueryResult` with the exact-accuracy envelope + the +/// `data_source: gorilla_archive` info line. Pulled out so tests +/// can pin the wrapping shape independently of the executor. +pub fn wrap_result(plan: &QueryPlan, outcome: ExecutionOutcome) -> QueryResult { + // Result timestamp is the right edge of the requested range — + // mirrors `SimpleEngine`'s convention for instant-vector queries + // against a closed time window. + let result_ts = plan.time_range_ms.1.max(0) as u64; + + let labels = KeyByLabelValues::new_with_labels(Vec::new()); + let element = InstantVectorElement::new(labels, outcome.value); + let envelope = AccuracyEnvelope::single(AccuracyProfile::exact()); + QueryResult::vector(vec![element], result_ts).with_accuracy(envelope) + // Window is the requested range, expressed in u64 ms. + .with_window_used(( + plan.time_range_ms.0.max(0) as u64, + plan.time_range_ms.1.max(0) as u64, + )) +} + +/// Output of an executed plan. Kept narrow on purpose — Phase 4's +/// MVP returns a single scalar per query. Higher-cardinality +/// (per-group) shapes will land in Phase 5+ once capability +/// routing decides which engine answers grouped queries. +#[derive(Debug, Clone, PartialEq)] +pub struct ExecutionOutcome { + /// Final scalar (e.g. `sum_over_time` total, `quantile_over_time` + /// φ-quantile). NaN when the time range carries no samples. + pub value: f64, + /// Number of raw samples that contributed to `value`. + pub samples_scanned: usize, + /// Number of chunks the executor fetched from the cold store. + pub chunks_fetched: usize, +} + +impl ExecutionOutcome { + /// "no data" sentinel — used when the time range is empty. + pub fn empty() -> Self { + Self { + value: f64::NAN, + samples_scanned: 0, + chunks_fetched: 0, + } + } + + /// Build the `infos` array surfaced on the wire response. + /// Pulled out so tests can pin the exact strings. + pub fn info_lines(&self) -> Vec { + vec![ + AccuracyProfile::exact().summary(), + DATA_SOURCE_GORILLA_ARCHIVE.to_string(), + format!("samples_scanned: {}", self.samples_scanned), + format!("chunks_fetched: {}", self.chunks_fetched), + ] + } +} diff --git a/asap-query-engine/src/engines/gorilla_engine/query_planner.rs b/asap-query-engine/src/engines/gorilla_engine/query_planner.rs new file mode 100644 index 00000000..6334f6fd --- /dev/null +++ b/asap-query-engine/src/engines/gorilla_engine/query_planner.rs @@ -0,0 +1,307 @@ +//! PromQL → [`QueryPlan`] translator for the Phase-4 Gorilla engine. +//! +//! The Phase-4 surface is intentionally narrow: instant-vector +//! queries that wrap a single matrix selector with one of the +//! supported `*_over_time` / `rate` / `increase` functions, OR +//! a top-level `quantile_over_time(φ, m[range])` / +//! `topk(k, m[range])`-style aggregation. +//! +//! Time range is `(now - lookback_ms, now)` where `now` is the +//! caller-supplied "query time" — fixed to `chrono::Utc::now()` +//! when not specified, so callers that don't care about backdating +//! a query don't need to thread a clock through. + +use std::time::SystemTime; + +use chrono::Utc; +use promql_parser::parser::{ + AggregateExpr, Call, Expr, FunctionArgs, MatrixSelector, NumberLiteral, ParenExpr, + VectorSelector, +}; + +/// Statistic to compute, alongside any extra parameters +/// (quantile φ, top-k k). +#[derive(Debug, Clone, PartialEq)] +pub enum QueryStatistic { + /// `sum_over_time(m[range])` + SumOverTime, + /// `count_over_time(m[range])` + CountOverTime, + /// `avg_over_time(m[range])` (= sum / count) + AvgOverTime, + /// `min_over_time(m[range])` + MinOverTime, + /// `max_over_time(m[range])` + MaxOverTime, + /// `rate(m[range])` — `(last - first) / range_seconds` + Rate, + /// `increase(m[range])` — `last - first` + Increase, + /// `quantile_over_time(φ, m[range])` + QuantileOverTime { phi: f64 }, + /// `topk(k, sum_over_time(m[range]))`-style aggregation. The + /// MVP Phase 4 implementation returns the sum of the top-`k` + /// sample values in the range — once Phase 5 adds spatial + /// grouping the executor will return a per-group vector. + TopK { k: usize }, +} + +impl QueryStatistic { + /// True iff the executor can answer this statistic via the + /// streaming-additive path; false → buffered path (everything + /// has to be in memory before producing the answer). + pub fn is_streaming_additive(&self) -> bool { + matches!( + self, + Self::SumOverTime + | Self::CountOverTime + | Self::AvgOverTime + | Self::MinOverTime + | Self::MaxOverTime + | Self::Rate + | Self::Increase + ) + } +} + +/// Output of [`plan_query`]. +#[derive(Debug, Clone, PartialEq)] +pub struct QueryPlan { + pub metric: String, + /// Half-open `[start_ms, end_ms)` request window. Computed as + /// `(now_ms - range_ms, now_ms)` from the matrix selector's + /// `[range]` duration. + pub time_range_ms: (i64, i64), + pub statistic: QueryStatistic, +} + +/// Parse `query` and produce a [`QueryPlan`]. `now` defaults to +/// the system clock; the [`plan_query_at`] variant lets tests pin +/// a deterministic timestamp. +pub fn plan_query(query: &str) -> Result { + let now_ms = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or_else(|_| Utc::now().timestamp_millis()); + plan_query_at(query, now_ms) +} + +/// As [`plan_query`], with a caller-supplied `now_ms`. +pub fn plan_query_at(query: &str, now_ms: i64) -> Result { + let ast = promql_parser::parser::parse(query).map_err(|e| format!("parse: {e}"))?; + plan_from_ast(&ast, now_ms) +} + +fn plan_from_ast(ast: &Expr, now_ms: i64) -> Result { + match ast { + Expr::Paren(ParenExpr { expr }) => plan_from_ast(expr, now_ms), + Expr::Call(call) => plan_from_call(call, now_ms), + Expr::Aggregate(agg) => plan_from_aggregate(agg, now_ms), + other => Err(format!( + "unsupported top-level expression: {:?}; the Gorilla engine \ + expects a single function call (rate/increase/*_over_time) \ + or topk(k, ...) aggregation", + std::mem::discriminant(other) + )), + } +} + +fn plan_from_call(call: &Call, now_ms: i64) -> Result { + let name = call.func.name.to_lowercase(); + match name.as_str() { + "sum_over_time" + | "count_over_time" + | "avg_over_time" + | "min_over_time" + | "max_over_time" + | "rate" + | "increase" => { + let ms = expect_single_matrix_arg(&call.args, &name)?; + let (metric, range_ms) = matrix_metric_and_range_ms(ms); + let stat = match name.as_str() { + "sum_over_time" => QueryStatistic::SumOverTime, + "count_over_time" => QueryStatistic::CountOverTime, + "avg_over_time" => QueryStatistic::AvgOverTime, + "min_over_time" => QueryStatistic::MinOverTime, + "max_over_time" => QueryStatistic::MaxOverTime, + "rate" => QueryStatistic::Rate, + "increase" => QueryStatistic::Increase, + _ => unreachable!(), + }; + Ok(QueryPlan { + metric, + time_range_ms: (now_ms - range_ms, now_ms), + statistic: stat, + }) + } + "quantile_over_time" => { + // quantile_over_time(φ, m[range]) + if call.args.args.len() != 2 { + return Err(format!( + "quantile_over_time expects 2 args, got {}", + call.args.args.len() + )); + } + let phi = expect_number(&call.args.args[0], "quantile_over_time φ")?; + let ms = expect_matrix_selector(&call.args.args[1], "quantile_over_time")?; + let (metric, range_ms) = matrix_metric_and_range_ms(ms); + Ok(QueryPlan { + metric, + time_range_ms: (now_ms - range_ms, now_ms), + statistic: QueryStatistic::QuantileOverTime { phi }, + }) + } + other => Err(format!( + "unsupported PromQL function: {other}; the Gorilla engine \ + supports rate/increase/*_over_time/quantile_over_time" + )), + } +} + +fn plan_from_aggregate(agg: &AggregateExpr, now_ms: i64) -> Result { + // PromQL grammar requires aggregation operators to take a + // vector — so the legal Phase-4 spellings are e.g. + // `topk(2, sum_over_time(m[10s]))`. We strip the outer + // aggregation, recurse into the inner call to recover the + // `(metric, range)` pair, then overlay the TopK statistic. + let op_str = format!("{}", agg.op); + if !op_str.eq_ignore_ascii_case("topk") { + return Err(format!( + "unsupported top-level aggregation: {op_str}; only `topk(k, ...)` \ + is supported in Phase 4" + )); + } + let k_expr = agg + .param + .as_deref() + .ok_or_else(|| "topk requires a numeric parameter (k)".to_string())?; + let k = expect_number(k_expr, "topk k")?; + if !k.is_finite() || k <= 0.0 { + return Err(format!("topk k must be positive, got {k}")); + } + // Recurse into the inner expression — it can be a matrix + // selector (handled by [`matrix_metric_and_range_ms`] + // directly) OR a vector-returning function call (the legal + // PromQL spelling). Either way we end up with a + // `(metric, range_ms)` pair we can overlay TopK on. + let inner_plan = match &*agg.expr { + Expr::MatrixSelector(ms) => { + let (metric, range_ms) = matrix_metric_and_range_ms(ms); + QueryPlan { + metric, + time_range_ms: (now_ms - range_ms, now_ms), + statistic: QueryStatistic::SumOverTime, // overlay below + } + } + _ => plan_from_ast(&agg.expr, now_ms)?, + }; + Ok(QueryPlan { + metric: inner_plan.metric, + time_range_ms: inner_plan.time_range_ms, + statistic: QueryStatistic::TopK { k: k as usize }, + }) +} + +fn expect_single_matrix_arg<'a>( + args: &'a FunctionArgs, + fname: &str, +) -> Result<&'a MatrixSelector, String> { + if args.args.len() != 1 { + return Err(format!( + "{fname} expects 1 matrix-selector arg, got {}", + args.args.len() + )); + } + expect_matrix_selector(&args.args[0], fname) +} + +fn expect_matrix_selector<'a>(expr: &'a Expr, ctx: &str) -> Result<&'a MatrixSelector, String> { + match expr { + Expr::MatrixSelector(ms) => Ok(ms), + Expr::Paren(ParenExpr { expr }) => expect_matrix_selector(expr, ctx), + other => Err(format!( + "{ctx}: expected matrix selector `metric[range]`, got {:?}", + std::mem::discriminant(other) + )), + } +} + +fn expect_number(expr: &Expr, ctx: &str) -> Result { + match expr { + Expr::NumberLiteral(NumberLiteral { val }) => Ok(*val), + Expr::Paren(ParenExpr { expr }) => expect_number(expr, ctx), + other => Err(format!( + "{ctx}: expected numeric literal, got {:?}", + std::mem::discriminant(other) + )), + } +} + +fn matrix_metric_and_range_ms(ms: &MatrixSelector) -> (String, i64) { + let metric = vector_selector_metric(&ms.vs); + let range_ms = ms.range.as_millis() as i64; + (metric, range_ms) +} + +fn vector_selector_metric(vs: &VectorSelector) -> String { + if let Some(name) = &vs.name { + return name.clone(); + } + // Fallback: inspect matchers for an `__name__` exact match. + for m in vs.matchers.matchers.iter() { + if m.name == "__name__" { + return m.value.clone(); + } + } + String::new() +} + +#[cfg(test)] +mod tests { + use super::*; + + const NOW: i64 = 1_715_000_000_000; + + #[test] + fn plans_sum_over_time() { + let plan = plan_query_at("sum_over_time(http_requests_total[5m])", NOW).unwrap(); + assert_eq!(plan.metric, "http_requests_total"); + assert_eq!(plan.statistic, QueryStatistic::SumOverTime); + assert_eq!(plan.time_range_ms, (NOW - 5 * 60_000, NOW)); + } + + #[test] + fn plans_quantile_over_time() { + let plan = plan_query_at("quantile_over_time(0.99, latency_ms[1m])", NOW).unwrap(); + assert_eq!(plan.metric, "latency_ms"); + assert!(matches!( + plan.statistic, + QueryStatistic::QuantileOverTime { phi } if (phi - 0.99).abs() < 1e-12 + )); + } + + #[test] + fn plans_topk() { + // Legal PromQL spelling: aggregation wraps a vector-returning + // function call. The Phase-4 planner peels off the outer + // `topk` and recovers the `(metric, range)` pair from the + // inner `sum_over_time(...)`. + let plan = plan_query_at("topk(3, sum_over_time(m[10s]))", NOW).unwrap(); + assert!(matches!(plan.statistic, QueryStatistic::TopK { k } if k == 3)); + assert_eq!(plan.metric, "m"); + assert_eq!(plan.time_range_ms, (NOW - 10_000, NOW)); + } + + #[test] + fn rejects_binary_expression() { + assert!(plan_query_at("foo + bar", NOW).is_err()); + } + + #[test] + fn streaming_classification() { + assert!(QueryStatistic::SumOverTime.is_streaming_additive()); + assert!(QueryStatistic::Rate.is_streaming_additive()); + assert!(!QueryStatistic::QuantileOverTime { phi: 0.5 }.is_streaming_additive()); + assert!(!QueryStatistic::TopK { k: 1 }.is_streaming_additive()); + } +} diff --git a/asap-query-engine/src/engines/gorilla_engine/tests.rs b/asap-query-engine/src/engines/gorilla_engine/tests.rs new file mode 100644 index 00000000..70b68d15 --- /dev/null +++ b/asap-query-engine/src/engines/gorilla_engine/tests.rs @@ -0,0 +1,615 @@ +//! Phase-4 unit tests for the Gorilla query engine. +//! +//! Tests exercise the engine end-to-end via a `MockColdStore` +//! injected in place of the production `GorillaS3ColdStore`. The +//! mock is intentionally minimal: it owns a `Vec<(ChunkRef, +//! Vec)>` and answers `list_chunks` / `read_chunk` +//! straight off it, with optional latency injection for the +//! timeout test. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use tokio::time::sleep; + +use crate::drivers::query::fallback::cold_store::{ + ChunkRef, ColdStore, ColdStoreError, RawSample, +}; +use crate::engines::query_result::QueryResult; +use crate::stores::sketch_db::accuracy::{AccuracyKind, AccuracyProfile}; + +use super::query_planner::{plan_query_at, QueryStatistic}; +use super::{ + wrap_result, EngineError, ExactExecutor, ExecutionOutcome, GorillaEngineConfig, + GorillaQueryEngine, DATA_SOURCE_GORILLA_ARCHIVE, +}; + +// ───────────────────────────────────────────────────────────────────── +// Mock cold store +// ───────────────────────────────────────────────────────────────────── + +/// In-process mock that satisfies the [`ColdStore`] trait without +/// any S3 / disk roundtrip. Built once in each test from a list of +/// `(ChunkRef, samples)` pairs. +#[derive(Default)] +struct MockColdStore { + chunks: Vec<(ChunkRef, Vec)>, + /// If set, every `read_chunk` call sleeps for this duration — + /// used by the timeout test. + read_delay: Option, +} + +impl MockColdStore { + fn new(chunks: Vec<(ChunkRef, Vec)>) -> Self { + Self { + chunks, + read_delay: None, + } + } + + fn with_read_delay(mut self, d: Duration) -> Self { + self.read_delay = Some(d); + self + } +} + +#[async_trait] +impl ColdStore for MockColdStore { + async fn scan( + &self, + metric: &str, + start_ms: i64, + end_ms: i64, + ) -> Result, ColdStoreError> { + let mut out = Vec::new(); + let chunks = self.list_chunks(metric, start_ms, end_ms).await?; + for c in chunks { + for s in self.read_chunk(&c).await? { + if s.ts_ms >= start_ms && s.ts_ms < end_ms { + out.push(s); + } + } + } + Ok(out) + } + + async fn list_chunks( + &self, + metric: &str, + start_ms: i64, + end_ms: i64, + ) -> Result, ColdStoreError> { + Ok(self + .chunks + .iter() + .filter(|(c, _)| { + c.metric == metric + && c.time_range_ms.0 < end_ms + && c.time_range_ms.1 >= start_ms + }) + .map(|(c, _)| c.clone()) + .collect()) + } + + async fn read_chunk(&self, chunk: &ChunkRef) -> Result, ColdStoreError> { + if let Some(d) = self.read_delay { + sleep(d).await; + } + for (c, samples) in &self.chunks { + if c.key == chunk.key { + return Ok(samples.clone()); + } + } + Err(ColdStoreError::Backend(format!( + "mock: no such chunk {}", + chunk.key + ))) + } +} + +// ───────────────────────────────────────────────────────────────────── +// Fixture helpers +// ───────────────────────────────────────────────────────────────────── + +const NOW_MS: i64 = 1_715_000_000_000; +const METRIC: &str = "http_requests_total"; + +fn raw(ts_ms: i64, value: f64) -> RawSample { + RawSample { + ts_ms, + labels: BTreeMap::new(), + value, + } +} + +/// One chunk covering `[start, start + n*step]` with a +/// monotonically-increasing value column (`base + i*step_v`). +fn linear_chunk( + key: &str, + start_ms: i64, + step_ms: i64, + n: usize, + base: f64, + step_v: f64, +) -> (ChunkRef, Vec) { + let samples: Vec = (0..n) + .map(|i| raw(start_ms + (i as i64) * step_ms, base + (i as f64) * step_v)) + .collect(); + let last_ts = samples.last().map(|s| s.ts_ms).unwrap_or(start_ms); + let chunk = ChunkRef { + key: key.to_string(), + metric: METRIC.to_string(), + time_range_ms: (start_ms, last_ts + 1), + label_hash: 0, + sample_count: n as u32, + size_bytes: 0, + }; + (chunk, samples) +} + +fn cfg() -> GorillaEngineConfig { + GorillaEngineConfig { + max_buffered_samples: 1_000_000, + query_timeout_secs: 30, + } +} + +fn engine_with(chunks: Vec<(ChunkRef, Vec)>) -> GorillaQueryEngine { + GorillaQueryEngine::new(Arc::new(MockColdStore::new(chunks)), cfg()) +} + +fn engine_with_config( + chunks: Vec<(ChunkRef, Vec)>, + config: GorillaEngineConfig, +) -> GorillaQueryEngine { + GorillaQueryEngine::new(Arc::new(MockColdStore::new(chunks)), config) +} + +// ───────────────────────────────────────────────────────────────────── +// Streaming-additive happy paths +// ───────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn execute_sum_over_time_streaming() { + // 60 samples × value 2.0 = 120.0 + let chunks = vec![linear_chunk( + "c1", + NOW_MS - 60_000, + 1_000, + 60, + 2.0, + 0.0, + )]; + let engine = engine_with(chunks); + let plan = plan_query_at(&format!("sum_over_time({METRIC}[5m])"), NOW_MS).unwrap(); + assert_eq!(plan.statistic, QueryStatistic::SumOverTime); + + let exec = ExactExecutor::new( + Arc::new(MockColdStore::new(vec![linear_chunk( + "c1", + NOW_MS - 60_000, + 1_000, + 60, + 2.0, + 0.0, + )])), + cfg(), + ); + let outcome = exec.execute_plan(&plan).await.unwrap(); + assert_eq!(outcome.value, 120.0); + assert_eq!(outcome.samples_scanned, 60); + assert_eq!(outcome.chunks_fetched, 1); + + // Also verify via the high-level engine. + let result = engine + .execute_at(&format!("sum_over_time({METRIC}[5m])"), NOW_MS) + .await + .unwrap(); + assert!(matches!(result, QueryResult::Vector(_))); +} + +#[tokio::test] +async fn execute_count_over_time() { + let chunks = vec![linear_chunk("c1", NOW_MS - 30_000, 1_000, 30, 0.0, 0.0)]; + let engine = engine_with(chunks); + let plan = plan_query_at(&format!("count_over_time({METRIC}[1m])"), NOW_MS).unwrap(); + let exec = ExactExecutor::new( + Arc::new(MockColdStore::new(vec![linear_chunk( + "c1", + NOW_MS - 30_000, + 1_000, + 30, + 0.0, + 0.0, + )])), + cfg(), + ); + let outcome = exec.execute_plan(&plan).await.unwrap(); + assert_eq!(outcome.value, 30.0); + + let result = engine + .execute_at(&format!("count_over_time({METRIC}[1m])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!(iv.values[0].value, 30.0); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_avg_over_time() { + // 4 samples: 1, 2, 3, 4 → avg = 2.5 + let chunks = vec![linear_chunk("c1", NOW_MS - 4_000, 1_000, 4, 1.0, 1.0)]; + let engine = engine_with(chunks); + let result = engine + .execute_at(&format!("avg_over_time({METRIC}[10s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert!((iv.values[0].value - 2.5).abs() < 1e-12); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_min_over_time() { + // values 5, 1, 3, 4 → min = 1 + let samples = vec![ + raw(NOW_MS - 4_000, 5.0), + raw(NOW_MS - 3_000, 1.0), + raw(NOW_MS - 2_000, 3.0), + raw(NOW_MS - 1_000, 4.0), + ]; + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 4_000, NOW_MS), + label_hash: 0, + sample_count: 4, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("min_over_time({METRIC}[10s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!(iv.values[0].value, 1.0); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_max_over_time() { + let samples = vec![ + raw(NOW_MS - 4_000, 5.0), + raw(NOW_MS - 3_000, 1.0), + raw(NOW_MS - 2_000, 3.0), + raw(NOW_MS - 1_000, 4.0), + ]; + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 4_000, NOW_MS), + label_hash: 0, + sample_count: 4, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("max_over_time({METRIC}[10s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!(iv.values[0].value, 5.0); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_rate_basic() { + // Counter goes from 100 at t=NOW-10s to 200 at t=NOW-1s. + // rate over 10s window = (200 - 100) / 10s = 10.0 + let samples = vec![ + raw(NOW_MS - 10_000, 100.0), + raw(NOW_MS - 5_000, 150.0), + raw(NOW_MS - 1_000, 200.0), + ]; + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 10_000, NOW_MS), + label_hash: 0, + sample_count: 3, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("rate({METRIC}[10s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert!((iv.values[0].value - 10.0).abs() < 1e-9); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_increase_basic() { + let samples = vec![ + raw(NOW_MS - 10_000, 100.0), + raw(NOW_MS - 5_000, 150.0), + raw(NOW_MS - 1_000, 250.0), + ]; + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 10_000, NOW_MS), + label_hash: 0, + sample_count: 3, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("increase({METRIC}[10s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert!((iv.values[0].value - 150.0).abs() < 1e-9); + } else { + panic!("expected Vector"); + } +} + +// ───────────────────────────────────────────────────────────────────── +// Buffered quantile + topk +// ───────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn execute_quantile_buffered_basic() { + // Values 0..100; q0.99 → index round((100-1)*0.99) = round(98.01) = 98 → value 98. + let mut samples: Vec = (0..100) + .map(|i| raw(NOW_MS - 100_000 + (i as i64) * 1_000, i as f64)) + .collect(); + // Shuffle the value order so the executor must sort. + samples.sort_by_key(|s| s.value as i64); + samples.reverse(); + + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 100_000, NOW_MS), + label_hash: 0, + sample_count: 100, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("quantile_over_time(0.99, {METRIC}[2m])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!(iv.values[0].value, 98.0); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_quantile_too_many_samples_errors() { + // Generate 100 samples but cap the buffered budget at 5. + let samples: Vec = (0..100) + .map(|i| raw(NOW_MS - 100_000 + (i as i64) * 1_000, i as f64)) + .collect(); + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 100_000, NOW_MS), + label_hash: 0, + sample_count: 100, + size_bytes: 0, + }; + let cfg = GorillaEngineConfig { + max_buffered_samples: 5, + query_timeout_secs: 30, + }; + let engine = engine_with_config(vec![(chunk, samples)], cfg); + let res = engine + .execute_at(&format!("quantile_over_time(0.5, {METRIC}[2m])"), NOW_MS) + .await; + match res { + Err(EngineError::TooManySamples { count, limit }) => { + assert_eq!(limit, 5); + assert!(count > limit); + } + other => panic!("expected TooManySamples, got {other:?}"), + } +} + +#[tokio::test] +async fn execute_topk_basic() { + // Values [1, 2, 3, 10, 20]; topk(2) → 30 + let samples = vec![ + raw(NOW_MS - 5_000, 1.0), + raw(NOW_MS - 4_000, 2.0), + raw(NOW_MS - 3_000, 3.0), + raw(NOW_MS - 2_000, 10.0), + raw(NOW_MS - 1_000, 20.0), + ]; + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 5_000, NOW_MS), + label_hash: 0, + sample_count: 5, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("topk(2, sum_over_time({METRIC}[10s]))"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!(iv.values[0].value, 30.0); + } else { + panic!("expected Vector"); + } +} + +// ───────────────────────────────────────────────────────────────────── +// Edge cases +// ───────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn execute_empty_chunks_returns_zero_or_nan() { + let engine = engine_with(Vec::new()); + let sum = engine + .execute_at(&format!("sum_over_time({METRIC}[5m])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = sum { + assert!(iv.values[0].value.is_nan(), "sum on empty should be NaN"); + } else { + panic!("expected Vector"); + } + let count = engine + .execute_at(&format!("count_over_time({METRIC}[5m])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = count { + assert_eq!(iv.values[0].value, 0.0); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_chunks_partially_outside_range_filtered() { + // Chunk has 100 samples spanning [NOW-100s, NOW]; request + // covers the latter half [NOW-50s, NOW] → exactly 50 samples + // contribute. + let samples: Vec = (0..100) + .map(|i| raw(NOW_MS - 100_000 + (i as i64) * 1_000, 1.0)) + .collect(); + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 100_000, NOW_MS), + label_hash: 0, + sample_count: 100, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("count_over_time({METRIC}[50s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!( + iv.values[0].value, 50.0, + "exactly 50 samples should match a 50s window" + ); + } else { + panic!("expected Vector"); + } +} + +// ───────────────────────────────────────────────────────────────────── +// Result wrapping +// ───────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn result_carries_exact_accuracy_envelope() { + let chunks = vec![linear_chunk("c", NOW_MS - 1_000, 100, 10, 1.0, 0.0)]; + let engine = engine_with(chunks); + let result = engine + .execute_at(&format!("sum_over_time({METRIC}[5s])"), NOW_MS) + .await + .unwrap(); + let env = result + .accuracy() + .expect("Gorilla engine result must carry an accuracy envelope"); + assert_eq!(env.profile.kind, AccuracyKind::Exact); + assert_eq!(env.profile.epsilon, 0.0); + assert_eq!(env.profile.delta, 0.0); + // And the summary string the dashboards parse: + assert_eq!(env.profile.summary(), AccuracyProfile::exact().summary()); +} + +#[tokio::test] +async fn result_includes_data_source_gorilla_archive() { + // The wrapping fn surfaces the data_source line on + // ExecutionOutcome::info_lines — pin both the marker constant + // and the assembled info strings. + let outcome = ExecutionOutcome { + value: 42.0, + samples_scanned: 7, + chunks_fetched: 2, + }; + let infos = outcome.info_lines(); + assert!( + infos.contains(&DATA_SOURCE_GORILLA_ARCHIVE.to_string()), + "infos must include `{DATA_SOURCE_GORILLA_ARCHIVE}`; got {infos:?}" + ); + assert!( + infos.iter().any(|i| i == "samples_scanned: 7"), + "infos must report the scanned-samples count" + ); + assert!( + infos.iter().any(|i| i == "chunks_fetched: 2"), + "infos must report the chunk-fetch count" + ); + + // And via the wrap_result path, the QueryResult itself carries + // the exact-accuracy envelope (data_source line is on the + // info-array which is assembled at the HTTP-driver layer; see + // wrap_result docs). + let plan = plan_query_at(&format!("sum_over_time({METRIC}[5s])"), NOW_MS).unwrap(); + let qr = wrap_result(&plan, outcome.clone()); + let env = qr.accuracy().expect("wrap_result must attach envelope"); + assert_eq!(env.profile.kind, AccuracyKind::Exact); +} + +// ───────────────────────────────────────────────────────────────────── +// Timeout +// ───────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn engine_respects_config_timeout() { + // 1 chunk + 250 ms read delay; engine timeout = 1 s ceil. We + // configure the timeout to 1s (the floor) and force the chunk + // count up so the cumulative read time > 1 s. + let mut chunks = Vec::new(); + for i in 0..10 { + let chunk = ChunkRef { + key: format!("k-{i}"), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 60_000, NOW_MS), + label_hash: 0, + sample_count: 1, + size_bytes: 0, + }; + chunks.push((chunk, vec![raw(NOW_MS - 1_000, 1.0)])); + } + let mock = MockColdStore::new(chunks).with_read_delay(Duration::from_millis(250)); + let cfg = GorillaEngineConfig { + max_buffered_samples: 1_000_000, + query_timeout_secs: 1, + }; + let engine = GorillaQueryEngine::new(Arc::new(mock), cfg); + let res = engine + .execute_at(&format!("sum_over_time({METRIC}[5m])"), NOW_MS) + .await; + match res { + Err(EngineError::Timeout(_)) => {} + other => panic!("expected Timeout, got {other:?}"), + } +} diff --git a/asap-query-engine/src/engines/mod.rs b/asap-query-engine/src/engines/mod.rs index 06ed0db1..8395ad2a 100644 --- a/asap-query-engine/src/engines/mod.rs +++ b/asap-query-engine/src/engines/mod.rs @@ -1,3 +1,4 @@ +pub mod gorilla_engine; pub mod logical; pub mod physical; pub mod query_result; @@ -5,6 +6,9 @@ pub mod simple_engine; pub mod timeline_dispatch; pub mod window_merger; +pub use gorilla_engine::{ + EngineError as GorillaEngineError, GorillaEngineConfig, GorillaQueryEngine, +}; pub use query_result::{InstantVector, QueryResult, RangeVector, RangeVectorElement, Sample}; pub use simple_engine::SimpleEngine; pub use timeline_dispatch::{combine_statistic, CombinedResult}; From b5f70129647d023083ba8a5e46f2b7155f1ee129 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Wed, 6 May 2026 01:31:35 -0400 Subject: [PATCH 3/3] =?UTF-8?q?feat(capability):=20Phase=205=20=E2=80=94?= =?UTF-8?q?=20QueryEngine=20trait=20+=20StorageBackend=20routing=20for=20G?= =?UTF-8?q?orilla-S3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the storage-backend axis to capability matching so queries against metrics configured for Gorilla-S3 dispatch to GorillaQueryEngine instead of the warm-tier SimpleEngine. Per docs/design-gorilla-s3-cold-engine.md §8. asap_types extensions: - StorageBackend enum (SketchWarmTier default + GorillaS3Archive + ColdJsonlFallback + DoubleWrite) with snake_case serde + #[serde(default)] so pre-Phase-5 configs decode unchanged. - AccuracyTarget enum (Exact / Approximate, default Approximate). - compatible_storage_backends(stat, accuracy, metric_storage_config) returns the ordered failover list the router walks. Exact-on-archive subsumes approximate-on-warm (Gorilla-only metrics still archive). - StreamingConfig.storage_backend field with serde-default + builder constructor (with_storage_backend) + getter. asap-query-engine extensions (additive): - engines::EngineError envelope (CapabilityMiss vs Backend) the trait returns. - engines::router::QueryEngine async trait + EngineCapabilities struct. - engines::router::EngineRouter dispatcher: consults compatible_storage_backends, dispatches to first registered engine, falls through on Backend / CapabilityMiss to the next compatible backend (typically ColdJsonlFallback). - impl QueryEngine for SimpleEngine (adapts handle_query → Result; None → CapabilityMiss). - impl QueryEngine for GorillaQueryEngine (Plan errors fold to CapabilityMiss; ColdStore / Timeout / TooManySamples to Backend). Tests added (16 new): - capability_matching: 6 routing tests + 1 source-of-truth agreement test enumerating (Statistic × AccuracyTarget × StorageBackend). - streaming_config: 2 serde tests pinning legacy + Phase-5 wire format. - engines::router: 6 tests (warm-tier dispatch, archive dispatch, JSONL fallback when archive+warm fail, no-engines clean error, all-failed surface, hot-swap re-registration). Test-utility struct-literal sites updated to include the new StreamingConfig.storage_backend field. No precompute_engine/, stores/, or precompute_operators/ touched. Verified clean: cargo build --release -p {asap_types, query_engine_rust}; cargo test -p asap_types; cargo test --release -p query_engine_rust router; cargo clippy --release -p asap_types --all-targets -- -D warnings. Pre-existing query_engine_rust clippy/test failures unchanged (verified via git-stash diff). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../rs/asap_types/src/capability_matching.rs | 321 +++++++++++++ .../dependencies/rs/asap_types/src/lib.rs | 4 +- .../rs/asap_types/src/streaming_config.rs | 51 +++ .../src/engines/gorilla_engine/mod.rs | 46 ++ asap-query-engine/src/engines/mod.rs | 59 +++ asap-query-engine/src/engines/router.rs | 428 ++++++++++++++++++ .../src/engines/simple_engine.rs | 43 ++ .../src/tests/capability_matching_tests.rs | 2 + .../src/tests/sql_pattern_matching_tests.rs | 1 + .../tests/test_utilities/config_builders.rs | 2 + .../tests/test_utilities/engine_factories.rs | 6 + .../tests/inference_yaml_pattern_coverage.rs | 1 + 12 files changed, 963 insertions(+), 1 deletion(-) create mode 100644 asap-query-engine/src/engines/router.rs diff --git a/asap-common/dependencies/rs/asap_types/src/capability_matching.rs b/asap-common/dependencies/rs/asap_types/src/capability_matching.rs index 2b31b2dc..17d8ae7d 100644 --- a/asap-common/dependencies/rs/asap_types/src/capability_matching.rs +++ b/asap-common/dependencies/rs/asap_types/src/capability_matching.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::Statistic; +use serde::{Deserialize, Serialize}; use tracing::{debug, warn}; use crate::aggregation_config::{AggregationConfig, AggregationIdInfo}; @@ -11,6 +12,81 @@ use crate::query_requirements::QueryRequirements; use crate::utils::normalize_spatial_filter; use promql_utilities::query_logics::enums::AggregationType; +// --------------------------------------------------------------------------- +// Phase-5: storage-backend capability axis +// +// Today's `find_compatible_aggregation` matches on +// `(metric, statistic, sub_type, window_size, grouping_labels, spatial_filter)` +// — there is no axis for "which storage tier serves this query." The Phase-5 +// `GorillaQueryEngine` (PR #85) introduces a parallel exact tier; the planner / +// router needs to disambiguate between warm-tier sketches and Gorilla-S3 +// chunks. See `docs/design-gorilla-s3-cold-engine.md` §8. +// --------------------------------------------------------------------------- + +/// Which physical storage tier a query (or a metric configuration) routes to. +/// +/// `SketchWarmTier` is the default — every existing `AggregationConfig` and +/// `StreamingConfig` decodes into this variant via `#[serde(default)]`, so +/// pre-Phase-5 deploys keep dispatching to `SimpleEngine` unchanged. +#[derive( + Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default, +)] +#[serde(rename_all = "snake_case")] +pub enum StorageBackend { + /// Warm-tier sketch DB (today's `SimpleMapStore` + accumulators). + /// Served by `SimpleEngine`. Default for unconfigured metrics. + #[default] + SketchWarmTier, + + /// NEW (Phase 5) — Gorilla-S3 archive. Served by `GorillaQueryEngine`, + /// reading per-hour Gorilla chunks from S3 / MinIO via the Phase-3 + /// `GorillaS3ColdStore`. + GorillaS3Archive, + + /// Local-FS JSONL fallback (PR #54 §5.2). Served by the existing + /// cold-fallback path; used when no warm-tier or archive aggregation + /// can answer the query. + ColdJsonlFallback, + + /// Double-write: the metric is written to both warm-tier sketches AND the + /// Gorilla-S3 archive. Capability matching surfaces both options and the + /// cost-aware dispatcher picks per query (typically warm-tier for low- + /// latency approximate, archive for exact). + DoubleWrite, +} + +impl StorageBackend { + /// String tag pinned for byte-comparable dispatch on the wire (mirrors + /// the `data_source: ` info-line on `QueryResult`). Engines + /// register themselves under these IDs in the router. + pub const fn data_source_id(self) -> &'static str { + match self { + StorageBackend::SketchWarmTier => "sketch_warm", + StorageBackend::GorillaS3Archive => "gorilla_archive", + StorageBackend::ColdJsonlFallback => "cold_jsonl", + StorageBackend::DoubleWrite => "double_write", + } + } +} + +/// Accuracy hint pushed by the controller at intent-binding time +/// (`controller/docs/design.md` §6 `core::workload`). The Phase-5 capability +/// router consults this to decide whether a metric configured for both warm- +/// tier and Gorilla-S3 should answer from the archive (Exact) or the +/// approximate warm-tier sketch. +#[derive( + Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default, +)] +#[serde(rename_all = "snake_case")] +pub enum AccuracyTarget { + /// Caller demands an exact answer; warm-tier sketches are not eligible + /// unless they happen to be exact accumulators (Sum, MinMax, Increase). + Exact, + /// Caller accepts ε/δ-bounded approximate answers. Default. + #[default] + Approximate, +} + // --------------------------------------------------------------------------- // Pure compatibility helpers // --------------------------------------------------------------------------- @@ -67,6 +143,67 @@ pub fn compatible_agg_types(stat: Statistic) -> &'static [AggregationType] { } } +/// Returns the storage backends that can serve a `(statistic, accuracy)` +/// query when the metric is configured for `metric_storage_config`. +/// +/// The returned list is **ordered by preference**: the router walks it in +/// order and dispatches to the first backend whose engine is registered. +/// `ColdJsonlFallback` is appended whenever the warm tier is in play so a +/// capability miss falls through to the §5.2 raw-store path before erroring. +/// +/// Routing rules (mirrors `docs/design-gorilla-s3-cold-engine.md` §8): +/// +/// * Metric configured for `GorillaS3Archive`: always +/// `[GorillaS3Archive]`. Exact-on-archive subsumes approximate-on-warm, +/// so even a `Statistic::Quantile` with an `Approximate` target still +/// routes to the archive when the metric is Gorilla-only — there is no +/// warm-tier sketch to fall back to in that deploy shape. +/// * Metric configured for `SketchWarmTier` (or unconfigured / default): +/// `[SketchWarmTier, ColdJsonlFallback]` — warm tier first, raw-store +/// fallback if no compatible aggregation exists. +/// * Metric configured for `DoubleWrite`: head depends on accuracy hint, +/// tail is the failover sequence (the cost-aware `EngineRouter` picks +/// the head, walks the tail on failure): +/// - `Exact` → `[GorillaS3Archive, SketchWarmTier, ColdJsonlFallback]` +/// - `Approximate` → `[SketchWarmTier, GorillaS3Archive, ColdJsonlFallback]` +/// * Metric explicitly configured for `ColdJsonlFallback`: just +/// `[ColdJsonlFallback]`. +pub fn compatible_storage_backends( + _stat: Statistic, + accuracy: AccuracyTarget, + metric_storage_config: StorageBackend, +) -> Vec { + match metric_storage_config { + StorageBackend::GorillaS3Archive => { + // Exact-on-archive subsumes approximate-on-warm: a Gorilla-only + // metric has no sketch to back-fall to, and the archive can + // always answer exact (and therefore also approximate) queries. + vec![StorageBackend::GorillaS3Archive] + } + StorageBackend::SketchWarmTier => { + vec![ + StorageBackend::SketchWarmTier, + StorageBackend::ColdJsonlFallback, + ] + } + StorageBackend::DoubleWrite => match accuracy { + AccuracyTarget::Exact => vec![ + StorageBackend::GorillaS3Archive, + StorageBackend::SketchWarmTier, + StorageBackend::ColdJsonlFallback, + ], + AccuracyTarget::Approximate => vec![ + StorageBackend::SketchWarmTier, + StorageBackend::GorillaS3Archive, + StorageBackend::ColdJsonlFallback, + ], + }, + StorageBackend::ColdJsonlFallback => { + vec![StorageBackend::ColdJsonlFallback] + } + } +} + /// Returns the required aggregation_sub_type for this statistic, if any. /// `Min` requires `"min"`, `Max` requires `"max"`. All others are unconstrained. pub fn required_sub_type(stat: Statistic) -> Option<&'static str> { @@ -967,4 +1104,188 @@ mod tests { ); assert_eq!(info.aggregation_id_for_key, 43); } + + // ----------------------------------------------------------------------- + // Phase-5: storage-backend routing + // + // The Phase-5 `EngineRouter` (see `asap-query-engine/src/engines/router.rs`) + // consults `compatible_storage_backends(stat, accuracy, metric_storage)` + // to pick a backend. These tests pin the routing matrix so the dispatcher + // stays in lock-step with the design doc §8. + // ----------------------------------------------------------------------- + + #[test] + fn gorilla_s3_metric_routes_to_archive() { + let backends = compatible_storage_backends( + Statistic::Sum, + AccuracyTarget::Exact, + StorageBackend::GorillaS3Archive, + ); + assert_eq!(backends, vec![StorageBackend::GorillaS3Archive]); + } + + #[test] + fn sketch_warm_tier_metric_routes_to_simple_engine() { + let backends = compatible_storage_backends( + Statistic::Quantile, + AccuracyTarget::Approximate, + StorageBackend::SketchWarmTier, + ); + assert_eq!( + backends, + vec![ + StorageBackend::SketchWarmTier, + StorageBackend::ColdJsonlFallback, + ] + ); + } + + #[test] + fn double_write_metric_returns_both_options() { + // Exact: archive head, warm-tier failover, then JSONL. + let exact = compatible_storage_backends( + Statistic::Sum, + AccuracyTarget::Exact, + StorageBackend::DoubleWrite, + ); + assert_eq!( + exact, + vec![ + StorageBackend::GorillaS3Archive, + StorageBackend::SketchWarmTier, + StorageBackend::ColdJsonlFallback, + ] + ); + // Approximate: warm-tier head (cheaper for ε/δ-bounded answers), + // archive failover, then JSONL. + let approx = compatible_storage_backends( + Statistic::Quantile, + AccuracyTarget::Approximate, + StorageBackend::DoubleWrite, + ); + assert_eq!( + approx, + vec![ + StorageBackend::SketchWarmTier, + StorageBackend::GorillaS3Archive, + StorageBackend::ColdJsonlFallback, + ] + ); + } + + /// Exact-on-archive subsumes approximate-on-warm: a metric configured + /// only for Gorilla-S3 still routes to the archive even when the caller + /// asks for an approximate answer (no warm-tier sketch exists to back- + /// fall to in that deploy shape). + #[test] + fn gorilla_s3_with_non_exact_accuracy_still_archives() { + let backends = compatible_storage_backends( + Statistic::Quantile, + AccuracyTarget::Approximate, + StorageBackend::GorillaS3Archive, + ); + assert_eq!(backends, vec![StorageBackend::GorillaS3Archive]); + } + + #[test] + fn cold_jsonl_only_metric_routes_to_jsonl() { + let backends = compatible_storage_backends( + Statistic::Sum, + AccuracyTarget::Exact, + StorageBackend::ColdJsonlFallback, + ); + assert_eq!(backends, vec![StorageBackend::ColdJsonlFallback]); + } + + #[test] + fn storage_backend_default_is_warm_tier() { + // `#[serde(default)]` on `StreamingConfig.storage_backend` (and on + // `StorageBackend::default()`) MUST be `SketchWarmTier` so pre-Phase-5 + // configs decode without bumping deploys onto the archive. + assert_eq!(StorageBackend::default(), StorageBackend::SketchWarmTier); + } + + #[test] + fn storage_backend_data_source_id_is_pinned() { + // The router registers engines by these strings; dashboards + // byte-compare them. Pin to catch accidental rename. + assert_eq!(StorageBackend::SketchWarmTier.data_source_id(), "sketch_warm"); + assert_eq!( + StorageBackend::GorillaS3Archive.data_source_id(), + "gorilla_archive", + ); + assert_eq!( + StorageBackend::ColdJsonlFallback.data_source_id(), + "cold_jsonl", + ); + } + + /// Source-of-truth agreement check, mirrors + /// `capability_canonical_map_agreement` for the storage axis. + /// + /// For every `(Statistic, AccuracyTarget, StorageBackend)` triple: + /// 1. The returned backend list is non-empty. + /// 2. The first element matches the expected head per the routing matrix + /// in `compatible_storage_backends`'s docstring (kept sync'd by hand). + /// 3. Every list ends in something the router can dispatch — either the + /// archive (Gorilla-only deploys) or `ColdJsonlFallback` (every + /// other deploy shape). + #[test] + fn capability_storage_backend_agreement() { + let stats = [ + Statistic::Count, + Statistic::Sum, + Statistic::Cardinality, + Statistic::Increase, + Statistic::Rate, + Statistic::Min, + Statistic::Max, + Statistic::Quantile, + Statistic::Topk, + ]; + let accuracies = [AccuracyTarget::Exact, AccuracyTarget::Approximate]; + let configs = [ + StorageBackend::SketchWarmTier, + StorageBackend::GorillaS3Archive, + StorageBackend::ColdJsonlFallback, + StorageBackend::DoubleWrite, + ]; + + for &stat in &stats { + for &acc in &accuracies { + for &cfg in &configs { + let backends = compatible_storage_backends(stat, acc, cfg); + assert!( + !backends.is_empty(), + "compatible_storage_backends({stat:?}, {acc:?}, {cfg:?}) returned empty \ + — every metric configuration must route to at least one backend", + ); + let last = *backends.last().unwrap(); + assert!( + last == StorageBackend::ColdJsonlFallback + || last == StorageBackend::GorillaS3Archive, + "backend list for ({stat:?}, {acc:?}, {cfg:?}) must terminate in a \ + dispatchable failover (ColdJsonlFallback or GorillaS3Archive); got {last:?}", + ); + // The expected head is determined by `(metric_storage_config, accuracy)`: + let expected_head = match (cfg, acc) { + (StorageBackend::GorillaS3Archive, _) => StorageBackend::GorillaS3Archive, + (StorageBackend::SketchWarmTier, _) => StorageBackend::SketchWarmTier, + (StorageBackend::ColdJsonlFallback, _) => StorageBackend::ColdJsonlFallback, + (StorageBackend::DoubleWrite, AccuracyTarget::Exact) => { + StorageBackend::GorillaS3Archive + } + (StorageBackend::DoubleWrite, AccuracyTarget::Approximate) => { + StorageBackend::SketchWarmTier + } + }; + assert_eq!( + backends[0], expected_head, + "head mismatch for ({stat:?}, {acc:?}, {cfg:?}): expected {expected_head:?}, got {:?}", + backends[0], + ); + } + } + } + } } diff --git a/asap-common/dependencies/rs/asap_types/src/lib.rs b/asap-common/dependencies/rs/asap_types/src/lib.rs index 2ae3471e..dc633d28 100644 --- a/asap-common/dependencies/rs/asap_types/src/lib.rs +++ b/asap-common/dependencies/rs/asap_types/src/lib.rs @@ -12,7 +12,9 @@ pub mod utils; pub use aggregation_config::*; pub use aggregation_reference::*; -pub use capability_matching::find_compatible_aggregation; +pub use capability_matching::{ + compatible_storage_backends, find_compatible_aggregation, AccuracyTarget, StorageBackend, +}; pub use enums::*; pub use inference_config::*; pub use promql_schema::*; diff --git a/asap-common/dependencies/rs/asap_types/src/streaming_config.rs b/asap-common/dependencies/rs/asap_types/src/streaming_config.rs index 5a7b800c..aae8a82b 100644 --- a/asap-common/dependencies/rs/asap_types/src/streaming_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/streaming_config.rs @@ -8,6 +8,7 @@ use std::ops::Index; use crate::aggregation_config::{AggregationConfig, AggregationIdInfo}; use crate::capability_matching::find_compatible_aggregation as common_find_compatible; +use crate::capability_matching::StorageBackend; use crate::enums::QueryLanguage; use crate::inference_config::{InferenceConfig, SchemaConfig}; use crate::query_requirements::QueryRequirements; @@ -15,15 +16,43 @@ use crate::query_requirements::QueryRequirements; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct StreamingConfig { pub aggregation_configs: HashMap, + /// Phase-5 capability-routing axis: which storage tier serves this + /// per-metric runtime config. The controller pushes this when planning + /// (see `docs/design-gorilla-s3-cold-engine.md` §8); pre-Phase-5 + /// configs decode with `#[serde(default)]` to `SketchWarmTier` so + /// existing deploys keep dispatching to `SimpleEngine`. + #[serde(default)] + pub storage_backend: StorageBackend, } impl StreamingConfig { pub fn new(aggregation_configs: HashMap) -> Self { Self { aggregation_configs, + storage_backend: StorageBackend::default(), } } + /// Phase-5 constructor: build with an explicit storage-backend pin. + /// Used by the controller-driven plan-push path; tests typically + /// stay on `Self::new(...)` and let the default land. + pub fn with_storage_backend( + aggregation_configs: HashMap, + storage_backend: StorageBackend, + ) -> Self { + Self { + aggregation_configs, + storage_backend, + } + } + + /// Read-only access to the storage backend pinned at construction + /// time. The Phase-5 router consults this to pick which engine + /// answers a query. + pub fn storage_backend(&self) -> StorageBackend { + self.storage_backend + } + pub fn get_aggregation_config(&self, aggregation_id: u64) -> Option<&AggregationConfig> { self.aggregation_configs.get(&aggregation_id) } @@ -129,3 +158,25 @@ impl Default for StreamingConfig { Self::new(HashMap::new()) } } + +#[cfg(test)] +mod tests { + use super::*; + + /// Pre-Phase-5 deploys serialize `StreamingConfig` without the + /// `storage_backend` field; deserialize must default to `SketchWarmTier` + /// so the router keeps dispatching to `SimpleEngine` unchanged. + #[test] + fn deserialize_legacy_yaml_defaults_to_warm_tier() { + let yaml = "{\"aggregation_configs\":{}}"; + let cfg: StreamingConfig = serde_json::from_str(yaml).expect("legacy decode"); + assert_eq!(cfg.storage_backend(), StorageBackend::SketchWarmTier); + } + + #[test] + fn deserialize_with_explicit_archive_pin() { + let yaml = "{\"aggregation_configs\":{},\"storage_backend\":\"gorilla_s3_archive\"}"; + let cfg: StreamingConfig = serde_json::from_str(yaml).expect("Phase-5 decode"); + assert_eq!(cfg.storage_backend(), StorageBackend::GorillaS3Archive); + } +} diff --git a/asap-query-engine/src/engines/gorilla_engine/mod.rs b/asap-query-engine/src/engines/gorilla_engine/mod.rs index bc2c3ce2..8c5ee3b2 100644 --- a/asap-query-engine/src/engines/gorilla_engine/mod.rs +++ b/asap-query-engine/src/engines/gorilla_engine/mod.rs @@ -264,3 +264,49 @@ impl ExecutionOutcome { ] } } + +// --------------------------------------------------------------------------- +// Phase-5: `QueryEngine` trait impl. +// +// Wraps `GorillaQueryEngine::execute` with the EngineError envelope the +// router speaks. Plan-time / parse-time failures fold into +// `EngineError::CapabilityMiss` (the engine cannot serve this query +// shape; router should fall through). Cold-store / timeout / buffer-budget +// failures fold into `EngineError::Backend` (the engine could have served +// the query but its backend transiently failed; router should also fall +// through, typically to `ColdJsonlFallback`). +// --------------------------------------------------------------------------- + +#[async_trait::async_trait] +impl crate::engines::router::QueryEngine for GorillaQueryEngine { + async fn execute( + &self, + query: &str, + ) -> Result { + match GorillaQueryEngine::execute(self, query).await { + Ok(result) => Ok(result), + Err(EngineError::Plan(msg)) => Err(crate::engines::EngineError::capability_miss( + asap_types::StorageBackend::GorillaS3Archive.data_source_id(), + msg, + )), + Err(other) => Err(crate::engines::EngineError::backend( + asap_types::StorageBackend::GorillaS3Archive.data_source_id(), + other, + )), + } + } + + fn capabilities(&self) -> crate::engines::router::EngineCapabilities { + crate::engines::router::EngineCapabilities { + data_source_id: asap_types::StorageBackend::GorillaS3Archive.data_source_id(), + storage_backend: asap_types::StorageBackend::GorillaS3Archive, + // The buffered-aggregate budget gives a natural ceiling: each + // sample is ~16 B (i64 ts + f64 value), so the byte budget is + // ~16 × max_buffered_samples. + supports_streams_above_bytes: self + .config + .max_buffered_samples + .saturating_mul(16), + } + } +} diff --git a/asap-query-engine/src/engines/mod.rs b/asap-query-engine/src/engines/mod.rs index 8395ad2a..aad49a0f 100644 --- a/asap-query-engine/src/engines/mod.rs +++ b/asap-query-engine/src/engines/mod.rs @@ -2,6 +2,7 @@ pub mod gorilla_engine; pub mod logical; pub mod physical; pub mod query_result; +pub mod router; pub mod simple_engine; pub mod timeline_dispatch; pub mod window_merger; @@ -10,6 +11,64 @@ pub use gorilla_engine::{ EngineError as GorillaEngineError, GorillaEngineConfig, GorillaQueryEngine, }; pub use query_result::{InstantVector, QueryResult, RangeVector, RangeVectorElement, Sample}; +pub use router::{EngineCapabilities, EngineRouter, EngineRouterError, QueryEngine}; pub use simple_engine::SimpleEngine; pub use timeline_dispatch::{combine_statistic, CombinedResult}; pub use window_merger::{create_window_merger, NaiveMerger, WindowMerger}; + +// --------------------------------------------------------------------------- +// Phase-5: shared `EngineError` surface returned by the `QueryEngine` trait. +// +// The trait is engine-agnostic, so its `execute` must return an error type +// that can wrap *any* concrete engine's failure mode. Today's two engines +// — `SimpleEngine` (capability-miss → `None`) and `GorillaQueryEngine` +// (rich `EngineError`) — fold into this common envelope. New engines +// can plug in by adding a `Backend(String)` arm or a typed conversion. +// --------------------------------------------------------------------------- + +use thiserror::Error; + +/// Top-level error returned by any [`QueryEngine`] impl. +/// +/// Concrete engines convert their internal error types into this envelope +/// via the conversions in this file (or via `?` for `GorillaEngineError`). +/// The router uses the variant to decide whether a failover is sensible +/// (e.g. `Backend(_)` falls through to the next compatible backend; +/// `CapabilityMiss` does not — the caller should escalate). +#[derive(Debug, Error)] +pub enum EngineError { + /// The engine has no aggregation that can answer this query. Mirrors + /// `SimpleEngine::handle_query` returning `None`. The router treats + /// this as a "hard miss" and falls through to the next backend in the + /// `compatible_storage_backends` list (typically `ColdJsonlFallback`). + #[error("no compatible aggregation in {engine_id}: {detail}")] + CapabilityMiss { + engine_id: &'static str, + detail: String, + }, + + /// The engine's backend (cold-store, S3, planner, …) failed during + /// execution. Wraps the underlying engine's error as a string so the + /// router doesn't take a hard dep on every engine's error type. + #[error("backend failure in {engine_id}: {message}")] + Backend { + engine_id: &'static str, + message: String, + }, +} + +impl EngineError { + pub fn capability_miss(engine_id: &'static str, detail: impl Into) -> Self { + Self::CapabilityMiss { + engine_id, + detail: detail.into(), + } + } + + pub fn backend(engine_id: &'static str, source: impl ToString) -> Self { + Self::Backend { + engine_id, + message: source.to_string(), + } + } +} diff --git a/asap-query-engine/src/engines/router.rs b/asap-query-engine/src/engines/router.rs new file mode 100644 index 00000000..ae345134 --- /dev/null +++ b/asap-query-engine/src/engines/router.rs @@ -0,0 +1,428 @@ +//! Phase-5 capability router: dispatches a `(query, metric_storage)` pair +//! to the engine that owns the chosen storage tier. +//! +//! The router holds a small map keyed by +//! [`asap_types::StorageBackend::data_source_id`] +//! and walks the ordered backend list returned by +//! [`asap_types::compatible_storage_backends`]. The first registered +//! engine answers; on a recoverable backend failure (`EngineError::Backend`), +//! the router falls through to the next compatible backend if the list +//! still has options. A hard capability miss in the head engine likewise +//! falls through (matching the pre-Phase-5 §5.2 fallback contract). +//! +//! See `docs/design-gorilla-s3-cold-engine.md` §8 for the design rationale +//! and the routing matrix the router walks. + +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use thiserror::Error; +use tracing::{debug, warn}; + +use asap_types::{compatible_storage_backends, AccuracyTarget, StorageBackend}; +use promql_utilities::query_logics::enums::Statistic; + +use super::{EngineError, QueryResult}; + +// --------------------------------------------------------------------------- +// `QueryEngine` trait — the abstraction the router holds. +// +// The trait is intentionally narrow: a single `execute(&str)` method (so it +// integrates with both `SimpleEngine::handle_query` and +// `GorillaQueryEngine::execute` without forcing either side to refactor its +// public surface), plus a `capabilities()` accessor the router consults at +// registration time. +// --------------------------------------------------------------------------- + +/// What a [`QueryEngine`] can serve. The router uses this to key its +/// internal map (`data_source_id`) and to estimate cost when several +/// compatible engines are registered. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct EngineCapabilities { + /// Stable identifier for the storage tier this engine answers from. + /// Must equal `self.storage_backend().data_source_id()`. Pinned so the + /// router can byte-compare and so dashboards parsing the wire response's + /// `data_source: ` info-line stay in sync. + pub data_source_id: &'static str, + /// Which physical tier this engine owns. The router asks + /// `compatible_storage_backends(...)` for an ordered list of + /// `StorageBackend`s and looks each up via this field. + pub storage_backend: StorageBackend, + /// Memory budget (in bytes) the engine is willing to buffer for + /// streaming-aggregate queries. Used by the cost-aware dispatcher + /// when several engines are eligible for the same query (today it's + /// purely informational; the Phase-6 cost model will consume it). + pub supports_streams_above_bytes: usize, +} + +/// The Phase-5 dispatch boundary. Concrete engines implement this trait so +/// the [`EngineRouter`] can hold them as `Arc` rather +/// than case-on-concrete. +/// +/// `execute` takes the query as `&str` (matching `GorillaQueryEngine`'s +/// existing surface) and returns a wire-ready [`QueryResult`]. Internal +/// engine signatures (e.g. `SimpleEngine::handle_query`'s `Option<...>`) +/// are translated by the impl so callers can program against the trait. +#[async_trait] +pub trait QueryEngine: Send + Sync { + /// Answer `query` against this engine's storage tier. + async fn execute(&self, query: &str) -> Result; + + /// What this engine can serve. Cheap; the router calls it on every + /// `register` and may re-call to refresh cost estimates. + fn capabilities(&self) -> EngineCapabilities; +} + +// --------------------------------------------------------------------------- +// `EngineRouter` — the dispatcher. +// --------------------------------------------------------------------------- + +/// Errors surfaced by [`EngineRouter::execute`] when neither the head engine +/// nor any failover backend can answer the query. +#[derive(Debug, Error)] +pub enum EngineRouterError { + /// The router has no engine registered for any of the compatible + /// backends. This is a configuration bug — the deploy didn't + /// register an engine for a tier the metric is supposed to use. + #[error( + "no engine registered for any compatible backend; \ + tried {tried:?}, registered={registered:?}" + )] + NoEngineRegistered { + tried: Vec, + registered: Vec<&'static str>, + }, + + /// The router walked the failover sequence and every engine errored + /// or missed. Returns the *last* error so callers see the deepest + /// failure (typically the cold-tier fallback's error, which is the + /// most informative). + #[error("all compatible engines failed; last error: {last}")] + AllFailed { last: EngineError }, +} + +/// Dispatches PromQL queries to the engine that owns the chosen storage +/// tier. Built once at startup, cloned cheaply (the inner map holds +/// `Arc` so engines are shared, not duplicated). +#[derive(Default, Clone)] +pub struct EngineRouter { + /// Keyed by [`EngineCapabilities::data_source_id`] for O(1) lookup. + engines: HashMap<&'static str, Arc>, +} + +impl EngineRouter { + /// Build an empty router. Use [`Self::register`] to plug engines in. + pub fn new() -> Self { + Self { + engines: HashMap::new(), + } + } + + /// Register an engine. The router keys it by + /// `engine.capabilities().data_source_id`. If two engines claim the + /// same `data_source_id` the later registration wins (matches + /// `HashMap::insert` semantics; tests assume this for hot-swap). + pub fn register(&mut self, engine: Arc) { + let caps = engine.capabilities(); + debug!( + data_source_id = caps.data_source_id, + backend = ?caps.storage_backend, + "router: registering engine", + ); + self.engines.insert(caps.data_source_id, engine); + } + + /// Number of engines registered. Test-only convenience. + pub fn len(&self) -> usize { + self.engines.len() + } + + /// Whether no engines are registered. Test-only convenience. + pub fn is_empty(&self) -> bool { + self.engines.is_empty() + } + + /// Walk the compatible-backend list for `(stat, accuracy, metric_storage)`, + /// dispatch to the first registered engine, and (on + /// [`EngineError::Backend`] or [`EngineError::CapabilityMiss`]) fall + /// through to the next compatible backend. + /// + /// On exhaustion returns either: + /// - [`EngineRouterError::NoEngineRegistered`] if zero of the + /// compatible backends had an engine registered, or + /// - [`EngineRouterError::AllFailed`] if every registered engine in + /// the failover sequence returned an error. + pub async fn execute( + &self, + query: &str, + stat: Statistic, + accuracy: AccuracyTarget, + metric_storage: StorageBackend, + ) -> Result { + let backends = compatible_storage_backends(stat, accuracy, metric_storage); + debug!( + query = query, + stat = ?stat, + accuracy = ?accuracy, + metric_storage = ?metric_storage, + backends = ?backends, + "router: dispatching", + ); + + let mut last_err: Option = None; + let mut any_engine_tried = false; + + for backend in &backends { + let id = backend.data_source_id(); + let Some(engine) = self.engines.get(id) else { + debug!( + backend = ?backend, + data_source_id = id, + "router: no engine registered, trying next failover", + ); + continue; + }; + any_engine_tried = true; + match engine.execute(query).await { + Ok(result) => { + debug!( + backend = ?backend, + "router: dispatch succeeded", + ); + return Ok(result); + } + Err(e) => { + warn!( + backend = ?backend, + error = %e, + "router: engine failed, falling through to next backend", + ); + last_err = Some(e); + } + } + } + + if !any_engine_tried { + return Err(EngineRouterError::NoEngineRegistered { + tried: backends, + registered: self.engines.keys().copied().collect(), + }); + } + + Err(EngineRouterError::AllFailed { + last: last_err.expect("at least one engine ran (any_engine_tried=true)"), + }) + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::engines::query_result::QueryResult; + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// Stub engine that records call counts and returns either a canned + /// vector result or a configured error. Keeps the router tests + /// hermetic — no SimpleEngine / GorillaQueryEngine wire-up needed. + struct StubEngine { + caps: EngineCapabilities, + calls: Arc, + outcome: Outcome, + } + + enum Outcome { + Ok, + Backend, + CapabilityMiss, + } + + impl StubEngine { + fn new(backend: StorageBackend, outcome: Outcome) -> (Arc, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + let engine = Arc::new(Self { + caps: EngineCapabilities { + data_source_id: backend.data_source_id(), + storage_backend: backend, + supports_streams_above_bytes: 1024 * 1024, + }, + calls: calls.clone(), + outcome, + }); + (engine, calls) + } + } + + #[async_trait] + impl QueryEngine for StubEngine { + async fn execute(&self, _query: &str) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + match self.outcome { + Outcome::Ok => Ok(QueryResult::vector(Vec::new(), 0)), + Outcome::Backend => Err(EngineError::backend( + self.caps.data_source_id, + "simulated backend failure", + )), + Outcome::CapabilityMiss => Err(EngineError::capability_miss( + self.caps.data_source_id, + "no compatible aggregation", + )), + } + } + fn capabilities(&self) -> EngineCapabilities { + self.caps + } + } + + #[tokio::test] + async fn router_dispatches_to_warm_tier_for_sketch_metrics() { + let mut router = EngineRouter::new(); + let (warm, warm_calls) = StubEngine::new(StorageBackend::SketchWarmTier, Outcome::Ok); + let (jsonl, jsonl_calls) = + StubEngine::new(StorageBackend::ColdJsonlFallback, Outcome::Ok); + router.register(warm); + router.register(jsonl); + + let result = router + .execute( + "sum_over_time(foo[5m])", + Statistic::Sum, + AccuracyTarget::Approximate, + StorageBackend::SketchWarmTier, + ) + .await; + assert!(result.is_ok()); + assert_eq!(warm_calls.load(Ordering::SeqCst), 1); + assert_eq!(jsonl_calls.load(Ordering::SeqCst), 0, "JSONL must not run when warm-tier succeeds"); + } + + #[tokio::test] + async fn router_dispatches_to_gorilla_for_archive_metrics() { + let mut router = EngineRouter::new(); + let (warm, warm_calls) = StubEngine::new(StorageBackend::SketchWarmTier, Outcome::Ok); + let (gorilla, gorilla_calls) = + StubEngine::new(StorageBackend::GorillaS3Archive, Outcome::Ok); + router.register(warm); + router.register(gorilla); + + let result = router + .execute( + "sum_over_time(audit_events[1h])", + Statistic::Sum, + AccuracyTarget::Exact, + StorageBackend::GorillaS3Archive, + ) + .await; + assert!(result.is_ok()); + assert_eq!(gorilla_calls.load(Ordering::SeqCst), 1); + assert_eq!( + warm_calls.load(Ordering::SeqCst), + 0, + "warm-tier must not run for an archive-only metric", + ); + } + + #[tokio::test] + async fn router_falls_back_to_jsonl_when_archive_fails() { + // Double-write deploy: archive fails, warm-tier fails too, JSONL answers. + let mut router = EngineRouter::new(); + let (gorilla, gorilla_calls) = + StubEngine::new(StorageBackend::GorillaS3Archive, Outcome::Backend); + let (warm, warm_calls) = + StubEngine::new(StorageBackend::SketchWarmTier, Outcome::CapabilityMiss); + let (jsonl, jsonl_calls) = + StubEngine::new(StorageBackend::ColdJsonlFallback, Outcome::Ok); + router.register(gorilla); + router.register(warm); + router.register(jsonl); + + let result = router + .execute( + "sum_over_time(foo[5m])", + Statistic::Sum, + AccuracyTarget::Exact, + StorageBackend::DoubleWrite, + ) + .await; + assert!(result.is_ok(), "router must reach JSONL on archive+warm failure"); + assert_eq!(gorilla_calls.load(Ordering::SeqCst), 1); + assert_eq!(warm_calls.load(Ordering::SeqCst), 1); + assert_eq!(jsonl_calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn router_with_no_engines_errors_cleanly() { + let router = EngineRouter::new(); + let result = router + .execute( + "sum_over_time(foo[5m])", + Statistic::Sum, + AccuracyTarget::Approximate, + StorageBackend::SketchWarmTier, + ) + .await; + match result { + Err(EngineRouterError::NoEngineRegistered { tried, registered }) => { + assert_eq!( + tried, + vec![ + StorageBackend::SketchWarmTier, + StorageBackend::ColdJsonlFallback, + ] + ); + assert!(registered.is_empty()); + } + other => panic!("expected NoEngineRegistered, got {other:?}"), + } + } + + #[tokio::test] + async fn router_returns_all_failed_when_every_engine_errors() { + let mut router = EngineRouter::new(); + let (warm, _) = StubEngine::new(StorageBackend::SketchWarmTier, Outcome::Backend); + let (jsonl, _) = StubEngine::new(StorageBackend::ColdJsonlFallback, Outcome::Backend); + router.register(warm); + router.register(jsonl); + + let result = router + .execute( + "sum_over_time(foo[5m])", + Statistic::Sum, + AccuracyTarget::Approximate, + StorageBackend::SketchWarmTier, + ) + .await; + match result { + Err(EngineRouterError::AllFailed { last }) => { + // The deepest failure (JSONL) is what the caller sees. + assert!(matches!(last, EngineError::Backend { .. })); + } + other => panic!("expected AllFailed, got {other:?}"), + } + } + + #[tokio::test] + async fn register_overwrites_same_data_source_id() { + let mut router = EngineRouter::new(); + let (first, first_calls) = StubEngine::new(StorageBackend::SketchWarmTier, Outcome::Ok); + let (second, second_calls) = StubEngine::new(StorageBackend::SketchWarmTier, Outcome::Ok); + router.register(first); + router.register(second); + assert_eq!(router.len(), 1, "two registrations under same id collapse to one"); + + let _ = router + .execute( + "sum_over_time(foo[5m])", + Statistic::Sum, + AccuracyTarget::Approximate, + StorageBackend::SketchWarmTier, + ) + .await; + assert_eq!(first_calls.load(Ordering::SeqCst), 0); + assert_eq!(second_calls.load(Ordering::SeqCst), 1, "later registration wins"); + } +} diff --git a/asap-query-engine/src/engines/simple_engine.rs b/asap-query-engine/src/engines/simple_engine.rs index 530917ef..8dbbdb06 100644 --- a/asap-query-engine/src/engines/simple_engine.rs +++ b/asap-query-engine/src/engines/simple_engine.rs @@ -4088,6 +4088,49 @@ impl SimpleEngine { } } +// --------------------------------------------------------------------------- +// Phase-5: `QueryEngine` trait impl. +// +// Adapter only — does NOT change `handle_query` or any other existing +// surface. The trait's `execute(&str)` walks the same `handle_query` code +// path the binary's HTTP driver uses today; `None` (capability miss) is +// translated to `EngineError::CapabilityMiss` so the router can fall through +// to the next compatible backend. +// --------------------------------------------------------------------------- + +#[async_trait::async_trait] +impl crate::engines::router::QueryEngine for SimpleEngine { + async fn execute( + &self, + query: &str, + ) -> Result { + // `handle_query` is sync + needs a `time: f64` (epoch millis as float). + // The router doesn't pass a query time, so we use wall-clock now — + // matches `GorillaQueryEngine::execute`'s convention. + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::SystemTime::UNIX_EPOCH) + .map(|d| d.as_millis() as f64) + .unwrap_or(0.0); + match self.handle_query(query.to_string(), now_ms) { + Some((_labels, result)) => Ok(result), + None => Err(crate::engines::EngineError::capability_miss( + asap_types::StorageBackend::SketchWarmTier.data_source_id(), + format!("SimpleEngine has no compatible aggregation for `{query}`"), + )), + } + } + + fn capabilities(&self) -> crate::engines::router::EngineCapabilities { + crate::engines::router::EngineCapabilities { + data_source_id: asap_types::StorageBackend::SketchWarmTier.data_source_id(), + storage_backend: asap_types::StorageBackend::SketchWarmTier, + // Warm-tier sketches are O(sketch-size); call it 16 MiB ceiling + // for buffered ops (KLL with k=200 is well below this). + supports_streams_above_bytes: 16 * 1024 * 1024, + } + } +} + #[cfg(test)] mod range_query_tests { use crate::data_model::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; diff --git a/asap-query-engine/src/tests/capability_matching_tests.rs b/asap-query-engine/src/tests/capability_matching_tests.rs index 789064c4..c1f20c51 100644 --- a/asap-query-engine/src/tests/capability_matching_tests.rs +++ b/asap-query-engine/src/tests/capability_matching_tests.rs @@ -68,6 +68,7 @@ fn engine_no_query_configs( } let streaming_config = Arc::new(StreamingConfig { aggregation_configs: agg_map, + storage_backend: Default::default(), }); let store = Arc::new(SimpleMapStore::new( streaming_config.clone(), @@ -129,6 +130,7 @@ fn engine_with_query_config( agg_map.insert(agg_id, agg_config.clone()); let streaming_config = Arc::new(StreamingConfig { aggregation_configs: agg_map, + storage_backend: Default::default(), }); let store = Arc::new(SimpleMapStore::new( streaming_config.clone(), diff --git a/asap-query-engine/src/tests/sql_pattern_matching_tests.rs b/asap-query-engine/src/tests/sql_pattern_matching_tests.rs index d0b20d12..abe70fee 100644 --- a/asap-query-engine/src/tests/sql_pattern_matching_tests.rs +++ b/asap-query-engine/src/tests/sql_pattern_matching_tests.rs @@ -77,6 +77,7 @@ mod tests { agg_configs.insert(agg_id, agg_config); let streaming_config = Arc::new(StreamingConfig { aggregation_configs: agg_configs, + storage_backend: Default::default(), }); let store = Arc::new(SimpleMapStore::new( diff --git a/asap-query-engine/src/tests/test_utilities/config_builders.rs b/asap-query-engine/src/tests/test_utilities/config_builders.rs index ec691272..cde37b1b 100644 --- a/asap-query-engine/src/tests/test_utilities/config_builders.rs +++ b/asap-query-engine/src/tests/test_utilities/config_builders.rs @@ -211,6 +211,7 @@ impl TestConfigBuilder { // Create StreamingConfig let streaming_config = StreamingConfig { aggregation_configs: self.streaming_configs, + storage_backend: Default::default(), }; (inference_config, Arc::new(streaming_config)) @@ -252,6 +253,7 @@ impl TestConfigBuilder { // Create StreamingConfig let streaming_config = StreamingConfig { aggregation_configs: self.streaming_configs, + storage_backend: Default::default(), }; ( diff --git a/asap-query-engine/src/tests/test_utilities/engine_factories.rs b/asap-query-engine/src/tests/test_utilities/engine_factories.rs index cbaaeae6..f971dd87 100644 --- a/asap-query-engine/src/tests/test_utilities/engine_factories.rs +++ b/asap-query-engine/src/tests/test_utilities/engine_factories.rs @@ -95,6 +95,7 @@ pub fn create_engine_single_pop_with_aggregated( let streaming_config = Arc::new(StreamingConfig { aggregation_configs, + storage_backend: Default::default(), }); let store = Arc::new(SimpleMapStore::new( @@ -215,6 +216,7 @@ pub fn create_engine_dual_input( let streaming_config = Arc::new(StreamingConfig { aggregation_configs, + storage_backend: Default::default(), }); let store = Arc::new(SimpleMapStore::new( @@ -330,6 +332,7 @@ pub fn create_engine_two_metrics( let streaming_config = Arc::new(StreamingConfig { aggregation_configs, + storage_backend: Default::default(), }); let store = Arc::new(SimpleMapStore::new( @@ -434,6 +437,7 @@ pub fn create_engine_three_metrics( let streaming_config = Arc::new(StreamingConfig { aggregation_configs, + storage_backend: Default::default(), }); let store = Arc::new(SimpleMapStore::new( @@ -514,6 +518,7 @@ pub fn create_engine_multi_timestamp( let streaming_config = Arc::new(StreamingConfig { aggregation_configs, + storage_backend: Default::default(), }); let store = Arc::new(SimpleMapStore::new( @@ -594,6 +599,7 @@ pub fn create_engine_multi_timestamp_with_window( let streaming_config = Arc::new(StreamingConfig { aggregation_configs, + storage_backend: Default::default(), }); let store = Arc::new(SimpleMapStore::new( diff --git a/asap-query-engine/tests/inference_yaml_pattern_coverage.rs b/asap-query-engine/tests/inference_yaml_pattern_coverage.rs index ee233948..4ef9f609 100644 --- a/asap-query-engine/tests/inference_yaml_pattern_coverage.rs +++ b/asap-query-engine/tests/inference_yaml_pattern_coverage.rs @@ -188,6 +188,7 @@ fn build_engine( ); let streaming_config = Arc::new(StreamingConfig { aggregation_configs, + storage_backend: Default::default(), }); let store = Arc::new(SimpleMapStore::new( streaming_config.clone(),