refactor(flight): own outbound streams directly

Author: dqhl76

src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs

@@ -1874,11 +1874,12 @@ mod tests {
     use databend_common_base::runtime::Runtime;
     use databend_common_exception::ErrorCode;
     use futures::StreamExt;
+    use tokio::sync::Semaphore;
 
     use super::QueryCoordinator;
     use crate::servers::flight::NewFlightStream;
-    use crate::servers::flight::v1::exchange::exchange_packet_sender::ExchangePacketSender;
     use crate::servers::flight::v1::packets::DataPacket;
+    use crate::servers::flight::v1::transport::OutboundStream;
     use crate::servers::flight::v1::transport::StreamSendOutcome;
     use crate::servers::flight::v1::transport::frame_lane;
     use crate::servers::flight::v1::transport::reliable::DoExchangeConnector;
@@ -1965,20 +1966,20 @@ mod tests {
         let pending = attached.remove("target").unwrap();
         assert!(attached.is_empty());
 
-        let sender = ExchangePacketSender::from_reliable(vec![pending], 1, &runtime);
+        let stream = pending.start(Arc::new(Semaphore::new(64)), Some(256 * 1024), &runtime);
         let payload = FlightData {
             app_metadata: vec![0, 0].into(),
             data_body: vec![1, 2, 3].into(),
             ..Default::default()
         };
         assert_eq!(
-            sender.send(0, 0, payload).await.unwrap(),
+            stream.send(0, payload).await.unwrap(),
             StreamSendOutcome::Accepted
         );
-        sender.finish_producer().await.unwrap();
+        stream.finish().await.unwrap();
 
         assert!(finish_seen.load(Ordering::SeqCst));
-        assert!(sender.destination(0).is_closed());
+        assert!(stream.is_closed());
     }
 
     #[tokio::test]

src/query/service/src/servers/flight/v1/exchange/exchange_packet_sender.rs

@@ -28,25 +28,22 @@ use databend_common_pipeline::sinks::AsyncSink;
 use databend_common_pipeline::sinks::AsyncSinker;
 
 use super::serde::ExchangeSerializeMeta;
-use crate::servers::flight::v1::exchange::exchange_packet_sender::ExchangePacketSender;
+use crate::servers::flight::v1::transport::OutboundStreamRef;
 use crate::servers::flight::v1::transport::StreamSendOutcome;
 
 pub struct ExchangePacketSink {
-    sender: Arc<ExchangePacketSender>,
-    destination: usize,
+    stream: OutboundStreamRef,
     ignore_exchange: bool,
 }
 
 impl ExchangePacketSink {
     fn create(
         input: Arc<InputPort>,
-        sender: Arc<ExchangePacketSender>,
-        destination: usize,
+        stream: OutboundStreamRef,
         ignore_exchange: bool,
     ) -> ProcessorPtr {
         ProcessorPtr::create(AsyncSinker::create(input, Self {
-            sender,
-            destination,
+            stream,
             ignore_exchange,
         }))
     }
@@ -57,7 +54,7 @@ impl AsyncSink for ExchangePacketSink {
     const NAME: &'static str = "ExchangePacketSink";
 
     async fn on_finish(&mut self) -> Result<()> {
-        self.sender.finish_producer().await
+        self.stream.finish().await
     }
 
     async fn consume(&mut self, mut data_block: DataBlock) -> Result<bool> {
@@ -77,9 +74,7 @@ impl AsyncSink for ExchangePacketSink {
             bytes += packet.bytes_size();
             let flight_data = FlightData::try_from(packet)?;
 
-            if self.sender.send(0, self.destination, flight_data).await?
-                == StreamSendOutcome::ConsumerClosed
-            {
+            if self.stream.send(0, flight_data).await? == StreamSendOutcome::ConsumerClosed {
                 return Ok(true);
             }
         }
@@ -89,14 +84,10 @@ impl AsyncSink for ExchangePacketSink {
     }
 }
 
-pub fn create_packet_writer_item(
-    sender: Arc<ExchangePacketSender>,
-    destination: usize,
-    ignore_exchange: bool,
-) -> PipeItem {
+pub fn create_packet_writer_item(stream: OutboundStreamRef, ignore_exchange: bool) -> PipeItem {
     let input = InputPort::create();
     PipeItem::create(
-        ExchangePacketSink::create(input.clone(), sender, destination, ignore_exchange),
+        ExchangePacketSink::create(input.clone(), stream, ignore_exchange),
         vec![input],
         vec![],
     )