[PATCH rest-server] rest-server: stop a json-seq stream once its client is gone

[Dietmar Maurer <dietmar@proxmox.com>]

The task forwarding the records of a streamed response only noticed
a closed connection when sending the next record failed. A stream
that stays silent for a long time, such as an API handler waiting on
a slow backend, thus kept running after its client went away, and so
did the work behind it, as the producer only learns of the gone
client once the stream is dropped.

Stop forwarding as soon as the client is gone, which drops the stream
right away.

Signed-off-by: Dietmar Maurer <dietmar@proxmox.com>
---
 proxmox-rest-server/src/rest.rs | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/proxmox-rest-server/src/rest.rs b/proxmox-rest-server/src/rest.rs
index 779851eb..f47761bd 100644
--- a/proxmox-rest-server/src/rest.rs
+++ b/proxmox-rest-server/src/rest.rs
@@ -896,7 +896,16 @@ fn handle_stream_as_json_seq(stream: proxmox_router::Stream) -> Result<Response<
         use futures::StreamExt;
 
         let mut stream = stream.into_inner();
-        while let Some(record) = stream.next().await {
+        loop {
+            // A stream may stay silent for a long time. Stop once the client is gone, which drops
+            // the stream, so its producer notices without having to send another record.
+            let record = tokio::select! {
+                record = stream.next() => record,
+                () = send.closed() => break,
+            };
+            let Some(record) = record else {
+                break;
+            };
             if send.send(Ok(record.to_bytes())).await.is_err() {
                 break;
             }
@@ -1475,6 +1484,21 @@ impl Unformatted {
 mod tests {
     use super::*;
 
+    #[tokio::test]
+    async fn json_seq_stream_stops_once_the_client_is_gone() {
+        let (send, receive) = tokio::sync::mpsc::unbounded_channel::<proxmox_router::Record>();
+        let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(receive);
+        let response = handle_stream_as_json_seq(proxmox_router::Stream::from(stream)).unwrap();
+        tokio::task::yield_now().await;
+        assert!(!send.is_closed());
+
+        // the stream sends nothing, yet it is dropped with the response
+        drop(response);
+        tokio::time::timeout(std::time::Duration::from_secs(5), send.closed())
+            .await
+            .expect("stream of a gone client still running");
+    }
+
     fn headers(name: &'static str, value: &str) -> HeaderMap {
         let mut headers = HeaderMap::new();
         headers.insert(name, value.parse().unwrap());
-- 
2.47.3