refactor(query): add lock for meta semaphore request (#18579)

Author: zhang2014

src/common/base/src/runtime/workload_group.rs

@@ -19,6 +19,7 @@ use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
 use std::time::Duration;
 
+use tokio::sync::Mutex;
 use tokio::sync::Semaphore;
 
 use crate::runtime::MemStat;
@@ -80,6 +81,7 @@ pub struct WorkloadGroupResource {
     pub meta: WorkloadGroup,
     pub queue_key: String,
     pub permits: usize,
+    pub mutex: Arc<Mutex<()>>,
     pub semaphore: Arc<Semaphore>,
     pub mem_stat: Arc<MemStat>,
     pub max_memory_usage: Arc<AtomicUsize>,

src/query/management/src/workload/workload_resource.rs

@@ -20,6 +20,7 @@ use std::sync::Mutex;
 use std::sync::PoisonError;
 use std::sync::Weak;
 
+use databend_common_base::base::tokio::sync::Mutex as TokioMutex;
 use databend_common_base::runtime::workload_group::QuotaValue;
 use databend_common_base::runtime::workload_group::WorkloadGroupResource;
 use databend_common_base::runtime::workload_group::MAX_CONCURRENCY_QUOTA_KEY;
@@ -153,6 +154,7 @@ impl WorkloadGroupResourceManagerInner {
                 }
             })),
             permits: Semaphore::MAX_PERMITS,
+            mutex: Arc::new(TokioMutex::new(())),
             semaphore: Arc::new(Semaphore::new(Semaphore::MAX_PERMITS)),
         });
 

src/query/service/src/sessions/queue_mgr.rs

@@ -29,6 +29,7 @@ use std::time::SystemTime;
 
 use databend_common_ast::ast::ExplainKind;
 use databend_common_base::base::escape_for_key;
+use databend_common_base::base::tokio::sync::Mutex as TokioMutex;
 use databend_common_base::base::GlobalInstance;
 use databend_common_base::base::WatchNotify;
 use databend_common_base::runtime::workload_group::QuotaValue;
@@ -100,6 +101,7 @@ pub(crate) struct Inner<Data: QueueData> {
 pub struct QueueManager<Data: QueueData> {
     permits: usize,
     meta_store: MetaStore,
+    tokio_mutex: Arc<TokioMutex<()>>,
     semaphore: Arc<Semaphore>,
     global_statement_queue: bool,
     queue: Mutex<HashMap<Data::Key, Inner<Data>>>,
@@ -152,6 +154,7 @@ impl<Data: QueueData> QueueManager<Data> {
             global_statement_queue,
             queue: Mutex::new(HashMap::new()),
             semaphore: Arc::new(Semaphore::new(permits)),
+            tokio_mutex: Arc::new(TokioMutex::new(())),
         })
     }
 
@@ -241,6 +244,9 @@ impl<Data: QueueData> QueueManager<Data> {
 
                 timeout -= instant.elapsed();
 
+                // Prevent concurrent access to meta and serialize the submission of acquire requests.
+                // This ensures that at most permits + nodes acquirers will be in the queue at any given time.
+                let _guard = workload_group.mutex.clone().lock_owned().await;
                 data.set_status("[QUERY-QUEUE] Waiting for global workload semaphore");
 
                 let workload_queue_guard = self
@@ -331,6 +337,9 @@ impl<Data: QueueData> QueueManager<Data> {
                     data.lock_ttl(),
                 );
 
+                // Prevent concurrent access to meta and serialize the submission of acquire requests.
+                // This ensures that at most permits + nodes acquirers will be in the queue at any given time.
+                let _guard = self.tokio_mutex.clone().lock_owned().await;
                 let acquire = tokio::time::timeout(timeout, semaphore_acquire);
                 let queue_future = AcquireQueueFuture::create(data, acquire, self.clone());
                 queue_future.await

src/query/service/tests/it/sessions/queue_mgr.rs

@@ -19,6 +19,7 @@ use std::time::Instant;
 use std::time::SystemTime;
 use std::time::UNIX_EPOCH;
 
+use databend_common_base::base::tokio::sync::Mutex;
 use databend_common_base::base::WatchNotify;
 use databend_common_base::runtime::workload_group::QuotaValue;
 use databend_common_base::runtime::workload_group::WorkloadGroup;
@@ -680,6 +681,7 @@ async fn test_workload_group_concurrency_control() -> Result<()> {
         },
         queue_key: "test_concurrency_queue".to_string(),
         permits: 2,
+        mutex: Arc::new(Mutex::new(())),
         semaphore: Arc::new(Semaphore::new(2)),
         mem_stat: MemStat::create(String::new()),
         max_memory_usage: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
@@ -736,6 +738,7 @@ async fn test_workload_group_concurrent_queue_acquisition() -> Result<()> {
         },
         queue_key: "test_concurrent_acquisition".to_string(),
         permits: 1,
+        mutex: Arc::new(Default::default()),
         semaphore: Arc::new(Semaphore::new(1)),
         mem_stat: MemStat::create(String::new()),
         max_memory_usage: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
@@ -798,6 +801,7 @@ async fn test_workload_group_multilevel_queue_guards() -> Result<()> {
         },
         queue_key: "test_multilevel_queue".to_string(),
         permits: 3,
+        mutex: Arc::new(Default::default()),
         semaphore: Arc::new(Semaphore::new(3)),
         mem_stat: MemStat::create(String::new()),
         max_memory_usage: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
@@ -849,6 +853,7 @@ async fn test_workload_group_zero_concurrency() -> Result<()> {
         },
         queue_key: "test_zero_queue".to_string(),
         permits: 0,
+        mutex: Arc::new(Default::default()),
         semaphore: Arc::new(Semaphore::new(0)),
         mem_stat: MemStat::create(String::new()),
         max_memory_usage: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
@@ -897,6 +902,7 @@ async fn test_workload_group_with_timeout() -> Result<()> {
         },
         queue_key: "test_timeout_queue".to_string(),
         permits: 1,
+        mutex: Arc::new(Default::default()),
         semaphore: Arc::new(Semaphore::new(1)),
         mem_stat: MemStat::create(String::new()),
         max_memory_usage: Arc::new(std::sync::atomic::AtomicUsize::new(0)),