starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-05 00:09:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support grpc search partition (#5457)

Browse Source 
close #5458
- [x] add a grpc search partition function
- [x] add a function for create grpc SearchClient
This commit is contained in:
Huaijin
2024-12-26 15:52:53 +08:00
committed by GitHub
parent 593b1ea787
commit 4dbdf651ec
7 changed files with 369 additions and 264 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/handler/grpc/request/search/mod.rs
+41 -4
View File
@@ -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<Response<SearchResponse>, Status> {
let req = req.into_inner();
let request = json::from_slice::<search::Request>(&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<SearchPartitionRequest>,
) -> Result<Response<SearchPartitionResponse>, Status> {
let req = req.into_inner();
let request =
json::from_slice::<search::SearchPartitionRequest>(&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,
src/proto/proto/cluster/search.proto
+42 -23
View File
@@ -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;
}
src/proto/src/generated/cluster.rs
+132 -32
View File
@@ -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<QueryStatus>,
}
#[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<u8>,
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<u8>,
}
#[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<u8>,
#[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<u8>,
}
#[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<u8>,
}
/// 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<QueryStatus>,
}
#[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<super::SearchPartitionRequest>,
) -> std::result::Result<
tonic::Response<super::SearchPartitionResponse>,
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<super::GetResultRequest>,
@@ -1336,6 +1383,13 @@ pub mod search_server {
&self,
request: tonic::Request<super::SearchRequest>,
) -> std::result::Result<tonic::Response<super::SearchResponse>, tonic::Status>;
async fn search_partition(
&self,
request: tonic::Request<super::SearchPartitionRequest>,
) -> std::result::Result<
tonic::Response<super::SearchPartitionResponse>,
tonic::Status,
>;
async fn get_result(
&self,
request: tonic::Request<super::GetResultRequest>,
@@ -1612,6 +1666,52 @@ pub mod search_server {
};
Box::pin(fut)
}
"/cluster.Search/SearchPartition" => {
#[allow(non_camel_case_types)]
struct SearchPartitionSvc<T: Search>(pub Arc<T>);
impl<
T: Search,
> tonic::server::UnaryService<super::SearchPartitionRequest>
for SearchPartitionSvc<T> {
type Response = super::SearchPartitionResponse;
type Future = BoxFuture<
tonic::Response<Self::Response>,
tonic::Status,
>;
fn call(
&mut self,
request: tonic::Request<super::SearchPartitionRequest>,
) -> Self::Future {
let inner = Arc::clone(&self.0);
let fut = async move {
<T as Search>::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<T: Search>(pub Arc<T>);
src/service/grpc.rs
+62 -6
View File
@@ -13,13 +13,22 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
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<RwAHashMap<String, Channel>> = Lazy::new(Default::default);
@@ -98,3 +107,50 @@ async fn create_channel(grpc_addr: &str) -> Result<Channel, tonic::Status> {
})?;
Ok(channel)
}
pub async fn make_grpc_search_client<T>(
request: &mut Request<T>,
node: &Arc<dyn NodeInfo>,
) -> Result<
cluster_rpc::search_client::SearchClient<
InterceptedService<Channel, impl Fn(Request<()>) -> Result<Request<()>, 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))
}
src/service/search/grpc_search.rs
+79 -46
View File
@@ -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<String>,
in_req: &search::Request,
) -> Result<search::Response, Error> {
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<search::SearchPartitionResponse, Error> {
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::<search::SearchPartitionResponse>(&response.response)?;
Ok(response)
}
src/service/search/mod.rs
+8 -75
View File
@@ -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<search::QueryStatusResponse, Error> {
// 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<search::QueryStatusResponse, Error> {
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 {
src/service/search_jobs.rs
+5 -78
View File
@@ -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<Response, anyhow::Error> {
let cfg = config::get_config();
if *cluster == config::get_cluster_name() {
let buf = storage::get(path).await?;
let res: Response = json::from_slice::<Response>(&buf)?;
@@ -443,41 +439,7 @@ pub async fn get_result(path: &str, cluster: &str) -> Result<Response, anyhow::E
let path = path.to_string();
let task = tokio::task::spawn(async move {
let mut request = tonic::Request::new(proto::cluster_rpc::GetResultRequest { path });
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.get_result(request).await {
Ok(res) => res.into_inner(),
Err(err) => {
@@ -526,44 +488,9 @@ pub async fn delete_result(paths: Vec<String>) -> 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) => {