use dora_core::{
config::{DataId, NodeId},
descriptor::{Descriptor, OperatorDefinition},
};
use dora_node_api::{DoraArray, EncodedSample, Event, MetadataParameters, SampleAllocator};
use eyre::Result;
use std::any::Any;
use std::sync::{Arc, OnceLock};
use tokio::sync::{mpsc::Sender, oneshot};
pub type SharedAllocator = Arc<OnceLock<SampleAllocator>>;
#[derive(Clone)]
pub struct RuntimeHandle {
events_tx: Sender<OperatorEvent>,
allocator: SharedAllocator,
}
impl RuntimeHandle {
pub fn new(events_tx: Sender<OperatorEvent>, allocator: SharedAllocator) -> Self {
Self {
events_tx,
allocator,
}
}
pub fn send_output(
&self,
output_id: DataId,
parameters: MetadataParameters,
array: &DoraArray,
) -> Result<()> {
let allocator = self
.allocator
.get()
.ok_or_else(|| eyre::eyre!("cannot send an output before the node is initialized"))?;
let encoded = allocator.encode_arrow(array)?;
self.events_tx
.blocking_send(OperatorEvent::Output {
output_id,
parameters,
encoded,
})
.map_err(|_| eyre::eyre!("failed to send output to runtime"))
}
pub fn report(&self, event: OperatorEvent) {
let _ = self.events_tx.blocking_send(event);
}
}
pub trait OperatorRunner {
fn run_operator(
&self,
node_id: &NodeId,
operator: OperatorDefinition,
incoming_events: flume::Receiver<Event>,
handle: RuntimeHandle,
init_done: oneshot::Sender<Result<()>>,
dataflow_descriptor: &Descriptor,
) -> eyre::Result<RunnerGuard>;
}
pub type RunnerGuard = Option<Box<dyn Any>>;
#[derive(Debug)]
#[allow(clippy::large_enum_variant)]
pub enum OperatorEvent {
Output {
output_id: DataId,
parameters: MetadataParameters,
encoded: EncodedSample,
},
Error(eyre::Error),
Panic(Box<dyn Any + Send>),
Finished {
reason: StopReason,
},
}
#[derive(Debug)]
pub enum StopReason {
InputsClosed,
ExplicitStop,
ExplicitStopAll,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn send_output_releases_the_operator_payload_before_it_crosses_the_channel() {
use arrow::buffer::Buffer;
use std::ptr::NonNull;
use std::sync::atomic::{AtomicBool, Ordering};
struct ForeignOwner {
released: Arc<AtomicBool>,
_backing: Vec<u8>,
}
impl Drop for ForeignOwner {
fn drop(&mut self) {
self.released.store(true, Ordering::SeqCst);
}
}
let backing = vec![0xCDu8; 4096];
let ptr = NonNull::new(backing.as_ptr() as *mut u8).expect("non-null");
let released = Arc::new(AtomicBool::new(false));
let owner = Arc::new(ForeignOwner {
released: released.clone(),
_backing: backing,
});
let buffer = unsafe { Buffer::from_custom_allocation(ptr, 4096, owner) };
let array = dora_node_api::DoraArray::from_array(arrow::array::make_array(
arrow::array::ArrayData::builder(dora_node_api::arrow_v59::datatypes::DataType::UInt8)
.len(4096)
.add_buffer(buffer)
.build()
.expect("valid UInt8 array"),
));
let (events_tx, mut events_rx) = tokio::sync::mpsc::channel(1);
let allocator: SharedAllocator = Arc::new(OnceLock::new());
allocator
.set(SampleAllocator::heap())
.expect("fresh OnceLock");
let handle = RuntimeHandle::new(events_tx, allocator);
handle
.send_output(
DataId::from("out".to_string()),
MetadataParameters::default(),
&array,
)
.expect("send_output");
drop(array);
assert!(
released.load(Ordering::SeqCst),
"the queued event must not keep the operator's payload alive"
);
match events_rx.try_recv().expect("an output event was queued") {
OperatorEvent::Output { encoded, .. } => {
assert!(
encoded.as_bytes().len() > 4096,
"the sample must hold the encoded payload, not a reference to it"
);
}
other => panic!("expected an Output event, got {other:?}"),
}
}
#[test]
fn send_output_before_node_init_is_a_clear_error() {
let (events_tx, _events_rx) = tokio::sync::mpsc::channel(1);
let handle = RuntimeHandle::new(events_tx, SharedAllocator::default());
let array = dora_node_api::DoraArray::from_array(arrow::array::NullArray::new(1));
let err = handle
.send_output(
DataId::from("out".to_string()),
MetadataParameters::default(),
&array,
)
.expect_err("no allocator yet");
assert!(
err.to_string().contains("before the node is initialized"),
"unexpected error: {err}"
);
}
}