dora_node_api/event_stream/
mod.rs

1use std::{
2    collections::{BTreeMap, HashMap, VecDeque},
3    pin::pin,
4    sync::Arc,
5    time::Duration,
6};
7
8use dora_message::{
9    daemon_to_node::{DaemonCommunication, DaemonReply, DataMessage, NodeEvent},
10    id::DataId,
11    node_to_daemon::{DaemonRequest, Timestamped},
12    DataflowId,
13};
14pub use event::{Event, MappedInputData, RawData};
15use futures::{
16    future::{select, Either},
17    Stream, StreamExt,
18};
19use futures_timer::Delay;
20use scheduler::{Scheduler, NON_INPUT_EVENT};
21
22use self::{
23    event::SharedMemoryData,
24    thread::{EventItem, EventStreamThreadHandle},
25};
26use crate::daemon_connection::DaemonChannel;
27use dora_core::{
28    config::{Input, NodeId},
29    uhlc,
30};
31use eyre::{eyre, Context};
32
33mod event;
34pub mod merged;
35mod scheduler;
36mod thread;
37
38pub struct EventStream {
39    node_id: NodeId,
40    receiver: flume::r#async::RecvStream<'static, EventItem>,
41    _thread_handle: EventStreamThreadHandle,
42    close_channel: DaemonChannel,
43    clock: Arc<uhlc::HLC>,
44    scheduler: Scheduler,
45}
46
47impl EventStream {
48    #[tracing::instrument(level = "trace", skip(clock))]
49    pub(crate) fn init(
50        dataflow_id: DataflowId,
51        node_id: &NodeId,
52        daemon_communication: &DaemonCommunication,
53        input_config: BTreeMap<DataId, Input>,
54        clock: Arc<uhlc::HLC>,
55    ) -> eyre::Result<Self> {
56        let channel = match daemon_communication {
57            DaemonCommunication::Shmem {
58                daemon_events_region_id,
59                ..
60            } => unsafe { DaemonChannel::new_shmem(daemon_events_region_id) }.wrap_err_with(
61                || format!("failed to create shmem event stream for node `{node_id}`"),
62            )?,
63            DaemonCommunication::Tcp { socket_addr } => DaemonChannel::new_tcp(*socket_addr)
64                .wrap_err_with(|| format!("failed to connect event stream for node `{node_id}`"))?,
65            #[cfg(unix)]
66            DaemonCommunication::UnixDomain { socket_file } => {
67                DaemonChannel::new_unix_socket(socket_file).wrap_err_with(|| {
68                    format!("failed to connect event stream for node `{node_id}`")
69                })?
70            }
71        };
72
73        let close_channel = match daemon_communication {
74            DaemonCommunication::Shmem {
75                daemon_events_close_region_id,
76                ..
77            } => unsafe { DaemonChannel::new_shmem(daemon_events_close_region_id) }.wrap_err_with(
78                || format!("failed to create shmem event close channel for node `{node_id}`"),
79            )?,
80            DaemonCommunication::Tcp { socket_addr } => DaemonChannel::new_tcp(*socket_addr)
81                .wrap_err_with(|| {
82                    format!("failed to connect event close channel for node `{node_id}`")
83                })?,
84            #[cfg(unix)]
85            DaemonCommunication::UnixDomain { socket_file } => {
86                DaemonChannel::new_unix_socket(socket_file).wrap_err_with(|| {
87                    format!("failed to connect event close channel for node `{node_id}`")
88                })?
89            }
90        };
91
92        let mut queue_size_limit: HashMap<DataId, (usize, VecDeque<EventItem>)> = input_config
93            .iter()
94            .map(|(input, config)| {
95                (
96                    input.clone(),
97                    (config.queue_size.unwrap_or(1), VecDeque::new()),
98                )
99            })
100            .collect();
101
102        queue_size_limit.insert(
103            DataId::from(NON_INPUT_EVENT.to_string()),
104            (1_000, VecDeque::new()),
105        );
106
107        let scheduler = Scheduler::new(queue_size_limit);
108
109        Self::init_on_channel(
110            dataflow_id,
111            node_id,
112            channel,
113            close_channel,
114            clock,
115            scheduler,
116        )
117    }
118
119    pub(crate) fn init_on_channel(
120        dataflow_id: DataflowId,
121        node_id: &NodeId,
122        mut channel: DaemonChannel,
123        mut close_channel: DaemonChannel,
124        clock: Arc<uhlc::HLC>,
125        scheduler: Scheduler,
126    ) -> eyre::Result<Self> {
127        channel.register(dataflow_id, node_id.clone(), clock.new_timestamp())?;
128        let reply = channel
129            .request(&Timestamped {
130                inner: DaemonRequest::Subscribe,
131                timestamp: clock.new_timestamp(),
132            })
133            .map_err(|e| eyre!(e))
134            .wrap_err("failed to create subscription with dora-daemon")?;
135
136        match reply {
137            DaemonReply::Result(Ok(())) => {}
138            DaemonReply::Result(Err(err)) => {
139                eyre::bail!("subscribe failed: {err}")
140            }
141            other => eyre::bail!("unexpected subscribe reply: {other:?}"),
142        }
143
144        close_channel.register(dataflow_id, node_id.clone(), clock.new_timestamp())?;
145
146        let (tx, rx) = flume::bounded(100_000_000);
147
148        let thread_handle = thread::init(node_id.clone(), tx, channel, clock.clone())?;
149
150        Ok(EventStream {
151            node_id: node_id.clone(),
152            receiver: rx.into_stream(),
153            _thread_handle: thread_handle,
154            close_channel,
155            clock,
156            scheduler,
157        })
158    }
159
160    /// wait for the next event on the events stream.
161    pub fn recv(&mut self) -> Option<Event> {
162        futures::executor::block_on(self.recv_async())
163    }
164
165    /// wait for the next event on the events stream until timeout
166    pub fn recv_timeout(&mut self, dur: Duration) -> Option<Event> {
167        futures::executor::block_on(self.recv_async_timeout(dur))
168    }
169
170    pub async fn recv_async(&mut self) -> Option<Event> {
171        loop {
172            if self.scheduler.is_empty() {
173                if let Some(event) = self.receiver.next().await {
174                    self.scheduler.add_event(event);
175                } else {
176                    break;
177                }
178            } else {
179                match select(Delay::new(Duration::from_micros(300)), self.receiver.next()).await {
180                    Either::Left((_elapsed, _)) => break,
181                    Either::Right((Some(event), _)) => self.scheduler.add_event(event),
182                    Either::Right((None, _)) => break,
183                };
184            }
185        }
186        let event = self.scheduler.next();
187        event.map(Self::convert_event_item)
188    }
189
190    pub async fn recv_async_timeout(&mut self, dur: Duration) -> Option<Event> {
191        match select(Delay::new(dur), pin!(self.recv_async())).await {
192            Either::Left((_elapsed, _)) => Some(Self::convert_event_item(EventItem::TimeoutError(
193                eyre!("Receiver timed out"),
194            ))),
195            Either::Right((event, _)) => event,
196        }
197    }
198
199    fn convert_event_item(item: EventItem) -> Event {
200        match item {
201            EventItem::NodeEvent { event, ack_channel } => match event {
202                NodeEvent::Stop => Event::Stop,
203                NodeEvent::Reload { operator_id } => Event::Reload { operator_id },
204                NodeEvent::InputClosed { id } => Event::InputClosed { id },
205                NodeEvent::Input { id, metadata, data } => {
206                    let data = match data {
207                        None => Ok(None),
208                        Some(DataMessage::Vec(v)) => Ok(Some(RawData::Vec(v))),
209                        Some(DataMessage::SharedMemory {
210                            shared_memory_id,
211                            len,
212                            drop_token: _, // handled in `event_stream_loop`
213                        }) => unsafe {
214                            MappedInputData::map(&shared_memory_id, len).map(|data| {
215                                Some(RawData::SharedMemory(SharedMemoryData {
216                                    data,
217                                    _drop: ack_channel,
218                                }))
219                            })
220                        },
221                    };
222                    let data = data.and_then(|data| {
223                        let raw_data = data.unwrap_or(RawData::Empty);
224                        raw_data
225                            .into_arrow_array(&metadata.type_info)
226                            .map(arrow::array::make_array)
227                    });
228                    match data {
229                        Ok(data) => Event::Input {
230                            id,
231                            metadata,
232                            data: data.into(),
233                        },
234                        Err(err) => Event::Error(format!("{err:?}")),
235                    }
236                }
237                NodeEvent::AllInputsClosed => {
238                    let err = eyre!(
239                        "received `AllInputsClosed` event, which should be handled by background task"
240                    );
241                    tracing::error!("{err:?}");
242                    Event::Error(err.wrap_err("internal error").to_string())
243                }
244            },
245
246            EventItem::FatalError(err) => {
247                Event::Error(format!("fatal event stream error: {err:?}"))
248            }
249            EventItem::TimeoutError(err) => {
250                Event::Error(format!("Timeout event stream error: {err:?}"))
251            }
252        }
253    }
254}
255
256impl Stream for EventStream {
257    type Item = Event;
258
259    fn poll_next(
260        mut self: std::pin::Pin<&mut Self>,
261        cx: &mut std::task::Context<'_>,
262    ) -> std::task::Poll<Option<Self::Item>> {
263        self.receiver
264            .poll_next_unpin(cx)
265            .map(|item| item.map(Self::convert_event_item))
266    }
267}
268
269impl Drop for EventStream {
270    #[tracing::instrument(skip(self), fields(%self.node_id))]
271    fn drop(&mut self) {
272        let request = Timestamped {
273            inner: DaemonRequest::EventStreamDropped,
274            timestamp: self.clock.new_timestamp(),
275        };
276        let result = self
277            .close_channel
278            .request(&request)
279            .map_err(|e| eyre!(e))
280            .wrap_err("failed to signal event stream closure to dora-daemon")
281            .and_then(|r| match r {
282                DaemonReply::Result(Ok(())) => Ok(()),
283                DaemonReply::Result(Err(err)) => Err(eyre!("EventStreamClosed failed: {err}")),
284                other => Err(eyre!("unexpected EventStreamClosed reply: {other:?}")),
285            });
286        if let Err(err) = result {
287            tracing::warn!("{err:?}")
288        }
289    }
290}
dora_node_api/event_stream/mod.rs

dora_node_api/event_stream/
mod.rs
