Merge pull request #9091 from youngsofun/http

refactor(rest api): drop ctx after query finished.

Author: mergify[bot]

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

@@ -116,7 +116,6 @@ pub struct ExecuteRunning {
 }
 
 pub struct ExecuteStopped {
-    pub ctx: Arc<QueryContext>,
     pub stats: Progresses,
     pub affect: Option<QueryAffect>,
     pub reason: Result<()>,
@@ -183,7 +182,6 @@ impl Executor {
                         .unwrap_or_else(|e| error!("fail to write query_log {:?}", e));
                 }
                 guard.state = Stopped(Box::new(ExecuteStopped {
-                    ctx: s.ctx.clone(),
                     stats: Default::default(),
                     reason,
                     stop_time: Instant::now(),
@@ -203,7 +201,6 @@ impl Executor {
                 }
 
                 guard.state = Stopped(Box::new(ExecuteStopped {
-                    ctx: r.ctx.clone(),
                     stats: Progresses::from_context(&r.ctx),
                     reason,
                     stop_time: Instant::now(),
@@ -227,7 +224,7 @@ impl ExecuteState {
         session: Arc<Session>,
         ctx: Arc<QueryContext>,
         block_sender: SizedChannelSender<DataBlock>,
-    ) -> Result<ExecuteRunning> {
+    ) -> Result<()> {
         let mut planner = Planner::new(ctx.clone());
         let (plan, _, _) = planner.plan_sql(sql).await?;
         ctx.attach_query_str(plan.to_string(), sql);
@@ -238,30 +235,31 @@ impl ExecuteState {
             ctx: ctx.clone(),
         };
 
+        info!("http query {}, change state to Running", &ctx.get_id());
+        Executor::start_to_running(&executor, Running(running_state)).await;
+
         let executor_clone = executor.clone();
         let ctx_clone = ctx.clone();
         let block_sender_closer = block_sender.closer();
 
-        ctx.try_spawn(async move {
-            let res = execute(interpreter, ctx_clone, block_sender, executor_clone.clone());
-            match AssertUnwindSafe(res).catch_unwind().await {
-                Ok(Err(err)) => {
-                    Executor::stop(&executor_clone, Err(err), false).await;
-                    block_sender_closer.close();
-                }
-                Err(_) => {
-                    Executor::stop(
-                        &executor_clone,
-                        Err(ErrorCode::PanicError("interpreter panic!")),
-                        false,
-                    )
-                    .await;
-                    block_sender_closer.close();
-                }
-                _ => {}
+        let res = execute(interpreter, ctx_clone, block_sender, executor_clone.clone());
+        match AssertUnwindSafe(res).catch_unwind().await {
+            Ok(Err(err)) => {
+                Executor::stop(&executor_clone, Err(err), false).await;
+                block_sender_closer.close();
             }
-        })?;
-        Ok(running_state)
+            Err(_) => {
+                Executor::stop(
+                    &executor_clone,
+                    Err(ErrorCode::PanicError("interpreter panic!")),
+                    false,
+                )
+                .await;
+                block_sender_closer.close();
+            }
+            _ => {}
+        }
+        Ok(())
     }
 }
 

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

@@ -239,42 +239,30 @@ impl HttpQuery {
         let block_sender_closer = block_sender.closer();
         let state_clone = state.clone();
         let ctx_clone = ctx.clone();
+        let ctx_clone2 = ctx.clone();
         let sql = request.sql.clone();
         let query_id = id.clone();
         let query_id_clone = id.clone();
         ctx.try_spawn(async move {
             let state = state_clone.clone();
-            let running_state = ExecuteState::try_start_query(
-                state,
-                &sql,
-                session,
-                ctx_clone.clone(),
-                block_sender,
-            )
-            .await;
-            match running_state {
-                Ok(s) => {
-                    tracing::info!("http query {}, change state to Running", &query_id);
-                    Executor::start_to_running(&state_clone, ExecuteState::Running(s)).await;
-                }
-                Err(e) => {
-                    InterpreterQueryLog::fail_to_start(ctx_clone.clone(), e.clone());
-                    let state = ExecuteStopped {
-                        ctx: ctx_clone.clone(),
-                        stats: Progresses::default(),
-                        reason: Err(e.clone()),
-                        stop_time: Instant::now(),
-                        affect: ctx_clone.get_affect(),
-                    };
-                    tracing::info!(
-                        "http query {}, change state to Stopped, fail to start {:?}",
-                        &query_id,
-                        e
-                    );
-                    Executor::start_to_stop(&state_clone, ExecuteState::Stopped(Box::new(state)))
-                        .await;
-                    block_sender_closer.close();
-                }
+            if let Err(e) =
+                ExecuteState::try_start_query(state, &sql, session, ctx_clone.clone(), block_sender)
+                    .await
+            {
+                InterpreterQueryLog::fail_to_start(ctx_clone.clone(), e.clone());
+                let state = ExecuteStopped {
+                    stats: Progresses::default(),
+                    reason: Err(e.clone()),
+                    stop_time: Instant::now(),
+                    affect: ctx_clone.get_affect(),
+                };
+                tracing::info!(
+                    "http query {}, change state to Stopped, fail to start {:?}",
+                    &query_id,
+                    e
+                );
+                Executor::start_to_stop(&state_clone, ExecuteState::Stopped(Box::new(state))).await;
+                block_sender_closer.close();
             }
         })?;
 
@@ -284,6 +272,7 @@ impl HttpQuery {
             request.pagination.max_rows_per_page,
             block_receiver,
             format_settings,
+            ctx_clone2,
         )));
         let query = HttpQuery {
             id,

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

@@ -17,6 +17,8 @@ use std::sync::Arc;
 use std::time::Instant;
 
 use common_base::base::tokio;
+use common_base::base::GlobalIORuntime;
+use common_base::base::TrySpawn;
 use common_datablocks::DataBlock;
 use common_datavalues::DataSchema;
 use common_datavalues::DataSchemaRef;
@@ -29,6 +31,7 @@ use tracing::info;
 use crate::servers::http::v1::json_block::block_to_json_value;
 use crate::servers::http::v1::query::sized_spsc::SizedChannelReceiver;
 use crate::servers::http::v1::JsonBlock;
+use crate::sessions::QueryContext;
 
 #[derive(Debug, PartialEq, Eq)]
 pub enum Wait {
@@ -59,6 +62,7 @@ pub struct PageManager {
     row_buffer: VecDeque<Vec<JsonValue>>,
     block_receiver: SizedChannelReceiver<DataBlock>,
     format_settings: FormatSettings,
+    query_ctx_ref: Option<Arc<QueryContext>>,
 }
 
 impl PageManager {
@@ -67,6 +71,7 @@ impl PageManager {
         max_rows_per_page: usize,
         block_receiver: SizedChannelReceiver<DataBlock>,
         format_settings: FormatSettings,
+        query_ctx_ref: Arc<QueryContext>,
     ) -> PageManager {
         PageManager {
             query_id,
@@ -80,6 +85,7 @@ impl PageManager {
             block_receiver,
             max_rows_per_page,
             format_settings,
+            query_ctx_ref: Some(query_ctx_ref),
         }
     }
 
@@ -194,6 +200,12 @@ impl PageManager {
         if !self.block_end {
             self.block_end = self.block_receiver.is_empty();
         }
+        if self.block_end {
+            let ctx = self.query_ctx_ref.take();
+            GlobalIORuntime::instance().spawn(async move {
+                drop(ctx);
+            });
+        }
         let end = self.block_end && self.row_buffer.is_empty();
         Ok((block, end))
     }