dora/apis/python/node/src/lib.rs

use dora_node_api::config::{DataId, NodeId};
use dora_node_api::{DoraNode, Input};
use eyre::{Context, Result};
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use std::sync::Arc;
use std::thread;
use tokio::sync::mpsc;
use tokio::sync::mpsc::{Receiver, Sender};

#[pyclass]
pub struct Node {
    id: NodeId,
    pub rx_input: Receiver<Input>,
    pub tx_output: Sender<(String, Vec<u8>)>,
}

pub struct PyInput(Input);

impl IntoPy<PyObject> for PyInput {
    fn into_py(self, py: Python) -> PyObject {
        (self.0.id.to_string(), PyBytes::new(py, &self.0.data)).into_py(py)
    }
}

#[pymethods]
impl Node {
    #[new]
    pub fn new() -> Result<Self> {
        let id = {
            let raw =
                std::env::var("DORA_NODE_ID").wrap_err("env variable DORA_NODE_ID must be set")?;
            serde_yaml::from_str(&raw).context("failed to deserialize operator config")?
        };

        let (tx_input, rx_input) = mpsc::channel(1);
        let (tx_output, mut rx_output) = mpsc::channel::<(String, Vec<u8>)>(1);

        // Dispatching a tokio threadpool enables us to conveniently use Dora Future stream
        // through tokio channel.
        // It would have been difficult to expose the FutureStream of Dora directly.
        thread::spawn(move || -> Result<()> {
            let rt = tokio::runtime::Builder::new_multi_thread().build()?;

            let node = Arc::new(DoraNode::init_from_env()?);
            let _node = node.clone();
            let receive_handle = tokio::task::spawn_blocking(async move {
                let mut inputs = _node.inputs().unwrap();
                while let Ok(input) = inputs.recv() {
                    tx_input.blocking_send(input)?
                }
                Result::<_, eyre::Error>::Ok(())
            });
            let send_handle = tokio::task::spawn_blocking(async move {
                while let Some((output_str, data)) = rx_output.recv().await {
                    let output_id = DataId::from(output_str);
                    node.send_output(&output_id, data.as_slice())?
                }
                Result::<_, eyre::Error>::Ok(())
            });
            let (receiver, sender) = tokio::join!(receive_handle, send_handle);
            receiver
                .wrap_err("Handle to the receiver failed")?
                .wrap_err("Receiving messages from receiver channel failed")?;
            sender
                .wrap_err("Handle to the sender failed")?
                .wrap_err("Sending messages using sender channel failed")?;
            Ok(())
        });

        Ok(Node {
            id,
            rx_input,
            tx_output,
        })
    }

    #[allow(clippy::should_implement_trait)]
    pub fn next(&mut self) -> PyResult<Option<PyInput>> {
        self.__next__()
    }

    pub fn __next__(&mut self) -> PyResult<Option<PyInput>> {
        Ok(self.rx_input.blocking_recv().map(PyInput))
    }

    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    pub fn send_output(&self, output_str: String, data: &PyBytes) -> Result<()> {
        self.tx_output
            .blocking_send((output_str, data.as_bytes().to_vec()))
            .wrap_err("Could not send output")
    }

    pub fn id(&self) -> String {
        self.id.to_string()
    }
}

#[pymodule]
fn dora(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<Node>().unwrap();
    Ok(())
}