refactor dedup request

Author: baojinri

Cargo.lock

@@ -47,6 +47,7 @@ lazy_static = { workspace = true }
 log = { workspace = true }
 macros = { workspace = true }
 meta_client = { workspace = true }
+notifier = { workspace = true }
 paste = { workspace = true }
 prom-remote-api = { workspace = true, features = ["warp"] }
 prometheus = { workspace = true }

proxy/Cargo.toml

@@ -18,7 +18,6 @@
 #![feature(trait_alias)]
 
 pub mod context;
-pub mod dedup_requests;
 pub mod error;
 mod error_util;
 pub mod forward;
@@ -82,16 +81,14 @@ use table_engine::{
     PARTITION_TABLE_ENGINE_TYPE,
 };
 use time_ext::{current_time_millis, parse_duration};
-use tokio::sync::mpsc::Sender;
 use tonic::{transport::Channel, IntoRequest};
 
 use crate::{
-    dedup_requests::RequestNotifiers,
     error::{ErrNoCause, ErrWithCause, Error, Internal, Result},
     forward::{ForwardRequest, ForwardResult, Forwarder, ForwarderRef},
     hotspot::HotspotRecorder,
     instance::InstanceRef,
-    read::SqlResponse,
+    read::ReadRequestNotifiers,
     schema_config_provider::SchemaConfigProviderRef,
 };
 
@@ -125,7 +122,7 @@ pub struct Proxy {
     engine_runtimes: Arc<EngineRuntimes>,
     cluster_with_meta: bool,
     sub_table_access_perm: SubTableAccessPerm,
-    request_notifiers: Option<Arc<RequestNotifiers<String, Sender<Result<SqlResponse>>>>>,
+    request_notifiers: Option<ReadRequestNotifiers>,
 }
 
 impl Proxy {
@@ -142,7 +139,7 @@ impl Proxy {
         engine_runtimes: Arc<EngineRuntimes>,
         cluster_with_meta: bool,
         sub_table_access_perm: SubTableAccessPerm,
-        request_notifiers: Option<Arc<RequestNotifiers<String, Sender<Result<SqlResponse>>>>>,
+        request_notifiers: Option<ReadRequestNotifiers>,
     ) -> Self {
         let forwarder = Arc::new(Forwarder::new(
             forward_config,

proxy/src/dedup_requests.rs

@@ -24,6 +24,7 @@ use generic_error::BoxError;
 use http::StatusCode;
 use interpreters::interpreter::Output;
 use log::{error, info, warn};
+use notifier::notifier::{ExecutionGuard, RequestNotifiers, RequestResult};
 use query_frontend::{
     frontend,
     frontend::{Context as SqlContext, Frontend},
@@ -36,13 +37,15 @@ use tokio::sync::mpsc::{self, Sender};
 use tonic::{transport::Channel, IntoRequest};
 
 use crate::{
-    dedup_requests::{ExecutionGuard, RequestNotifiers, RequestResult},
     error::{ErrNoCause, ErrWithCause, Error, Internal, InternalNoCause, Result},
     forward::{ForwardRequest, ForwardResult},
     metrics::GRPC_HANDLER_COUNTER_VEC,
     Context, Proxy,
 };
 
+const DEDUP_READ_CHANNEL_LEN: usize = 1;
+pub type ReadRequestNotifiers = Arc<RequestNotifiers<String, Sender<Result<SqlResponse>>>>;
+
 pub enum SqlResponse {
     Forwarded(SqlQueryResponse),
     Local(Output),
@@ -77,10 +80,10 @@ impl Proxy {
         ctx: &Context,
         schema: &str,
         sql: &str,
-        request_notifiers: Arc<RequestNotifiers<String, Sender<Result<SqlResponse>>>>,
+        request_notifiers: ReadRequestNotifiers,
         enable_partition_table_access: bool,
     ) -> Result<SqlResponse> {
-        let (tx, mut rx) = mpsc::channel(1);
+        let (tx, mut rx) = mpsc::channel(DEDUP_READ_CHANNEL_LEN);
         let mut guard = match request_notifiers.insert_notifier(sql.to_string(), tx) {
             RequestResult::First => ExecutionGuard::new(|| {
                 request_notifiers.take_notifiers(&sql.to_string());

proxy/src/lib.rs

@@ -25,9 +25,9 @@ use interpreters::table_manipulator::TableManipulatorRef;
 use log::{info, warn};
 use logger::RuntimeLevel;
 use macros::define_result;
+use notifier::notifier::RequestNotifiers;
 use partition_table_engine::PartitionTableEngine;
 use proxy::{
-    dedup_requests::RequestNotifiers,
     hotspot::HotspotRecorder,
     instance::{DynamicConfig, Instance, InstanceRef},
     limiter::Limiter,
@@ -441,12 +441,6 @@ impl Builder {
             timeout: self.server_config.timeout.map(|v| v.0),
         };
 
-        let request_notifiers = if self.server_config.query_dedup.enable {
-            Some(Arc::new(RequestNotifiers::default()))
-        } else {
-            None
-        };
-
         let proxy = Arc::new(Proxy::new(
             router.clone(),
             instance.clone(),
@@ -459,7 +453,10 @@ impl Builder {
             engine_runtimes.clone(),
             self.cluster.is_some(),
             self.server_config.sub_table_access_perm,
-            request_notifiers,
+            self.server_config
+                .query_dedup
+                .enable
+                .then(|| Arc::new(RequestNotifiers::default())),
         ));
 
         let http_service = http::Builder::new(http_config)