fix(s3select): cancel queries after client disconnect (#5560)

rustfs/src/app/select_object.rs

@@ -129,7 +129,13 @@ async fn send_select_events(
         return;
     }
 
-    while let Some(result) = output.next().await {
+    let receiver_closed = tx.closed();
+    tokio::pin!(receiver_closed);
+    while let Some(result) = tokio::select! {
+        biased;
+        _ = &mut receiver_closed => return,
+        result = output.next() => result,
+    } {
         let batch = match result {
             Ok(batch) => batch,
             Err(err) => {
@@ -822,6 +828,110 @@ mod tests {
         assert!(rx.recv().await.is_none());
     }
 
+    #[tokio::test(start_paused = true)]
+    async fn producer_preserves_finite_stream_terminal_events() {
+        let schema = Arc::new(datafusion::arrow::datatypes::Schema::new(vec![datafusion::arrow::datatypes::Field::new(
+            "value",
+            datafusion::arrow::datatypes::DataType::Utf8,
+            false,
+        )]));
+        let batch = |value| {
+            RecordBatch::try_new(schema.clone(), vec![Arc::new(datafusion::arrow::array::StringArray::from(vec![value]))])
+                .expect("test record batch should be valid")
+        };
+        let batches = [Ok(batch("a")), Ok(batch("b"))];
+        let output = Box::pin(RecordBatchStreamAdapter::new(schema, futures::stream::iter(batches)));
+        let (tx, mut rx) = mpsc::channel(8);
+        let validation = SelectValidation {
+            output_format: SelectOutputFormat::Csv(CSVOutput::default()),
+            progress_enabled: false,
+        };
+
+        send_select_events(output, &tx, validation).await;
+        drop(tx);
+
+        assert!(matches!(rx.recv().await, Some(Ok(SelectObjectContentEvent::Cont(_)))));
+        for expected in [b"a\n".as_slice(), b"b\n".as_slice()] {
+            let Some(Ok(SelectObjectContentEvent::Records(records))) = rx.recv().await else {
+                panic!("producer should emit a records event for each batch");
+            };
+            assert_eq!(records.payload.as_deref(), Some(expected));
+        }
+        let Some(Ok(SelectObjectContentEvent::Stats(stats))) = rx.recv().await else {
+            panic!("producer should emit final stats");
+        };
+        assert_eq!(stats.details.and_then(|details| details.bytes_returned), Some(4));
+        assert!(matches!(rx.recv().await, Some(Ok(SelectObjectContentEvent::End(_)))));
+        assert!(rx.recv().await.is_none());
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn producer_drops_query_stream_when_receiver_closes() {
+        let (stream_dropped_tx, stream_dropped_rx) = tokio::sync::oneshot::channel::<()>();
+        let output = Box::pin(RecordBatchStreamAdapter::new(
+            Arc::new(Schema::empty()),
+            futures::stream::once(async move {
+                let _stream_dropped = stream_dropped_tx;
+                futures::future::pending::<Result<RecordBatch, DataFusionError>>().await
+            }),
+        ));
+        let (tx, mut rx) = mpsc::channel(2);
+        let validation = SelectValidation {
+            output_format: SelectOutputFormat::Csv(CSVOutput::default()),
+            progress_enabled: false,
+        };
+        let producer = send_select_events(output, &tx, validation);
+        tokio::pin!(producer);
+
+        assert!(futures::poll!(producer.as_mut()).is_pending());
+        assert!(matches!(rx.recv().await, Some(Ok(SelectObjectContentEvent::Cont(_)))));
+        drop(rx);
+
+        assert!(
+            futures::poll!(producer.as_mut()).is_ready(),
+            "producer should observe the closed receiver"
+        );
+        assert!(
+            stream_dropped_rx.await.is_err(),
+            "query stream should be dropped when the receiver closes"
+        );
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn producer_prefers_closed_receiver_over_ready_query_stream() {
+        let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<()>();
+        let (stream_polled_tx, stream_polled_rx) = tokio::sync::oneshot::channel::<()>();
+        let output = Box::pin(RecordBatchStreamAdapter::new(
+            Arc::new(Schema::empty()),
+            futures::stream::once(async move {
+                ready_rx.await.expect("test should release the query stream");
+                let _ = stream_polled_tx.send(());
+                Ok(RecordBatch::new_empty(Arc::new(Schema::empty())))
+            }),
+        ));
+        let (tx, mut rx) = mpsc::channel(2);
+        let validation = SelectValidation {
+            output_format: SelectOutputFormat::Csv(CSVOutput::default()),
+            progress_enabled: false,
+        };
+        let producer = send_select_events(output, &tx, validation);
+        tokio::pin!(producer);
+
+        assert!(futures::poll!(producer.as_mut()).is_pending());
+        assert!(matches!(rx.recv().await, Some(Ok(SelectObjectContentEvent::Cont(_)))));
+        drop(rx);
+        ready_tx.send(()).expect("test should make the query stream ready");
+
+        assert!(
+            futures::poll!(producer.as_mut()).is_ready(),
+            "producer should prioritize the closed receiver"
+        );
+        assert!(
+            stream_polled_rx.await.is_err(),
+            "closed receiver should win before the ready query stream is consumed"
+        );
+    }
+
     #[test]
     fn validate_defaults_csv_header_and_compression() {
         let mut input = base_input();