Expose AllInputClosed message as a Stop message (dora-rs#1026)

This PR makes it possible for nodes to stop when all input is closed.
This is necessary because when merging two different event stream and we
want to stop when all event from dora finishes we can rely on the stop
messages to stop.

It also removes the complexity of the event stream code.

Author: haixuanTao

apis/c++/node/src/lib.rs

@@ -144,7 +144,7 @@ pub struct DoraEvent(Option<Event>);
 fn event_type(event: &DoraEvent) -> ffi::DoraEventType {
     match &event.0 {
         Some(event) => match event {
-            Event::Stop => ffi::DoraEventType::Stop,
+            Event::Stop(_) => ffi::DoraEventType::Stop,
             Event::Input { .. } => ffi::DoraEventType::Input,
             Event::InputClosed { .. } => ffi::DoraEventType::InputClosed,
             Event::Error(_) => ffi::DoraEventType::Error,

apis/c/node/src/lib.rs

@@ -91,7 +91,7 @@ pub unsafe extern "C" fn dora_next_event(context: *mut c_void) -> *mut c_void {
 pub unsafe extern "C" fn read_dora_event_type(event: *const ()) -> EventType {
     let event: &Event = unsafe { &*event.cast() };
     match event {
-        Event::Stop => EventType::Stop,
+        Event::Stop(_) => EventType::Stop,
         Event::Input { .. } => EventType::Input,
         Event::InputClosed { .. } => EventType::InputClosed,
         Event::Error(_) => EventType::Error,

apis/python/operator/src/lib.rs

@@ -6,7 +6,7 @@ use std::{
 use arrow::pyarrow::ToPyArrow;
 use dora_node_api::{
     merged::{MergeExternalSend, MergedEvent},
-    DoraNode, Event, EventStream, Metadata, MetadataParameters, Parameter,
+    DoraNode, Event, EventStream, Metadata, MetadataParameters, Parameter, StopCause,
 };
 use eyre::{Context, Result};
 use futures::{Stream, StreamExt};
@@ -146,7 +146,7 @@ impl PyEvent {
 
     fn ty(event: &Event) -> &str {
         match event {
-            Event::Stop => "STOP",
+            Event::Stop(_) => "STOP",
             Event::Input { .. } => "INPUT",
             Event::InputClosed { .. } => "INPUT_CLOSED",
             Event::Error(_) => "ERROR",
@@ -158,6 +158,11 @@ impl PyEvent {
         match event {
             Event::Input { id, .. } => Some(id),
             Event::InputClosed { id } => Some(id),
+            Event::Stop(cause) => match cause {
+                StopCause::Manual => Some("MANUAL"),
+                StopCause::AllInputsClosed => Some("ALL_INPUTS_CLOSED"),
+                &_ => None,
+            },
             _ => None,
         }
     }

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

@@ -10,7 +10,7 @@ use shared_memory_extended::{Shmem, ShmemConf};
 #[derive(Debug)]
 #[non_exhaustive]
 pub enum Event {
-    Stop,
+    Stop(StopCause),
     Reload {
         operator_id: Option<OperatorId>,
     },
@@ -25,6 +25,13 @@ pub enum Event {
     Error(String),
 }
 
+#[derive(Debug, Clone)]
+#[non_exhaustive]
+pub enum StopCause {
+    Manual,
+    AllInputsClosed,
+}
+
 pub enum RawData {
     Empty,
     Vec(AVec<u8, ConstAlign<128>>),

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

@@ -11,7 +11,7 @@ use dora_message::{
     node_to_daemon::{DaemonRequest, Timestamped},
     DataflowId,
 };
-pub use event::{Event, MappedInputData, RawData};
+pub use event::{Event, MappedInputData, RawData, StopCause};
 use futures::{
     future::{select, Either},
     Stream, StreamExt,
@@ -199,7 +199,7 @@ impl EventStream {
     fn convert_event_item(item: EventItem) -> Event {
         match item {
             EventItem::NodeEvent { event, ack_channel } => match event {
-                NodeEvent::Stop => Event::Stop,
+                NodeEvent::Stop => Event::Stop(event::StopCause::Manual),
                 NodeEvent::Reload { operator_id } => Event::Reload { operator_id },
                 NodeEvent::InputClosed { id } => Event::InputClosed { id },
                 NodeEvent::Input { id, metadata, data } => {
@@ -234,13 +234,7 @@ impl EventStream {
                         Err(err) => Event::Error(format!("{err:?}")),
                     }
                 }
-                NodeEvent::AllInputsClosed => {
-                    let err = eyre!(
-                        "received `AllInputsClosed` event, which should be handled by background task"
-                    );
-                    tracing::error!("{err:?}");
-                    Event::Error(err.wrap_err("internal error").to_string())
-                }
+                NodeEvent::AllInputsClosed => Event::Stop(event::StopCause::AllInputsClosed),
             },
 
             EventItem::FatalError(err) => {

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

@@ -92,6 +92,7 @@ fn event_stream_loop(
     clock: Arc<uhlc::HLC>,
 ) {
     let mut tx = Some(tx);
+    let mut close_tx = false;
     let mut pending_drop_tokens: Vec<(DropToken, flume::Receiver<()>, Instant, u64)> = Vec::new();
     let mut drop_tokens = Vec::new();
 
@@ -135,10 +136,8 @@ fn event_stream_loop(
                     data: Some(data), ..
                 } => data.drop_token(),
                 NodeEvent::AllInputsClosed => {
-                    // close the event stream
-                    tx = None;
-                    // skip this internal event
-                    continue;
+                    close_tx = true;
+                    None
                 }
                 _ => None,
             };
@@ -166,6 +165,10 @@ fn event_stream_loop(
             } else {
                 tracing::warn!("dropping event because event `tx` was already closed: `{inner:?}`");
             }
+
+            if close_tx {
+                tx = None;
+            };
         }
     };
     if let Err(err) = result {

apis/rust/node/src/lib.rs

@@ -20,7 +20,7 @@ pub use dora_message::{
     metadata::{Metadata, MetadataParameters, Parameter},
     DataflowId,
 };
-pub use event_stream::{merged, Event, EventStream, MappedInputData, RawData};
+pub use event_stream::{merged, Event, EventStream, MappedInputData, RawData, StopCause};
 pub use flume::Receiver;
 pub use node::{arrow_utils, DataSample, DoraNode, ZERO_COPY_THRESHOLD};
 

binaries/daemon/src/spawn.rs

@@ -540,7 +540,7 @@ pub async fn spawn_node(
             // If log is an output, we're sending the logs to the dataflow
             if let Some(stdout_output_name) = &send_stdout_to {
                 // Convert logs to DataMessage
-                let array = message.into_arrow();
+                let array = message.as_str().into_arrow();
 
                 let array: ArrayData = array.into();
                 let total_len = required_data_size(&array);

binaries/runtime/src/lib.rs

@@ -232,10 +232,10 @@ async fn run(
                     }
                 }
             }
-            RuntimeEvent::Event(Event::Stop) => {
+            RuntimeEvent::Event(Event::Stop(cause)) => {
                 // forward stop event to all operators and close the event channels
                 for (_, channel) in operator_channels.drain() {
-                    let _ = channel.send_async(Event::Stop).await;
+                    let _ = channel.send_async(Event::Stop(cause.clone())).await;
                 }
             }
             RuntimeEvent::Event(Event::Reload {

binaries/runtime/src/operator/shared_lib.rs

@@ -182,7 +182,7 @@ impl<'lib> SharedLibraryOperator<'lib> {
             }
 
             let mut operator_event = match event {
-                Event::Stop => dora_operator_api_types::RawEvent {
+                Event::Stop(_) => dora_operator_api_types::RawEvent {
                     input: None,
                     input_closed: None,
                     stop: true,