diff --git a/Cargo.lock b/Cargo.lock index 41854c3c2a..1682ef95ae 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -199,7 +199,7 @@ dependencies = [ "serde_urlencoded", "smallvec", "socket2", - "time 0.3.20", + "time 0.3.21", "url", ] @@ -706,6 +706,17 @@ dependencies = [ "zstd-safe 5.0.2+zstd.1.5.2", ] +[[package]] +name = "async-recursion" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e97ce7de6cf12de5d7226c73f5ba9811622f4db3a5b91b55c53e987e5f91cba" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.15", +] + [[package]] name = "async-stream" version = "0.3.5" @@ -807,7 +818,7 @@ dependencies = [ "http", "hyper", "ring", - "time 0.3.20", + "time 0.3.21", "tokio", "tower", "tracing", @@ -939,7 +950,7 @@ dependencies = [ "percent-encoding", "regex", "sha2", - "time 0.3.20", + "time 0.3.21", "tracing", ] @@ -1043,7 +1054,7 @@ dependencies = [ "itoa", "num-integer", "ryu", - "time 0.3.20", + "time 0.3.21", ] [[package]] @@ -1645,7 +1656,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e859cd57d0710d9e06c381b550c06e76992472a8c6d527aecd2fc673dcc231fb" dependencies = [ "percent-encoding", - "time 0.3.20", + "time 0.3.21", "version_check", ] @@ -2164,9 +2175,9 @@ checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10" [[package]] name = "dotenv_config" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc8db7cd099cfe5d7b33f170b3d19054045425fe6d0d876defb7912c31d2fb67" +checksum = "72ef519d2aabc15ca3dba4729976066fa23f80187bf2b19d623d24fe1a0ec3ea" dependencies = [ "anyhow", "askama", @@ -3000,9 +3011,9 @@ dependencies = [ [[package]] name = "keccak" -version = "0.1.3" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3afef3b6eff9ce9d8ff9b3601125eec7f0c8cbac7abd14f355d053fa56c98768" +checksum = "8f6d5ed8676d904364de097082f4e7d240b571b67989ced0240f08b7f966f940" dependencies = [ "cpufeatures", ] @@ -3113,9 +3124,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.142" +version = "0.2.144" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a987beff54b60ffa6d51982e1aa1146bc42f19bd26be28b0586f252fccf5317" +checksum = "2b00cc1c228a6782d0f076e7b232802e0c5689d41bb5df366f2a6b6621cfdfe1" [[package]] name = "libm" @@ -4395,9 +4406,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.26" +version = "1.0.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4424af4bf778aae2051a77b60283332f386554255d722233d09fbfc7e30da2fc" +checksum = "8f4f29d145265ec1c483c7c654450edde0bfe043d3938d6972630663356d9500" dependencies = [ "proc-macro2", ] @@ -4894,7 +4905,7 @@ dependencies = [ "serde", "serde_json", "thiserror", - "time 0.3.20", + "time 0.3.21", ] [[package]] @@ -5022,7 +5033,7 @@ dependencies = [ "serde", "serde_json", "thiserror", - "time 0.3.20", + "time 0.3.21", "url", "uuid", ] @@ -5035,18 +5046,18 @@ checksum = "e6b44e8fc93a14e66336d230954dda83d18b4605ccace8fe09bc7514a71ad0bc" [[package]] name = "serde" -version = "1.0.160" +version = "1.0.162" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bb2f3770c8bce3bcda7e149193a069a0f4365bda1fa5cd88e03bca26afc1216c" +checksum = "71b2f6e1ab5c2b98c05f0f35b236b22e8df7ead6ffbf51d7808da7f8817e7ab6" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.160" +version = "1.0.162" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "291a097c63d8497e00160b166a967a4a79c64f3facdd01cbd7502231688d77df" +checksum = "a2a0814352fd64b58489904a44ea8d90cb1a91dcb6b4f5ebabc32c8318e93cb6" dependencies = [ "proc-macro2", "quote", @@ -5123,9 +5134,9 @@ dependencies = [ [[package]] name = "sha3" -version = "0.10.7" +version = "0.10.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54c2bb1a323307527314a36bfb73f24febb08ce2b8a554bf4ffd6f51ad15198c" +checksum = "75872d278a8f37ef87fa0ddbda7802605cb18344497949862c0d4dcb291eba60" dependencies = [ "digest", "keccak", @@ -5489,9 +5500,9 @@ dependencies = [ [[package]] name = "time" -version = "0.3.20" +version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd0cbfecb4d19b5ea75bb31ad904eb5b9fa13f21079c3b92017ebdf4999a5890" +checksum = "8f3403384eaacbca9923fa06940178ac13e4edb725486d70e8e15881d0c836cc" dependencies = [ "itoa", "serde", @@ -5501,15 +5512,15 @@ dependencies = [ [[package]] name = "time-core" -version = "0.1.0" +version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e153e1f1acaef8acc537e68b44906d2db6436e2b35ac2c6b42640fff91f00fd" +checksum = "7300fbefb4dadc1af235a9cef3737cea692a9d97e1b9cbcd4ebdae6f8868e6fb" [[package]] name = "time-macros" -version = "0.2.8" +version = "0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd80a657e71da814b8e5d60d3374fc6d35045062245d80224748ae522dd76f36" +checksum = "372950940a5f07bf38dbe211d7283c9e6d7327df53794992d293e534c733d09b" dependencies = [ "time-core", ] @@ -6075,7 +6086,7 @@ dependencies = [ "getset", "rustversion", "thiserror", - "time 0.3.20", + "time 0.3.21", ] [[package]] @@ -6677,6 +6688,7 @@ dependencies = [ "argon2", "arrow", "arrow-schema", + "async-recursion", "async-trait", "async_once", "awc", @@ -6737,7 +6749,7 @@ dependencies = [ "strum", "sys-info", "thiserror", - "time 0.3.20", + "time 0.3.21", "tokio", "tonic", "tonic-build", @@ -6757,9 +6769,9 @@ dependencies = [ [[package]] name = "zip" -version = "0.6.4" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0445d0fbc924bb93539b4316c11afb121ea39296f99a3c4c9edad09e3658cdef" +checksum = "7e92305c174683d78035cbf1b70e18db6329cc0f1b9cae0a52ca90bf5bfe7125" dependencies = [ "byteorder", "crc32fast", diff --git a/Cargo.toml b/Cargo.toml index 79b2c2f353..954b2a1723 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -48,6 +48,7 @@ arrow = { version = "36", features = ["simd", "ipc_compression"] } arrow-schema = { version = "36", features = ["serde"] } async-trait = "0.1" async_once = "0.2" +async-recursion = "1.0" awc = "3.1" base64 = "0.21" blake3 = "1.3" @@ -117,15 +118,17 @@ reqwest = { version = "0.11", default-features = false, features = [ "rustls-tls", ] } rs-snowflake = "0.6" -rustc-hash = "1.1.0" rust-embed-for-web = "11.1" +rustc-hash = "1.1.0" segment = "0.2" +sentry = { version = "0.31", features = ["log"] } serde = { version = "1", features = ["derive"] } serde_json = "1" simd-json = "0.7" sled = "0.34" snap = "1" sqlparser = { version = "0.32", features = ["serde"] } +strum = { version = "0.24", features = ["derive"] } sys-info = "0.9" thiserror = "1.0" time = "0.3" @@ -146,8 +149,6 @@ vrl-value = { package = "value", git = "https://github.com/zinclabs/vrl", rev = vrl-stdlib = { package = "vrl-stdlib", git = "https://github.com/zinclabs/vrl", rev = "e723f96c9305d77c853d8f95a1955ff393c8aff1" } walkdir = "2" zstd = "0.12" -sentry = { version = "0.31" ,features = ["log"] } -strum = { version = "0.24", features = ["derive"] } [build-dependencies] chrono = { version = "0.4", default-features = false, features = ["clock"] } diff --git a/build.rs b/build.rs index 0f66ace33e..25d934407c 100644 --- a/build.rs +++ b/build.rs @@ -37,10 +37,17 @@ fn main() -> Result<()> { .type_attribute("SearchAggRequest", "#[derive(serde::Serialize)]") .type_attribute("SearchAggResponse", "#[derive(Eq)]") .type_attribute("SearchAggResponse", "#[derive(serde::Serialize)]") + .type_attribute("Series", "#[derive(serde::Serialize)]") + .type_attribute("Label", "#[derive(serde::Serialize)]") + .type_attribute("Sample", "#[derive(serde::Serialize)]") + .type_attribute("MetricsQueryStmt", "#[derive(serde::Serialize)]") + .type_attribute("MetricsQueryRequest", "#[derive(serde::Serialize)]") + .type_attribute("MetricsQueryResponse", "#[derive(serde::Serialize)]") .compile( &[ "proto/cluster/common.proto", "proto/cluster/event.proto", + "proto/cluster/metrics.proto", "proto/cluster/search.proto", ], &["proto"], diff --git a/deploy/build/buildspec-test.yml b/deploy/build/buildspec-test.yml index 95ced75704..14edaf4fb6 100644 --- a/deploy/build/buildspec-test.yml +++ b/deploy/build/buildspec-test.yml @@ -5,7 +5,9 @@ phases: commands: - apt-get update - apt-get install -y bc - - apt-get install -y protobuf-compiler + - apt-get install -y zip unzip wget + - wget https://github.com/protocolbuffers/protobuf/releases/download/v21.12/protoc-21.12-linux-x86_64.zip + - unzip protoc-21.12-linux-x86_64.zip -d protoc && cp protoc/bin/protoc /usr/local/bin/ && cp -r protoc/include/google /usr/local/include/ - curl https://sh.rustup.rs -sSf | sh -s -- -y - . "$HOME/.cargo/env" - rustup toolchain install nightly-2023-03-28 diff --git a/deploy/build/buildspec.yml b/deploy/build/buildspec.yml index 054fd87f84..d10eab9eff 100644 --- a/deploy/build/buildspec.yml +++ b/deploy/build/buildspec.yml @@ -32,12 +32,12 @@ phases: - docker push public.ecr.aws/zinclabs/zincobserve-dev:latest-amd64 # build zincobserve-mimalloc - - docker pull public.ecr.aws/zinclabs/rust:bullseye-sccache - - docker build -t zincobserve:latest-amd64-mimalloc -f deploy/build/Dockerfile.mimalloc.amd64 . - - docker tag zincobserve:latest-amd64-mimalloc public.ecr.aws/zinclabs/zincobserve-dev:$GIT_TAG-$GIT_HASH-amd64-mimalloc - - docker tag zincobserve:latest-amd64-mimalloc public.ecr.aws/zinclabs/zincobserve-dev:latest-amd64-mimalloc - - docker push public.ecr.aws/zinclabs/zincobserve-dev:$GIT_TAG-$GIT_HASH-amd64-mimalloc - - docker push public.ecr.aws/zinclabs/zincobserve-dev:latest-amd64-mimalloc + # - docker pull public.ecr.aws/zinclabs/rust:bullseye-sccache + # - docker build -t zincobserve:latest-amd64-mimalloc -f deploy/build/Dockerfile.mimalloc.amd64 . + # - docker tag zincobserve:latest-amd64-mimalloc public.ecr.aws/zinclabs/zincobserve-dev:$GIT_TAG-$GIT_HASH-amd64-mimalloc + # - docker tag zincobserve:latest-amd64-mimalloc public.ecr.aws/zinclabs/zincobserve-dev:latest-amd64-mimalloc + # - docker push public.ecr.aws/zinclabs/zincobserve-dev:$GIT_TAG-$GIT_HASH-amd64-mimalloc + # - docker push public.ecr.aws/zinclabs/zincobserve-dev:latest-amd64-mimalloc cache: paths: diff --git a/proto/cluster/common.proto b/proto/cluster/common.proto index ce19427f63..75088fc727 100644 --- a/proto/cluster/common.proto +++ b/proto/cluster/common.proto @@ -1,7 +1,7 @@ syntax = "proto3"; option java_multiple_files = true; -option java_package = "org.zinc.oxide.cluster"; +option java_package = "org.zinc.observe.cluster"; option java_outer_classname = "clusterProto"; package cluster; @@ -27,4 +27,19 @@ enum StreamType { FILELIST = 4; } -message EmptyResponse {} \ No newline at end of file +// Search request type +enum SearchType { + USER = 0; // user input search request + CLUSTER = 1; // cluster dispatch search request + CACHE_ONLY = 2; // ingester node just search local cache +} + +// Job information for a request +message Job { + string session_id = 1; + string job = 2; + int32 stage = 3; + int32 partition = 4; +} + +message EmptyResponse {} diff --git a/proto/cluster/event.proto b/proto/cluster/event.proto index bd749e6f06..1c00597777 100644 --- a/proto/cluster/event.proto +++ b/proto/cluster/event.proto @@ -1,7 +1,7 @@ syntax = "proto3"; option java_multiple_files = true; -option java_package = "org.zinc.oxide.cluster"; +option java_package = "org.zinc.observe.cluster"; option java_outer_classname = "clusterProto"; package cluster; @@ -20,6 +20,4 @@ message FileKey { string key = 1; FileMeta meta = 2; bool deleted = 3; - } - \ No newline at end of file diff --git a/proto/cluster/metrics.proto b/proto/cluster/metrics.proto new file mode 100644 index 0000000000..5a80df65b6 --- /dev/null +++ b/proto/cluster/metrics.proto @@ -0,0 +1,57 @@ +syntax = "proto3"; + +option java_multiple_files = true; +option java_package = "org.zinc.observe.cluster"; +option java_outer_classname = "clusterProto"; + +package cluster; + +import "cluster/common.proto"; + +service Metrics { + rpc Query (MetricsQueryRequest) returns (MetricsQueryResponse) {} +} + +// Search request query +message MetricsQueryStmt { + string query = 1; + bool is_range_query = 2; + int64 start = 3; + int64 end = 4; + int64 step = 5; +} + +// Search request +message MetricsQueryRequest { + Job job = 1; + string org_id = 2; + SearchType stype = 3; + MetricsQueryStmt query = 4; +} + +// The response message containing the greetings +message MetricsQueryResponse { + Job job = 1; + int32 took = 2; + int32 file_count = 3; + int32 scan_size = 4; // unit: mb + string result_type = 5; // vector, matrix, scalar + repeated Series result = 6; +} + +message Series { + repeated Label metric = 1; + repeated Sample values = 2; + optional Sample value = 3; + optional double scalar = 4; +} + +message Label { + string name = 1; + string value = 2; +} + +message Sample { + int64 time = 1; + double value = 2; +} diff --git a/proto/cluster/search.proto b/proto/cluster/search.proto index 9405917533..2dacdb327f 100644 --- a/proto/cluster/search.proto +++ b/proto/cluster/search.proto @@ -1,31 +1,17 @@ syntax = "proto3"; option java_multiple_files = true; -option java_package = "org.zinc.oxide.cluster"; +option java_package = "org.zinc.observe.cluster"; option java_outer_classname = "clusterProto"; package cluster; +import "cluster/common.proto"; service Search { rpc Search (SearchRequest) returns (SearchResponse) {} } -// Job information for a request -message Job { - string session_id = 1; - string job = 2; - int32 stage = 3; - int32 partition = 4; -} - -// Search request type -enum SearchType { - USER = 0; // user input search request - CLUSTER = 1; // cluster dispatch search request - CACHE_ONLY = 2; // ingester node just search local cache -} - // Search request query message SearchQuery { string sql = 1; @@ -43,13 +29,13 @@ message SearchQuery { // Search request message SearchRequest { - Job job = 1; - string org_id = 2; - string stream_type = 3; - SearchType stype = 4; - SearchQuery query = 5; - repeated string file_list = 6; - repeated SearchAggRequest aggs = 7; + Job job = 1; + string org_id = 2; + string stream_type = 3; + SearchType stype = 4; + SearchQuery query = 5; + repeated string file_list = 6; + repeated SearchAggRequest aggs = 7; } // The response message containing the greetings diff --git a/src/handler/grpc/mod.rs b/src/handler/grpc/mod.rs index 6825f2352c..53496496b4 100644 --- a/src/handler/grpc/mod.rs +++ b/src/handler/grpc/mod.rs @@ -15,6 +15,7 @@ use uuid::Uuid; use crate::meta; +use crate::service::promql; pub mod auth; pub mod request; @@ -97,6 +98,67 @@ impl From<&meta::common::FileKey> for cluster_rpc::FileKey { } } +impl From for cluster_rpc::MetricsQueryRequest { + fn from(req: promql::MetricsQueryRequest) -> Self { + let req_query = cluster_rpc::MetricsQueryStmt { + query: req.query.to_owned(), + is_range_query: req.is_range_query, + start: req.start, + end: req.end, + step: req.step, + }; + + let job = cluster_rpc::Job { + session_id: Uuid::new_v4().to_string(), + job: "".to_string(), + stage: 0, + partition: 0, + }; + + cluster_rpc::MetricsQueryRequest { + job: Some(job), + org_id: "".to_string(), + stype: cluster_rpc::SearchType::User.into(), + query: Some(req_query), + } + } +} + +impl From<&cluster_rpc::Label> for promql::value::Label { + fn from(req: &cluster_rpc::Label) -> Self { + promql::value::Label { + name: req.name.to_owned(), + value: req.value.to_owned(), + } + } +} + +impl From<&promql::value::Label> for cluster_rpc::Label { + fn from(req: &promql::value::Label) -> Self { + cluster_rpc::Label { + name: req.name.to_owned(), + value: req.value.to_owned(), + } + } +} + +impl From<&cluster_rpc::Sample> for promql::value::Sample { + fn from(req: &cluster_rpc::Sample) -> Self { + promql::value::Sample { + timestamp: req.time, + value: req.value, + } + } +} +impl From<&promql::value::Sample> for cluster_rpc::Sample { + fn from(req: &promql::value::Sample) -> Self { + cluster_rpc::Sample { + time: req.timestamp, + value: req.value, + } + } +} + #[cfg(test)] mod test { use std::collections::HashMap; diff --git a/src/handler/grpc/request/event.rs b/src/handler/grpc/request/event.rs index 852737dd0a..54812bb04a 100644 --- a/src/handler/grpc/request/event.rs +++ b/src/handler/grpc/request/event.rs @@ -13,7 +13,6 @@ // limitations under the License. use opentelemetry::global; -use std::time::Instant; use tonic::{Request, Response, Status}; use tracing_opentelemetry::OpenTelemetrySpanExt; @@ -34,7 +33,7 @@ impl Event for Eventer { &self, req: Request, ) -> Result, Status> { - let start = Instant::now(); + let start = std::time::Instant::now(); let parent_cx = global::get_text_map_propagator(|prop| { prop.extract(&super::MetadataMap(req.metadata())) }); diff --git a/src/handler/grpc/request/metrics.rs b/src/handler/grpc/request/metrics.rs new file mode 100644 index 0000000000..5f016fad3f --- /dev/null +++ b/src/handler/grpc/request/metrics.rs @@ -0,0 +1,75 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use opentelemetry::global; +use tonic::{Request, Response, Status}; +use tracing_opentelemetry::OpenTelemetrySpanExt; + +use crate::handler::grpc::cluster_rpc::metrics_server::Metrics; +use crate::handler::grpc::cluster_rpc::MetricsQueryRequest; +use crate::handler::grpc::cluster_rpc::MetricsQueryResponse; +use crate::infra::errors; +use crate::infra::metrics; +use crate::meta; +use crate::service::promql::search as SearchService; + +#[derive(Default)] +pub struct Querier {} + +#[tonic::async_trait] +impl Metrics for Querier { + #[tracing::instrument(name = "grpc:metrics:enter", skip_all)] + async fn query( + &self, + req: Request, + ) -> Result, Status> { + let start = std::time::Instant::now(); + let parent_cx = global::get_text_map_propagator(|prop| { + prop.extract(&super::MetadataMap(req.metadata())) + }); + tracing::Span::current().set_parent(parent_cx); + + let req = req.get_ref(); + let org_id = req.org_id.clone(); + let stream_type = meta::StreamType::Metrics.to_string(); + let result = match SearchService::grpc::search(req).await { + Ok(res) => res, + Err(err) => { + // metrics + let time = start.elapsed().as_secs_f64(); + metrics::GRPC_RESPONSE_TIME + .with_label_values(&["/prom/v1/query", "500", &org_id, "", &stream_type]) + .observe(time); + metrics::GRPC_INCOMING_REQUESTS + .with_label_values(&["/prom/v1/query", "500", &org_id, "", &stream_type]) + .inc(); + return Err(match err { + errors::Error::ErrorCode(code) => Status::internal(code.to_json()), + _ => Status::internal(err.to_string()), + }); + } + }; + + // metrics + let time = start.elapsed().as_secs_f64(); + metrics::GRPC_RESPONSE_TIME + .with_label_values(&["/prom/v1/query", "200", &org_id, "", &stream_type]) + .observe(time); + metrics::GRPC_INCOMING_REQUESTS + .with_label_values(&["/prom/v1/query", "200", &org_id, "", &stream_type]) + .inc(); + + Ok(Response::new(result)) + } +} diff --git a/src/handler/grpc/request/mod.rs b/src/handler/grpc/request/mod.rs index 148961ff2b..fa1c70e072 100644 --- a/src/handler/grpc/request/mod.rs +++ b/src/handler/grpc/request/mod.rs @@ -15,6 +15,7 @@ use opentelemetry::propagation::Extractor; pub mod event; +pub mod metrics; pub mod search; pub mod traces; diff --git a/src/handler/grpc/request/search.rs b/src/handler/grpc/request/search.rs index c2571c569a..b0077a49b7 100644 --- a/src/handler/grpc/request/search.rs +++ b/src/handler/grpc/request/search.rs @@ -13,7 +13,6 @@ // limitations under the License. use opentelemetry::global; -use std::time::Instant; use tonic::{Request, Response, Status}; use tracing_opentelemetry::OpenTelemetrySpanExt; @@ -29,12 +28,12 @@ pub struct Searcher {} #[tonic::async_trait] impl Search for Searcher { - #[tracing::instrument(name = "grpc:search:enter", skip(self, req))] + #[tracing::instrument(name = "grpc:search:enter", skip_all)] async fn search( &self, req: Request, ) -> Result, Status> { - let start = Instant::now(); + let start = std::time::Instant::now(); let parent_cx = global::get_text_map_propagator(|prop| { prop.extract(&super::MetadataMap(req.metadata())) }); @@ -43,7 +42,7 @@ impl Search for Searcher { let req = req.get_ref(); let org_id = req.org_id.clone(); let stream_type = req.stream_type.as_str(); - let result = match SearchService::exec::search(req).await { + let result = match SearchService::grpc::search(req).await { Ok(res) => res, Err(err) => { // metrics diff --git a/src/handler/http/request/prom/mod.rs b/src/handler/http/request/prom/mod.rs index 448beaacd3..0acb5f4310 100644 --- a/src/handler/http/request/prom/mod.rs +++ b/src/handler/http/request/prom/mod.rs @@ -15,10 +15,11 @@ use actix_web::{get, http, post, web, HttpRequest, HttpResponse}; use std::io::Error; -use crate::{ - meta::{self, http::HttpResponse as MetaHttpResponse}, - service::metrics, -}; +use crate::common::time::{parse_milliseconds, parse_str_to_timestamp_micros}; +use crate::infra::errors; +use crate::meta::{self, http::HttpResponse as MetaHttpResponse}; +use crate::service::metrics; +use crate::service::promql; /** prometheus remote-write endpoint for metrics */ #[utoipa::path( @@ -105,8 +106,54 @@ pub async fn query( query: web::Query, ) -> Result { let query = query.into_inner(); - tracing::trace!(%org_id, ?query); - Ok(HttpResponse::NotImplemented().body("Not Implemented")) + let start = match query.time { + None => chrono::Utc::now().timestamp_micros(), + Some(v) => match parse_str_to_timestamp_micros(&v) { + Ok(v) => v, + Err(e) => { + log::error!("parse time error: {}", e); + return Ok(HttpResponse::BadRequest().json(promql::QueryResponse { + status: promql::Status::Error, + data: None, + error_type: Some("bad_data".to_string()), + error: Some(e.to_string()), + })); + } + }, + }; + let end = start; + + let req = promql::MetricsQueryRequest { + query: query.query, + is_range_query: false, + start, + end, + step: 300_000_000, + }; + + match promql::search::search(&org_id, &req).await { + Ok(data) => Ok(HttpResponse::Ok().json(promql::QueryResponse { + status: promql::Status::Success, + data: Some(promql::QueryResult { + result_type: data.get_type().to_string(), + result: data, + }), + error_type: None, + error: None, + })), + Err(err) => { + let err = match err { + errors::Error::ErrorCode(code) => code.get_error_detail(), + _ => err.to_string(), + }; + Ok(HttpResponse::BadRequest().json(promql::QueryResponse { + status: promql::Status::Error, + data: None, + error_type: Some("bad_data".to_string()), + error: Some(err), + })) + } + } } /** prometheus range queries */ @@ -168,8 +215,77 @@ pub async fn query_range( range_query: web::Query, ) -> Result { let range_query = range_query.into_inner(); - tracing::trace!(%org_id, ?range_query); - Ok(HttpResponse::NotImplemented().body("Not Implemented")) + let start = match parse_str_to_timestamp_micros(&range_query.start) { + Ok(v) => v, + Err(e) => { + log::error!("parse time error: {}", e); + return Ok(HttpResponse::BadRequest().json(promql::QueryResponse { + status: promql::Status::Error, + data: None, + error_type: Some("bad_data".to_string()), + error: Some(e.to_string()), + })); + } + }; + let end = match parse_str_to_timestamp_micros(&range_query.end) { + Ok(v) => v, + Err(e) => { + log::error!("parse time error: {}", e); + return Ok(HttpResponse::BadRequest().json(promql::QueryResponse { + status: promql::Status::Error, + data: None, + error_type: Some("bad_data".to_string()), + error: Some(e.to_string()), + })); + } + }; + let step = match range_query.step { + None => 300_000_000, + Some(v) => match parse_milliseconds(&v) { + Ok(v) => (v * 1_000) as i64, + Err(e) => { + log::error!("parse time error: {}", e); + return Ok(HttpResponse::BadRequest().json(promql::QueryResponse { + status: promql::Status::Error, + data: None, + error_type: Some("bad_data".to_string()), + error: Some(e.to_string()), + })); + } + }, + }; + + let req = promql::MetricsQueryRequest { + query: range_query.query, + is_range_query: true, + start, + end, + step, + }; + + match promql::search::search(&org_id, &req).await { + Ok(data) => Ok(HttpResponse::Ok().json(promql::QueryResponse { + status: promql::Status::Success, + data: Some(promql::QueryResult { + result_type: data.get_type().to_string(), + result: data, + }), + error_type: None, + error: None, + })), + Err(err) => { + let err = match err { + errors::Error::ErrorCode(code) => code.get_error_detail(), + _ => err.to_string(), + }; + Ok(HttpResponse::BadRequest().json(promql::QueryResponse { + status: promql::Status::Error, + data: None, + error_type: Some("bad_data".to_string()), + error: Some(err), + })) + } + } } /** prometheus query metric metadata */ diff --git a/src/handler/http/request/search/mod.rs b/src/handler/http/request/search/mod.rs index 7907e2d08c..754b45983c 100644 --- a/src/handler/http/request/search/mod.rs +++ b/src/handler/http/request/search/mod.rs @@ -17,7 +17,6 @@ use actix_web::{get, post, web, HttpRequest, HttpResponse}; use chrono::Duration; use std::collections::HashMap; use std::io::Error; -use std::time::Instant; use crate::common::base64; use crate::common::functions; @@ -100,7 +99,7 @@ pub async fn search( in_req: HttpRequest, body: actix_web::web::Bytes, ) -> Result { - let start = Instant::now(); + let start = std::time::Instant::now(); let org_id = org_id.into_inner(); let query = web::Query::>::from_query(in_req.query_string()).unwrap(); let stream_type = match get_stream_type_from_request(&query) { @@ -237,7 +236,7 @@ pub async fn around( path: web::Path<(String, String)>, in_req: HttpRequest, ) -> Result { - let start = Instant::now(); + let start = std::time::Instant::now(); let mut uses_fn = false; let (org_id, stream_name) = path.into_inner(); let query = web::Query::>::from_query(in_req.query_string()).unwrap(); @@ -465,7 +464,7 @@ pub async fn values( path: web::Path<(String, String)>, in_req: HttpRequest, ) -> Result { - let start = Instant::now(); + let start = std::time::Instant::now(); let mut uses_fn = false; let (org_id, stream_name) = path.into_inner(); let query = web::Query::>::from_query(in_req.query_string()).unwrap(); diff --git a/src/infra/config.rs b/src/infra/config.rs index 2162259280..21ad26c078 100644 --- a/src/infra/config.rs +++ b/src/infra/config.rs @@ -37,7 +37,7 @@ pub static BUILD_DATE: &str = env!("GIT_BUILD_DATE"); pub static HAS_FUNCTIONS: bool = true; #[cfg(not(feature = "zo_functions"))] pub static HAS_FUNCTIONS: bool = false; -pub static SEARCHING_IN_CACHE: AtomicU8 = AtomicU8::new(0); +pub static SEARCHING_IN_WAL: AtomicU8 = AtomicU8::new(0); pub static CONFIG: Lazy = Lazy::new(init); pub static INSTANCE_ID: Lazy> = Lazy::new(DashMap::new); diff --git a/src/job/files/disk.rs b/src/job/files/disk.rs index 50bf1629dd..f84ac2c0ba 100644 --- a/src/job/files/disk.rs +++ b/src/job/files/disk.rs @@ -151,7 +151,7 @@ async fn move_files_to_storage() -> Result<(), anyhow::Error> { match db::file_list::local::set(&key, meta, false).await { Ok(_) => { loop { - let searching = config::SEARCHING_IN_CACHE.load(Ordering::Relaxed); + let searching = config::SEARCHING_IN_WAL.load(Ordering::Relaxed); if searching == 0 { break; } else { diff --git a/src/main.rs b/src/main.rs index 88cad57940..ff4544713f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -30,8 +30,11 @@ use tracing_subscriber::prelude::*; use tracing_subscriber::Registry; use zincobserve::handler::grpc::auth::check_auth; use zincobserve::handler::grpc::cluster_rpc::event_server::EventServer; +use zincobserve::handler::grpc::cluster_rpc::metrics_server::MetricsServer; use zincobserve::handler::grpc::cluster_rpc::search_server::SearchServer; -use zincobserve::handler::grpc::request::{event::Eventer, search::Searcher, traces::TraceServer}; +use zincobserve::handler::grpc::request::{ + event::Eventer, metrics::Querier, search::Searcher, traces::TraceServer, +}; use zincobserve::handler::http::router::{ get_basic_routes, get_other_service_routes, get_service_routes, }; @@ -53,6 +56,7 @@ async fn main() -> Result<(), anyhow::Error> { if cli().await? { return Ok(()); } + let _guard; if CONFIG.common.tracing_enabled { let service_name = format!("zo-{}", CONFIG.common.instance_name); @@ -110,12 +114,13 @@ async fn main() -> Result<(), anyhow::Error> { rx.await?; if !router::is_router() { let gaddr: SocketAddr = format!("0.0.0.0:{}", CONFIG.grpc.port).parse()?; - let searcher = Searcher::default(); - let search_svc = SearchServer::new(searcher) + let event_svc = EventServer::new(Eventer::default()) .send_compressed(CompressionEncoding::Gzip) .accept_compressed(CompressionEncoding::Gzip); - let eventer = Eventer::default(); - let event_svc = EventServer::new(eventer) + let search_svc = SearchServer::new(Searcher::default()) + .send_compressed(CompressionEncoding::Gzip) + .accept_compressed(CompressionEncoding::Gzip); + let metrics_svc = MetricsServer::new(Querier::default()) .send_compressed(CompressionEncoding::Gzip) .accept_compressed(CompressionEncoding::Gzip); let tracer = TraceServer::default(); @@ -125,8 +130,9 @@ async fn main() -> Result<(), anyhow::Error> { log::info!("starting gRPC server at {}", gaddr); tonic::transport::Server::builder() .layer(tonic::service::interceptor(check_auth)) - .add_service(search_svc) .add_service(event_svc) + .add_service(search_svc) + .add_service(metrics_svc) .add_service(trace_svc) .serve(gaddr) .await diff --git a/src/service/logs/mod.rs b/src/service/logs/mod.rs index 0c784ae596..07914c7c21 100644 --- a/src/service/logs/mod.rs +++ b/src/service/logs/mod.rs @@ -59,13 +59,6 @@ fn parse_bulk_index(v: &Value) -> Option<(String, String, String)> { None } -// generate partition key for query -pub fn get_partition_key_query(s: &str) -> String { - let mut s = s.replace(['/', '.'], "_"); - s.truncate(100); - s -} - pub fn cast_to_type(mut value: Value, delta: Vec) -> (Option, Option) { let local_map = value.as_object_mut().unwrap(); //let mut error_msg = String::new(); diff --git a/src/service/mod.rs b/src/service/mod.rs index 559b9680ff..1d4341c0a9 100644 --- a/src/service/mod.rs +++ b/src/service/mod.rs @@ -24,6 +24,7 @@ pub mod kv; pub mod logs; pub mod metrics; pub mod organization; +pub mod promql; pub mod router; pub mod schema; pub mod search; @@ -31,3 +32,10 @@ pub mod stream; pub mod traces; pub mod triggers; pub mod users; + +// generate partition key for query +pub fn get_partition_key_query(s: &str) -> String { + let mut s = s.replace(['/', '.'], "_"); + s.truncate(100); + s +} diff --git a/src/service/promql/aggregations/avg.rs b/src/service/promql/aggregations/avg.rs new file mode 100644 index 0000000000..e65b3e0cbc --- /dev/null +++ b/src/service/promql/aggregations/avg.rs @@ -0,0 +1,37 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; +use promql_parser::parser::LabelModifier; + +use crate::service::promql::value::{InstantValue, Sample, Value}; + +pub fn avg(timestamp: i64, param: &Option, data: &Value) -> Result { + let score_values = super::eval_arithmetic(param, data, "avg", |total, val| total + val)?; + if score_values.is_none() { + return Ok(Value::None); + } + let values = score_values + .unwrap() + .values() + .map(|v| InstantValue { + labels: v.labels.clone(), + value: Sample { + timestamp, + value: v.value / v.num as f64, + }, + }) + .collect::>(); + Ok(Value::Vector(values)) +} diff --git a/src/service/promql/aggregations/bottomk.rs b/src/service/promql/aggregations/bottomk.rs new file mode 100644 index 0000000000..5c08a30de1 --- /dev/null +++ b/src/service/promql/aggregations/bottomk.rs @@ -0,0 +1,23 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; +use promql_parser::parser::Expr as PromExpr; + +use super::QueryEngine; +use crate::service::promql::value::Value; + +pub async fn bottomk(ctx: &mut QueryEngine, param: Box, data: &Value) -> Result { + super::eval_top(ctx, param, data, true).await +} diff --git a/src/service/promql/aggregations/count.rs b/src/service/promql/aggregations/count.rs new file mode 100644 index 0000000000..1aff5adb3f --- /dev/null +++ b/src/service/promql/aggregations/count.rs @@ -0,0 +1,37 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; +use promql_parser::parser::LabelModifier; + +use crate::service::promql::value::{InstantValue, Sample, Value}; + +pub fn count(timestamp: i64, param: &Option, data: &Value) -> Result { + let score_values = super::eval_arithmetic(param, data, "count", |_prev, _val| 0.0)?; + if score_values.is_none() { + return Ok(Value::None); + } + let values = score_values + .unwrap() + .values() + .map(|v| InstantValue { + labels: v.labels.clone(), + value: Sample { + timestamp, + value: v.num as f64, + }, + }) + .collect::>(); + Ok(Value::Vector(values)) +} diff --git a/src/service/promql/aggregations/max.rs b/src/service/promql/aggregations/max.rs new file mode 100644 index 0000000000..b17771b190 --- /dev/null +++ b/src/service/promql/aggregations/max.rs @@ -0,0 +1,49 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; +use promql_parser::parser::LabelModifier; + +use crate::service::promql::value::{InstantValue, Sample, Value}; + +pub fn max(timestamp: i64, param: &Option, data: &Value) -> Result { + let score_values = + super::eval_arithmetic( + param, + data, + "max", + |prev, val| { + if prev >= val { + prev + } else { + val + } + }, + )?; + if score_values.is_none() { + return Ok(Value::None); + } + let values = score_values + .unwrap() + .values() + .map(|v| InstantValue { + labels: v.labels.clone(), + value: Sample { + timestamp, + value: v.value, + }, + }) + .collect::>(); + Ok(Value::Vector(values)) +} diff --git a/src/service/promql/aggregations/min.rs b/src/service/promql/aggregations/min.rs new file mode 100644 index 0000000000..ffa8b6efc4 --- /dev/null +++ b/src/service/promql/aggregations/min.rs @@ -0,0 +1,43 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; +use promql_parser::parser::LabelModifier; + +use crate::service::promql::value::{InstantValue, Sample, Value}; + +pub fn min(timestamp: i64, param: &Option, data: &Value) -> Result { + let score_values = super::eval_arithmetic(param, data, "min", |prev, val| { + if prev > 0.0 && prev <= val { + prev + } else { + val + } + })?; + if score_values.is_none() { + return Ok(Value::None); + } + let values = score_values + .unwrap() + .values() + .map(|v| InstantValue { + labels: v.labels.clone(), + value: Sample { + timestamp, + value: v.value, + }, + }) + .collect::>(); + Ok(Value::Vector(values)) +} diff --git a/src/service/promql/aggregations/mod.rs b/src/service/promql/aggregations/mod.rs new file mode 100644 index 0000000000..99959b9b9c --- /dev/null +++ b/src/service/promql/aggregations/mod.rs @@ -0,0 +1,174 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::{DataFusionError, Result}; +use promql_parser::parser::Expr as PromExpr; +use promql_parser::parser::LabelModifier; +use rustc_hash::FxHashMap; + +use super::QueryEngine; +use crate::service::promql::value::{signature, Labels, Signature, Value}; + +mod avg; +mod bottomk; +mod count; +mod max; +mod min; +mod sum; +mod topk; + +pub(crate) use avg::avg; +pub(crate) use bottomk::bottomk; +pub(crate) use count::count; +pub(crate) use max::max; +pub(crate) use min::min; +pub(crate) use sum::sum; +pub(crate) use topk::topk; + +pub(crate) struct ArithmeticItem { + pub(crate) labels: Labels, + pub(crate) value: f64, + pub(crate) num: usize, +} + +#[derive(Debug, Clone)] +pub(crate) struct TopItem { + pub(crate) index: usize, + pub(crate) value: f64, +} + +pub(crate) fn eval_arithmetic( + param: &Option, + data: &Value, + f_name: &str, + f_handler: fn(total: f64, val: f64) -> f64, +) -> Result>> { + let data = match data { + Value::Vector(v) => v, + Value::None => return Ok(None), + _ => { + return Err(DataFusionError::Plan(format!( + "[{f_name}] function only accept vector values" + ))) + } + }; + + let mut score_values = FxHashMap::default(); + match param { + Some(v) => match v { + LabelModifier::Include(labels) => { + for item in data.iter() { + let mut sum_labels = Labels::default(); + for label in item.labels.iter() { + if labels.contains(&label.name) { + sum_labels.push(label.clone()); + } + } + let sum_hash = signature(&sum_labels); + let entry = score_values.entry(sum_hash).or_insert(ArithmeticItem { + labels: sum_labels, + value: 0.0, + num: 0, + }); + entry.value = f_handler(entry.value, item.value.value); + entry.num += 1; + } + } + LabelModifier::Exclude(labels) => { + for item in data.iter() { + let mut sum_labels = Labels::default(); + for label in item.labels.iter() { + if !labels.contains(&label.name) { + sum_labels.push(label.clone()); + } + } + let sum_hash = signature(&sum_labels); + let entry = score_values.entry(sum_hash).or_insert(ArithmeticItem { + labels: sum_labels, + value: 0.0, + num: 0, + }); + entry.value = f_handler(entry.value, item.value.value); + entry.num += 1; + } + } + }, + None => { + for item in data.iter() { + let entry = score_values + .entry(Signature::default()) + .or_insert(ArithmeticItem { + labels: Labels::default(), + value: 0.0, + num: 0, + }); + entry.value = f_handler(entry.value, item.value.value); + entry.num += 1; + } + } + } + Ok(Some(score_values)) +} + +pub async fn eval_top( + ctx: &mut QueryEngine, + param: Box, + data: &Value, + is_bottom: bool, +) -> Result { + let fn_name = if is_bottom { "bottomk" } else { "topk" }; + + let param = ctx.exec_expr(¶m).await?; + let n = match param { + Value::Float(v) => v as usize, + _ => { + return Err(DataFusionError::Plan(format!( + "[{fn_name}] param must be NumberLiteral" + ))) + } + }; + + let data = match data { + Value::Vector(v) => v, + Value::None => return Ok(Value::None), + _ => { + return Err(DataFusionError::Plan(format!( + "[{fn_name}] function only accept vector values" + ))) + } + }; + + let mut score_values = Vec::new(); + for (i, item) in data.iter().enumerate() { + if item.value.value.is_nan() { + continue; + } + score_values.push(TopItem { + index: i, + value: item.value.value, + }); + } + + if is_bottom { + score_values.sort_by(|a, b| a.value.partial_cmp(&b.value).unwrap()); + } else { + score_values.sort_by(|a, b| b.value.partial_cmp(&a.value).unwrap()); + } + let values = score_values + .iter() + .take(n) + .map(|v| data[v.index].clone()) + .collect(); + Ok(Value::Vector(values)) +} diff --git a/src/service/promql/aggregations/sum.rs b/src/service/promql/aggregations/sum.rs new file mode 100644 index 0000000000..ac8bb48bff --- /dev/null +++ b/src/service/promql/aggregations/sum.rs @@ -0,0 +1,37 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; +use promql_parser::parser::LabelModifier; + +use crate::service::promql::value::{InstantValue, Sample, Value}; + +pub fn sum(timestamp: i64, param: &Option, data: &Value) -> Result { + let score_values = super::eval_arithmetic(param, data, "sum", |total, val| total + val)?; + if score_values.is_none() { + return Ok(Value::None); + } + let values = score_values + .unwrap() + .values() + .map(|v| InstantValue { + labels: v.labels.clone(), + value: Sample { + timestamp, + value: v.value, + }, + }) + .collect::>(); + Ok(Value::Vector(values)) +} diff --git a/src/service/promql/aggregations/topk.rs b/src/service/promql/aggregations/topk.rs new file mode 100644 index 0000000000..4c1b24ad9a --- /dev/null +++ b/src/service/promql/aggregations/topk.rs @@ -0,0 +1,23 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; +use promql_parser::parser::Expr as PromExpr; + +use super::QueryEngine; +use crate::service::promql::value::Value; + +pub async fn topk(ctx: &mut QueryEngine, param: Box, data: &Value) -> Result { + super::eval_top(ctx, param, data, false).await +} diff --git a/src/service/promql/engine.rs b/src/service/promql/engine.rs new file mode 100644 index 0000000000..d1d8ea9ae8 --- /dev/null +++ b/src/service/promql/engine.rs @@ -0,0 +1,805 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use async_recursion::async_recursion; +use async_trait::async_trait; +use datafusion::{ + arrow::array::{Float64Array, Int64Array, StringArray}, + error::{DataFusionError, Result}, + prelude::{col, lit, SessionContext}, +}; +use promql_parser::{ + label::MatchOp, + parser::{ + token, AggregateExpr, Call, EvalStmt, Expr as PromExpr, Function, FunctionArgs, + LabelModifier, MatrixSelector, NumberLiteral, ParenExpr, TokenType, UnaryExpr, + VectorSelector, + }, +}; +use rustc_hash::FxHashMap; +use std::{ + str::FromStr, + sync::Arc, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use super::{aggregations, functions, value::*}; + +#[async_trait] +pub trait TableProvider: Sync + Send + 'static { + async fn create_context( + &self, + stream_name: &str, + time_range: (i64, i64), + filters: &[(String, String)], + ) -> Result; +} + +pub struct QueryEngine { + table_provider: Box, + /// The time boundaries for the evaluation. If start equals end an instant + /// is evaluated. + start: i64, + end: i64, + /// Time between two evaluated instants for the range [start:end]. + interval: i64, + /// Default look back from sample search. + lookback_delta: i64, + /// The index of the current time window. Used when evaluating a [range query]. + /// + /// [range query]: https://promlabs.com/blog/2020/06/18/the-anatomy-of-a-promql-query/#range-queries + time_window_idx: i64, + /// key — metric name; value — time series data + metrics_cache: FxHashMap, +} + +impl QueryEngine { + pub fn new

(provider: P) -> Self + where + P: TableProvider, + { + let now = micros_since_epoch(SystemTime::now()); + let five_min = micros(Duration::from_secs(300)); + Self { + table_provider: Box::new(provider), + start: now, + end: now, + interval: five_min, + lookback_delta: five_min, + time_window_idx: 0, + metrics_cache: FxHashMap::default(), + } + } + + pub async fn exec(&mut self, stmt: EvalStmt) -> Result { + self.start = micros_since_epoch(stmt.start); + self.end = micros_since_epoch(stmt.end); + if stmt.interval > Duration::ZERO { + self.interval = micros(stmt.interval); + } + if stmt.lookback_delta > Duration::ZERO { + self.lookback_delta = micros(stmt.lookback_delta); + } + + if self.start == self.end { + // Instant query + let mut value = self.exec_expr(&stmt.expr).await?; + if let Value::Float(v) = value { + value = Value::Sample(Sample { + timestamp: self.end, + value: v, + }); + } + value.sort(); + return Ok(value); + } + + // Range query + // See https://promlabs.com/blog/2020/06/18/the-anatomy-of-a-promql-query/#range-queries + let mut instant_vectors = Vec::new(); + let nr_steps = ((self.end - self.start) / self.interval) + 1; + for i in 0..nr_steps { + self.time_window_idx = i; + match self.exec_expr(&stmt.expr).await? { + Value::Instant(v) => instant_vectors.push(RangeValue { + labels: v.labels.to_owned(), + time_range: None, + values: vec![v.value], + }), + Value::Vector(vs) => instant_vectors.extend(vs.into_iter().map(|v| RangeValue { + labels: v.labels.to_owned(), + time_range: None, + values: vec![v.value], + })), + Value::Range(v) => instant_vectors.push(v), + Value::Matrix(v) => instant_vectors.extend(v), + Value::Sample(v) => instant_vectors.push(RangeValue { + labels: Labels::default(), + time_range: None, + values: vec![v], + }), + Value::Float(v) => instant_vectors.push(RangeValue { + labels: Labels::default(), + time_range: None, + values: vec![Sample { + timestamp: self.start + (self.interval * self.time_window_idx), + value: v, + }], + }), + Value::None => continue, + }; + } + + // empty result quick return + if instant_vectors.is_empty() { + return Ok(Value::None); + } + + // merge data + let mut merged_data = FxHashMap::default(); + let mut merged_metrics = FxHashMap::default(); + for value in instant_vectors { + merged_data + .entry(signature(&value.labels)) + .or_insert_with(Vec::new) + .extend(value.values); + merged_metrics.insert(signature(&value.labels), value.labels); + } + let merged_data = merged_data + .into_iter() + .map(|(sig, values)| RangeValue { + labels: merged_metrics.get(&sig).unwrap().to_owned(), + time_range: None, + values, + }) + .collect::>(); + + // sort data + let mut value = Value::Matrix(merged_data); + value.sort(); + Ok(value) + } + + #[async_recursion] + pub async fn exec_expr(&mut self, prom_expr: &PromExpr) -> Result { + Ok(match &prom_expr { + PromExpr::Aggregate(AggregateExpr { + op, + expr, + param, + modifier, + }) => self.aggregate_exprs(op, expr, param, modifier).await?, + PromExpr::Unary(UnaryExpr { expr }) => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Unary: {:?}", + expr + ))); + } + PromExpr::Binary(expr) => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Paren: {:?}", + expr + ))); + } + PromExpr::Paren(ParenExpr { expr }) => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Paren: {:?}", + expr + ))); + } + PromExpr::Subquery(expr) => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Subquery: {:?}", + expr + ))); + } + PromExpr::NumberLiteral(NumberLiteral { val }) => Value::Float(*val), + PromExpr::StringLiteral(expr) => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported StringLiteral: {:?}", + expr + ))); + } + PromExpr::VectorSelector(v) => { + let data = self.eval_vector_selector(v).await?; + if data.is_empty() { + Value::None + } else { + Value::Vector(data) + } + } + PromExpr::MatrixSelector(MatrixSelector { + vector_selector, + range, + }) => { + let data = self.eval_matrix_selector(vector_selector, *range).await?; + if data.is_empty() { + Value::None + } else { + Value::Matrix(data) + } + } + PromExpr::Call(Call { func, args }) => self.call_expr(func, args).await?, + PromExpr::Extension(expr) => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Extension: {:?}", + expr + ))); + } + }) + } + + /// MatrixSelector is a special case of VectorSelector that returns a matrix of samples. + async fn eval_vector_selector( + &mut self, + selector: &VectorSelector, + ) -> Result> { + let metrics_name = selector.name.as_ref().unwrap(); + if !self.metrics_cache.contains_key(metrics_name) { + self.selector_load_data(selector, None).await?; + } + let metrics_cache = match self.metrics_cache.get(metrics_name) { + Some(v) => match v.get_ref_matrix_values() { + Some(v) => v, + None => return Ok(vec![]), + }, + None => return Ok(vec![]), + }; + + let mut values = vec![]; + for metric in metrics_cache { + let mut value = match metric.values.last() { + Some(v) => *v, + None => continue, // have no sample + }; + if value.timestamp != self.end && metric.values.len() > 1 { + let mut metric = metric.clone(); + metric.time_range = Some((metric.values.first().unwrap().timestamp, self.end)); + if let Some(extra) = metric.extrapolate() { + value = extra.1; + } + } + values.push( + // See https://promlabs.com/blog/2020/06/18/the-anatomy-of-a-promql-query/#instant-queries + InstantValue { + labels: metric.labels.clone(), + value, + }, + ); + } + Ok(values) + } + + /// MatrixSelector is a special case of VectorSelector that returns a matrix of samples. + async fn eval_matrix_selector( + &mut self, + selector: &VectorSelector, + range: Duration, + ) -> Result> { + let metrics_name = selector.name.as_ref().unwrap(); + if !self.metrics_cache.contains_key(metrics_name) { + self.selector_load_data(selector, Some(range)).await?; + } + let metrics_cache = match self.metrics_cache.get(metrics_name) { + Some(v) => match v.get_ref_matrix_values() { + Some(v) => v, + None => return Ok(vec![]), + }, + None => return Ok(vec![]), + }; + + let end = self.start + (self.interval * self.time_window_idx); // 15s + let start = end - micros(range); // 5m + + let mut values = Vec::with_capacity(metrics_cache.len()); + for metric in metrics_cache { + let metric_data = metric + .values + .iter() + .filter(|v| v.timestamp > start && v.timestamp <= end) + .cloned() + .collect::>(); + values.push(RangeValue { + labels: metric.labels.clone(), + time_range: Some((start, end)), + values: metric_data, + }); + } + Ok(values) + } + + async fn selector_load_data( + &mut self, + selector: &VectorSelector, + range: Option, + ) -> Result<()> { + let start = { + // https://promlabs.com/blog/2020/07/02/selecting-data-in-promql/#lookback-delta + let lookback_delta = range.map_or(self.lookback_delta, micros); + self.start + (self.interval * self.time_window_idx) - lookback_delta + }; + let end = self.end; // 30 minutes + 5m = 35m + + // 1. Group by metrics (sets of label name-value pairs) + let table_name = selector.name.as_ref().unwrap(); + let filters: Vec<(String, String)> = selector + .matchers + .matchers + .iter() + .filter(|mat| mat.op == MatchOp::Equal) + .map(|mat| (mat.name.clone(), mat.value.clone())) + .collect(); + let ctx = self + .table_provider + .create_context(table_name, (start, end), &filters) + .await?; + let table = match ctx.table(table_name).await { + Ok(v) => v, + Err(_) => { + self.metrics_cache + .insert(table_name.to_string(), Value::None); + return Ok(()); // table not found + } + }; + + let mut df_group = table.clone().filter( + col(FIELD_TIME) + .gt(lit(start)) + .and(col(FIELD_TIME).lt_eq(lit(end))), + )?; + let regexp_match_udf = + crate::service::search::datafusion::regexp_udf::REGEX_MATCH_UDF.clone(); + let regexp_not_match_udf = + crate::service::search::datafusion::regexp_udf::REGEX_NOT_MATCH_UDF.clone(); + for mat in selector.matchers.matchers.iter() { + match &mat.op { + MatchOp::Equal => { + df_group = df_group.filter(col(mat.name.clone()).eq(lit(mat.value.clone())))? + } + MatchOp::NotEqual => { + df_group = + df_group.filter(col(mat.name.clone()).not_eq(lit(mat.value.clone())))? + } + MatchOp::Re(_re) => { + df_group = df_group.filter( + regexp_match_udf.call(vec![col(mat.name.clone()), lit(mat.value.clone())]), + )? + } + MatchOp::NotRe(_re) => { + df_group = df_group.filter( + regexp_not_match_udf + .call(vec![col(mat.name.clone()), lit(mat.value.clone())]), + )? + } + } + } + let batches = df_group + .sort(vec![col(FIELD_TIME).sort(true, true)])? + .collect() + .await?; + + let mut metrics: FxHashMap = FxHashMap::default(); + for batch in &batches { + let hash_values = batch + .column_by_name(FIELD_HASH) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let time_values = batch + .column_by_name(FIELD_TIME) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let value_values = batch + .column_by_name(FIELD_VALUE) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + let hash = hash_values.value(i).to_string(); + let entry = metrics.entry(hash).or_insert_with(|| { + let mut labels = Vec::with_capacity(batch.num_columns()); + for (k, v) in batch.schema().fields().iter().zip(batch.columns()) { + let name = k.name(); + if name == FIELD_HASH || name == FIELD_TIME || name == FIELD_VALUE { + continue; + } + let value = v.as_any().downcast_ref::().unwrap(); + labels.push(Arc::new(Label { + name: name.to_string(), + value: value.value(i).to_string(), + })); + } + labels.sort_by(|a, b| a.name.cmp(&b.name)); + RangeValue { + labels, + time_range: None, + values: Vec::with_capacity(20), + } + }); + entry.values.push(Sample { + timestamp: time_values.value(i), + value: value_values.value(i), + }); + } + } + + // We don't need the primary key (FIELD_HASH) any more + let mut metric_values = metrics.into_values().collect::>(); + + // TODO fix Extrapolation, it needs reduce by distance + // https://ihac.xyz/2018/12/11/Prometheus-Extrapolation%E5%8E%9F%E7%90%86%E8%A7%A3%E6%9E%90/ + + // Fix data about app restart + for series in metric_values.iter_mut() { + if let Some(metric_type) = labels_value(&series.labels, FIELD_TYPE) { + if metric_type != TYPE_COUNTER { + continue; + } + } + let mut delta: f64 = 0.0; + let mut last_value = 0.0; + for sample in series.values.iter_mut() { + if last_value > sample.value { + delta += last_value; + } + last_value = sample.value; + if delta > 0.0 { + sample.value += delta; + } + } + } + + // cache data + let values = if metric_values.is_empty() { + Value::None + } else { + Value::Matrix(metric_values) + }; + self.metrics_cache.insert(table_name.to_string(), values); + Ok(()) + } + + async fn aggregate_exprs( + &mut self, + op: &TokenType, + expr: &PromExpr, + param: &Option>, + modifier: &Option, + ) -> Result { + let sample_time = self.start + (self.interval * self.time_window_idx); + let input = self.exec_expr(expr).await?; + + Ok(match op.id() { + token::T_SUM => aggregations::sum(sample_time, modifier, &input)?, + token::T_AVG => aggregations::avg(sample_time, modifier, &input)?, + token::T_COUNT => aggregations::count(sample_time, modifier, &input)?, + token::T_MIN => aggregations::min(sample_time, modifier, &input)?, + token::T_MAX => aggregations::max(sample_time, modifier, &input)?, + token::T_GROUP => Value::None, + token::T_STDDEV => Value::None, + token::T_STDVAR => Value::None, + token::T_TOPK => aggregations::topk(self, param.clone().unwrap(), &input).await?, + token::T_BOTTOMK => aggregations::bottomk(self, param.clone().unwrap(), &input).await?, + token::T_COUNT_VALUES => Value::None, + token::T_QUANTILE => Value::None, + _ => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Aggregate: {:?}", + op + ))); + } + }) + } + + async fn call_expr(&mut self, func: &Function, args: &FunctionArgs) -> Result { + use crate::service::promql::functions::Func; + + let func_name = Func::from_str(func.name).map_err(|_| { + DataFusionError::NotImplemented(format!("Unsupported function: {}", func.name)) + })?; + + let last_arg = args + .last() + .expect("BUG: promql-parser should have validated function arguments"); + let input = self.exec_expr(&last_arg).await?; + + Ok(match func_name { + Func::Abs => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Absent => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::AbsentOverTime => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::AvgOverTime => functions::avg_over_time(&input)?, + Func::Ceil => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Changes => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Clamp => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::ClampMax => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::ClampMin => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::CountOverTime => functions::count_over_time(&input)?, + Func::DayOfMonth => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::DayOfWeek => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::DayOfYear => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::DaysInMonth => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Delta => functions::delta(&input)?, + Func::Deriv => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Exp => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Floor => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::HistogramCount => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::HistogramFraction => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::HistogramQuantile => { + let args = &args.args; + if args.len() != 2 { + return Err(DataFusionError::Plan(format!( + "{}: expected 2 arguments, got {}", + func.name, + args.len() + ))); + } + let phi = { + match *args[0] { + PromExpr::NumberLiteral(ref num) => num.val, + _ => { + return Err(DataFusionError::Plan(format!( + "{}: the first argument must be a number", + func.name + ))) + } + } + }; + let sample_time = self.start + (self.interval * self.time_window_idx); + functions::histogram_quantile(sample_time, phi, input)? + } + Func::HistogramSum => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::HoltWinters => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Hour => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Idelta => functions::idelta(&input)?, + Func::Increase => functions::increase(&input)?, + Func::Irate => functions::irate(&input)?, + Func::LabelJoin => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::LabelReplace => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Ln => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Log10 => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Log2 => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::MaxOverTime => functions::max_over_time(&input)?, + Func::MinOverTime => functions::min_over_time(&input)?, + Func::Minute => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Month => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::PredictLinear => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::QuantileOverTime => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Rate => functions::rate(&input)?, + Func::Resets => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Round => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Scalar => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Sgn => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Sort => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::SortDesc => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::SumOverTime => functions::sum_over_time(&input)?, + Func::Time => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Timestamp => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Vector => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + Func::Year => { + return Err(DataFusionError::NotImplemented(format!( + "Unsupported Function: {:?}", + func_name + ))); + } + }) + } +} + +/// Converts `t` to the number of microseconds elapsed since the beginning of the Unix epoch. +fn micros_since_epoch(t: SystemTime) -> i64 { + micros( + t.duration_since(UNIX_EPOCH) + .expect("BUG: {t} is earlier than Unix epoch"), + ) +} + +fn micros(t: Duration) -> i64 { + t.as_micros() + .try_into() + .expect("BUG: time value is too large to fit in i64") +} diff --git a/src/service/promql/functions/avg_over_time.rs b/src/service/promql/functions/avg_over_time.rs new file mode 100644 index 0000000000..12aaece940 --- /dev/null +++ b/src/service/promql/functions/avg_over_time.rs @@ -0,0 +1,28 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; + +use crate::service::promql::value::{RangeValue, Value}; + +pub(crate) fn avg_over_time(data: &Value) -> Result { + super::eval_idelta(data, "avg_over_time", exec) +} + +fn exec(data: &RangeValue) -> Option { + if data.values.is_empty() { + return None; + } + Some(data.values.iter().map(|s| s.value).sum::() / data.values.len() as f64) +} diff --git a/src/service/promql/functions/count_over_time.rs b/src/service/promql/functions/count_over_time.rs new file mode 100644 index 0000000000..92f457bc86 --- /dev/null +++ b/src/service/promql/functions/count_over_time.rs @@ -0,0 +1,28 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; + +use crate::service::promql::value::{RangeValue, Value}; + +pub(crate) fn count_over_time(data: &Value) -> Result { + super::eval_idelta(data, "count_over_time", exec) +} + +fn exec(data: &RangeValue) -> Option { + if data.values.is_empty() { + return None; + } + Some(data.values.len() as f64) +} diff --git a/src/service/promql/functions/delta.rs b/src/service/promql/functions/delta.rs new file mode 100644 index 0000000000..23b5ffcf24 --- /dev/null +++ b/src/service/promql/functions/delta.rs @@ -0,0 +1,27 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; + +use crate::service::promql::value::{RangeValue, Value}; + +pub(crate) fn delta(data: &Value) -> Result { + super::eval_idelta(data, "delta", exec) +} + +fn exec(range: &RangeValue) -> Option { + range + .extrapolate() + .map(|(first, last)| last.value - first.value) +} diff --git a/src/service/promql/functions/histogram.rs b/src/service/promql/functions/histogram.rs new file mode 100644 index 0000000000..79b783913e --- /dev/null +++ b/src/service/promql/functions/histogram.rs @@ -0,0 +1,302 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::{DataFusionError, Result}; +use rustc_hash::FxHashMap; + +use crate::service::promql::value::{self, InstantValue, Labels, Signature, Value}; + +// https://github.com/prometheus/prometheus/blob/cf1bea344a3c390a90c35ea8764c4a468b345d5e/promql/quantile.go#L33 +#[derive(Debug, Clone, Copy, PartialEq)] +struct Bucket { + upper_bound: f64, + count: f64, +} + +// https://github.com/prometheus/prometheus/blob/cf1bea344a3c390a90c35ea8764c4a468b345d5e/promql/quantile.go#L45 +#[derive(Debug)] +struct MetricWithBuckets { + labels: Labels, + buckets: Vec, +} + +/// TODO: support native histograms; see [`histogramQuantile`] +/// +/// [`histogramQuantile`]: https://github.com/prometheus/prometheus/blob/f7c6130ff27a2a12412c02cce223f7a8abc59e49/promql/quantile.go#L146 +pub(crate) fn histogram_quantile(sample_time: i64, phi: f64, data: Value) -> Result { + let in_vec = match data { + Value::Vector(v) => v, + Value::None => return Ok(Value::None), + _ => { + return Err(DataFusionError::Plan( + "histogram_quantile: vector argument expected".to_owned(), + )) + } + }; + + let mut metrics_with_buckets: FxHashMap = FxHashMap::default(); + for InstantValue { + mut labels, + value: sample, + } in in_vec + { + // [https://prometheus.io/docs/prometheus/latest/querying/functions/#histogram_quantile]: + // + // The conventional float samples in `in_vec` are considered the counts + // of observations in each bucket of one or more conventional + // histograms. Each float sample must have a label `le` where the label + // value denotes the inclusive upper bound of the bucket. *Float samples + // without such a label are silently ignored.* + let upper_bound: f64 = match labels + .iter() + .find(|v| v.name == "le") + .map(|s| s.value.parse()) + { + Some(Ok(u)) => u, + None | Some(Err(_)) => continue, + }; + + let sig = value::signature_without_labels( + &labels, + &[value::FIELD_HASH, value::FIELD_NAME, value::FIELD_BUCKET], + ); + let entry = metrics_with_buckets.entry(sig).or_insert_with(|| { + labels.retain(|l| { + l.name != value::FIELD_HASH + && l.name != value::FIELD_NAME + && l.name != value::FIELD_BUCKET + }); + MetricWithBuckets { + labels, + buckets: Vec::new(), + } + }); + entry.buckets.push(Bucket { + upper_bound, + count: sample.value, + }); + } + + let values = metrics_with_buckets + .into_values() + .map(|mb| value::InstantValue { + labels: mb.labels, + value: value::Sample { + timestamp: sample_time, + value: bucket_quantile(phi, mb.buckets), + }, + }) + .collect(); + + Ok(Value::Vector(values)) +} + +// cf. https://github.com/prometheus/prometheus/blob/cf1bea344a3c390a90c35ea8764c4a468b345d5e/promql/quantile.go#L76 +fn bucket_quantile(phi: f64, mut buckets: Vec) -> f64 { + if phi.is_nan() { + return f64::NAN; + } + if phi < 0.0 { + return f64::NEG_INFINITY; + } + if phi > 1.0 { + return f64::INFINITY; + } + buckets.sort_by(|a, b| a.upper_bound.partial_cmp(&b.upper_bound).unwrap()); + // The caller of `bucket_quantile` guarantees that `buckets` is non-empty. + let highest_bucket = buckets[buckets.len() - 1]; + if !(highest_bucket.upper_bound.is_infinite() && highest_bucket.upper_bound.is_sign_positive()) + { + return f64::NAN; + } + let mut buckets = coalesce_buckets(buckets); + ensure_monotonic(&mut buckets); + let buckets = buckets; + if buckets.len() < 2 { + return f64::NAN; + } + let observations = highest_bucket.count; + if observations == 0.0 { + return f64::NAN; + } + let mut rank = phi * observations; + let b = match buckets[..buckets.len() - 1] + .binary_search_by(|b| b.count.partial_cmp(&rank).unwrap()) + { + Ok(b) => b, + Err(b) => b, + }; + if b == buckets.len() - 1 { + return buckets[buckets.len() - 2].upper_bound; + } + if b == 0 && buckets[0].upper_bound <= 0.0 { + return buckets[0].upper_bound; + } + let bucket_end = buckets[b].upper_bound; + let mut count = buckets[b].count; + let bucket_start = if b == 0 { + 0.0 + } else { + count -= buckets[b - 1].count; + rank -= buckets[b - 1].count; + buckets[b - 1].upper_bound + }; + bucket_start + (bucket_end - bucket_start) * rank / count +} + +/// `coalesce_buckets` merges buckets with the same upper bound. +/// The input buckets must be sorted. +fn coalesce_buckets(buckets: Vec) -> Vec { + let mut st = None; + let mut buckets = buckets + .into_iter() + .filter_map(|b| match st { + None => { + st = Some(b); + None + } + Some(last) => { + if b.upper_bound == last.upper_bound { + st = Some(Bucket { + count: last.count + b.count, + ..last + }); + None + } else { + st = Some(b); + Some(last) + } + } + }) + .collect::>(); + if let Some(last) = st { + buckets.push(last); + } + buckets +} + +// For the rationale behind this function, see +// https://github.com/prometheus/prometheus/blob/0bf707e288eaa8694105e53c81a102017529793d/promql/quantile.go#L314-L347 +fn ensure_monotonic(buckets: &mut [Bucket]) { + let mut max = buckets[0].count; + for bucket in &mut buckets[1..] { + if bucket.count > max { + max = bucket.count; + } else if bucket.count < max { + bucket.count = max; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use expect_test::expect; + + #[test] + fn test_coalesce_buckets() { + let buckets = vec![ + Bucket { + upper_bound: 1.0, + count: 2.0, + }, + Bucket { + upper_bound: 1.0, + count: 3.0, + }, + Bucket { + upper_bound: 2.0, + count: 4.0, + }, + Bucket { + upper_bound: 3.0, + count: 1.0, + }, + Bucket { + upper_bound: 3.0, + count: 1.0, + }, + ]; + + expect![[r#" + [ + Bucket { + upper_bound: 1.0, + count: 5.0, + }, + Bucket { + upper_bound: 2.0, + count: 4.0, + }, + Bucket { + upper_bound: 3.0, + count: 2.0, + }, + ] + "#]] + .assert_debug_eq(&coalesce_buckets(buckets)); + } + + #[test] + fn test_ensure_monotonic() { + let mut buckets = vec![ + Bucket { + upper_bound: 1.0, + count: 2.0, + }, + Bucket { + upper_bound: 2.0, + count: 1.0, + }, + Bucket { + upper_bound: 3.0, + count: 4.0, + }, + Bucket { + upper_bound: 4.0, + count: 3.0, + }, + Bucket { + upper_bound: 5.0, + count: 5.0, + }, + ]; + ensure_monotonic(&mut buckets); + expect![[r#" + [ + Bucket { + upper_bound: 1.0, + count: 2.0, + }, + Bucket { + upper_bound: 2.0, + count: 2.0, + }, + Bucket { + upper_bound: 3.0, + count: 4.0, + }, + Bucket { + upper_bound: 4.0, + count: 4.0, + }, + Bucket { + upper_bound: 5.0, + count: 5.0, + }, + ] + "#]] + .assert_debug_eq(&buckets); + } +} diff --git a/src/service/promql/functions/idelta.rs b/src/service/promql/functions/idelta.rs new file mode 100644 index 0000000000..f235488acf --- /dev/null +++ b/src/service/promql/functions/idelta.rs @@ -0,0 +1,32 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; + +use crate::service::promql::value::{RangeValue, Value}; + +pub(crate) fn idelta(data: &Value) -> Result { + super::eval_idelta(data, "idelta", exec) +} + +fn exec(data: &RangeValue) -> Option { + if data.values.is_empty() { + return None; + } else if data.values.len() == 1 { + return Some(0.0); + } + let (last, data) = data.values.split_last().unwrap(); + let previous = data.last().unwrap(); + Some(last.value - previous.value) +} diff --git a/src/service/promql/functions/increase.rs b/src/service/promql/functions/increase.rs new file mode 100644 index 0000000000..a0a1187393 --- /dev/null +++ b/src/service/promql/functions/increase.rs @@ -0,0 +1,27 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; + +use crate::service::promql::value::{RangeValue, Value}; + +pub(crate) fn increase(data: &Value) -> Result { + super::eval_idelta(data, "increase", exec) +} + +fn exec(range: &RangeValue) -> Option { + range + .extrapolate() + .map(|(first, last)| last.value - first.value) +} diff --git a/src/service/promql/functions/irate.rs b/src/service/promql/functions/irate.rs new file mode 100644 index 0000000000..04366be0f1 --- /dev/null +++ b/src/service/promql/functions/irate.rs @@ -0,0 +1,36 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; + +use crate::service::promql::value::{RangeValue, Value}; + +pub(crate) fn irate(data: &Value) -> Result { + super::eval_idelta(data, "irate", exec) +} + +fn exec(data: &RangeValue) -> Option { + if data.values.is_empty() { + return None; + } else if data.values.len() == 1 { + return Some(0.0); + } + let (last, data) = data.values.split_last().unwrap(); + let previous = data.last().unwrap(); + let dt_seconds = ((last.timestamp - previous.timestamp) / 1_000_000) as f64; + if dt_seconds == 0.0 { + return Some(0.0); + } + Some((last.value - previous.value) / dt_seconds) +} diff --git a/src/service/promql/functions/max_over_time.rs b/src/service/promql/functions/max_over_time.rs new file mode 100644 index 0000000000..26611d2f1d --- /dev/null +++ b/src/service/promql/functions/max_over_time.rs @@ -0,0 +1,34 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; + +use crate::service::promql::value::{RangeValue, Value}; + +pub(crate) fn max_over_time(data: &Value) -> Result { + super::eval_idelta(data, "max_over_time", exec) +} + +fn exec(data: &RangeValue) -> Option { + if data.values.is_empty() { + return None; + } + Some( + data.values + .iter() + .map(|s| s.value) + .max_by(|a, b| a.partial_cmp(b).unwrap()) + .unwrap(), + ) +} diff --git a/src/service/promql/functions/min_over_time.rs b/src/service/promql/functions/min_over_time.rs new file mode 100644 index 0000000000..2b1b3c1991 --- /dev/null +++ b/src/service/promql/functions/min_over_time.rs @@ -0,0 +1,34 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; + +use crate::service::promql::value::{RangeValue, Value}; + +pub(crate) fn min_over_time(data: &Value) -> Result { + super::eval_idelta(data, "min_over_time", exec) +} + +fn exec(data: &RangeValue) -> Option { + if data.values.is_empty() { + return None; + } + Some( + data.values + .iter() + .map(|s| s.value) + .min_by(|a, b| a.partial_cmp(b).unwrap()) + .unwrap(), + ) +} diff --git a/src/service/promql/functions/mod.rs b/src/service/promql/functions/mod.rs new file mode 100644 index 0000000000..5a79bc22e4 --- /dev/null +++ b/src/service/promql/functions/mod.rs @@ -0,0 +1,133 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::{DataFusionError, Result}; + +use crate::service::promql::value::{InstantValue, RangeValue, Sample, Value}; + +mod avg_over_time; +mod count_over_time; +mod delta; +mod histogram; +mod idelta; +mod increase; +mod irate; +mod max_over_time; +mod min_over_time; +mod rate; +mod sum_over_time; + +pub(crate) use avg_over_time::avg_over_time; +pub(crate) use count_over_time::count_over_time; +pub(crate) use delta::delta; +pub(crate) use histogram::histogram_quantile; +pub(crate) use idelta::idelta; +pub(crate) use increase::increase; +pub(crate) use irate::irate; +pub(crate) use max_over_time::max_over_time; +pub(crate) use min_over_time::min_over_time; +pub(crate) use rate::rate; +pub(crate) use sum_over_time::sum_over_time; + +use strum::EnumString; + +/// Reference: https://prometheus.io/docs/prometheus/latest/querying/functions/ +#[derive(Debug, Clone, Copy, PartialEq, EnumString)] +#[strum(serialize_all = "snake_case")] +pub(crate) enum Func { + Abs, + Absent, + AbsentOverTime, + AvgOverTime, + Ceil, + Changes, + Clamp, + ClampMax, + ClampMin, + CountOverTime, + DayOfMonth, + DayOfWeek, + DayOfYear, + DaysInMonth, + Delta, + Deriv, + Exp, + Floor, + HistogramCount, + HistogramFraction, + HistogramQuantile, + HistogramSum, + HoltWinters, + Hour, + Idelta, + Increase, + Irate, + LabelJoin, + LabelReplace, + Ln, + Log10, + Log2, + MaxOverTime, + MinOverTime, + Minute, + Month, + PredictLinear, + QuantileOverTime, + Rate, + Resets, + Round, + Scalar, + Sgn, + Sort, + SortDesc, + SumOverTime, + Time, + Timestamp, + Vector, + Year, +} + +pub(crate) fn eval_idelta( + data: &Value, + fn_name: &str, + fn_handler: fn(&RangeValue) -> Option, +) -> Result { + let data = match data { + Value::Matrix(v) => v, + Value::None => return Ok(Value::None), + _ => { + return Err(DataFusionError::Plan(format!( + "{fn_name}: matrix argument expected" + ))) + } + }; + + let mut rate_values = Vec::with_capacity(data.len()); + for metric in data.iter() { + let mut labels = metric.labels.clone(); + labels.retain(|l| l.name != super::value::FIELD_NAME); + let value = fn_handler(metric); + if value.is_none() { + continue; + } + rate_values.push(InstantValue { + labels, + value: Sample { + timestamp: metric.time_range.unwrap().1, + value: value.unwrap(), + }, + }); + } + Ok(Value::Vector(rate_values)) +} diff --git a/src/service/promql/functions/rate.rs b/src/service/promql/functions/rate.rs new file mode 100644 index 0000000000..856dd124bc --- /dev/null +++ b/src/service/promql/functions/rate.rs @@ -0,0 +1,32 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; + +use crate::service::promql::value::{RangeValue, Value}; + +pub(crate) fn rate(data: &Value) -> Result { + super::eval_idelta(data, "rate", exec) +} + +fn exec(range: &RangeValue) -> Option { + range.extrapolate().map(|(first, last)| { + let dt_seconds = ((last.timestamp - first.timestamp) / 1_000_000) as f64; + if dt_seconds == 0.0 { + 0.0 + } else { + (last.value - first.value) / dt_seconds + } + }) +} diff --git a/src/service/promql/functions/sum_over_time.rs b/src/service/promql/functions/sum_over_time.rs new file mode 100644 index 0000000000..0e9aa6e459 --- /dev/null +++ b/src/service/promql/functions/sum_over_time.rs @@ -0,0 +1,28 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::error::Result; + +use crate::service::promql::value::{RangeValue, Value}; + +pub(crate) fn sum_over_time(data: &Value) -> Result { + super::eval_idelta(data, "sum_over_time", exec) +} + +fn exec(data: &RangeValue) -> Option { + if data.values.is_empty() { + return None; + } + Some(data.values.iter().map(|s| s.value).sum()) +} diff --git a/src/service/promql/mod.rs b/src/service/promql/mod.rs new file mode 100644 index 0000000000..b7944cd634 --- /dev/null +++ b/src/service/promql/mod.rs @@ -0,0 +1,59 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use serde::{Deserialize, Serialize}; +use utoipa::ToSchema; + +mod aggregations; +mod engine; +mod functions; +pub mod search; +pub mod value; + +pub use engine::QueryEngine; + +#[derive(Clone, Debug, Serialize, Deserialize, ToSchema)] +pub struct MetricsQueryRequest { + pub query: String, + pub is_range_query: bool, + pub start: i64, + pub end: i64, + pub step: i64, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum Status { + Success, + Error, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct QueryResult { + pub result_type: String, // vector, matrix, scalar, string + pub result: value::Value, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct QueryResponse { + pub status: Status, + #[serde(skip_serializing_if = "Option::is_none")] + pub data: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error_type: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} diff --git a/src/service/promql/search/grpc/mod.rs b/src/service/promql/search/grpc/mod.rs new file mode 100644 index 0000000000..85015b3582 --- /dev/null +++ b/src/service/promql/search/grpc/mod.rs @@ -0,0 +1,432 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::arrow::datatypes::Schema; +use datafusion::datasource::file_format::file_type::{FileType, GetExt}; +use datafusion::datasource::file_format::json::JsonFormat; +use datafusion::datasource::file_format::parquet::ParquetFormat; +use datafusion::datasource::listing::{ListingOptions, ListingTable}; +use datafusion::datasource::listing::{ListingTableConfig, ListingTableUrl}; +use datafusion::datasource::TableProvider; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::DataFusionError; +use std::sync::Arc; +use tracing::{info_span, Instrument}; + +use crate::common::str::find; +use crate::handler::grpc::cluster_rpc; +use crate::infra::cluster; +use crate::infra::config::CONFIG; +use crate::infra::errors::{Error, ErrorCodes, Result}; +use crate::meta::{self, StreamType}; +use crate::service::promql::value; +use crate::service::search; +use crate::service::{file_list, get_partition_key_query}; + +mod storage; +mod wal; + +pub async fn search( + req: &cluster_rpc::MetricsQueryRequest, +) -> Result { + let start = std::time::Instant::now(); + let session_id = req.job.as_ref().unwrap().session_id.to_string(); + let session_id = Arc::new(session_id); + let is_range_query = req.query.as_ref().unwrap().is_range_query; + + // result + let mut results = Vec::new(); + + let span1 = info_span!("service:promql:search:grpc:in_cache"); + + // search in cache + let session_id1 = session_id.clone(); + let org_id1 = req.org_id.clone(); + let req1: cluster_rpc::MetricsQueryStmt = req.query.as_ref().unwrap().clone(); + let task1 = tokio::task::spawn( + async move { + if cluster::is_ingester(&cluster::LOCAL_NODE_ROLE) { + wal::search(&session_id1, &org_id1, &req1).await + } else { + Ok(value::Value::None) + } + } + .instrument(span1), + ); + + let span2 = info_span!("service:promql:search:grpc:in_storage"); + + // search in object storage + let req_stype = req.stype; + let session_id2 = session_id.clone(); + let org_id2 = req.org_id.clone(); + let req2 = req.query.as_ref().unwrap().clone(); + let task2 = tokio::task::spawn( + async move { + if req_stype == cluster_rpc::SearchType::CacheOnly as i32 { + Ok(value::Value::None) + } else { + storage::search(&session_id2, &org_id2, &req2).await + } + } + .instrument(span2), + ); + + // merge local wal + let value1 = match task1.await { + Ok(result) => match result { + Ok(val) => val, + Err(err) => { + log::error!("datafusion execute error: {}", err); + return Err(handle_datafusion_error(err)); + } + }, + Err(err) => { + return Err(Error::ErrorCode(ErrorCodes::ServerInternalError( + err.to_string(), + ))) + } + }; + + match value1 { + value::Value::None => {} + _ => { + results.push(value1); + } + } + + // merge object storage search + let value2 = match task2.await { + Ok(result) => match result { + Ok(val) => val, + Err(err) => { + log::error!("datafusion execute error: {}", err); + return Err(handle_datafusion_error(err)); + } + }, + Err(err) => { + return Err(Error::ErrorCode(ErrorCodes::ServerInternalError( + err.to_string(), + ))) + } + }; + + match value2 { + value::Value::None => {} + _ => { + results.push(value2); + } + } + + let mut result_type = if is_range_query { "matrix" } else { "vector" }; + + let mut resp = cluster_rpc::MetricsQueryResponse { + job: req.job.clone(), + took: start.elapsed().as_millis() as i32, + result_type: result_type.to_string(), + ..Default::default() + }; + + if results.is_empty() { + return Ok(resp); + } + + for result in results.iter() { + result_type = result.get_type(); + match result { + value::Value::None => { + continue; + } + value::Value::Instant(v) => { + resp.result.push(cluster_rpc::Series { + metric: v.labels.iter().map(|l| l.as_ref().into()).collect(), + values: vec![], + value: Some((&v.value).into()), + scalar: None, + }); + } + value::Value::Range(v) => { + resp.result.push(cluster_rpc::Series { + metric: v.labels.iter().map(|l| l.as_ref().into()).collect(), + values: v.values.iter().map(|l| l.into()).collect(), + value: None, + scalar: None, + }); + } + value::Value::Vector(v) => { + v.iter().for_each(|v| { + resp.result.push(cluster_rpc::Series { + metric: v.labels.iter().map(|l| l.as_ref().into()).collect(), + values: vec![], + value: Some((&v.value).into()), + scalar: None, + }); + }); + } + value::Value::Matrix(v) => { + v.iter().for_each(|v| { + resp.result.push(cluster_rpc::Series { + metric: v.labels.iter().map(|l| l.as_ref().into()).collect(), + values: v.values.iter().map(|l| l.into()).collect(), + value: None, + scalar: None, + }); + }); + } + value::Value::Sample(v) => { + resp.result.push(cluster_rpc::Series { + metric: vec![], + values: vec![], + value: Some(v.into()), + scalar: None, + }); + } + value::Value::Float(v) => { + resp.result.push(cluster_rpc::Series { + metric: vec![], + values: vec![], + value: None, + scalar: Some(*v), + }); + } + } + } + + // reset result type + resp.result_type = result_type.to_string(); + + Ok(resp) +} + +pub fn handle_datafusion_error(err: DataFusionError) -> Error { + let err = err.to_string(); + if err.contains("Schema error: No field named") { + let pos = err.find("Schema error: No field named").unwrap(); + return match get_key_from_error(&err, pos) { + Some(key) => Error::ErrorCode(ErrorCodes::SearchFieldNotFound(key)), + None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)), + }; + } + if err.contains("parquet not found") { + return Error::ErrorCode(ErrorCodes::SearchParquetFileNotFound); + } + if err.contains("Invalid function ") { + let pos = err.find("Invalid function ").unwrap(); + return match get_key_from_error(&err, pos) { + Some(key) => Error::ErrorCode(ErrorCodes::SearchFunctionNotDefined(key)), + None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)), + }; + } + if err.contains("Incompatible data types") { + let pos = err.find("for field").unwrap(); + let pos_start = err[pos..].find(' ').unwrap(); + let pos_end = err[pos + pos_start + 1..].find('.').unwrap(); + let field = err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string(); + return Error::ErrorCode(ErrorCodes::SearchFieldHasNoCompatibleDataType(field)); + } + Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)) +} + +fn get_key_from_error(err: &str, pos: usize) -> Option { + for ponct in ['\'', '"'] { + let pos_start = err[pos..].find(ponct); + if pos_start.is_none() { + continue; + } + let pos_start = pos_start.unwrap(); + let pos_end = err[pos + pos_start + 1..].find(ponct); + if pos_end.is_none() { + continue; + } + let pos_end = pos_end.unwrap(); + return Some(err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string()); + } + None +} + +/// match a source is a valid file or not +pub async fn match_source( + org_id: &str, + stream_name: &str, + time_range: Option<(i64, i64)>, + filters: &[(String, String)], + source: &str, + is_wal: bool, + match_min_ts_only: bool, +) -> bool { + let stream_type = StreamType::Metrics; + // match org_id & table + if !source.starts_with(format!("files/{}/{}/{}/", org_id, stream_type, stream_name).as_str()) { + return false; + } + + // check partition key + if !filter_source_by_partition_key(source, filters) { + return false; + } + + if is_wal { + return true; + } + + // check time range + let file_meta = file_list::get_file_meta(source).await.unwrap_or_default(); + if file_meta.min_ts == 0 || file_meta.max_ts == 0 { + return true; + } + log::trace!( + "time range: {:?}, file time: {}-{}, {}", + time_range, + file_meta.min_ts, + file_meta.max_ts, + source + ); + + // match partation clause + if time_range.is_some() { + let (time_min, time_max) = time_range.unwrap(); + if match_min_ts_only && time_min > 0 { + return file_meta.min_ts >= time_min && file_meta.min_ts < time_max; + } + if time_min > 0 && time_min > file_meta.max_ts { + return false; + } + if time_max > 0 && time_max < file_meta.min_ts { + return false; + } + } + true +} + +fn filter_source_by_partition_key(source: &str, filters: &[(String, String)]) -> bool { + !filters.iter().any(|(k, v)| { + let field = get_partition_key_query(&format!("{k}=")); + let value = get_partition_key_query(&format!("{k}={v}")); + find(source, &format!("/{field}")) && !find(source, &format!("/{value}/")) + }) +} + +pub async fn register_table( + session: &meta::search::Session, + org_id: &str, + stream_name: &str, + stream_type: StreamType, + schema: Option>, + files: &Vec, + file_type: FileType, +) -> datafusion::error::Result { + if files.is_empty() { + return Ok(SessionContext::new()); + } + + let start = std::time::Instant::now(); + let runtime_env = search::datafusion::exec::create_runtime_env()?; + let session_config = SessionConfig::new() + .with_information_schema(schema.is_none()) + .with_batch_size(8192); + let ctx = SessionContext::with_config_rt(session_config.clone(), Arc::new(runtime_env)); + + // Configure listing options + let listing_options = match file_type { + FileType::PARQUET => { + let file_format = ParquetFormat::default().with_enable_pruning(Some(false)); + ListingOptions::new(Arc::new(file_format)) + .with_file_extension(FileType::PARQUET.get_ext()) + .with_target_partitions(CONFIG.limit.cpu_num) + } + FileType::JSON => { + let file_format = JsonFormat::default(); + ListingOptions::new(Arc::new(file_format)) + .with_file_extension(FileType::JSON.get_ext()) + .with_target_partitions(CONFIG.limit.cpu_num) + } + _ => { + return Err(DataFusionError::Execution(format!( + "Unsupported file type scheme {file_type:?}", + ))); + } + }; + + let prefix = if session + .data_type + .eq(&search::datafusion::storage::file_list::SessionType::Wal) + { + format!( + "{}files/{}/{stream_type}/{}/", + &CONFIG.common.data_wal_dir, org_id, stream_name + ) + } else { + search::datafusion::storage::file_list::set(&session.id, files) + .await + .unwrap(); + format!("mem:///{}/", session.id) + }; + let prefix = match ListingTableUrl::parse(prefix) { + Ok(url) => url, + Err(e) => { + return Err(datafusion::error::DataFusionError::Execution(format!( + "ListingTableUrl error: {e}", + ))); + } + }; + let prefixes = vec![prefix]; + log::info!( + "Prepare table took {:.3} seconds.", + start.elapsed().as_secs_f64() + ); + + let mut config = + ListingTableConfig::new_with_multi_paths(prefixes).with_listing_options(listing_options); + let schema = if schema.is_none() + || session + .data_type + .eq(&search::datafusion::storage::file_list::SessionType::Wal) + { + config = config.infer_schema(&ctx.state()).await.unwrap(); + let table = ListingTable::try_new(config.clone())?; + let infered_schema = table.schema(); + if schema.is_none() { + infered_schema + } else { + match Schema::try_merge(vec![ + schema.unwrap().as_ref().to_owned(), + infered_schema.as_ref().to_owned(), + ]) { + Ok(schema) => Arc::new(schema), + Err(e) => { + return Err(datafusion::error::DataFusionError::Execution(format!( + "ListingTable Merge schema error: {e}" + ))); + } + } + } + } else { + schema.as_ref().unwrap().clone() + }; + config = config.with_schema(schema); + log::info!( + "infer schema took {:.3} seconds.", + start.elapsed().as_secs_f64() + ); + + let table = ListingTable::try_new(config)?; + ctx.register_table(stream_name, Arc::new(table))?; + + log::info!( + "Register table took {:.3} seconds.", + start.elapsed().as_secs_f64() + ); + + Ok(ctx) +} diff --git a/src/service/promql/search/grpc/storage.rs b/src/service/promql/search/grpc/storage.rs new file mode 100644 index 0000000000..cfdf01dbe3 --- /dev/null +++ b/src/service/promql/search/grpc/storage.rs @@ -0,0 +1,203 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use async_trait::async_trait; +use datafusion::datasource::file_format::file_type::FileType; +use datafusion::error::{DataFusionError, Result}; +use datafusion::prelude::SessionContext; +use promql_parser::parser; +use std::sync::Arc; +use std::time::{Duration, UNIX_EPOCH}; +use tokio::sync::Semaphore; + +use crate::handler::grpc::cluster_rpc; +use crate::infra::cache::file_data; +use crate::infra::config::CONFIG; +use crate::meta; +use crate::service::promql::{engine, value}; +use crate::service::search::datafusion::storage::file_list::SessionType; +use crate::service::{db, file_list, promql, search}; + +struct StorageProvider { + org_id: String, + session_id: String, +} + +#[async_trait] +impl engine::TableProvider for StorageProvider { + async fn create_context( + &self, + stream_name: &str, + time_range: (i64, i64), + filters: &[(String, String)], + ) -> Result { + // get file list + let files = get_file_list(&self.org_id, stream_name, time_range, filters).await?; + let file_count = files.len(); + + // load files to local cache + let mut tasks = Vec::new(); + let semaphore = std::sync::Arc::new(Semaphore::new(CONFIG.limit.query_thread_num)); + for file in files.iter() { + let file = file.clone(); + let permit = semaphore.clone().acquire_owned().await.unwrap(); + let task: tokio::task::JoinHandle> = + tokio::task::spawn(async move { + if !file_data::exist(&file).unwrap_or_default() { + if let Err(e) = file_data::download(&file).await { + log::error!("storage->search: load file {}, err: {}", &file, e); + } + }; + drop(permit); + Ok(()) + }); + tasks.push(task); + } + + for task in tasks { + match task.await { + Ok(ret) => { + if let Err(err) = ret { + return Err(DataFusionError::Execution(err.to_string())); + } + } + Err(err) => { + return Err(DataFusionError::Execution(err.to_string())); + } + }; + } + log::info!( + "[TRACE] promql->search->storage: load files {} done", + file_count + ); + + // fetch all schema versions, get latest schema + let stream_type = meta::StreamType::Metrics; + let schema = match db::schema::get(&self.org_id, stream_name, Some(stream_type)).await { + Ok(schema) => schema, + Err(err) => { + log::error!("get schema error: {}", err); + return Err(datafusion::error::DataFusionError::Execution( + err.to_string(), + )); + } + }; + let schema = Arc::new( + schema + .to_owned() + .with_metadata(std::collections::HashMap::new()), + ); + let session = meta::search::Session { + id: self.session_id.clone(), + data_type: SessionType::Storage, + }; + + super::register_table( + &session, + &self.org_id, + stream_name, + stream_type, + Some(schema), + &files, + FileType::PARQUET, + ) + .await + } +} + +/// search in remote object storage + +/// search in local wal, which haven't been sync to object storage +pub async fn search( + session_id: &str, + org_id: &str, + query: &cluster_rpc::MetricsQueryStmt, +) -> Result { + let prom_expr = match parser::parse(&query.query) { + Ok(expr) => expr, + Err(e) => { + return Err(DataFusionError::Execution(e)); + } + }; + + let eval_stmt = parser::EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH + .checked_add(Duration::from_micros(query.start as u64)) + .unwrap(), + end: UNIX_EPOCH + .checked_add(Duration::from_micros(query.end as u64)) + .unwrap(), + interval: Duration::from_micros(query.step as u64), + lookback_delta: Duration::from_secs(300), + }; + + let mut engine = promql::QueryEngine::new(StorageProvider { + org_id: org_id.to_string(), + session_id: session_id.to_string(), + }); + let data = engine.exec(eval_stmt).await?; + + // clear session + search::datafusion::storage::file_list::clear(session_id) + .await + .unwrap(); + + Ok(data) +} + +#[inline] +async fn get_file_list( + org_id: &str, + stream_name: &str, + time_range: (i64, i64), + filters: &[(String, String)], +) -> Result> { + let (time_min, time_max) = time_range; + let results = match file_list::get_file_list( + org_id, + stream_name, + Some(meta::StreamType::Metrics), + time_min, + time_max, + ) + .await + { + Ok(results) => results, + Err(err) => { + log::error!("get file list error: {}", err); + return Err(DataFusionError::Execution( + "get file list error".to_string(), + )); + } + }; + + let mut files = Vec::new(); + for file in results { + if super::match_source( + org_id, + stream_name, + Some(time_range), + filters, + &file, + false, + false, + ) + .await + { + files.push(file.clone()); + } + } + Ok(files) +} diff --git a/src/service/promql/search/grpc/wal.rs b/src/service/promql/search/grpc/wal.rs new file mode 100644 index 0000000000..7d464470d0 --- /dev/null +++ b/src/service/promql/search/grpc/wal.rs @@ -0,0 +1,188 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use async_trait::async_trait; +use datafusion::datasource::file_format::file_type::FileType; +use datafusion::error::{DataFusionError, Result}; +use datafusion::prelude::SessionContext; +use promql_parser::parser; +use std::path::Path; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::{Duration, UNIX_EPOCH}; + +use crate::common::file::scan_files; +use crate::handler::grpc::cluster_rpc; +use crate::infra::config::{self, CONFIG}; +use crate::meta; +use crate::service::db; +use crate::service::promql; +use crate::service::promql::engine; +use crate::service::promql::value; +use crate::service::search::datafusion::storage::file_list::SessionType; + +struct WalProvider { + org_id: String, + session_id: String, +} + +#[async_trait] +impl engine::TableProvider for WalProvider { + async fn create_context( + &self, + stream_name: &str, + time_range: (i64, i64), + filters: &[(String, String)], + ) -> Result { + // get file list + let files = get_file_list(&self.org_id, stream_name, time_range, filters).await?; + + // fetch all schema versions, get latest schema + let stream_type = meta::StreamType::Metrics; + let schema = match db::schema::get(&self.org_id, stream_name, Some(stream_type)).await { + Ok(schema) => schema, + Err(err) => { + log::error!("get schema error: {}", err); + return Err(datafusion::error::DataFusionError::Execution( + err.to_string(), + )); + } + }; + let schema = Arc::new( + schema + .to_owned() + .with_metadata(std::collections::HashMap::new()), + ); + let session = meta::search::Session { + id: self.session_id.clone(), + data_type: SessionType::Wal, + }; + + super::register_table( + &session, + &self.org_id, + stream_name, + stream_type, + Some(schema), + &files, + FileType::JSON, + ) + .await + } +} + +/// search in local wal, which haven't been sync to object storage +pub async fn search( + session_id: &str, + org_id: &str, + query: &cluster_rpc::MetricsQueryStmt, +) -> Result { + // mark searching in wal + let searching = Searching::new(); + + let prom_expr = match parser::parse(&query.query) { + Ok(expr) => expr, + Err(e) => { + return Err(DataFusionError::Execution(e)); + } + }; + + let eval_stmt = parser::EvalStmt { + expr: prom_expr, + start: UNIX_EPOCH + .checked_add(Duration::from_micros(query.start as u64)) + .unwrap(), + end: UNIX_EPOCH + .checked_add(Duration::from_micros(query.end as u64)) + .unwrap(), + interval: Duration::from_micros(query.step as u64), + lookback_delta: Duration::from_secs(300), + }; + + let mut engine = promql::QueryEngine::new(WalProvider { + org_id: org_id.to_string(), + session_id: session_id.to_string(), + }); + let data = engine.exec(eval_stmt).await?; + + // searching done. + drop(searching); + + Ok(data) +} + +/// get file list from local cache, no need match_source, each file will be searched +#[inline] +async fn get_file_list( + org_id: &str, + stream_name: &str, + time_range: (i64, i64), + filters: &[(String, String)], +) -> Result> { + let pattern = format!( + "{}/files/{}/{}/{}/*.json", + &CONFIG.common.data_wal_dir, + org_id, + meta::StreamType::Metrics, + stream_name + ); + let files = scan_files(&pattern); + + let mut result = Vec::new(); + let data_dir = match Path::new(&CONFIG.common.data_wal_dir).canonicalize() { + Ok(path) => path, + Err(_) => { + return Ok(result); + } + }; + for file in files { + let file = Path::new(&file).canonicalize().unwrap(); + let file = file.strip_prefix(&data_dir).unwrap(); + let local_file = file.to_str().unwrap(); + let file_path = file.parent().unwrap().to_str().unwrap().replace('\\', "/"); + let file_name = file.file_name().unwrap().to_str().unwrap(); + let file_name = file_name.replace('_', "/"); + let source_file = format!("{file_path}/{file_name}"); + if super::match_source( + org_id, + stream_name, + Some(time_range), + filters, + &source_file, + false, + true, + ) + .await + { + result.push(format!("{}{local_file}", &CONFIG.common.data_wal_dir)); + } + } + Ok(result) +} + +/// Searching for marking searching in wal +struct Searching; + +impl Searching { + pub fn new() -> Self { + config::SEARCHING_IN_WAL.store(1, Ordering::Relaxed); + Searching + } +} + +impl Drop for Searching { + fn drop(&mut self) { + config::SEARCHING_IN_WAL.store(0, Ordering::Relaxed); + } +} diff --git a/src/service/promql/search/mod.rs b/src/service/promql/search/mod.rs new file mode 100644 index 0000000000..f6b06f597f --- /dev/null +++ b/src/service/promql/search/mod.rs @@ -0,0 +1,350 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use rustc_hash::FxHashMap; +use std::sync::Arc; +use std::{cmp::min, time::Duration}; +use tonic::{codec::CompressionEncoding, metadata::MetadataValue, transport::Channel, Request}; +use tracing::{info_span, instrument, Instrument}; +use tracing_opentelemetry::OpenTelemetrySpanExt; +use uuid::Uuid; + +use super::value::{InstantValue, Labels, Value}; +use crate::handler::grpc::cluster_rpc; +use crate::infra::cluster::{self, get_internal_grpc_token}; +use crate::infra::config::CONFIG; +use crate::infra::db::etcd; +use crate::infra::errors::{Error, ErrorCodes, Result}; +use crate::service::promql::value::*; + +pub mod grpc; + +pub async fn search(org_id: &str, req: &super::MetricsQueryRequest) -> Result { + let root_span = info_span!("service:promql:search:enter"); + let mut req: cluster_rpc::MetricsQueryRequest = req.to_owned().into(); + req.org_id = org_id.to_string(); + req.stype = cluster_rpc::SearchType::User as i32; + search_in_cluster(req).instrument(root_span).await +} + +#[inline(always)] +async fn get_queue_lock() -> Result { + let mut lock = etcd::Locker::new("search/cluster_queue"); + lock.lock(0).await.map_err(server_internal_error)?; + Ok(lock) +} + +#[instrument(name = "service:promql:search:cluster", skip(req))] +async fn search_in_cluster(req: cluster_rpc::MetricsQueryRequest) -> Result { + let start = std::time::Instant::now(); + + // get a cluster search queue lock + let locker = if CONFIG.common.local_mode { + None + } else { + Some(get_queue_lock().await?) + }; + let took_wait = start.elapsed().as_millis() as usize; + + // get nodes from cluster + let mut nodes = cluster::get_cached_online_query_nodes().unwrap(); + // sort nodes by node_id this will improve hit cache ratio + nodes.sort_by_key(|x| x.id); + let nodes = nodes; + + let querier_num = match nodes + .iter() + .filter(|node| cluster::is_querier(&node.role)) + .count() + { + 0 => 1, + n => n as i64, + }; + + let is_range_query = req.query.as_ref().unwrap().is_range_query; + let req_start: i64 = req.query.as_ref().unwrap().start; + let req_end: i64 = req.query.as_ref().unwrap().end; + let req_step: i64 = req.query.as_ref().unwrap().step; + let point_num = if !is_range_query || req_start == req_end { + 1 + } else { + (req_end - req_start) / req_step + }; + let offset = if querier_num >= point_num { + req_step + } else { + ((point_num / querier_num) + 1) * req_step + }; + + // partition request, here plus 1 second, because division is integer, maybe lose some precision + let mut session_id = Uuid::new_v4().to_string(); + let job = cluster_rpc::Job { + session_id: session_id.clone(), + job: session_id.split_off(30), // take the last 6 characters as job id + stage: 0, + partition: 0, + }; + + // make cluster request + let mut tasks = Vec::new(); + let mut offset_start = req_start; + for (partition_no, node) in nodes.iter().cloned().enumerate() { + let mut req = req.clone(); + let mut req_query = req.query.as_mut().unwrap(); + let mut job = job.clone(); + job.partition = partition_no as i32; + req.job = Some(job); + req.stype = cluster_rpc::SearchType::CacheOnly as i32; + let is_querier = cluster::is_querier(&node.role); + if is_querier { + if offset_start <= req_end { + req.stype = cluster_rpc::SearchType::Cluster as i32; + req_query.start = offset_start; + req_query.end = min(req_end, offset_start + offset); + offset_start += offset; + } else if !cluster::is_ingester(&node.role) { + continue; // no need more querier + } + } + log::info!( + "[TRACE] promql->search->partition: node: {}, is_querier: {}, start: {}, end: {}", + node.id, + is_querier, + req_query.start, + req_query.end, + ); + + let grpc_span = info_span!("service:promql:search:cluster:grpc_search"); + + let node_addr = node.grpc_addr.clone(); + let task = tokio::task::spawn( + async move { + let org_id: MetadataValue<_> = match req.org_id.parse() { + Ok(org_id) => org_id, + Err(_) => return Err(Error::Message("invalid org_id".to_string())), + }; + let mut request = tonic::Request::new(req); + request.set_timeout(Duration::from_secs(CONFIG.grpc.timeout)); + + opentelemetry::global::get_text_map_propagator(|propagator| { + propagator.inject_context( + &tracing::Span::current().context(), + &mut MetadataMap(request.metadata_mut()), + ) + }); + + let token: MetadataValue<_> = match get_internal_grpc_token().parse() { + Ok(token) => token, + Err(_) => return Err(Error::Message("invalid token".to_string())), + }; + let channel = match Channel::from_shared(node_addr).unwrap().connect().await { + Ok(channel) => channel, + Err(err) => { + log::error!( + "[TRACE] promql->search->grpc: node: {}, connect err: {:?}", + node.id, + err + ); + return Err(server_internal_error("connect search node error")); + } + }; + let mut client = cluster_rpc::metrics_client::MetricsClient::with_interceptor( + channel, + move |mut req: Request<()>| { + req.metadata_mut().insert("authorization", token.clone()); + req.metadata_mut() + .insert(CONFIG.grpc.org_header_key.as_str(), org_id.clone()); + Ok(req) + }, + ); + client = client + .send_compressed(CompressionEncoding::Gzip) + .accept_compressed(CompressionEncoding::Gzip); + let response: cluster_rpc::MetricsQueryResponse = match client.query(request).await { + Ok(res) => res.into_inner(), + Err(err) => { + log::error!( + "[TRACE] promql->search->grpc: node: {}, search err: {:?}", + node.id, + err + ); + if err.code() == tonic::Code::Internal { + let err = ErrorCodes::from_json(err.message())?; + return Err(Error::ErrorCode(err)); + } + return Err(server_internal_error("search node error")); + } + }; + + log::info!( + "[TRACE] promql->search->grpc: result node: {}, is_querier: {}, took: {}, files: {}", + node.id, + is_querier, + response.took, + response.file_count + ); + + Ok(response) + } + .instrument(grpc_span), + ); + tasks.push(task); + } + + let mut results = Vec::new(); + for task in tasks { + let result = task + .await + .map_err(|err| Error::ErrorCode(ErrorCodes::ServerInternalError(err.to_string())))?; + match result { + Ok(res) => results.push(res), + Err(err) => { + // search done, release lock + if locker.is_some() { + if let Err(e) = locker.unwrap().unlock().await { + log::error!("search in cluster unlock error: {}", e); + } + } + return Err(err); + } + } + } + + // search done, release lock + if locker.is_some() { + if let Err(e) = locker.unwrap().unlock().await { + log::error!("search in cluster unlock error: {}", e); + } + } + + // merge multiple instances data + let mut file_count = 0; + let mut scan_size = 0; + let mut result_type = String::new(); + let mut series_data: Vec = Vec::new(); + for resp in results { + file_count += resp.file_count; + scan_size += resp.scan_size; + if result_type.is_empty() { + result_type = resp.result_type.clone(); + } + resp.result.into_iter().for_each(|series| { + series_data.push(series); + }); + } + + // merge result + let values = if result_type == "matrix" { + merge_matrix_query(&series_data) + } else if result_type == "vector" { + merge_vector_query(&series_data) + } else if result_type == "scalar" { + merge_scalar_query(&series_data) + } else { + return Err(server_internal_error("invalid result type")); + }; + log::info!( + "[TRACE] promql->search->result: took: {}, took_wait: {}, file_count: {}, scan_size: {}", + start.elapsed().as_millis(), + took_wait, + file_count, + scan_size, + ); + + Ok(values) +} + +fn merge_matrix_query(series: &[cluster_rpc::Series]) -> Value { + let mut merged_data = FxHashMap::default(); + let mut merged_metrics = FxHashMap::default(); + series.iter().for_each(|v| { + let labels: Labels = v.metric.iter().map(|v| Arc::new(Label::from(v))).collect(); + let values: Vec = v.values.iter().map(Sample::from).collect(); + merged_data + .entry(signature(&labels)) + .or_insert_with(Vec::new) + .extend(values); + merged_metrics.insert(signature(&labels), labels); + }); + let merged_data = merged_data + .into_iter() + .map(|(sig, values)| RangeValue { + labels: merged_metrics.get(&sig).unwrap().to_owned(), + values, + time_range: None, + }) + .collect::>(); + + // sort data + let mut value = Value::Matrix(merged_data); + value.sort(); + value +} + +fn merge_vector_query(series: &[cluster_rpc::Series]) -> Value { + let mut merged_data = FxHashMap::default(); + let mut merged_metrics: std::collections::HashMap< + Signature, + Vec>, + std::hash::BuildHasherDefault, + > = FxHashMap::default(); + series.iter().for_each(|v| { + let labels: Labels = v.metric.iter().map(|v| Arc::new(Label::from(v))).collect(); + let value: Sample = v.value.as_ref().unwrap().into(); + merged_data.insert(signature(&labels), value); + merged_metrics.insert(signature(&labels), labels); + }); + let merged_data = merged_data + .into_iter() + .map(|(sig, value)| InstantValue { + labels: merged_metrics.get(&sig).unwrap().to_owned(), + value, + }) + .collect::>(); + + // sort data + let mut value = Value::Vector(merged_data); + value.sort(); + value +} + +fn merge_scalar_query(series: &[cluster_rpc::Series]) -> Value { + let mut sample: Sample = Default::default(); + series.iter().for_each(|v| { + if v.scalar.is_some() { + sample.value = v.scalar.unwrap(); + } + if v.value.is_some() { + sample = v.value.as_ref().unwrap().into(); + } + }); + Value::Sample(sample) +} + +struct MetadataMap<'a>(&'a mut tonic::metadata::MetadataMap); + +impl<'a> opentelemetry::propagation::Injector for MetadataMap<'a> { + /// Set a key and value in the MetadataMap. Does nothing if the key or value are not valid inputs + fn set(&mut self, key: &str, value: String) { + if let Ok(key) = tonic::metadata::MetadataKey::from_bytes(key.as_bytes()) { + if let Ok(val) = tonic::metadata::MetadataValue::try_from(&value) { + self.0.insert(key, val); + } + } + } +} + +fn server_internal_error(error: impl ToString) -> Error { + Error::ErrorCode(ErrorCodes::ServerInternalError(error.to_string())) +} diff --git a/src/service/promql/value.rs b/src/service/promql/value.rs new file mode 100644 index 0000000000..96c87f4188 --- /dev/null +++ b/src/service/promql/value.rs @@ -0,0 +1,325 @@ +// Copyright 2022 Zinc Labs Inc. and Contributors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use serde::{ + ser::{SerializeSeq, SerializeStruct, Serializer}, + Serialize, +}; +use std::{cmp::Ordering, sync::Arc}; + +pub const FIELD_HASH: &str = "__hash__"; +pub const FIELD_NAME: &str = "__name__"; +pub const FIELD_TYPE: &str = "metric_type"; +pub const FIELD_TIME: &str = "_timestamp"; +pub const FIELD_VALUE: &str = "value"; +pub const FIELD_BUCKET: &str = "le"; + +pub const TYPE_COUNTER: &str = "counter"; +pub const TYPE_GAUGE: &str = "gauge"; +pub const TYPE_HISTOGRAM: &str = "histogram"; +pub const TYPE_SUMMARY: &str = "summary"; + +// See https://docs.rs/indexmap/latest/indexmap/#alternate-hashers +type FxIndexMap = + indexmap::IndexMap>; + +pub type Labels = Vec>; + +#[derive(Debug, Clone)] +pub struct Label { + pub name: String, + pub value: String, +} + +#[derive(Debug, Default, Clone, Copy)] +pub struct Sample { + /// Time in microseconds + pub timestamp: i64, + pub value: f64, +} + +impl Serialize for Sample { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let mut seq = serializer.serialize_seq(Some(2))?; + seq.serialize_element(&(self.timestamp / 1_000_000))?; + seq.serialize_element(&self.value.to_string())?; + seq.end() + } +} + +#[derive(Debug, Clone)] +pub struct InstantValue { + pub labels: Labels, + pub value: Sample, +} + +impl Serialize for InstantValue { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let mut seq = serializer.serialize_struct("instant_value", 2)?; + let labels_map = self + .labels + .iter() + .map(|l| (l.name.as_str(), l.value.as_str())) + .collect::>(); + seq.serialize_field("metric", &labels_map)?; + seq.serialize_field("value", &self.value)?; + seq.end() + } +} + +#[derive(Debug, Clone)] +pub struct RangeValue { + pub labels: Labels, + pub time_range: Option<(i64, i64)>, // start, end + pub values: Vec, +} + +impl Serialize for RangeValue { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let mut seq = serializer.serialize_struct("range_value", 2)?; + let labels_map = self + .labels + .iter() + .map(|l| (l.name.as_str(), l.value.as_str())) + .collect::>(); + seq.serialize_field("metric", &labels_map)?; + seq.serialize_field("values", &self.values)?; + seq.end() + } +} + +impl RangeValue { + /// Returns first and last data points, [extrapolated] to the time window + /// boundaries. + /// + /// [extrapolated]: https://promlabs.com/blog/2021/01/29/how-exactly-does-promql-calculate-rates/#extrapolation-of-data + pub(crate) fn extrapolate(&self) -> Option<(Sample, Sample)> { + let samples = &self.values; + if samples.len() < 2 { + return None; + } + let first = samples.first().unwrap(); + let last = samples.last().unwrap(); + + let (t_start, t_end) = self.time_range.unwrap(); + assert!(t_start < t_end); + assert!(t_start <= first.timestamp); + assert!(first.timestamp <= last.timestamp); + assert!(last.timestamp <= t_end); + + Some(( + if first.timestamp == t_start { + *first + } else { + extrapolated_sample(first, last, t_start) + }, + if last.timestamp == t_end { + *last + } else { + extrapolated_sample(first, last, t_end) + }, + )) + } +} + +// https://promlabs.com/blog/2021/01/29/how-exactly-does-promql-calculate-rates/#extrapolation-of-data +fn extrapolated_sample(p1: &Sample, p2: &Sample, t: i64) -> Sample { + let dt = p2.timestamp - p1.timestamp; + let dv = p2.value - p1.value; + let dt2 = t - p1.timestamp; + let dv2 = dv * dt2 as f64 / dt as f64; + Sample { + timestamp: t, + value: p1.value + dv2, + } +} + +pub fn labels_value(labels: &Labels, name: &str) -> Option { + labels + .binary_search_by_key(&name, |label| label.name.as_str()) + .ok() + .map(|index| labels[index].value.clone()) +} + +#[derive(Debug, Clone, Serialize)] +#[serde(untagged)] +pub enum Value { + Instant(InstantValue), + Range(RangeValue), + Vector(Vec), + Matrix(Vec), + Sample(Sample), // only used for return literal value + Float(f64), + None, +} + +impl Value { + pub(crate) fn get_ref_matrix_values(&self) -> Option<&Vec> { + match self { + Value::Matrix(values) => Some(values), + _ => None, + } + } + + pub fn get_type(&self) -> &str { + match self { + Value::Instant(_) => "vector", + Value::Range(_) => "matrix", + Value::Vector(_) => "vector", + Value::Matrix(_) => "matrix", + Value::Sample(_) => "scalar", + Value::Float(_) => "scalar", + Value::None => "scalar", + } + } + + pub fn sort(&mut self) { + match self { + Value::Vector(v) => { + v.sort_by(|a, b| { + b.value + .value + .partial_cmp(&a.value.value) + .unwrap_or(Ordering::Equal) + }); + } + Value::Matrix(v) => { + v.sort_by(|a, b| { + let a = a.values.iter().map(|x| x.value).sum::(); + let b = b.values.iter().map(|x| x.value).sum::(); + b.partial_cmp(&a).unwrap_or(Ordering::Equal) + }); + } + _ => {} + } + } +} + +#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)] +pub struct Signature([u8; 32]); + +impl From for String { + fn from(sig: Signature) -> Self { + hex::encode(sig.0) + } +} + +// REFACTORME: make this a method of `Metric` +pub fn signature(labels: &Labels) -> Signature { + signature_without_labels(labels, &[]) +} + +/// `signature_without_labels` is just as [`signature`], but only for labels not matching `names`. +// REFACTORME: make this a method of `Metric` +pub fn signature_without_labels(labels: &Labels, exclude_names: &[&str]) -> Signature { + let mut hasher = blake3::Hasher::new(); + labels + .iter() + .filter(|item| !exclude_names.contains(&item.name.as_str())) + .for_each(|item| { + hasher.update(item.name.as_bytes()); + hasher.update(item.value.as_bytes()); + }); + Signature(hasher.finalize().into()) +} + +#[cfg(test)] +mod tests { + use super::*; + use expect_test::expect; + + #[test] + fn test_signature_without_labels() { + let mut labels: Labels = Default::default(); + labels.push(Arc::new(Label { + name: "a".to_owned(), + value: "1".to_owned(), + })); + labels.push(Arc::new(Label { + name: "b".to_owned(), + value: "2".to_owned(), + })); + labels.push(Arc::new(Label { + name: "c".to_owned(), + value: "3".to_owned(), + })); + labels.push(Arc::new(Label { + name: "d".to_owned(), + value: "4".to_owned(), + })); + + let sig = signature(&labels); + expect![[r#" + "f287fde2994111abd7740b5c7c28b0eeabe3f813ae65397bb6acb684e2ab6b22" + "#]] + .assert_debug_eq(&String::from(sig)); + + let sig: String = signature_without_labels(&labels, &["a", "c"]).into(); + expect![[r#" + "ec9c3a0c9c03420d330ab62021551cffe993c07b20189c5ed831dad22f54c0c7" + "#]] + .assert_debug_eq(&sig); + assert_eq!(sig.len(), 64); + + assert_eq!(signature(&labels), signature_without_labels(&labels, &[])); + } + + #[test] + fn test_extrapolated_sample() { + let p1 = Sample { + timestamp: 100, + value: 10.0, + }; + let p2 = Sample { + timestamp: 200, + value: 20.0, + }; + let p3 = extrapolated_sample(&p1, &p2, 300); + assert_eq!(p3.timestamp, 300); + assert_eq!(p3.value, 30.0); + + let p1 = Sample { + timestamp: 225, + value: 1.0, + }; + let p2 = Sample { + timestamp: 675, + value: 2.0, + }; + let p3 = extrapolated_sample(&p1, &p2, 750); + let p4 = extrapolated_sample(&p1, &p2, 150); + assert_eq!(format!("{:.2}", p3.value - p4.value), "1.33"); + + let p1 = Sample { + timestamp: 375, + value: 1.0, + }; + let p2 = Sample { + timestamp: 675, + value: 2.0, + }; + let p3 = extrapolated_sample(&p1, &p2, 750); + let p4 = extrapolated_sample(&p1, &p2, 300); + assert_eq!(format!("{:.2}", p3.value - p4.value), "1.50"); + } +} diff --git a/src/service/router/mod.rs b/src/service/router/mod.rs index c79ebb983d..1ea54d6ac8 100644 --- a/src/service/router/mod.rs +++ b/src/service/router/mod.rs @@ -21,7 +21,17 @@ use rand::thread_rng; use crate::infra::cluster; use crate::infra::config::CONFIG; -const QUERIER_ROUTES: [&str; 3] = ["/_search", "/_around", "/_values"]; +const QUERIER_ROUTES: [&str; 9] = [ + "/_search", + "/_around", + "/_values", + "/prometheus/v1/series", + "/prometheus/v1/query_range", + "/prometheus/v1/query", + "/prometheus/v1/metadata", + "/prometheus/v1/labels", + "/prometheus/v1/label/", +]; #[inline] pub fn is_router() -> bool { diff --git a/src/service/search/datafusion/exec.rs b/src/service/search/datafusion/exec.rs index 70318cdf84..bc78960942 100644 --- a/src/service/search/datafusion/exec.rs +++ b/src/service/search/datafusion/exec.rs @@ -89,7 +89,7 @@ pub async fn sql( } }; - let prefix = if session.data_type.eq(&file_list::SessionType::Cache) { + let prefix = if session.data_type.eq(&file_list::SessionType::Wal) { format!( "{}files/{}/{stream_type}/{}/", &CONFIG.common.data_wal_dir, sql.org_id, sql.stream_name @@ -114,7 +114,7 @@ pub async fn sql( let mut config = ListingTableConfig::new_with_multi_paths(prefixes).with_listing_options(listing_options); - let schema = if schema.is_none() || session.data_type.eq(&file_list::SessionType::Cache) { + let schema = if schema.is_none() || session.data_type.eq(&file_list::SessionType::Wal) { config = config.infer_schema(&ctx.state()).await.unwrap(); let table = ListingTable::try_new(config.clone())?; let infered_schema = table.schema(); @@ -886,7 +886,7 @@ pub async fn merge_parquet_files( Ok(file_meta) } -fn create_runtime_env() -> Result { +pub fn create_runtime_env() -> Result { let object_store_registry = DefaultObjectStoreRegistry::new(); let mem = super::storage::memory::InMemory::new(); diff --git a/src/service/search/datafusion/mod.rs b/src/service/search/datafusion/mod.rs index 924c895513..0c382c8f77 100644 --- a/src/service/search/datafusion/mod.rs +++ b/src/service/search/datafusion/mod.rs @@ -16,8 +16,8 @@ use crate::meta::functions::ZoFunction; mod date_format_udf; pub mod exec; -mod match_udf; -mod regexp_udf; +pub mod match_udf; +pub mod regexp_udf; pub mod storage; mod time_range_udf; #[cfg(feature = "zo_functions")] diff --git a/src/service/search/datafusion/regexp_udf.rs b/src/service/search/datafusion/regexp_udf.rs index 9436577232..b0d9ae8711 100644 --- a/src/service/search/datafusion/regexp_udf.rs +++ b/src/service/search/datafusion/regexp_udf.rs @@ -81,7 +81,7 @@ pub fn regex_match_expr_impl(matches: bool) -> ScalarFunctionImplementation { } ColumnarValue::Scalar(ScalarValue::Utf8(pattern)) => pattern, ColumnarValue::Scalar(arg) => { - return Err(DataFusionError::Internal(format!( + return Err(DataFusionError::Plan(format!( "Expected string pattern to regex_match({matches}), got: {arg:?}" ))) } @@ -97,9 +97,8 @@ pub fn regex_match_expr_impl(matches: bool) -> ScalarFunctionImplementation { // the golang regexp library which is different than Rust's regexp let pattern = clean_non_meta_escapes(pattern); - let pattern = regex::Regex::new(&pattern).map_err(|e| { - DataFusionError::Internal(format!("error compiling regex pattern: {e}")) - })?; + let pattern = regex::Regex::new(&pattern) + .map_err(|e| DataFusionError::Plan(format!("error compiling regex pattern: {e}")))?; match &args[0] { ColumnarValue::Array(arr) => { @@ -118,7 +117,7 @@ pub fn regex_match_expr_impl(matches: bool) -> ScalarFunctionImplementation { let res = row.as_ref().map(|v| pattern.is_match(v) == matches); Ok(ColumnarValue::Scalar(ScalarValue::Boolean(res))) } - ColumnarValue::Scalar(v) => Err(DataFusionError::Internal(format!( + ColumnarValue::Scalar(v) => Err(DataFusionError::Plan(format!( "regex_match({matches}) expected first argument to be utf8, got ('{v}')" ))), } diff --git a/src/service/search/datafusion/storage/file_list.rs b/src/service/search/datafusion/storage/file_list.rs index d1885b2f57..bc83a86018 100644 --- a/src/service/search/datafusion/storage/file_list.rs +++ b/src/service/search/datafusion/storage/file_list.rs @@ -24,9 +24,8 @@ lazy_static! { #[derive(Clone, Debug, Eq, PartialEq)] pub enum SessionType { - Cache, - Local, - Remote, + Wal, + Storage, } pub async fn get(session_id: &str) -> Result, anyhow::Error> { diff --git a/src/service/search/cache.rs b/src/service/search/grpc/cache.rs similarity index 90% rename from src/service/search/cache.rs rename to src/service/search/grpc/cache.rs index 2dcd4d6584..3619274fbf 100644 --- a/src/service/search/cache.rs +++ b/src/service/search/grpc/cache.rs @@ -19,14 +19,14 @@ use std::sync::atomic::Ordering; use std::sync::Arc; use tracing::info_span; -use super::datafusion::storage::file_list::SessionType; -use super::sql::Sql; use crate::common::file::scan_files; use crate::infra::config::{self, CONFIG}; use crate::infra::errors::{Error, ErrorCodes}; use crate::meta; use crate::service::db; use crate::service::file_list::calculate_local_files_size; +use crate::service::search::datafusion::storage::file_list::SessionType; +use crate::service::search::sql::Sql; /// search in local cache, which haven't been sync to object storage #[tracing::instrument( @@ -92,7 +92,7 @@ pub async fn search( let session = meta::search::Session { id: session_id.to_string(), - data_type: SessionType::Cache, + data_type: SessionType::Wal, }; let result = match super::datafusion::exec::sql( &session, @@ -118,7 +118,7 @@ pub async fn search( Ok((result, file_count, scan_size as usize)) } -/// get file list from local cache, no need match_source, each file will be searched +/// get file list from local wal, no need match_source, each file will be searched #[inline] async fn get_file_list(sql: &Sql, stream_type: meta::StreamType) -> Result, Error> { let pattern = format!( @@ -142,25 +142,28 @@ async fn get_file_list(sql: &Sql, stream_type: meta::StreamType) -> Result Self { - config::SEARCHING_IN_CACHE.store(1, Ordering::Relaxed); + config::SEARCHING_IN_WAL.store(1, Ordering::Relaxed); Searching } } impl Drop for Searching { fn drop(&mut self) { - config::SEARCHING_IN_CACHE.store(0, Ordering::Relaxed); + config::SEARCHING_IN_WAL.store(0, Ordering::Relaxed); } } diff --git a/src/service/search/exec.rs b/src/service/search/grpc/mod.rs similarity index 68% rename from src/service/search/exec.rs rename to src/service/search/grpc/mod.rs index 2a3734e6b0..71641b330b 100644 --- a/src/service/search/exec.rs +++ b/src/service/search/grpc/mod.rs @@ -12,16 +12,24 @@ // See the License for the specific language governing permissions and // limitations under the License. -use ::datafusion::arrow::ipc; +use ::datafusion::arrow::{ipc, record_batch::RecordBatch}; use ahash::AHashMap as HashMap; +use datafusion_common::DataFusionError; use std::sync::Arc; use tracing::{info_span, Instrument}; +use super::datafusion; use crate::handler::grpc::cluster_rpc; +use crate::infra::cluster; use crate::infra::errors::{Error, ErrorCodes}; use crate::meta::StreamType; -#[tracing::instrument(name = "service:search:exec", skip(req))] +mod cache; +mod storage; + +pub type SearchResult = Result<(HashMap>, usize, usize), Error>; + +#[tracing::instrument(name = "service:search:grpc:search", skip(req))] pub async fn search( req: &cluster_rpc::SearchRequest, ) -> Result { @@ -36,17 +44,23 @@ pub async fn search( let mut file_count = 0; let mut scan_size = 0; - let span1 = info_span!("service:search:exec:in_cache"); + let span1 = info_span!("service:search:grpc:in_cache"); // search in cache let session_id1 = session_id.clone(); let sql1 = sql.clone(); let task1 = tokio::task::spawn( - async move { super::cache::search(&session_id1, sql1, stream_type).await } - .instrument(span1), + async move { + if cluster::is_ingester(&cluster::LOCAL_NODE_ROLE) { + cache::search(&session_id1, sql1, stream_type).await + } else { + Ok((HashMap::new(), 0, 0)) + } + } + .instrument(span1), ); - let span2 = info_span!("service:search:exec:in_storage"); + let span2 = info_span!("service:search:grpc:in_storage"); // search in object storage let req_stype = req.stype; @@ -58,7 +72,7 @@ pub async fn search( if req_stype == cluster_rpc::SearchType::CacheOnly as i32 { Ok((HashMap::new(), 0, 0)) } else { - super::storage::search(&session_id2, sql2, file_list.as_slice(), stream_type).await + storage::search(&session_id2, sql2, file_list.as_slice(), stream_type).await } } .instrument(span2), @@ -112,7 +126,7 @@ pub async fn search( file_count += file_count2; scan_size += scan_size2; - let span3 = info_span!("service:search:exec:merge"); + let span3 = info_span!("service:search:grpc:merge"); let guard3 = span3.enter(); // merge all batches @@ -134,18 +148,18 @@ pub async fn search( Ok(res) => res, Err(err) => { log::error!("datafusion merge error: {}", err); - return Err(super::handle_datafusion_error(err)); + return Err(handle_datafusion_error(err)); } }; } // clear session data - super::datafusion::storage::file_list::clear(&session_id) + datafusion::storage::file_list::clear(&session_id) .await .unwrap(); drop(guard3); - let span4 = info_span!("service:search:exec:response"); + let span4 = info_span!("service:search:grpc:response"); let _guard4 = span4.enter(); // final result @@ -210,3 +224,49 @@ pub async fn search( Ok(result) } + +pub fn handle_datafusion_error(err: DataFusionError) -> Error { + let err = err.to_string(); + if err.contains("Schema error: No field named") { + let pos = err.find("Schema error: No field named").unwrap(); + return match get_key_from_error(&err, pos) { + Some(key) => Error::ErrorCode(ErrorCodes::SearchFieldNotFound(key)), + None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)), + }; + } + if err.contains("parquet not found") { + return Error::ErrorCode(ErrorCodes::SearchParquetFileNotFound); + } + if err.contains("Invalid function ") { + let pos = err.find("Invalid function ").unwrap(); + return match get_key_from_error(&err, pos) { + Some(key) => Error::ErrorCode(ErrorCodes::SearchFunctionNotDefined(key)), + None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)), + }; + } + if err.contains("Incompatible data types") { + let pos = err.find("for field").unwrap(); + let pos_start = err[pos..].find(' ').unwrap(); + let pos_end = err[pos + pos_start + 1..].find('.').unwrap(); + let field = err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string(); + return Error::ErrorCode(ErrorCodes::SearchFieldHasNoCompatibleDataType(field)); + } + Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)) +} + +fn get_key_from_error(err: &str, pos: usize) -> Option { + for punctuation in ['\'', '"'] { + let pos_start = err[pos..].find(punctuation); + if pos_start.is_none() { + continue; + } + let pos_start = pos_start.unwrap(); + let pos_end = err[pos + pos_start + 1..].find(punctuation); + if pos_end.is_none() { + continue; + } + let pos_end = pos_end.unwrap(); + return Some(err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string()); + } + None +} diff --git a/src/service/search/storage.rs b/src/service/search/grpc/storage.rs similarity index 97% rename from src/service/search/storage.rs rename to src/service/search/grpc/storage.rs index a952860f46..1030d54941 100644 --- a/src/service/search/storage.rs +++ b/src/service/search/grpc/storage.rs @@ -19,12 +19,12 @@ use std::sync::Arc; use tokio::sync::Semaphore; use tracing::{info_span, Instrument}; -use super::datafusion::storage::file_list::SessionType; -use super::sql::Sql; use crate::infra::cache::file_data; use crate::infra::config::CONFIG; use crate::infra::errors::{Error, ErrorCodes}; use crate::meta; +use crate::service::search::datafusion::storage::file_list::SessionType; +use crate::service::search::sql::Sql; use crate::service::{db, file_list}; /// search in remote object storage @@ -171,7 +171,7 @@ pub async fn search( let sql = sql.clone(); let session = meta::search::Session { id: format!("{session_id}-{ver}"), - data_type: SessionType::Remote, + data_type: SessionType::Storage, }; // cacluate the diff between latest schema and group schema let mut diff_fields = HashMap::new(); @@ -252,7 +252,7 @@ async fn get_file_list(sql: &Sql, stream_type: meta::StreamType) -> Result>, usize, usize), Error>; pub async fn search( org_id: &str, stream_type: StreamType, req: &meta::search::Request, ) -> Result { - let root_span = info_span!("srv:search:enter"); + let root_span = info_span!("service:search:enter"); let mut req: cluster_rpc::SearchRequest = req.to_owned().into(); req.org_id = org_id.to_string(); @@ -57,6 +51,7 @@ pub async fn search( search_in_cluster(req).instrument(root_span).await } +#[inline(always)] async fn get_queue_lock() -> Result { let mut lock = etcd::Locker::new("search/cluster_queue"); lock.lock(0).await.map_err(server_internal_error)?; @@ -84,7 +79,7 @@ async fn get_times(sql: &sql::Sql, stream_type: StreamType) -> (i64, i64) { } #[instrument(skip_all)] -async fn partition_by_file_list(sql: &sql::Sql, stream_type: StreamType) -> Vec { +async fn get_file_list(sql: &sql::Sql, stream_type: StreamType) -> Vec { let (time_min, time_max) = get_times(sql, stream_type).await; match file_list::get_file_list( &sql.org_id, @@ -108,7 +103,7 @@ async fn partition_by_file_list(sql: &sql::Sql, stream_type: StreamType) -> Vec< } else { // less than 1 hour, use file meta reduce file list for file in file_list { - if sql.match_source(&file, false, stream_type).await { + if sql.match_source(&file, false, false, stream_type).await { files.push(file.clone()); } } @@ -119,7 +114,7 @@ async fn partition_by_file_list(sql: &sql::Sql, stream_type: StreamType) -> Vec< } } -#[instrument(name = "srv:search:cluster", skip(req))] +#[instrument(name = "service:search:cluster", skip(req))] async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result { let start = std::time::Instant::now(); @@ -132,7 +127,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result Result= file_num { 1 @@ -183,7 +178,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result Result Result Result>> = HashMap::new(); let sql = Arc::new(meta); - if !results.is_empty() { - for resp in results { - file_count += resp.file_count; - scan_size += resp.scan_size; - // handle hits - let value = batches.entry("query".to_string()).or_default(); - if !resp.hits.is_empty() { - let buf = Cursor::new(resp.hits); + for resp in results { + file_count += resp.file_count; + scan_size += resp.scan_size; + // handle hits + let value = batches.entry("query".to_string()).or_default(); + if !resp.hits.is_empty() { + let buf = Cursor::new(resp.hits); + let reader = ipc::reader::FileReader::try_new(buf, None).unwrap(); + log::info!( + "search_in_cluster: query num_batches: {:?}", + reader.num_batches() + ); + let batch = reader.into_iter().map(Result::unwrap).collect::>(); + value.push(batch); + } + // handle aggs + for agg in resp.aggs { + let value = batches.entry(format!("agg_{}", agg.name)).or_default(); + if !agg.hits.is_empty() { + let buf = Cursor::new(agg.hits); let reader = ipc::reader::FileReader::try_new(buf, None).unwrap(); log::info!( - "search_in_cluster: query num_batches: {:?}", + "search_in_cluster: agg:{} num_batches: {:?}", + agg.name, reader.num_batches() ); let batch = reader.into_iter().map(Result::unwrap).collect::>(); value.push(batch); } - // handle aggs - for agg in resp.aggs { - let value = batches.entry(format!("agg_{}", agg.name)).or_default(); - if !agg.hits.is_empty() { - let buf = Cursor::new(agg.hits); - let reader = ipc::reader::FileReader::try_new(buf, None).unwrap(); - log::info!( - "search_in_cluster: agg:{} num_batches: {:?}", - agg.name, - reader.num_batches() - ); - let batch = reader.into_iter().map(Result::unwrap).collect::>(); - value.push(batch); - } - } } // merge all batches @@ -404,7 +397,7 @@ async fn search_in_cluster(req: cluster_rpc::SearchRequest) -> Result) -> Vec { new_sources } -pub fn handle_datafusion_error(err: DataFusionError) -> Error { - let err = err.to_string(); - if err.contains("Schema error: No field named") { - let pos = err.find("Schema error: No field named").unwrap(); - return match get_key_from_error(&err, pos) { - Some(key) => Error::ErrorCode(ErrorCodes::SearchFieldNotFound(key)), - None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)), - }; - } - if err.contains("parquet not found") { - return Error::ErrorCode(ErrorCodes::SearchParquetFileNotFound); - } - if err.contains("Invalid function ") { - let pos = err.find("Invalid function ").unwrap(); - return match get_key_from_error(&err, pos) { - Some(key) => Error::ErrorCode(ErrorCodes::SearchFunctionNotDefined(key)), - None => Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)), - }; - } - if err.contains("Incompatible data types") { - let pos = err.find("for field").unwrap(); - let pos_start = err[pos..].find(' ').unwrap(); - let pos_end = err[pos + pos_start + 1..].find('.').unwrap(); - let field = err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string(); - return Error::ErrorCode(ErrorCodes::SearchFieldHasNoCompatibleDataType(field)); - } - Error::ErrorCode(ErrorCodes::SearchSQLExecuteError(err)) -} - -fn get_key_from_error(err: &str, pos: usize) -> Option { - for ponct in ['\'', '"'] { - let pos_start = err[pos..].find(ponct); - if pos_start.is_none() { - continue; - } - let pos_start = pos_start.unwrap(); - let pos_end = err[pos + pos_start + 1..].find(ponct); - if pos_end.is_none() { - continue; - } - let pos_end = pos_end.unwrap(); - return Some(err[pos + pos_start + 1..pos + pos_start + 1 + pos_end].to_string()); - } - None -} - struct MetadataMap<'a>(&'a mut tonic::metadata::MetadataMap); impl<'a> opentelemetry::propagation::Injector for MetadataMap<'a> { diff --git a/src/service/search/sql.rs b/src/service/search/sql.rs index ee1ba63917..7844c90495 100644 --- a/src/service/search/sql.rs +++ b/src/service/search/sql.rs @@ -27,7 +27,7 @@ use crate::infra::errors::{Error, ErrorCodes}; use crate::meta::sql::Sql as MetaSql; use crate::meta::StreamType; use crate::service::stream::get_stream_setting_fts_fields; -use crate::service::{db, file_list, logs}; +use crate::service::{db, file_list, get_partition_key_query}; const SQL_DELIMITERS: [u8; 12] = [ b' ', b'*', b'(', b')', b'<', b'>', b',', b';', b'=', b'!', b'\r', b'\n', @@ -587,6 +587,7 @@ impl Sql { &self, source: &str, match_min_ts_only: bool, + is_wal: bool, stream_type: StreamType, ) -> bool { // match org_id & table @@ -605,6 +606,10 @@ impl Sql { return false; } + if is_wal { + return true; + } + // check time range let file_meta = file_list::get_file_meta(source).await.unwrap_or_default(); if file_meta.min_ts == 0 || file_meta.max_ts == 0 { @@ -797,8 +802,8 @@ fn path_matches_by_partition_key<'a>( partitions: impl IntoIterator, ) -> bool { !partitions.into_iter().any(|(k, v)| { - let field = logs::get_partition_key_query(&format!("{k}=")); - let value = logs::get_partition_key_query(&format!("{k}={v}")); + let field = get_partition_key_query(&format!("{k}=")); + let value = get_partition_key_query(&format!("{k}={v}")); find(path, &format!("/{field}")) && !find(path, &format!("/{value}/")) }) }