fix(query): Support the release of the query queue after the client cancels the query (#18558)

* fix(query): Support the release of the query queue after the client cancels the query

* fix(query): Support the release of the query queue after the client cancels the query

* Update interpreter.rs

Author: zhang2014

src/query/catalog/src/table_context.rs

@@ -24,6 +24,7 @@ use std::time::SystemTime;
 use dashmap::DashMap;
 use databend_common_base::base::Progress;
 use databend_common_base::base::ProgressValues;
+use databend_common_base::base::WatchNotify;
 use databend_common_base::runtime::ExecutorStatsSnapshot;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
@@ -199,6 +200,7 @@ pub trait TableContext: Send + Sync {
     fn get_id(&self) -> String;
     fn get_current_catalog(&self) -> String;
     fn check_aborting(&self) -> Result<(), ContextError>;
+    fn get_abort_notify(&self) -> Arc<WatchNotify>;
     fn get_abort_checker(self: Arc<Self>) -> AbortChecker
     where Self: 'static {
         struct Checker<S> {

src/query/service/src/servers/http/v1/query/execute_state.rs

@@ -282,6 +282,8 @@ impl Executor {
                     query_id = guard.query_id, reason:? = reason;
                     "[HTTP-QUERY] Query state transitioning from Starting to Stopped"
                 );
+
+                s.ctx.get_abort_notify().notify_waiters();
                 if let Err(e) = &reason {
                     InterpreterQueryLog::log_finish(
                         &s.ctx,

src/query/service/src/sessions/query_ctx.rs

@@ -36,6 +36,7 @@ use dashmap::DashMap;
 use databend_common_base::base::Progress;
 use databend_common_base::base::ProgressValues;
 use databend_common_base::base::SpillProgress;
+use databend_common_base::base::WatchNotify;
 use databend_common_base::runtime::profile::Profile;
 use databend_common_base::runtime::profile::ProfileStatisticsName;
 use databend_common_base::runtime::ExecutorStatsSnapshot;
@@ -888,6 +889,10 @@ impl TableContext for QueryContext {
         self.shared.check_aborting()
     }
 
+    fn get_abort_notify(&self) -> Arc<WatchNotify> {
+        self.shared.abort_notify.clone()
+    }
+
     fn get_error(&self) -> Option<ErrorCode<ContextError>> {
         self.shared.get_error()
     }

src/query/service/src/sessions/query_ctx_shared.rs

@@ -30,6 +30,7 @@ use dashmap::DashMap;
 use databend_common_base::base::short_sql;
 use databend_common_base::base::Progress;
 use databend_common_base::base::SpillProgress;
+use databend_common_base::base::WatchNotify;
 use databend_common_base::runtime::drop_guard;
 use databend_common_base::runtime::ExecutorStatsSnapshot;
 use databend_common_base::runtime::MemStat;
@@ -117,6 +118,7 @@ pub struct QueryContextShared {
     pub(in crate::sessions) running_query_text_hash: Arc<RwLock<Option<String>>>,
     pub(in crate::sessions) running_query_parameterized_hash: Arc<RwLock<Option<String>>>,
     pub(in crate::sessions) aborting: Arc<AtomicBool>,
+    pub(in crate::sessions) abort_notify: Arc<WatchNotify>,
     pub(in crate::sessions) tables_refs: Arc<Mutex<HashMap<DatabaseAndTable, Arc<dyn Table>>>>,
     pub(in crate::sessions) streams_refs: Arc<RwLock<HashMap<DatabaseAndTable, bool>>>,
     pub(in crate::sessions) affect: Arc<Mutex<Option<QueryAffect>>>,
@@ -221,6 +223,7 @@ impl QueryContextShared {
             running_query_text_hash: Arc::new(RwLock::new(None)),
             running_query_parameterized_hash: Arc::new(RwLock::new(None)),
             aborting: Arc::new(AtomicBool::new(false)),
+            abort_notify: Arc::new(WatchNotify::new()),
             tables_refs: Arc::new(Mutex::new(HashMap::new())),
             streams_refs: Default::default(),
             affect: Arc::new(Mutex::new(None)),
@@ -368,7 +371,7 @@ impl QueryContextShared {
         }
 
         self.aborting.store(true, Ordering::Release);
-
+        self.abort_notify.notify_waiters();
         // TODO: Wait for the query to be processed (write out the last error)
     }
 

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

@@ -30,6 +30,7 @@ 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_base::base::WatchNotify;
 use databend_common_base::runtime::workload_group::QuotaValue;
 use databend_common_base::runtime::workload_group::MAX_CONCURRENCY_QUOTA_KEY;
 use databend_common_base::runtime::workload_group::QUERY_QUEUED_TIMEOUT_QUOTA_KEY;
@@ -54,6 +55,7 @@ use databend_common_sql::plans::ModifyColumnAction;
 use databend_common_sql::plans::ModifyTableColumnPlan;
 use databend_common_sql::plans::Plan;
 use databend_common_sql::PlanExtras;
+use futures_util::future::Either;
 use log::info;
 use parking_lot::Mutex;
 use pin_project_lite::pin_project;
@@ -82,6 +84,8 @@ pub trait QueueData: Send + Sync + 'static {
     fn enter_wait_pending(&self) {}
 
     fn exit_wait_pending(&self, _wait_time: Duration) {}
+
+    fn get_abort_notify(&self) -> Arc<WatchNotify>;
 }
 
 pub(crate) struct Inner<Data: QueueData> {
@@ -174,6 +178,20 @@ impl<Data: QueueData> QueueManager<Data> {
     }
 
     pub async fn acquire(self: &Arc<Self>, data: Data) -> Result<AcquireQueueGuard> {
+        let abort_notify = data.get_abort_notify();
+
+        let watch_abort_notify = Box::pin(async move { abort_notify.notified().await });
+
+        let acquire = Box::pin(self.acquire_inner(data));
+        match futures::future::select(acquire, watch_abort_notify).await {
+            Either::Left((left, _)) => left,
+            Either::Right((_, _)) => Err(ErrorCode::AbortedQuery(
+                "[QUERY-QUEUE] recv query abort notify.",
+            )),
+        }
+    }
+
+    async fn acquire_inner(self: &Arc<Self>, data: Data) -> Result<AcquireQueueGuard> {
         if !data.need_acquire_to_queue() {
             info!("[QUERY-QUEUE] Non-heavy queries skip the query queue and execute directly.");
             return Ok(AcquireQueueGuard::create(vec![]));
@@ -490,6 +508,7 @@ pub struct QueryEntry {
     pub timeout: Duration,
     pub lock_ttl: Duration,
     pub need_acquire_to_queue: bool,
+    pub abort_watch_notify: Arc<WatchNotify>,
 }
 
 impl QueryEntry {
@@ -511,6 +530,7 @@ impl QueryEntry {
                 timeout => Duration::from_secs(timeout),
             },
             lock_ttl: Duration::from_secs(settings.get_statement_queue_ttl_in_seconds()?),
+            abort_watch_notify: ctx.get_abort_notify(),
         })
     }
 
@@ -668,6 +688,10 @@ impl QueueData for QueryEntry {
         );
         self.ctx.set_query_queued_duration(wait_time)
     }
+
+    fn get_abort_notify(&self) -> Arc<WatchNotify> {
+        self.abort_watch_notify.clone()
+    }
 }
 
 pub type QueriesQueueManager = QueueManager<QueryEntry>;