feat(query): support cluster level concurrent limit (#17778)

zhang2014 · web-flow · commit 3dc7840d04af · 2025-04-15T18:16:57.000+08:00
* feat(query): support cluster level concurrent limit

* feat(query): support cluster level concurrent limit

* feat(query): support cluster level concurrent limit

* feat(query): support cluster level concurrent limit

* feat(query): support cluster level concurrent limit
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/meta/raft-store/src/mem_meta.rs b/src/meta/raft-store/src/mem_meta.rs
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
 use std::future;
 use std::io;
 use std::sync::Arc;
@@ -32,6 +33,7 @@ use futures_util::StreamExt;
 use futures_util::TryStreamExt;
 use log::debug;
 use tokio::sync::Mutex;
+use tokio::sync::Semaphore;
 
 use crate::applier::Applier;
 use crate::mem_state_machine::MemStateMachine;
@@ -40,6 +42,7 @@ use crate::state_machine_api_ext::StateMachineApiExt;
 #[derive(Clone, Default)]
 pub struct MemMeta {
     sm: Arc<Mutex<MemStateMachine>>,
+    pub locks: Arc<Mutex<HashMap<String, Arc<Semaphore>>>>,
 }
 
 impl MemMeta {
diff --git a/src/meta/semaphore/src/acquirer/permit.rs b/src/meta/semaphore/src/acquirer/permit.rs
@@ -38,7 +38,7 @@ use crate::storage::PermitKey;
 /// the oneshot sender signals the lease extending task to stop,
 /// allowing for proper cleanup of the [`PermitEntry`] in the meta-service.
 pub struct Permit {
-    pub(crate) fu: BoxFuture<'static, Result<(), ConnectionClosed>>,
+    pub fu: BoxFuture<'static, Result<(), ConnectionClosed>>,
 }
 
 impl std::fmt::Debug for Permit {
diff --git a/src/meta/store/Cargo.toml b/src/meta/store/Cargo.toml
@@ -12,10 +12,12 @@ io-uring = []
 
 [dependencies]
 async-trait = { workspace = true }
+databend-common-base = { workspace = true }
 databend-common-grpc = { workspace = true }
 databend-common-meta-client = { workspace = true }
 databend-common-meta-embedded = { workspace = true }
 databend-common-meta-kvapi = { workspace = true }
+databend-common-meta-semaphore = { workspace = true }
 databend-common-meta-types = { workspace = true }
 log = { workspace = true }
 tokio-stream = { workspace = true }
diff --git a/src/meta/store/src/lib.rs b/src/meta/store/src/lib.rs
@@ -12,11 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::hash_map::Entry;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::Context;
 use std::task::Poll;
+use std::time::Duration;
 
+use databend_common_base::base::tokio::sync::Semaphore as TokioSemaphore;
 use databend_common_grpc::RpcClientConf;
 use databend_common_meta_client::errors::CreationError;
 use databend_common_meta_client::ClientHandle;
@@ -25,6 +28,10 @@ use databend_common_meta_embedded::MemMeta;
 use databend_common_meta_kvapi::kvapi;
 use databend_common_meta_kvapi::kvapi::KVStream;
 use databend_common_meta_kvapi::kvapi::UpsertKVReply;
+use databend_common_meta_semaphore::acquirer::Permit;
+use databend_common_meta_semaphore::errors::AcquireError;
+use databend_common_meta_semaphore::errors::ConnectionClosed;
+use databend_common_meta_semaphore::Semaphore;
 use databend_common_meta_types::protobuf::WatchRequest;
 use databend_common_meta_types::protobuf::WatchResponse;
 use databend_common_meta_types::MetaError;
@@ -80,6 +87,42 @@ impl MetaStore {
             }
         }
     }
+
+    pub async fn new_acquired(
+        &self,
+        prefix: impl ToString,
+        capacity: u64,
+        id: impl ToString,
+        lease: Duration,
+    ) -> Result<Permit, AcquireError> {
+        match self {
+            MetaStore::L(v) => {
+                let mut local_lock_map = v.locks.lock().await;
+
+                let acquire_res = match local_lock_map.entry(prefix.to_string()) {
+                    Entry::Occupied(v) => v.get().clone(),
+                    Entry::Vacant(v) => v
+                        .insert(Arc::new(TokioSemaphore::new(capacity as usize)))
+                        .clone(),
+                };
+
+                match acquire_res.acquire_owned().await {
+                    Ok(guard) => Ok(Permit {
+                        fu: Box::pin(async move {
+                            let _guard = guard;
+                            Ok(())
+                        }),
+                    }),
+                    Err(_e) => Err(AcquireError::ConnectionClosed(ConnectionClosed::new_str(
+                        "",
+                    ))),
+                }
+            }
+            MetaStore::R(grpc_client) => {
+                Semaphore::new_acquired(grpc_client.clone(), prefix, capacity, id, lease).await
+            }
+        }
+    }
 }
 
 #[async_trait::async_trait]
