diff --git a/apis/rust/node/src/daemon_connection/control_channel.rs b/apis/rust/node/src/daemon_connection/control_channel.rs index d03cf746..a2bdddbf 100644 --- a/apis/rust/node/src/daemon_connection/control_channel.rs +++ b/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>, } 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, ) -> eyre::Result { 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<()> { diff --git a/apis/rust/node/src/daemon_connection/drop_stream.rs b/apis/rust/node/src/daemon_connection/drop_stream.rs new file mode 100644 index 00000000..dfa667b5 --- /dev/null +++ b/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> { + 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, 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; + } + } + } + } + } +} diff --git a/apis/rust/node/src/daemon_connection/event_stream.rs b/apis/rust/node/src/daemon_connection/event_stream.rs index ea697f25..bcf91af0 100644 --- a/apis/rust/node/src/daemon_connection/event_stream.rs +++ b/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, - _thread_handle: Arc, + _thread_handle: EventStreamThreadHandle, } impl EventStream { @@ -23,11 +20,7 @@ impl EventStream { dataflow_id: DataflowId, node_id: &NodeId, mut channel: DaemonChannel, - ) -> eyre::Result<( - Self, - Arc, - flume::Receiver, - )> { + ) -> eyre::Result { 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 { @@ -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, - mut channel: DaemonChannel, - finished_drop_tokens: flume::Sender, -) { +#[tracing::instrument(skip(tx, channel))] +fn event_stream_loop(node_id: NodeId, tx: flume::Sender, 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(()); } } diff --git a/apis/rust/node/src/daemon_connection/mod.rs b/apis/rust/node/src/daemon_connection/mod.rs index d62b3f43..2adfabd4 100644 --- a/apis/rust/node/src/daemon_connection/mod.rs +++ b/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 { - 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>); impl EventStreamThreadHandle { - fn new(join_handle: std::thread::JoinHandle<()>) -> Arc { + 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) } } diff --git a/binaries/daemon/src/lib.rs b/binaries/daemon/src/lib.rs index 22bfa1cc..1df167cb 100644 --- a/binaries/daemon/src/lib.rs +++ b/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, Result<(), String>)>, subscribe_channels: HashMap>, + drop_channels: HashMap>, mappings: HashMap>, timers: BTreeMap>, open_inputs: BTreeMap>, @@ -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, reply_sender: oneshot::Sender, }, + SubscribeDrop { + event_sender: UnboundedSender, + reply_sender: oneshot::Sender, + }, CloseOutputs { outputs: Vec, reply_sender: oneshot::Sender, diff --git a/binaries/daemon/src/listener/mod.rs b/binaries/daemon/src/listener/mod.rs index 31ce5056..73695988 100644 --- a/binaries/daemon/src/listener/mod.rs +++ b/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, subscribed_events: Option>, + subscribed_drop_events: Option>, queue: VecDeque>>, queue_sizes: BTreeMap, } @@ -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(()) } diff --git a/binaries/runtime/src/lib.rs b/binaries/runtime/src/lib.rs index 8ce40671..58eaa0b8 100644 --- a/binaries/runtime/src/lib.rs +++ b/binaries/runtime/src/lib.rs @@ -110,7 +110,7 @@ fn queue_sizes(config: &OperatorConfig) -> std::collections::BTreeMap, 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 { diff --git a/libraries/core/src/daemon_messages.rs b/libraries/core/src/daemon_messages.rs index 79fe7505..bedf7a3d 100644 --- a/libraries/core/src/daemon_messages.rs +++ b/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, }, + 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), + NextDropEvents(Vec), 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)]