Merge pull request #10718 from zhang2014/chore/read_finished

chore(cluster): close flight receiver in event

Author: BohuTANG

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

@@ -66,7 +66,8 @@ impl Processor for ExchangeSourceReader {
 
         if self.output.is_finished() {
             if !self.finished {
-                return Ok(Event::Async);
+                self.finished = true;
+                self.flight_receiver.close();
             }
 
             return Ok(Event::Finished);

src/query/service/src/api/rpc/flight_client.rs

@@ -81,6 +81,8 @@ impl FlightClient {
                         break;
                     }
                 }
+
+                tx.close();
             }
         });
 

src/query/service/src/api/rpc/flight_service.rs

@@ -146,6 +146,8 @@ impl FlightService for DatabendQueryFlightService {
                     .create_session(SessionType::FlightRPC)
                     .await?;
                 let ctx = session.create_query_context().await?;
+                // Keep query id
+                ctx.set_id(init_query_fragments_plan.executor_packet.query_id.clone());
 
                 let spawner = ctx.clone();
                 match_join_handle(spawner.spawn(async move {

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

@@ -200,6 +200,10 @@ impl QueryContext {
         self.shared.set_affect(affect)
     }
 
+    pub fn set_id(&self, id: String) {
+        *self.shared.init_query_id.write() = id;
+    }
+
     pub fn set_executor(&self, weak_ptr: Weak<PipelineExecutor>) {
         self.shared.set_executor(weak_ptr)
     }