Send drop tokens over separate channel

apis/rust/node/src/daemon_connection/control_channel.rs

@@ -1,31 +1,25 @@
-use super::{communication::DaemonChannel, EventStreamThreadHandle};
+use super::communication::DaemonChannel;
 use dora_core::{
     config::{DataId, NodeId},
     daemon_messages::{DaemonRequest, Data, DataflowId},
     message::Metadata,
 };
 use eyre::{bail, eyre, Context};
-use std::sync::Arc;
 
 pub(crate) struct ControlChannel {
     channel: DaemonChannel,
-    _event_stream_thread_handle: Option<Arc<EventStreamThreadHandle>>,
 }
 
 impl ControlChannel {
-    #[tracing::instrument(skip(channel, event_stream_thread_handle), level = "trace")]
+    #[tracing::instrument(skip(channel), level = "trace")]
     pub(crate) fn init(
         dataflow_id: DataflowId,
         node_id: &NodeId,
         mut channel: DaemonChannel,
-        event_stream_thread_handle: Arc<EventStreamThreadHandle>,
     ) -> eyre::Result<Self> {
         channel.register(dataflow_id, node_id.clone())?;
 
-        Ok(Self {
-            channel,
-            _event_stream_thread_handle: Some(event_stream_thread_handle),
-        })
+        Ok(Self { channel })
     }
 
     pub fn report_stop(&mut self) -> eyre::Result<()> {

apis/rust/node/src/daemon_connection/drop_stream.rs

@@ -0,0 +1,63 @@
+use super::DaemonChannel;
+use dora_core::{
+    config::NodeId,
+    daemon_messages::{DaemonReply, DaemonRequest, DataflowId, DropToken, NodeDropEvent},
+};
+use eyre::{eyre, Context};
+
+pub(crate) fn init(
+    dataflow_id: DataflowId,
+    node_id: &NodeId,
+    mut channel: DaemonChannel,
+) -> eyre::Result<flume::Receiver<DropToken>> {
+    channel.register(dataflow_id, node_id.clone())?;
+
+    channel
+        .request(&DaemonRequest::SubscribeDrop)
+        .map_err(|e| eyre!(e))
+        .wrap_err("failed to create subscription with dora-daemon")?;
+
+    let (tx, rx) = flume::bounded(0);
+    let node_id_cloned = node_id.clone();
+
+    std::thread::spawn(|| drop_stream_loop(node_id_cloned, tx, channel));
+
+    Ok(rx)
+}
+
+#[tracing::instrument(skip(tx, channel))]
+fn drop_stream_loop(node_id: NodeId, tx: flume::Sender<DropToken>, mut channel: DaemonChannel) {
+    'outer: loop {
+        let daemon_request = DaemonRequest::NextFinishedDropTokens;
+        let events = match channel.request(&daemon_request) {
+            Ok(DaemonReply::NextDropEvents(events)) if events.is_empty() => {
+                tracing::debug!("Drop stream closed for node ID `{node_id}`");
+                break;
+            }
+            Ok(DaemonReply::NextDropEvents(events)) => events,
+            Ok(other) => {
+                let err = eyre!("unexpected drop reply: {other:?}");
+                tracing::warn!("{err:?}");
+                continue;
+            }
+            Err(err) => {
+                let err = eyre!(err).wrap_err("failed to receive incoming drop event");
+                tracing::warn!("{err:?}");
+                continue;
+            }
+        };
+        for event in events {
+            match event {
+                NodeDropEvent::OutputDropped { drop_token } => {
+                    if tx.send(drop_token).is_err() {
+                        tracing::warn!(
+                            "drop channel was closed already, could not forward \
+                            drop token`{drop_token:?}`"
+                        );
+                        break 'outer;
+                    }
+                }
+            }
+        }
+    }
+}

apis/rust/node/src/daemon_connection/event_stream.rs

@@ -3,10 +3,7 @@ use dora_core::{
     daemon_messages::{self, DaemonReply, DaemonRequest, DataflowId, DropToken, NodeEvent},
 };
 use eyre::{eyre, Context};
-use std::{
-    sync::Arc,
-    time::{Duration, Instant},
-};
+use std::time::{Duration, Instant};
 
 use crate::{event::Data, Event, MappedInputData};
 
@@ -15,7 +12,7 @@ use super::{DaemonChannel, EventStreamThreadHandle};
 pub struct EventStream {
     node_id: NodeId,
     receiver: flume::Receiver<EventItem>,
-    _thread_handle: Arc<EventStreamThreadHandle>,
+    _thread_handle: EventStreamThreadHandle,
 }
 
 impl EventStream {
@@ -23,11 +20,7 @@ impl EventStream {
         dataflow_id: DataflowId,
         node_id: &NodeId,
         mut channel: DaemonChannel,
-    ) -> eyre::Result<(
-        Self,
-        Arc<EventStreamThreadHandle>,
-        flume::Receiver<DropToken>,
-    )> {
+    ) -> eyre::Result<Self> {
         channel.register(dataflow_id, node_id.clone())?;
 
         channel
@@ -37,22 +30,15 @@ impl EventStream {
 
         let (tx, rx) = flume::bounded(0);
         let node_id_cloned = node_id.clone();
-        let (finished_drop_tokens, finished_drop_tokens_rx) = flume::unbounded();
 
-        let join_handle = std::thread::spawn(|| {
-            event_stream_loop(node_id_cloned, tx, channel, finished_drop_tokens)
-        });
+        let join_handle = std::thread::spawn(|| event_stream_loop(node_id_cloned, tx, channel));
         let thread_handle = EventStreamThreadHandle::new(join_handle);
 
-        Ok((
-            EventStream {
-                node_id: node_id.clone(),
-                receiver: rx,
-                _thread_handle: thread_handle.clone(),
-            },
-            thread_handle,
-            finished_drop_tokens_rx,
-        ))
+        Ok(EventStream {
+            node_id: node_id.clone(),
+            receiver: rx,
+            _thread_handle: thread_handle,
+        })
     }
 
     pub fn recv(&mut self) -> Option<Event> {
@@ -108,13 +94,6 @@ impl EventStream {
                     tracing::error!("{err:?}");
                     Event::Error(err.wrap_err("internal error").to_string())
                 }
-                NodeEvent::OutputDropped { .. } => {
-                    let err = eyre!(
-                        "received OutputDrop event, which should be handled by background task"
-                    );
-                    tracing::error!("{err:?}");
-                    Event::Error(err.wrap_err("internal error").to_string())
-                }
             },
             EventItem::FatalError(err) => {
                 Event::Error(format!("fatal event stream error: {err:?}"))
@@ -125,13 +104,8 @@ impl EventStream {
     }
 }
 
-#[tracing::instrument(skip(tx, channel, finished_drop_tokens))]
-fn event_stream_loop(
-    node_id: NodeId,
-    tx: flume::Sender<EventItem>,
-    mut channel: DaemonChannel,
-    finished_drop_tokens: flume::Sender<DropToken>,
-) {
+#[tracing::instrument(skip(tx, channel))]
+fn event_stream_loop(node_id: NodeId, tx: flume::Sender<EventItem>, mut channel: DaemonChannel) {
     let mut tx = Some(tx);
     let mut pending_drop_tokens: Vec<(DropToken, flume::Receiver<()>, Instant, u64)> = Vec::new();
     let mut drop_tokens = Vec::new();
@@ -172,13 +146,6 @@ fn event_stream_loop(
                     // skip this internal event
                     continue;
                 }
-                NodeEvent::OutputDropped { drop_token } => {
-                    if let Err(flume::SendError(token)) = finished_drop_tokens.send(*drop_token) {
-                        tracing::error!("failed to report drop_token `{token:?}` to dora node");
-                    }
-                    // skip this internal event
-                    continue;
-                }
                 _ => None,
             };
 
@@ -191,14 +158,11 @@ fn event_stream_loop(
                     Ok(()) => {}
                     Err(send_error) => {
                         let event = send_error.into_inner();
-                        tracing::debug!(
+                        tracing::trace!(
                             "event channel was closed already, could not forward `{event:?}`"
                         );
-                        if finished_drop_tokens.is_disconnected() {
-                            // both the event stream and the dora node were dropped
-                            // -> break from the `event_stream_loop`
-                            break 'outer Ok(());
-                        }
+
+                        break 'outer Ok(());
                     }
                 }
 

apis/rust/node/src/daemon_connection/mod.rs

@@ -5,13 +5,14 @@ use dora_core::{
 };
 use eyre::Context;
 use flume::RecvTimeoutError;
-use std::{net::TcpStream, sync::Arc, time::Duration};
+use std::{net::TcpStream, time::Duration};
 
 pub(crate) use control_channel::ControlChannel;
 pub use event_stream::EventStream;
 
 mod communication;
 mod control_channel;
+mod drop_stream;
 mod event_stream;
 
 pub(crate) struct DaemonConnection {
@@ -26,34 +27,40 @@ impl DaemonConnection {
         node_id: &NodeId,
         daemon_communication: &DaemonCommunication,
     ) -> eyre::Result<Self> {
-        let (control, events) = match daemon_communication {
+        let (control, events, drop) = match daemon_communication {
             DaemonCommunication::Shmem {
                 daemon_control_region_id,
                 daemon_events_region_id,
+                daemon_drop_region_id,
             } => {
-                let control = unsafe { DaemonChannel::new_shmem(daemon_control_region_id) }
-                    .wrap_err("failed to create shmem control channel")?;
                 let events = unsafe { DaemonChannel::new_shmem(daemon_events_region_id) }
                     .wrap_err("failed to create shmem event channel")?;
-                (control, events)
+                let drop = unsafe { DaemonChannel::new_shmem(daemon_drop_region_id) }
+                    .wrap_err("failed to create shmem event channel")?;
+                let control = unsafe { DaemonChannel::new_shmem(daemon_control_region_id) }
+                    .wrap_err("failed to create shmem control channel")?;
+                (control, events, drop)
             }
             DaemonCommunication::Tcp { socket_addr } => {
-                let control = DaemonChannel::new_tcp(
-                    TcpStream::connect(socket_addr).wrap_err("failed to connect control stream")?,
-                )?;
                 let events = DaemonChannel::new_tcp(
                     TcpStream::connect(socket_addr).wrap_err("failed to connect event stream")?,
                 )?;
-                (control, events)
+                let drop = DaemonChannel::new_tcp(
+                    TcpStream::connect(socket_addr).wrap_err("failed to connect drop stream")?,
+                )?;
+                let control = DaemonChannel::new_tcp(
+                    TcpStream::connect(socket_addr).wrap_err("failed to connect control stream")?,
+                )?;
+                (control, events, drop)
             }
         };
 
-        let (event_stream, event_stream_thread_handle, finished_drop_tokens) =
-            EventStream::init(dataflow_id, node_id, events)
-                .wrap_err("failed to init event stream")?;
-        let control_channel =
-            ControlChannel::init(dataflow_id, node_id, control, event_stream_thread_handle)
-                .wrap_err("failed to init control stream")?;
+        let event_stream = EventStream::init(dataflow_id, node_id, events)
+            .wrap_err("failed to init event stream")?;
+        let control_channel = ControlChannel::init(dataflow_id, node_id, control)
+            .wrap_err("failed to init control stream")?;
+        let finished_drop_tokens =
+            drop_stream::init(dataflow_id, node_id, drop).wrap_err("failed to init drop stream")?;
 
         Ok(Self {
             control_channel,
@@ -65,12 +72,12 @@ impl DaemonConnection {
 
 pub(crate) struct EventStreamThreadHandle(flume::Receiver<std::thread::Result<()>>);
 impl EventStreamThreadHandle {
-    fn new(join_handle: std::thread::JoinHandle<()>) -> Arc<Self> {
+    fn new(join_handle: std::thread::JoinHandle<()>) -> Self {
         let (tx, rx) = flume::bounded(1);
         std::thread::spawn(move || {
             let _ = tx.send(join_handle.join());
         });
-        Arc::new(Self(rx))
+        Self(rx)
     }
 }
 

binaries/daemon/src/lib.rs

@@ -385,6 +385,16 @@ impl Daemon {
                     dataflow.start(&self.events_tx).await?;
                 }
             }
+            DaemonNodeEvent::SubscribeDrop {
+                event_sender,
+                reply_sender,
+            } => {
+                let dataflow = self.running.get_mut(&dataflow_id).wrap_err_with(|| {
+                    format!("failed to subscribe: no running dataflow with ID `{dataflow_id}`")
+                })?;
+                dataflow.drop_channels.insert(node_id, event_sender);
+                let _ = reply_sender.send(DaemonReply::Result(Ok(())));
+            }
             DaemonNodeEvent::CloseOutputs {
                 outputs,
                 reply_sender,
@@ -829,6 +839,7 @@ pub struct RunningDataflow {
     subscribe_replies: HashMap<NodeId, (oneshot::Sender<DaemonReply>, Result<(), String>)>,
 
     subscribe_channels: HashMap<NodeId, UnboundedSender<daemon_messages::NodeEvent>>,
+    drop_channels: HashMap<NodeId, UnboundedSender<daemon_messages::NodeDropEvent>>,
     mappings: HashMap<OutputId, BTreeSet<InputId>>,
     timers: BTreeMap<Duration, BTreeSet<InputId>>,
     open_inputs: BTreeMap<NodeId, BTreeSet<DataId>>,
@@ -853,6 +864,7 @@ impl RunningDataflow {
             pending_nodes: nodes.iter().map(|n| n.id.clone()).collect(),
             subscribe_replies: HashMap::new(),
             subscribe_channels: HashMap::new(),
+            drop_channels: HashMap::new(),
             mappings: HashMap::new(),
             timers: BTreeMap::new(),
             open_inputs: BTreeMap::new(),
@@ -929,9 +941,9 @@ impl RunningDataflow {
             std::collections::hash_map::Entry::Occupied(entry) => {
                 if entry.get().pending_nodes.is_empty() {
                     let (drop_token, info) = entry.remove_entry();
-                    let result = match self.subscribe_channels.get_mut(&info.owner) {
+                    let result = match self.drop_channels.get_mut(&info.owner) {
                         Some(channel) => channel
-                            .send(daemon_messages::NodeEvent::OutputDropped { drop_token })
+                            .send(daemon_messages::NodeDropEvent::OutputDropped { drop_token })
                             .wrap_err("send failed"),
                         None => Err(eyre!("no subscribe channel for node `{}`", &info.owner)),
                     };
@@ -994,6 +1006,10 @@ pub enum DaemonNodeEvent {
         event_sender: UnboundedSender<daemon_messages::NodeEvent>,
         reply_sender: oneshot::Sender<DaemonReply>,
     },
+    SubscribeDrop {
+        event_sender: UnboundedSender<daemon_messages::NodeDropEvent>,
+        reply_sender: oneshot::Sender<DaemonReply>,
+    },
     CloseOutputs {
         outputs: Vec<dora_core::config::DataId>,
         reply_sender: oneshot::Sender<DaemonReply>,

binaries/daemon/src/listener/mod.rs

@@ -3,7 +3,7 @@ use dora_core::{
     config::{DataId, NodeId},
     daemon_messages::{
         DaemonCommunication, DaemonCommunicationConfig, DaemonReply, DaemonRequest, DataflowId,
-        NodeEvent,
+        NodeDropEvent, NodeEvent,
     },
 };
 use eyre::{eyre, Context};
@@ -67,8 +67,13 @@ pub async fn spawn_listener_loop(
                 .size(4096)
                 .create()
                 .wrap_err("failed to allocate daemon_events_region")?;
+            let daemon_drop_region = ShmemConf::new()
+                .size(4096)
+                .create()
+                .wrap_err("failed to allocate daemon_drop_region")?;
             let daemon_control_region_id = daemon_control_region.get_os_id().to_owned();
             let daemon_events_region_id = daemon_events_region.get_os_id().to_owned();
+            let daemon_drop_region_id = daemon_drop_region.get_os_id().to_owned();
 
             {
                 let server = unsafe { ShmemServer::new(daemon_control_region) }
@@ -83,15 +88,28 @@ pub async fn spawn_listener_loop(
                     .wrap_err("failed to create events server")?;
                 let event_loop_node_id = format!("{dataflow_id}/{node_id}");
                 let daemon_tx = daemon_tx.clone();
+                let queue_sizes = queue_sizes.clone();
                 tokio::task::spawn(async move {
                     shmem::listener_loop(server, daemon_tx, queue_sizes).await;
                     tracing::debug!("event listener loop finished for `{event_loop_node_id}`");
                 });
             }
 
+            {
+                let server = unsafe { ShmemServer::new(daemon_drop_region) }
+                    .wrap_err("failed to create drop server")?;
+                let drop_loop_node_id = format!("{dataflow_id}/{node_id}");
+                let daemon_tx = daemon_tx.clone();
+                tokio::task::spawn(async move {
+                    shmem::listener_loop(server, daemon_tx, queue_sizes).await;
+                    tracing::debug!("drop listener loop finished for `{drop_loop_node_id}`");
+                });
+            }
+
             Ok(DaemonCommunication::Shmem {
                 daemon_control_region_id,
                 daemon_events_region_id,
+                daemon_drop_region_id,
             })
         }
     }
@@ -102,6 +120,7 @@ struct Listener {
     node_id: NodeId,
     daemon_tx: mpsc::Sender<Event>,
     subscribed_events: Option<UnboundedReceiver<NodeEvent>>,
+    subscribed_drop_events: Option<UnboundedReceiver<NodeDropEvent>>,
     queue: VecDeque<Box<Option<NodeEvent>>>,
     queue_sizes: BTreeMap<DataId, usize>,
 }
@@ -155,6 +174,7 @@ impl Listener {
                             node_id,
                             daemon_tx,
                             subscribed_events: None,
+                            subscribed_drop_events: None,
                             queue_sizes,
                             queue: VecDeque::new(),
                         };
@@ -337,6 +357,20 @@ impl Listener {
                 .await?;
                 self.subscribed_events = Some(rx);
             }
+            DaemonRequest::SubscribeDrop => {
+                let (tx, rx) = mpsc::unbounded_channel();
+                let (reply_sender, reply) = oneshot::channel();
+                self.process_daemon_event(
+                    DaemonNodeEvent::SubscribeDrop {
+                        event_sender: tx,
+                        reply_sender,
+                    },
+                    Some(reply),
+                    connection,
+                )
+                .await?;
+                self.subscribed_drop_events = Some(rx);
+            }
             DaemonRequest::NextEvent { drop_tokens } => {
                 self.report_drop_tokens(drop_tokens).await?;
 
@@ -362,9 +396,11 @@ impl Listener {
                     DaemonReply::NextEvents(queued_events)
                 };
 
-                self.send_reply(reply, connection)
+                tracing::trace!("sending NextEvent reply: {reply:?}");
+
+                self.send_reply(reply.clone(), connection)
                     .await
-                    .wrap_err("failed to send NextEvent reply")?;
+                    .wrap_err_with(|| format!("failed to send NextEvent reply: {reply:?}"))?;
             }
             DaemonRequest::ReportDropTokens { drop_tokens } => {
                 self.report_drop_tokens(drop_tokens).await?;
@@ -373,6 +409,22 @@ impl Listener {
                     .await
                     .wrap_err("failed to send ReportDropTokens reply")?;
             }
+            DaemonRequest::NextFinishedDropTokens => {
+                let reply = match self.subscribed_drop_events.as_mut() {
+                    // wait for next event
+                    Some(events) => match events.recv().await {
+                        Some(event) => DaemonReply::NextDropEvents(vec![event]),
+                        None => DaemonReply::NextEvents(vec![]),
+                    },
+                    None => DaemonReply::Result(Err("Ignoring event request because no drop \
+                        subscribe message was sent yet"
+                        .into())),
+                };
+
+                self.send_reply(reply.clone(), connection)
+                    .await
+                    .wrap_err_with(|| format!("failed to send NextEvent reply: {reply:?}"))?;
+            }
         }
         Ok(())
     }

binaries/runtime/src/lib.rs

@@ -110,7 +110,7 @@ fn queue_sizes(config: &OperatorConfig) -> std::collections::BTreeMap<DataId, us
     sizes
 }
 
-#[tracing::instrument(skip(operator_events, operator_channels), fields(node.id), level = "trace" )]
+#[tracing::instrument(skip(operator_events, operator_channels), level = "trace")]
 async fn run(
     operators: HashMap<OperatorId, OperatorConfig>,
     config: NodeConfig,
@@ -200,7 +200,8 @@ async fn run(
                         operator_channels.remove(&operator_id);
 
                         if operator_channels.is_empty() {
-                            break;
+                            tracing::warn!("operator channels is empty");
+                            // break;
                         }
                     }
                     OperatorEvent::Output {

libraries/core/src/daemon_messages.rs

@@ -20,6 +20,7 @@ pub enum DaemonCommunication {
     Shmem {
         daemon_control_region_id: SharedMemoryId,
         daemon_events_region_id: SharedMemoryId,
+        daemon_drop_region_id: SharedMemoryId,
     },
     Tcp {
         socket_addr: SocketAddr,
@@ -53,6 +54,8 @@ pub enum DaemonRequest {
     ReportDropTokens {
         drop_tokens: Vec<DropToken>,
     },
+    SubscribeDrop,
+    NextFinishedDropTokens,
 }
 
 impl DaemonRequest {
@@ -107,15 +110,16 @@ impl fmt::Debug for Data {
 
 type SharedMemoryId = String;
 
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub enum DaemonReply {
     Result(Result<(), String>),
     PreparedMessage { shared_memory_id: SharedMemoryId },
     NextEvents(Vec<NodeEvent>),
+    NextDropEvents(Vec<NodeDropEvent>),
     Empty,
 }
 
-#[derive(Debug, serde::Serialize, serde::Deserialize)]
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub enum NodeEvent {
     Stop,
     Reload {
@@ -130,9 +134,11 @@ pub enum NodeEvent {
         id: DataId,
     },
     AllInputsClosed,
-    OutputDropped {
-        drop_token: DropToken,
-    },
+}
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub enum NodeDropEvent {
+    OutputDropped { drop_token: DropToken },
 }
 
 #[derive(Debug, serde::Serialize, serde::Deserialize)]