Merge pull request #7679 from zhang2014/fix/runtime_no_schedule

fix(cluster): add statistics receiver runtime

Author: Xuanwo

src/query/service/src/api/rpc/exchange/exchange_manager.rs

@@ -284,7 +284,7 @@ impl DataExchangeManager {
                 let mut build_res = query_coordinator.subscribe_fragment(&ctx, fragment_id)?;
 
                 let exchanges = std::mem::take(&mut query_coordinator.statistics_exchanges);
-                let mut statistics_receiver = StatisticsReceiver::create(ctx.clone(), exchanges);
+                let mut statistics_receiver = StatisticsReceiver::create(ctx.clone(), exchanges)?;
                 statistics_receiver.start();
 
                 let statistics_receiver: Mutex<StatisticsReceiver> =

src/query/service/src/api/rpc/exchange/statistics_receiver.rs

@@ -22,7 +22,9 @@ use common_base::base::select3;
 use common_base::base::tokio;
 use common_base::base::tokio::sync::Notify;
 use common_base::base::tokio::task::JoinHandle;
+use common_base::base::Runtime;
 use common_base::base::Select3Output;
+use common_base::base::TrySpawn;
 use common_exception::ErrorCode;
 use common_exception::Result;
 
@@ -36,17 +38,25 @@ pub struct StatisticsReceiver {
     shutdown_flag: Arc<AtomicBool>,
     shutdown_notify: Arc<Notify>,
     exchange_handler: Vec<JoinHandle<Result<()>>>,
+    runtime: Arc<Runtime>,
 }
 
 impl StatisticsReceiver {
-    pub fn create(ctx: Arc<QueryContext>, exchanges: Vec<FlightExchange>) -> StatisticsReceiver {
-        StatisticsReceiver {
+    pub fn create(
+        ctx: Arc<QueryContext>,
+        exchanges: Vec<FlightExchange>,
+    ) -> Result<StatisticsReceiver> {
+        Ok(StatisticsReceiver {
             ctx,
             exchanges,
             shutdown_flag: Arc::new(AtomicBool::new(false)),
             shutdown_notify: Arc::new(Notify::new()),
             exchange_handler: vec![],
-        }
+            runtime: Arc::new(Runtime::with_worker_threads(
+                2,
+                Some(String::from("StatisticsReceiver")),
+            )?),
+        })
     }
 
     pub fn start(&mut self) {
@@ -55,7 +65,7 @@ impl StatisticsReceiver {
             let shutdown_flag = self.shutdown_flag.clone();
             let shutdown_notify = self.shutdown_notify.clone();
 
-            self.exchange_handler.push(tokio::spawn(async move {
+            self.exchange_handler.push(self.runtime.spawn(async move {
                 let mut recv = Box::pin(flight_exchange.recv());
                 let mut notified = Box::pin(shutdown_notify.notified());
 

tests/logictest/suites/base/03_dml/03_0028_copy_into_stage

@@ -0,0 +1,38 @@
+statement ok
+DROP DATABASE IF EXISTS db1;
+
+statement ok
+CREATE DATABASE db1;
+
+statement ok
+USE db1;
+
+statement ok
+CREATE TABLE test_table(id INTEGER, name VARCHAR, age INT);
+
+statement ok
+insert into test_table (id,name,age) values(1,'2',3), (4, '5', 6);
+
+statement ok
+CREATE STAGE IF NOT EXISTS test;
+
+statement ok
+copy into @test from test_table FILE_FORMAT = (type = 'CSV');
+
+statement ok
+copy into test_table from @test;
+
+statement query A
+SELECT COUNT() FROM test_table;
+
+----
+4
+
+statement ok
+drop table test_table all;
+
+statement ok
+drop stage test;
+
+statement ok
+DROP DATABASE db1;