diff --git a/src/query/service/Cargo.toml b/src/query/service/Cargo.toml
@@ -68,6 +68,7 @@ databend-common-meta-api = { workspace = true }
 databend-common-meta-app = { workspace = true }
 databend-common-meta-app-types = { workspace = true }
 databend-common-meta-kvapi = { workspace = true }
+databend-common-meta-semaphore = { workspace = true }
 databend-common-meta-store = { workspace = true }
 databend-common-meta-types = { workspace = true }
 databend-common-metrics = { workspace = true }
diff --git a/src/query/service/src/global_services.rs b/src/query/service/src/global_services.rs
@@ -111,7 +111,7 @@ impl GlobalServices {
             CatalogManager::init(config, Arc::new(default_catalog), catalog_creator).await?;
         }
 
-        QueriesQueueManager::init(config.query.max_running_queries as usize)?;
+        QueriesQueueManager::init(config.query.max_running_queries as usize, config).await?;
         HttpQueryManager::init(config).await?;
         ClientSessionManager::init(config).await?;
         DataExchangeManager::init()?;
diff --git a/src/query/service/src/sessions/queue_mgr.rs b/src/query/service/src/sessions/queue_mgr.rs
@@ -28,11 +28,17 @@ use std::time::Instant;
 use std::time::SystemTime;
 
 use databend_common_ast::ast::ExplainKind;
+use databend_common_base::base::escape_for_key;
 use databend_common_base::base::GlobalInstance;
 use databend_common_catalog::table_context::TableContext;
+use databend_common_config::InnerConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_meta_app::principal::UserInfo;
+use databend_common_meta_semaphore::acquirer::Permit;
+use databend_common_meta_semaphore::errors::AcquireError;
+use databend_common_meta_store::MetaStore;
+use databend_common_meta_store::MetaStoreProvider;
 use databend_common_metrics::session::dec_session_running_acquired_queries;
 use databend_common_metrics::session::inc_session_running_acquired_queries;
 use databend_common_metrics::session::incr_session_queue_abort_count;
@@ -47,18 +53,17 @@ use databend_common_sql::PlanExtras;
 use log::info;
 use parking_lot::Mutex;
 use pin_project_lite::pin_project;
-use tokio::sync::AcquireError;
-use tokio::sync::OwnedSemaphorePermit;
-use tokio::sync::Semaphore;
 use tokio::time::error::Elapsed;
 
 use crate::sessions::QueryContext;
 
 pub trait QueueData: Send + Sync + 'static {
-    type Key: Send + Sync + Eq + Hash + Display + Clone + 'static;
+    type Key: Send + Sync + Eq + Hash + Display + Clone + ToString + 'static;
 
     fn get_key(&self) -> Self::Key;
 
+    fn get_lock_key(&self) -> String;
+
     fn remove_error_message(key: Option<Self::Key>) -> ErrorCode;
 
     fn timeout(&self) -> Duration;
