refactor: add request id in context

proxy/src/context.rs

@@ -16,6 +16,7 @@
 
 use std::time::Duration;
 
+use common_types::request_id::RequestId;
 use macros::define_result;
 use snafu::{ensure, Backtrace, Snafu};
 
@@ -42,15 +43,16 @@ define_result!(Error);
 /// Context for request, may contains
 /// 1. Request context and options
 /// 2. Info from http headers
+#[derive(Debug)]
 pub struct RequestContext {
     /// Catalog of the request
     pub catalog: String,
     /// Schema of request
     pub schema: String,
-    /// Enable partition table_access flag
-    pub enable_partition_table_access: bool,
     /// Request timeout
     pub timeout: Option<Duration>,
+    /// Request id
+    pub request_id: RequestId,
 }
 
 impl RequestContext {
@@ -63,7 +65,6 @@ impl RequestContext {
 pub struct Builder {
     catalog: String,
     schema: String,
-    enable_partition_table_access: bool,
     timeout: Option<Duration>,
 }
 
@@ -78,11 +79,6 @@ impl Builder {
         self
     }
 
-    pub fn enable_partition_table_access(mut self, enable_partition_table_access: bool) -> Self {
-        self.enable_partition_table_access = enable_partition_table_access;
-        self
-    }
-
     pub fn timeout(mut self, timeout: Option<Duration>) -> Self {
         self.timeout = timeout;
         self
@@ -95,8 +91,8 @@ impl Builder {
         Ok(RequestContext {
             catalog: self.catalog,
             schema: self.schema,
-            enable_partition_table_access: self.enable_partition_table_access,
             timeout: self.timeout,
+            request_id: RequestId::next_id(),
         })
     }
 }

proxy/src/grpc/prom_query.rs

@@ -26,7 +26,6 @@ use ceresdbproto::{
 use common_types::{
     datum::DatumKind,
     record_batch::RecordBatch,
-    request_id::RequestId,
     schema::{RecordSchema, TSID_COLUMN},
 };
 use generic_error::BoxError;
@@ -76,7 +75,7 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
         ctx: Context,
         req: PrometheusQueryRequest,
     ) -> Result<PrometheusQueryResponse> {
-        let request_id = RequestId::next_id();
+        let request_id = ctx.request_id;
         let begin_instant = Instant::now();
         let deadline = ctx.timeout.map(|t| begin_instant + t);
         let req_ctx = req.context.context(ErrNoCause {

proxy/src/grpc/sql_query.rs

@@ -33,7 +33,7 @@ use log::{error, warn};
 use query_engine::{executor::Executor as QueryExecutor, physical_planner::PhysicalPlanner};
 use router::endpoint::Endpoint;
 use snafu::ResultExt;
-use tokio::sync::mpsc;
+use tokio::{runtime::Handle, sync::mpsc};
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::{transport::Channel, IntoRequest};
 
@@ -55,10 +55,11 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
         GRPC_HANDLER_COUNTER_VEC.incoming_query.inc();
 
         self.hotspot_recorder.inc_sql_query_reqs(&req).await;
-        match self.handle_sql_query_internal(ctx, req).await {
+        match self.handle_sql_query_internal(&ctx, &req).await {
             Err(e) => {
+                error!("Failed to handle sql query, ctx:{ctx:?}, err:{e}");
+
                 GRPC_HANDLER_COUNTER_VEC.query_failed.inc();
-                error!("Failed to handle sql query, err:{e}");
                 SqlQueryResponse {
                     header: Some(error::build_err_header(e)),
                     ..Default::default()
@@ -73,8 +74,8 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
 
     async fn handle_sql_query_internal(
         &self,
-        ctx: Context,
-        req: SqlQueryRequest,
+        ctx: &Context,
+        req: &SqlQueryRequest,
     ) -> Result<SqlQueryResponse> {
         if req.context.is_none() {
             return ErrNoCause {
@@ -86,7 +87,7 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
 
         let req_context = req.context.as_ref().unwrap();
         let schema = &req_context.database;
-        match self.handle_sql(ctx, schema, &req.sql).await? {
+        match self.handle_sql(ctx, schema, &req.sql, false).await? {
             SqlResponse::Forwarded(resp) => Ok(resp),
             SqlResponse::Local(output) => convert_output(&output, self.resp_compress_min_length),
         }
@@ -99,7 +100,7 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
     ) -> BoxStream<'static, SqlQueryResponse> {
         GRPC_HANDLER_COUNTER_VEC.stream_query.inc();
         self.hotspot_recorder.inc_sql_query_reqs(&req).await;
-        match self.clone().handle_stream_query_internal(ctx, req).await {
+        match self.clone().handle_stream_query_internal(&ctx, &req).await {
             Err(e) => stream::once(async {
                 error!("Failed to handle stream sql query, err:{e}");
                 GRPC_HANDLER_COUNTER_VEC.stream_query_failed.inc();
@@ -118,8 +119,8 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
 
     async fn handle_stream_query_internal(
         self: Arc<Self>,
-        ctx: Context,
-        req: SqlQueryRequest,
+        ctx: &Context,
+        req: &SqlQueryRequest,
     ) -> Result<BoxStream<'static, SqlQueryResponse>> {
         if req.context.is_none() {
             return ErrNoCause {
@@ -130,12 +131,8 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
         }
 
         let req_context = req.context.as_ref().unwrap();
-        let schema = req_context.database.clone();
-        let req = match self
-            .clone()
-            .maybe_forward_stream_sql_query(ctx.clone(), &req)
-            .await
-        {
+        let schema = &req_context.database;
+        let req = match self.clone().maybe_forward_stream_sql_query(ctx, req).await {
             Some(resp) => match resp {
                 ForwardResult::Forwarded(resp) => return resp,
                 ForwardResult::Local => req,
@@ -144,13 +141,12 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
         };
 
         let (tx, rx) = mpsc::channel(STREAM_QUERY_CHANNEL_LEN);
-        let runtime = ctx.runtime.clone();
         let resp_compress_min_length = self.resp_compress_min_length;
         let output = self
             .as_ref()
-            .fetch_sql_query_output(ctx, &schema, &req.sql)
+            .fetch_sql_query_output(ctx, schema, &req.sql, false)
             .await?;
-        runtime.spawn(async move {
+        Handle::current().spawn(async move {
             match output {
                 Output::AffectedRows(rows) => {
                     let resp =
@@ -189,7 +185,7 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
 
     async fn maybe_forward_stream_sql_query(
         self: Arc<Self>,
-        ctx: Context,
+        ctx: &Context,
         req: &SqlQueryRequest,
     ) -> Option<ForwardResult<BoxStream<'static, SqlQueryResponse>, Error>> {
         if req.tables.len() != 1 {
@@ -203,7 +199,7 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
             schema: req_ctx.database.clone(),
             table: req.tables[0].clone(),
             req: req.clone().into_request(),
-            forwarded_from: ctx.forwarded_from,
+            forwarded_from: ctx.forwarded_from.clone(),
         };
         let do_query = |mut client: StorageServiceClient<Channel>,
                         request: tonic::Request<SqlQueryRequest>,

proxy/src/http/prom.rs

@@ -31,13 +31,12 @@ use ceresdbproto::storage::{
 };
 use common_types::{
     datum::DatumKind,
-    request_id::RequestId,
     schema::{RecordSchema, TSID_COLUMN},
 };
 use generic_error::BoxError;
 use http::StatusCode;
 use interpreters::interpreter::Output;
-use log::{debug, error};
+use log::{error, info};
 use prom_remote_api::types::{
     Label, LabelMatcher, Query, QueryResult, RemoteStorage, Sample, TimeSeries, WriteRequest,
 };
@@ -85,12 +84,7 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
             }),
             table_requests: write_table_requests,
         };
-        let ctx = ProxyContext {
-            runtime: self.engine_runtimes.write_runtime.clone(),
-            timeout: ctx.timeout,
-            enable_partition_table_access: false,
-            forwarded_from: None,
-        };
+        let ctx = ProxyContext::new(ctx.timeout, None);
 
         match self.handle_write_internal(ctx, table_request).await {
             Ok(result) => {
@@ -127,15 +121,14 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
         metric: String,
         query: Query,
     ) -> Result<QueryResult> {
-        // Open partition table if needed.
-        self.maybe_open_partition_table_if_not_exist(&ctx.catalog, &ctx.schema, &metric)
-            .await?;
-
-        let request_id = RequestId::next_id();
+        let request_id = ctx.request_id;
         let begin_instant = Instant::now();
         let deadline = ctx.timeout.map(|t| begin_instant + t);
+        info!("Handle prom remote query begin, ctx:{ctx:?}, metric:{metric}, request:{query:?}");
 
-        debug!("Query handler try to process request, request_id:{request_id}, request:{query:?}");
+        // Open partition table if needed.
+        self.maybe_open_partition_table_if_not_exist(&ctx.catalog, &ctx.schema, &metric)
+            .await?;
 
         let provider = CatalogMetaProvider {
             manager: self.instance.catalog_manager.clone(),
@@ -171,7 +164,7 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
             .await?;
 
         let cost = begin_instant.saturating_elapsed().as_millis();
-        debug!("Query handler finished, request_id:{request_id}, cost:{cost}ms, query:{query:?}");
+        info!("Handle prom remote query sucess, ctx:{ctx:?}, cost:{cost}ms");
 
         convert_query_result(metric, timestamp_col_name, field_col_name, output)
     }

proxy/src/http/sql.rs

@@ -26,6 +26,7 @@ use common_types::{
 use generic_error::BoxError;
 use http::StatusCode;
 use interpreters::interpreter::Output;
+use log::error;
 use query_engine::{
     executor::{Executor as QueryExecutor, RecordBatchVec},
     physical_planner::PhysicalPlanner,
@@ -49,14 +50,16 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
         ctx: &RequestContext,
         req: Request,
     ) -> Result<Output> {
-        let context = Context {
-            timeout: ctx.timeout,
-            runtime: self.engine_runtimes.read_runtime.clone(),
-            enable_partition_table_access: true,
-            forwarded_from: None,
-        };
+        let schema = &ctx.schema;
+        let ctx = Context::new(ctx.timeout, None);
 
-        match self.handle_sql(context, &ctx.schema, &req.query).await? {
+        match self
+            .handle_sql(&ctx, schema, &req.query, true)
+            .await
+            .map_err(|e| {
+                error!("Handle sql query failed, ctx:{ctx:?}, req:{req:?}, err:{e}");
+                e
+            })? {
             SqlResponse::Forwarded(resp) => convert_sql_response_to_output(resp),
             SqlResponse::Local(output) => Ok(output),
         }

proxy/src/influxdb/mod.rs

@@ -23,7 +23,6 @@ use std::time::Instant;
 use ceresdbproto::storage::{
     RequestContext as GrpcRequestContext, WriteRequest as GrpcWriteRequest,
 };
-use common_types::request_id::RequestId;
 use generic_error::BoxError;
 use http::StatusCode;
 use interpreters::interpreter::Output;
@@ -80,12 +79,7 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
             }),
             table_requests: write_table_requests,
         };
-        let proxy_context = Context {
-            timeout: ctx.timeout,
-            runtime: self.engine_runtimes.write_runtime.clone(),
-            enable_partition_table_access: false,
-            forwarded_from: None,
-        };
+        let proxy_context = Context::new(ctx.timeout, None);
 
         match self
             .handle_write_internal(proxy_context, table_request)
@@ -126,7 +120,7 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
         ctx: RequestContext,
         req: InfluxqlRequest,
     ) -> Result<Output> {
-        let request_id = RequestId::next_id();
+        let request_id = ctx.request_id;
         let begin_instant = Instant::now();
         let deadline = ctx.timeout.map(|t| begin_instant + t);
 

proxy/src/lib.rs

@@ -72,7 +72,6 @@ use log::{error, info};
 use query_engine::{executor::Executor as QueryExecutor, physical_planner::PhysicalPlanner};
 use query_frontend::plan::Plan;
 use router::{endpoint::Endpoint, Router};
-use runtime::Runtime;
 use snafu::{OptionExt, ResultExt};
 use table_engine::{
     engine::{EngineRuntimes, TableState},
@@ -554,10 +553,19 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct Context {
-    pub timeout: Option<Duration>,
-    pub runtime: Arc<Runtime>,
-    pub enable_partition_table_access: bool,
-    pub forwarded_from: Option<String>,
+    request_id: RequestId,
+    timeout: Option<Duration>,
+    forwarded_from: Option<String>,
+}
+
+impl Context {
+    pub fn new(timeout: Option<Duration>, forwarded_from: Option<String>) -> Self {
+        Self {
+            request_id: RequestId::next_id(),
+            timeout,
+            forwarded_from,
+        }
+    }
 }

proxy/src/opentsdb/mod.rs

@@ -56,12 +56,7 @@ impl<Q: QueryExecutor + 'static, P: PhysicalPlanner> Proxy<Q, P> {
             }),
             table_requests: write_table_requests,
         };
-        let proxy_context = Context {
-            timeout: ctx.timeout,
-            runtime: self.engine_runtimes.write_runtime.clone(),
-            enable_partition_table_access: false,
-            forwarded_from: None,
-        };
+        let proxy_context = Context::new(ctx.timeout, None);
 
         match self
             .handle_write_internal(proxy_context, table_request)