From 4dbdf651ecb8f4dbf8c401e4e5c040a7849fb9a8 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Thu, 26 Dec 2024 15:52:53 +0800 Subject: [PATCH] feat: support grpc search partition (#5457) close #5458 - [x] add a grpc search partition function - [x] add a function for create grpc SearchClient --- src/handler/grpc/request/search/mod.rs | 45 ++++++- src/proto/proto/cluster/search.proto | 65 ++++++---- src/proto/src/generated/cluster.rs | 164 ++++++++++++++++++++----- src/service/grpc.rs | 68 +++++++++- src/service/search/grpc_search.rs | 125 ++++++++++++------- src/service/search/mod.rs | 83 ++----------- src/service/search_jobs.rs | 83 +------------ 7 files changed, 369 insertions(+), 264 deletions(-) diff --git a/src/handler/grpc/request/search/mod.rs b/src/handler/grpc/request/search/mod.rs index 3c167c5d48..268e98c217 100644 --- a/src/handler/grpc/request/search/mod.rs +++ b/src/handler/grpc/request/search/mod.rs @@ -24,7 +24,8 @@ use o2_enterprise::enterprise::search::{QueryManager, TaskStatus, WorkGroup}; use proto::cluster_rpc::{ search_server::Search, CancelQueryRequest, CancelQueryResponse, DeleteResultRequest, DeleteResultResponse, GetResultRequest, GetResultResponse, QueryStatusRequest, - QueryStatusResponse, SearchRequest, SearchResponse, + QueryStatusResponse, SearchPartitionRequest, SearchPartitionResponse, SearchRequest, + SearchResponse, }; use tonic::{Request, Response, Status}; @@ -132,7 +133,7 @@ impl Search for Searcher { ) -> Result, Status> { let req = req.into_inner(); let request = json::from_slice::(&req.request) - .map_err(|e| Status::internal(format!("failed to parse request: {e}")))?; + .map_err(|e| Status::internal(format!("failed to parse search request: {e}")))?; let stream_type = StreamType::from(req.stream_type.as_str()); let ret = SearchService::search( &req.trace_id, @@ -145,8 +146,9 @@ impl Search for Searcher { match ret { Ok(ret) => { - let response = json::to_vec(&ret) - .map_err(|e| Status::internal(format!("failed to serialize response: {e}")))?; + let response = json::to_vec(&ret).map_err(|e| { + Status::internal(format!("failed to serialize search response: {e}")) + })?; Ok(Response::new(SearchResponse { trace_id: req.trace_id, response, @@ -156,6 +158,41 @@ impl Search for Searcher { } } + async fn search_partition( + &self, + req: Request, + ) -> Result, Status> { + let req = req.into_inner(); + let request = + json::from_slice::(&req.request).map_err(|e| { + Status::internal(format!("failed to parse search partition request: {e}")) + })?; + let stream_type = StreamType::from(req.stream_type.as_str()); + let ret = SearchService::search_partition( + &req.trace_id, + &req.org_id, + stream_type, + &request, + req.skip_max_query_range, + ) + .await; + + match ret { + Ok(ret) => { + let response = json::to_vec(&ret).map_err(|e| { + Status::internal(format!( + "failed to serialize search partition response: {e}" + )) + })?; + Ok(Response::new(SearchPartitionResponse { + trace_id: req.trace_id, + response, + })) + } + Err(e) => Err(Status::internal(format!("search partition failed: {e}"))), + } + } + #[cfg(feature = "enterprise")] async fn get_result( &self, diff --git a/src/proto/proto/cluster/search.proto b/src/proto/proto/cluster/search.proto index 76e3af7219..83c6f7ba88 100644 --- a/src/proto/proto/cluster/search.proto +++ b/src/proto/proto/cluster/search.proto @@ -13,10 +13,52 @@ service Search { rpc CancelQuery(CancelQueryRequest) returns (CancelQueryResponse) {} rpc ClusterCancelQuery(CancelQueryRequest) returns (CancelQueryResponse) {} rpc Search(SearchRequest) returns (SearchResponse) {} + rpc SearchPartition(SearchPartitionRequest) returns (SearchPartitionResponse) {} rpc GetResult(GetResultRequest) returns (GetResultResponse) {} rpc DeleteResult(DeleteResultRequest) returns (DeleteResultResponse) {} } +message QueryStatusRequest {} + +message QueryStatusResponse { + repeated QueryStatus status = 1; +} + +message CancelQueryRequest { + string trace_id = 1; +} + +message CancelQueryResponse { + bool is_success = 1; +} + +message SearchRequest { + string trace_id = 1; + string org_id = 2; + string stream_type = 3; + optional string user_id = 4; + bytes request = 5; +} + +message SearchResponse { + string trace_id = 1; + bytes response = 2; +} + + +message SearchPartitionRequest { + string trace_id = 1; + string org_id = 2; + string stream_type = 3; + bytes request = 5; + bool skip_max_query_range = 6; +} + +message SearchPartitionResponse { + string trace_id = 1; + bytes response = 2; +} + message DeleteResultRequest { repeated string paths = 1; } @@ -31,19 +73,6 @@ message GetResultResponse { bytes response = 1; } -message SearchRequest { - string trace_id = 1; - string org_id = 2; - string stream_type = 3; - optional string user_id = 4; - bytes request = 5; -} - -message SearchResponse { - string trace_id = 1; - bytes response = 2; -} - // Search request query message SearchQuery { string sql = 1; @@ -60,11 +89,7 @@ message SearchQuery { bool skip_wal = 14; } -message QueryStatusRequest {} -message QueryStatusResponse { - repeated QueryStatus status = 1; -} message QueryStatus { string trace_id = 1; @@ -91,10 +116,4 @@ message Query { int64 end_time = 3; } -message CancelQueryRequest { - string trace_id = 1; -} -message CancelQueryResponse { - bool is_success = 1; -} diff --git a/src/proto/src/generated/cluster.rs b/src/proto/src/generated/cluster.rs index a451dca38c..9c534dd878 100644 --- a/src/proto/src/generated/cluster.rs +++ b/src/proto/src/generated/cluster.rs @@ -919,24 +919,24 @@ pub mod metrics_server { } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] -pub struct DeleteResultRequest { - #[prost(string, repeated, tag = "1")] - pub paths: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, +pub struct QueryStatusRequest {} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct QueryStatusResponse { + #[prost(message, repeated, tag = "1")] + pub status: ::prost::alloc::vec::Vec, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] -pub struct DeleteResultResponse {} -#[allow(clippy::derive_partial_eq_without_eq)] -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct GetResultRequest { +pub struct CancelQueryRequest { #[prost(string, tag = "1")] - pub path: ::prost::alloc::string::String, + pub trace_id: ::prost::alloc::string::String, } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] -pub struct GetResultResponse { - #[prost(bytes = "vec", tag = "1")] - pub response: ::prost::alloc::vec::Vec, +pub struct CancelQueryResponse { + #[prost(bool, tag = "1")] + pub is_success: bool, } #[derive(Eq)] #[allow(clippy::derive_partial_eq_without_eq)] @@ -963,6 +963,49 @@ pub struct SearchResponse { #[prost(bytes = "vec", tag = "2")] pub response: ::prost::alloc::vec::Vec, } +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SearchPartitionRequest { + #[prost(string, tag = "1")] + pub trace_id: ::prost::alloc::string::String, + #[prost(string, tag = "2")] + pub org_id: ::prost::alloc::string::String, + #[prost(string, tag = "3")] + pub stream_type: ::prost::alloc::string::String, + #[prost(bytes = "vec", tag = "5")] + pub request: ::prost::alloc::vec::Vec, + #[prost(bool, tag = "6")] + pub skip_max_query_range: bool, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct SearchPartitionResponse { + #[prost(string, tag = "1")] + pub trace_id: ::prost::alloc::string::String, + #[prost(bytes = "vec", tag = "2")] + pub response: ::prost::alloc::vec::Vec, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct DeleteResultRequest { + #[prost(string, repeated, tag = "1")] + pub paths: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct DeleteResultResponse {} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct GetResultRequest { + #[prost(string, tag = "1")] + pub path: ::prost::alloc::string::String, +} +#[allow(clippy::derive_partial_eq_without_eq)] +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct GetResultResponse { + #[prost(bytes = "vec", tag = "1")] + pub response: ::prost::alloc::vec::Vec, +} /// Search request query #[derive(Eq)] #[allow(clippy::derive_partial_eq_without_eq)] @@ -995,15 +1038,6 @@ pub struct SearchQuery { } #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] -pub struct QueryStatusRequest {} -#[allow(clippy::derive_partial_eq_without_eq)] -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct QueryStatusResponse { - #[prost(message, repeated, tag = "1")] - pub status: ::prost::alloc::vec::Vec, -} -#[allow(clippy::derive_partial_eq_without_eq)] -#[derive(Clone, PartialEq, ::prost::Message)] pub struct QueryStatus { #[prost(string, tag = "1")] pub trace_id: ::prost::alloc::string::String, @@ -1038,18 +1072,6 @@ pub struct Query { #[prost(int64, tag = "3")] pub end_time: i64, } -#[allow(clippy::derive_partial_eq_without_eq)] -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct CancelQueryRequest { - #[prost(string, tag = "1")] - pub trace_id: ::prost::alloc::string::String, -} -#[allow(clippy::derive_partial_eq_without_eq)] -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct CancelQueryResponse { - #[prost(bool, tag = "1")] - pub is_success: bool, -} #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum WorkGroup { @@ -1255,6 +1277,31 @@ pub mod search_client { req.extensions_mut().insert(GrpcMethod::new("cluster.Search", "Search")); self.inner.unary(req, path, codec).await } + pub async fn search_partition( + &mut self, + request: impl tonic::IntoRequest, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + self.inner + .ready() + .await + .map_err(|e| { + tonic::Status::new( + tonic::Code::Unknown, + format!("Service was not ready: {}", e.into()), + ) + })?; + let codec = tonic::codec::ProstCodec::default(); + let path = http::uri::PathAndQuery::from_static( + "/cluster.Search/SearchPartition", + ); + let mut req = request.into_request(); + req.extensions_mut() + .insert(GrpcMethod::new("cluster.Search", "SearchPartition")); + self.inner.unary(req, path, codec).await + } pub async fn get_result( &mut self, request: impl tonic::IntoRequest, @@ -1336,6 +1383,13 @@ pub mod search_server { &self, request: tonic::Request, ) -> std::result::Result, tonic::Status>; + async fn search_partition( + &self, + request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + >; async fn get_result( &self, request: tonic::Request, @@ -1612,6 +1666,52 @@ pub mod search_server { }; Box::pin(fut) } + "/cluster.Search/SearchPartition" => { + #[allow(non_camel_case_types)] + struct SearchPartitionSvc(pub Arc); + impl< + T: Search, + > tonic::server::UnaryService + for SearchPartitionSvc { + type Response = super::SearchPartitionResponse; + type Future = BoxFuture< + tonic::Response, + tonic::Status, + >; + fn call( + &mut self, + request: tonic::Request, + ) -> Self::Future { + let inner = Arc::clone(&self.0); + let fut = async move { + ::search_partition(&inner, request).await + }; + Box::pin(fut) + } + } + let accept_compression_encodings = self.accept_compression_encodings; + let send_compression_encodings = self.send_compression_encodings; + let max_decoding_message_size = self.max_decoding_message_size; + let max_encoding_message_size = self.max_encoding_message_size; + let inner = self.inner.clone(); + let fut = async move { + let inner = inner.0; + let method = SearchPartitionSvc(inner); + let codec = tonic::codec::ProstCodec::default(); + let mut grpc = tonic::server::Grpc::new(codec) + .apply_compression_config( + accept_compression_encodings, + send_compression_encodings, + ) + .apply_max_message_size_config( + max_decoding_message_size, + max_encoding_message_size, + ); + let res = grpc.unary(method, req).await; + Ok(res) + }; + Box::pin(fut) + } "/cluster.Search/GetResult" => { #[allow(non_camel_case_types)] struct GetResultSvc(pub Arc); diff --git a/src/service/grpc.rs b/src/service/grpc.rs index 14e9fa4ad4..b97be36306 100644 --- a/src/service/grpc.rs +++ b/src/service/grpc.rs @@ -13,13 +13,22 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use config::{get_config, utils::rand::get_rand_element, RwAHashMap}; -use once_cell::sync::Lazy; -use tonic::{ - transport::{Certificate, Channel, ClientTlsConfig}, - Status, -}; +use std::sync::Arc; +use config::{get_config, meta::cluster::NodeInfo, utils::rand::get_rand_element, RwAHashMap}; +use infra::errors::{Error, ErrorCodes}; +use once_cell::sync::Lazy; +use proto::cluster_rpc; +use tonic::{ + codec::CompressionEncoding, + metadata::MetadataValue, + service::interceptor::InterceptedService, + transport::{Certificate, Channel, ClientTlsConfig}, + Request, Status, +}; +use tracing_opentelemetry::OpenTelemetrySpanExt; + +use super::search::MetadataMap; use crate::common::infra::cluster; static CHANNELS: Lazy> = Lazy::new(Default::default); @@ -98,3 +107,50 @@ async fn create_channel(grpc_addr: &str) -> Result { })?; Ok(channel) } + +pub async fn make_grpc_search_client( + request: &mut Request, + node: &Arc, +) -> Result< + cluster_rpc::search_client::SearchClient< + InterceptedService) -> Result, Status>>, + >, + Error, +> { + let cfg = get_config(); + request.set_timeout(std::time::Duration::from_secs(cfg.limit.query_timeout)); + + opentelemetry::global::get_text_map_propagator(|propagator| { + propagator.inject_context( + &tracing::Span::current().context(), + &mut MetadataMap(request.metadata_mut()), + ) + }); + + let token: MetadataValue<_> = node + .get_auth_token() + .parse() + .map_err(|_| Error::Message("invalid token".to_string()))?; + let channel = get_cached_channel(&node.get_grpc_addr()) + .await + .map_err(|err| { + log::error!( + "search->grpc: node: {}, connect err: {:?}", + &node.get_grpc_addr(), + err + ); + Error::ErrorCode(ErrorCodes::ServerInternalError(err.to_string())) + })?; + let client = cluster_rpc::search_client::SearchClient::with_interceptor( + channel, + move |mut req: Request<()>| { + req.metadata_mut().insert("authorization", token.clone()); + Ok(req) + }, + ); + Ok(client + .send_compressed(CompressionEncoding::Gzip) + .accept_compressed(CompressionEncoding::Gzip) + .max_decoding_message_size(cfg.grpc.max_message_size * 1024 * 1024) + .max_encoding_message_size(cfg.grpc.max_message_size * 1024 * 1024)) +} diff --git a/src/service/search/grpc_search.rs b/src/service/search/grpc_search.rs index cc08c961ef..2f26cf9b09 100644 --- a/src/service/search/grpc_search.rs +++ b/src/service/search/grpc_search.rs @@ -1,21 +1,16 @@ +use std::sync::Arc; + use config::{ - get_config, - meta::{cluster::NodeInfo, search, stream::StreamType}, + meta::{search, stream::StreamType}, utils::json, }; use infra::errors::{Error, ErrorCodes}; -use proto::cluster_rpc; use rand::{rngs::StdRng, seq::SliceRandom, SeedableRng}; -use tonic::{codec::CompressionEncoding, metadata::MetadataValue, Request}; use tracing::{info_span, Instrument}; -use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::{ common::infra::cluster as infra_cluster, - service::{ - grpc::get_cached_channel, - search::{server_internal_error, MetadataMap}, - }, + service::{grpc::make_grpc_search_client, search::server_internal_error}, }; #[tracing::instrument(name = "service:search:grpc_search", skip_all)] @@ -26,7 +21,6 @@ pub async fn grpc_search( user_id: Option, in_req: &search::Request, ) -> Result { - let cfg = get_config(); let mut nodes = match infra_cluster::get_cached_online_query_nodes(None).await { Some(nodes) => nodes, None => { @@ -61,46 +55,14 @@ pub async fn grpc_search( user_id: user_id.clone(), request: req, }); - request.set_timeout(std::time::Duration::from_secs(cfg.limit.query_timeout)); - - opentelemetry::global::get_text_map_propagator(|propagator| { - propagator.inject_context( - &tracing::Span::current().context(), - &mut MetadataMap(request.metadata_mut()), - ) - }); - - let token: MetadataValue<_> = node - .get_auth_token() - .parse() - .map_err(|_| Error::Message("invalid token".to_string()))?; - let channel = get_cached_channel(&node_addr).await.map_err(|err| { - log::error!( - "search->grpc: node: {}, connect err: {:?}", - &node.grpc_addr, - err - ); - server_internal_error("connect search node error") - })?; - let mut client = cluster_rpc::search_client::SearchClient::with_interceptor( - channel, - move |mut req: Request<()>| { - req.metadata_mut().insert("authorization", token.clone()); - Ok(req) - }, - ); - client = client - .send_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Gzip) - .max_decoding_message_size(cfg.grpc.max_message_size * 1024 * 1024) - .max_encoding_message_size(cfg.grpc.max_message_size * 1024 * 1024); + let node = Arc::new(node) as _; + let mut client = make_grpc_search_client(&mut request, &node).await?; let response = match client.search(request).await { Ok(res) => res.into_inner(), Err(err) => { log::error!( - "search->grpc: node: {}, search err: {:?}", - &node.grpc_addr, - err + "search->grpc: node: {}, search err: {err:?}", + &node.get_grpc_addr(), ); if err.code() == tonic::Code::Internal { let err = ErrorCodes::from_json(err.message())?; @@ -121,3 +83,74 @@ pub async fn grpc_search( Ok(response) } + +#[tracing::instrument(name = "service:search:grpc_search_partition", skip_all)] +pub async fn grpc_search_partition( + trace_id: &str, + org_id: &str, + stream_type: StreamType, + in_req: &search::SearchPartitionRequest, + skip_max_query_range: bool, +) -> Result { + let mut nodes = match infra_cluster::get_cached_online_query_nodes(None).await { + Some(nodes) => nodes, + None => { + log::error!("search->grpc: no querier node online"); + return Err(server_internal_error("no querier node online")); + } + }; + // sort nodes by node_id this will improve hit cache ratio + nodes.dedup_by(|a, b| a.grpc_addr == b.grpc_addr); + let mut rng = StdRng::from_entropy(); + let node = nodes.choose(&mut rng).unwrap().clone(); + + // make cluster request + let node_addr = node.grpc_addr.clone(); + let grpc_span = info_span!( + "service:search:cluster:grpc_search_partition", + node_id = node.id, + node_addr = node_addr.as_str(), + ); + + let trace_id = trace_id.to_string(); + let org_id = org_id.to_string(); + let stream_type = stream_type.to_string(); + let in_req = in_req.clone(); + let task = tokio::task::spawn( + async move { + let req = bytes::Bytes::from(json::to_string(&in_req)?).to_vec(); + let mut request = tonic::Request::new(proto::cluster_rpc::SearchPartitionRequest { + trace_id: trace_id.to_string(), + org_id: org_id.to_string(), + stream_type: stream_type.to_string(), + request: req, + skip_max_query_range, + }); + let node = Arc::new(node) as _; + let mut client = make_grpc_search_client(&mut request, &node).await?; + let response = match client.search_partition(request).await { + Ok(res) => res.into_inner(), + Err(err) => { + log::error!( + "search->grpc: node: {}, search err: {err:?}", + &node.get_grpc_addr(), + ); + 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")); + } + }; + Ok(response) + } + .instrument(grpc_span), + ); + + let response = task + .await + .map_err(|e| Error::ErrorCode(ErrorCodes::ServerInternalError(e.to_string())))??; + let response = json::from_slice::(&response.response)?; + + Ok(response) +} diff --git a/src/service/search/mod.rs b/src/service/search/mod.rs index 6d726ae4b4..cb0097b90a 100644 --- a/src/service/search/mod.rs +++ b/src/service/search/mod.rs @@ -53,12 +53,8 @@ use tracing::Instrument; use tracing_opentelemetry::OpenTelemetrySpanExt; #[cfg(feature = "enterprise")] use { - crate::service::grpc::get_cached_channel, - config::meta::cluster::get_internal_grpc_token, - o2_enterprise::enterprise::search::TaskStatus, - o2_enterprise::enterprise::search::WorkGroup, - std::collections::HashSet, - tonic::{codec::CompressionEncoding, metadata::MetadataValue, Request}, + crate::service::grpc::make_grpc_search_client, o2_enterprise::enterprise::search::TaskStatus, + o2_enterprise::enterprise::search::WorkGroup, std::collections::HashSet, std::str::FromStr, tracing::info_span, }; @@ -775,8 +771,6 @@ pub async fn search_partition( #[cfg(feature = "enterprise")] pub async fn query_status() -> Result { // get nodes from cluster - - use std::str::FromStr; let mut nodes = match infra_cluster::get_cached_online_query_nodes(None).await { Some(nodes) => nodes, None => { @@ -800,46 +794,15 @@ pub async fn query_status() -> Result { let task = tokio::task::spawn( async move { - let cfg = get_config(); let mut request = tonic::Request::new(proto::cluster_rpc::QueryStatusRequest {}); - request.set_timeout(std::time::Duration::from_secs(cfg.limit.query_timeout)); - - opentelemetry::global::get_text_map_propagator(|propagator| { - propagator.inject_context( - &tracing::Span::current().context(), - &mut MetadataMap(request.metadata_mut()), - ) - }); - - let token: MetadataValue<_> = get_internal_grpc_token() - .parse() - .map_err(|_| Error::Message("invalid token".to_string()))?; - let channel = get_cached_channel(&node_addr).await.map_err(|err| { - log::error!( - "search->grpc: node: {}, connect err: {:?}", - &node.grpc_addr, - err - ); - server_internal_error("connect search node error") - })?; - let mut client = cluster_rpc::search_client::SearchClient::with_interceptor( - channel, - move |mut req: Request<()>| { - req.metadata_mut().insert("authorization", token.clone()); - Ok(req) - }, - ); - client = client - .send_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Gzip) - .max_decoding_message_size(cfg.grpc.max_message_size * 1024 * 1024) - .max_encoding_message_size(cfg.grpc.max_message_size * 1024 * 1024); + let node = Arc::new(node) as _; + let mut client = make_grpc_search_client(&mut request, &node).await?; let response = match client.query_status(request).await { Ok(res) => res.into_inner(), Err(err) => { log::error!( "search->grpc: node: {}, search err: {:?}", - &node.grpc_addr, + &node.get_grpc_addr(), err ); if err.code() == tonic::Code::Internal { @@ -963,47 +926,17 @@ pub async fn cancel_query( let trace_id = trace_id.to_string(); let task = tokio::task::spawn( async move { - let cfg = get_config(); let mut request = tonic::Request::new(proto::cluster_rpc::CancelQueryRequest { trace_id }); - request.set_timeout(std::time::Duration::from_secs(cfg.limit.query_timeout)); - opentelemetry::global::get_text_map_propagator(|propagator| { - propagator.inject_context( - &tracing::Span::current().context(), - &mut MetadataMap(request.metadata_mut()), - ) - }); - - let token: MetadataValue<_> = get_internal_grpc_token() - .parse() - .map_err(|_| Error::Message("invalid token".to_string()))?; - let channel = get_cached_channel(&node_addr).await.map_err(|err| { - log::error!( - "grpc_cancel_query: node: {}, connect err: {:?}", - &node.grpc_addr, - err - ); - server_internal_error("connect search node error") - })?; - let mut client = cluster_rpc::search_client::SearchClient::with_interceptor( - channel, - move |mut req: Request<()>| { - req.metadata_mut().insert("authorization", token.clone()); - Ok(req) - }, - ); - client = client - .send_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Gzip) - .max_decoding_message_size(cfg.grpc.max_message_size * 1024 * 1024) - .max_encoding_message_size(cfg.grpc.max_message_size * 1024 * 1024); + let node = Arc::new(node) as _; + let mut client = make_grpc_search_client(&mut request, &node).await?; let response: cluster_rpc::CancelQueryResponse = match client.cancel_query(request).await { Ok(res) => res.into_inner(), Err(err) => { log::error!( "grpc_cancel_query: node: {}, search err: {:?}", - &node.grpc_addr, + &node.get_grpc_addr(), err ); if err.code() == tonic::Code::Internal { diff --git a/src/service/search_jobs.rs b/src/service/search_jobs.rs index 2f2bf67d07..65b7064187 100644 --- a/src/service/search_jobs.rs +++ b/src/service/search_jobs.rs @@ -33,15 +33,12 @@ use o2_enterprise::enterprise::{ search::{get_cluster_node_by_name, get_cluster_nodes}, }, }; -use proto::cluster_rpc; use tokio::sync::mpsc; -use tonic::{codec::CompressionEncoding, metadata::MetadataValue, Request}; -use tracing_opentelemetry::OpenTelemetrySpanExt; +use super::grpc::make_grpc_search_client; use crate::service::{ db::search_job::{search_job_partitions::*, search_job_results::*, search_jobs::*}, - grpc::get_cached_channel, - search::{self as SearchService, grpc_search::grpc_search, MetadataMap}, + search::grpc_search::{grpc_search, grpc_search_partition}, }; // 1. get the oldest job from `search_jobs` table @@ -165,7 +162,7 @@ async fn handle_search_partition(job: &Job) -> Result<(), anyhow::Error> { let stream_type = StreamType::from(job.stream_type.as_str()); let req: search::Request = json::from_str(&job.payload)?; let partition_req = SearchPartitionRequest::from(&req); - let res = SearchService::search_partition( + let res = grpc_search_partition( &job.trace_id, &job.org_id, stream_type, @@ -430,7 +427,6 @@ async fn merge_response( // get the response in this cluster or other cluster pub async fn get_result(path: &str, cluster: &str) -> Result { - let cfg = config::get_config(); if *cluster == config::get_cluster_name() { let buf = storage::get(path).await?; let res: Response = json::from_slice::(&buf)?; @@ -443,41 +439,7 @@ pub async fn get_result(path: &str, cluster: &str) -> Result = node - .get_auth_token() - .parse() - .map_err(|_| Error::Message("invalid token".to_string()))?; - let channel = get_cached_channel(&node.get_grpc_addr()) - .await - .map_err(|err| { - log::error!( - "search->grpc: node: {}, connect err: {:?}", - &node.get_grpc_addr(), - err - ); - Error::ErrorCode(ErrorCodes::ServerInternalError(err.to_string())) - })?; - let mut client = cluster_rpc::search_client::SearchClient::with_interceptor( - channel, - move |mut req: Request<()>| { - req.metadata_mut().insert("authorization", token.clone()); - Ok(req) - }, - ); - client = client - .send_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Gzip) - .max_decoding_message_size(cfg.grpc.max_message_size * 1024 * 1024) - .max_encoding_message_size(cfg.grpc.max_message_size * 1024 * 1024); + let mut client = make_grpc_search_client(&mut request, &node).await?; let response = match client.get_result(request).await { Ok(res) => res.into_inner(), Err(err) => { @@ -526,44 +488,9 @@ pub async fn delete_result(paths: Vec) -> Result<(), anyhow::Error> { } let paths = paths.clone(); let task = tokio::task::spawn(async move { - let cfg = config::get_config(); let mut request = tonic::Request::new(proto::cluster_rpc::DeleteResultRequest { paths }); - request.set_timeout(std::time::Duration::from_secs(cfg.limit.query_timeout)); - - opentelemetry::global::get_text_map_propagator(|propagator| { - propagator.inject_context( - &tracing::Span::current().context(), - &mut MetadataMap(request.metadata_mut()), - ) - }); - - let token: MetadataValue<_> = node - .get_auth_token() - .parse() - .map_err(|_| Error::Message("invalid token".to_string()))?; - let channel = get_cached_channel(&node.get_grpc_addr()) - .await - .map_err(|err| { - log::error!( - "search->grpc: node: {}, connect err: {:?}", - &node.get_grpc_addr(), - err - ); - Error::ErrorCode(ErrorCodes::ServerInternalError(err.to_string())) - })?; - let mut client = cluster_rpc::search_client::SearchClient::with_interceptor( - channel, - move |mut req: Request<()>| { - req.metadata_mut().insert("authorization", token.clone()); - Ok(req) - }, - ); - client = client - .send_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Gzip) - .max_decoding_message_size(cfg.grpc.max_message_size * 1024 * 1024) - .max_encoding_message_size(cfg.grpc.max_message_size * 1024 * 1024); + let mut client = make_grpc_search_client(&mut request, &node).await?; let response = match client.delete_result(request).await { Ok(res) => res.into_inner(), Err(err) => {