fix(flow): drain frontend probe response before selecting peer (#9082)

* fix(flow): drain frontend probe response before selecting peer

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* refactor(flow): reuse record batch collector for frontend probe

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

---------

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
(cherry picked from commit 2d78220a4f)
Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

src/flow/src/batching_mode/frontend_client.rs

@@ -183,13 +183,21 @@ impl DatabaseWithPeer {
     /// Try sending a "SELECT 1" to the database
     async fn try_select_one(&self) -> Result<(), Error> {
         // notice here use `sql` for `SELECT 1` return 1 row
-        let _ = self
+        let output = self
             .database
             .sql("SELECT 1")
             .await
             .with_context(|_| InvalidRequestSnafu {
                 context: format!("Failed to handle `SELECT 1` request at {:?}", self.peer),
             })?;
+
+        if let OutputData::Stream(stream) = output.data {
+            common_recordbatch::util::collect(stream)
+                .await
+                .map_err(BoxedError::new)
+                .context(ExternalSnafu)?;
+        }
+
         Ok(())
     }
 }
@@ -625,9 +633,11 @@ mod tests {
 
     use arrow_flight::flight_service_server::FlightServiceServer;
     use arrow_flight::{FlightData, Ticket};
+    use common_grpc::flight::FlightEncoder;
     use common_query::{Output, OutputData};
     use common_recordbatch::adapter::RecordBatchMetrics;
     use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream};
+    use datatypes::arrow::datatypes::Schema as ArrowSchema;
     use datatypes::prelude::{ConcreteDataType, VectorRef};
     use datatypes::schema::{ColumnSchema, Schema};
     use datatypes::vectors::Int32Vector;
@@ -702,6 +712,14 @@ mod tests {
     #[derive(Debug)]
     struct SlowFlight;
 
+    struct DelayedEofFlight {
+        schema_sent: Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
+        release: Arc<tokio::sync::Notify>,
+    }
+
+    #[derive(Debug)]
+    struct LateStreamErrorFlight;
+
     struct WaitForConcurrentFlight {
         barrier: Arc<tokio::sync::Barrier>,
     }
@@ -797,6 +815,45 @@ mod tests {
         }
     }
 
+    #[async_trait::async_trait]
+    impl FlightCraft for DelayedEofFlight {
+        async fn do_get(
+            &self,
+            _request: TonicRequest<Ticket>,
+        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
+            let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty());
+            let schema_sent = self.schema_sent.lock().unwrap().take();
+            let schema_stream = futures::stream::once(async move {
+                if let Some(schema_sent) = schema_sent {
+                    let _ = schema_sent.send(());
+                }
+                Ok(schema)
+            });
+            let release = self.release.clone();
+            let delayed_eof = futures::stream::unfold(release, |release| async move {
+                release.notified().await;
+                None::<(std::result::Result<FlightData, Status>, _)>
+            });
+
+            Ok(TonicResponse::new(Box::pin(
+                schema_stream.chain(delayed_eof),
+            )))
+        }
+    }
+
+    #[async_trait::async_trait]
+    impl FlightCraft for LateStreamErrorFlight {
+        async fn do_get(
+            &self,
+            _request: TonicRequest<Ticket>,
+        ) -> std::result::Result<TonicResponse<TonicStream<FlightData>>, Status> {
+            let schema = FlightEncoder::default().encode_schema(&ArrowSchema::empty());
+            let stream =
+                futures::stream::iter([Ok(schema), Err(Status::unavailable("late stream error"))]);
+            Ok(TonicResponse::new(Box::pin(stream)))
+        }
+    }
+
     #[async_trait::async_trait]
     impl FlightCraft for WaitForConcurrentFlight {
         async fn do_get(
@@ -983,6 +1040,71 @@ mod tests {
         assert!(format!("{err:?}").contains("Invalid value for flow.return_region_seq"));
     }
 
+    #[tokio::test]
+    async fn test_try_select_one_waits_for_stream_eof() {
+        let (schema_sent, schema_sent_rx) = tokio::sync::oneshot::channel();
+        let release = Arc::new(tokio::sync::Notify::new());
+        let (addr, server) = start_flight_server(DelayedEofFlight {
+            schema_sent: Mutex::new(Some(schema_sent)),
+            release: release.clone(),
+        })
+        .await;
+        let database = Database::new(
+            DEFAULT_CATALOG_NAME,
+            DEFAULT_SCHEMA_NAME,
+            Client::with_urls([addr.as_str()]),
+        );
+        let db = DatabaseWithPeer::new(
+            database,
+            Peer {
+                id: 1,
+                addr: addr.clone(),
+            },
+        );
+        let mut probe = tokio::spawn(async move { db.try_select_one().await });
+
+        timeout(Duration::from_secs(1), schema_sent_rx)
+            .await
+            .expect("server should send the schema")
+            .expect("schema signal should be sent");
+        assert!(
+            timeout(Duration::from_millis(100), &mut probe)
+                .await
+                .is_err(),
+            "SELECT 1 must wait for the delayed stream tail and EOF"
+        );
+
+        release.notify_one();
+        timeout(Duration::from_secs(1), &mut probe)
+            .await
+            .expect("SELECT 1 should complete after EOF")
+            .expect("probe task should not panic")
+            .expect("SELECT 1 should succeed after EOF");
+        server.abort();
+    }
+
+    #[tokio::test]
+    async fn test_try_select_one_propagates_late_stream_error() {
+        let (addr, server) = start_flight_server(LateStreamErrorFlight).await;
+        let database = Database::new(
+            DEFAULT_CATALOG_NAME,
+            DEFAULT_SCHEMA_NAME,
+            Client::with_urls([addr.as_str()]),
+        );
+        let db = DatabaseWithPeer::new(
+            database,
+            Peer {
+                id: 1,
+                addr: addr.clone(),
+            },
+        );
+
+        let err = db.try_select_one().await.unwrap_err();
+        server.abort();
+
+        assert!(format!("{err:?}").contains("late stream error"));
+    }
+
     #[tokio::test]
     async fn test_check_all_frontends_without_auth_fails_fast_on_unauthenticated_frontend() {
         let (addr, server) = start_flight_server(RejectUnauthenticatedFlight).await;