fix: Deadlock in new-streaming multiplexer (#21963)

Author: orlp

crates/polars-stream/src/nodes/multiplexer.rs

@@ -3,6 +3,7 @@ use std::collections::VecDeque;
 use tokio::sync::mpsc::{UnboundedSender, unbounded_channel};
 
 use super::compute_node_prelude::*;
+use crate::async_primitives::wait_group::WaitGroup;
 use crate::morsel::SourceToken;
 
 // TODO: replace this with an out-of-core buffering solution.
@@ -138,9 +139,10 @@ impl ComputeNode for MultiplexerNode {
             let buffered_source_token = buffered_source_token.clone();
             join_handles.push(scope.spawn_task(TaskPriority::High, async move {
                 loop {
-                    let Ok(morsel) = receiver.recv().await else {
+                    let Ok(mut morsel) = receiver.recv().await else {
                         break;
                     };
+                    drop(morsel.take_consume_token());
 
                     let mut anyone_interested = false;
                     let mut active_listener_interested = false;
@@ -154,11 +156,7 @@ impl ComputeNode for MultiplexerNode {
                                 Err(_) => *buf_sender = Listener::Inactive,
                             },
                             Listener::Buffering(b) => {
-                                // Make sure to count buffered morsels as
-                                // consumed to not block the source.
-                                let mut m = morsel.clone();
-                                m.take_consume_token();
-                                b.push_front(m);
+                                b.push_front(morsel.clone());
                                 anyone_interested = true;
                             },
                             Listener::Inactive => {},
@@ -185,23 +183,28 @@ impl ComputeNode for MultiplexerNode {
             if let Some((buf, mut rx)) = opt_buf_recv {
                 let mut sender = send_port.take().unwrap().serial();
 
+                let wait_group = WaitGroup::default();
                 let buffered_source_token = buffered_source_token.clone();
                 join_handles.push(scope.spawn_task(TaskPriority::High, async move {
                     // First we try to flush all the old buffered data.
                     while let Some(mut morsel) = buf.pop_back() {
                         morsel.replace_source_token(buffered_source_token.clone());
+                        morsel.set_consume_token(wait_group.token());
                         if sender.send(morsel).await.is_err()
                             || buffered_source_token.stop_requested()
                         {
                             break;
                         }
+                        wait_group.wait().await;
                     }
 
                     // Then send along data from the multiplexer.
-                    while let Some(morsel) = rx.recv().await {
+                    while let Some(mut morsel) = rx.recv().await {
+                        morsel.set_consume_token(wait_group.token());
                         if sender.send(morsel).await.is_err() {
                             break;
                         }
+                        wait_group.wait().await;
                     }
                     Ok(())
                 }));