v1.17: exits send_datagram_task if the connection is closed (backport…

… of #33836) (#34327) exits send_datagram_task if the connection is closed (#33836) Waiting on receiver.recv() can unnecessarily block while the connection is already closed. The commit exits send_datagram_task if the connection is closed. (cherry picked from commit 03fbe08) Co-authored-by: behzad nouri <[email protected]>
solana-labs · Dec 5, 2023 · a665996 · a665996
1 parent e3bd31c
commit a665996
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/turbine/src/quic_endpoint.rs b/turbine/src/quic_endpoint.rs
@@ -435,10 +435,21 @@ async fn send_datagram_task(
     connection: Connection,
     mut receiver: AsyncReceiver<Bytes>,
 ) -> Result<(), Error> {
-    while let Some(bytes) = receiver.recv().await {
-        connection.send_datagram(bytes)?;
+    tokio::pin! {
+        let connection_closed = connection.closed();
+    }
+    loop {
+        tokio::select! {
+            biased;
+            bytes = receiver.recv() => {
+                match bytes {
+                    None => return Ok(()),
+                    Some(bytes) => connection.send_datagram(bytes)?,
+                }
+            }
+            err = &mut connection_closed => return Err(Error::from(err)),
+        }
     }
-    Ok(())
 }
 
 async fn make_connection_task(