@@ -78,29 +83,39 @@ pub(crate) struct Inner<Data: QueueData> {
 }
 
 pub struct QueueManager<Data: QueueData> {
-    semaphore: Arc<Semaphore>,
+    permits: usize,
+    meta_store: MetaStore,
     queue: Mutex<HashMap<Data::Key, Inner<Data>>>,
 }
 
 impl<Data: QueueData> QueueManager<Data> {
-    pub fn init(permits: usize) -> Result<()> {
+    pub async fn init(permits: usize, conf: &InnerConfig) -> Result<()> {
+        let metastore = {
+            let provider = Arc::new(MetaStoreProvider::new(conf.meta.to_meta_grpc_client_conf()));
+
+            provider.create_meta_store().await.map_err(|e| {
+                ErrorCode::MetaServiceError(format!("Failed to create meta store: {}", e))
+            })?
+        };
+
         info!("queue manager permits: {:?}", permits);
-        GlobalInstance::set(Self::create(permits));
+        GlobalInstance::set(Self::create(permits, metastore));
         Ok(())
     }
 
     pub fn instance() -> Arc<Self> {
         GlobalInstance::get::<Arc<Self>>()
     }
 
-    pub fn create(mut permits: usize) -> Arc<QueueManager<Data>> {
+    pub fn create(mut permits: usize, meta_store: MetaStore) -> Arc<QueueManager<Data>> {
         if permits == 0 {
             permits = usize::MAX >> 4;
         }
 
         Arc::new(QueueManager {
+            permits,
+            meta_store,
             queue: Mutex::new(HashMap::new()),
-            semaphore: Arc::new(Semaphore::new(permits)),
         })
     }
 
@@ -139,9 +154,16 @@ impl<Data: QueueData> QueueManager<Data> {
             );
 
             let timeout = data.timeout();
+            let semaphore_acquire = self.meta_store.new_acquired(
+                data.get_lock_key(),
+                self.permits as u64,
+                data.get_key(), // ID of this acquirer
+                Duration::from_secs(3),
+            );
+
             let future = AcquireQueueFuture::create(
                 Arc::new(data),
-                tokio::time::timeout(timeout, self.semaphore.clone().acquire_owned()),
+                tokio::time::timeout(timeout, semaphore_acquire),
                 self.clone(),
             );
             let start_time = SystemTime::now();
@@ -209,9 +231,12 @@ impl<Data: QueueData> QueueManager<Data> {
 
 pub struct AcquireQueueGuard {
     #[allow(dead_code)]
-    permit: Option<OwnedSemaphorePermit>,
+    permit: Option<Permit>,
 }
 
+unsafe impl Send for AcquireQueueGuard {}
+unsafe impl Sync for AcquireQueueGuard {}
+
 impl Drop for AcquireQueueGuard {
     fn drop(&mut self) {
         if self.permit.is_some() {
@@ -221,14 +246,14 @@ impl Drop for AcquireQueueGuard {
 }
 
 impl AcquireQueueGuard {
-    pub fn create(permit: Option<OwnedSemaphorePermit>) -> Self {
+    pub fn create(permit: Option<Permit>) -> Self {
         AcquireQueueGuard { permit }
     }
 }
 
 pin_project! {
     pub struct AcquireQueueFuture<Data: QueueData, T>
-where T: Future<Output =  std::result::Result< std::result::Result<OwnedSemaphorePermit, AcquireError>, Elapsed>>
+where T: Future<Output =  std::result::Result<std::result::Result<Permit, AcquireError>, Elapsed>>
 {
     #[pin]
     inner: T,
@@ -243,12 +268,7 @@ where T: Future<Output =  std::result::Result< std::result::Result<OwnedSemaphor
 }
 
 impl<Data: QueueData, T> AcquireQueueFuture<Data, T>
-where T: Future<
-        Output = std::result::Result<
-            std::result::Result<OwnedSemaphorePermit, AcquireError>,
-            Elapsed,
-        >,
-    >
+where T: Future<Output = std::result::Result<std::result::Result<Permit, AcquireError>, Elapsed>>
 {
     pub fn create(data: Arc<Data>, inner: T, mgr: Arc<QueueManager<Data>>) -> Self {
         AcquireQueueFuture {
@@ -263,12 +283,7 @@ where T: Future<
 }
 
 impl<Data: QueueData, T> Future for AcquireQueueFuture<Data, T>
-where T: Future<
-        Output = std::result::Result<
-            std::result::Result<OwnedSemaphorePermit, AcquireError>,
-            Elapsed,
-        >,
-    >
+where T: Future<Output = std::result::Result<std::result::Result<Permit, AcquireError>, Elapsed>>
 {
     type Output = Result<AcquireQueueGuard>;
 
@@ -443,6 +458,22 @@ impl QueueData for QueryEntry {
         self.query_id.clone()
     }
 
+    fn get_lock_key(&self) -> String {
+        let cluster = self.ctx.get_cluster();
+        let local_id = escape_for_key(&cluster.local_id).unwrap();
+        let mut lock_key = format!("__fd_queries_queue/lost_contact/{}", local_id);
+
+        for node in &cluster.nodes {
+            if node.id == cluster.local_id {
+                let cluster_id = escape_for_key(&node.cluster_id).unwrap();
+                let warehouse_id = escape_for_key(&node.warehouse_id).unwrap();
+                lock_key = format!("__fd_queries_queue/queue/{}/{}", warehouse_id, cluster_id);
+            }
+        }
+
+        lock_key
+    }
+
     fn remove_error_message(key: Option<Self::Key>) -> ErrorCode {
         match key {
             None => ErrorCode::AbortedQuery("The query has be kill while in queries queue"),
diff --git a/src/query/service/tests/it/sessions/queue_mgr.rs b/src/query/service/tests/it/sessions/queue_mgr.rs

Original file line number	Diff line number	Diff line change
`@@ -38,7 +38,7 @@ use crate::storage::PermitKey;`
`38`	`38`	`/// the oneshot sender signals the lease extending task to stop,`
`39`	`39`	/// allowing for proper cleanup of the [`PermitEntry`] in the meta-service.
`40`	`40`	`pub struct Permit {`
`41`		`- pub(crate) fu: BoxFuture<'static, Result<(), ConnectionClosed>>,`
	`41`	`+ pub fu: BoxFuture<'static, Result<(), ConnectionClosed>>,`
`42`	`42`	`}`
`43`	`43`
`44`	`44`	`impl std::fmt::Debug for Permit {`
Original file line number	Diff line number	Diff line change
`@@ -111,7 +111,7 @@ impl GlobalServices {`
`111`	`111`	`CatalogManager::init(config, Arc::new(default_catalog), catalog_creator).await?;`
`112`	`112`	`}`
`113`	`113`
`114`		`- QueriesQueueManager::init(config.query.max_running_queries as usize)?;`
	`114`	`+ QueriesQueueManager::init(config.query.max_running_queries as usize, config).await?;`
`115`	`115`	`HttpQueryManager::init(config).await?;`
`116`	`116`	`ClientSessionManager::init(config).await?;`
`117`	`117`	`DataExchangeManager::init()?;`