Python node event mutex (dora-rs#1244)

Allows concurrent read/write access through a node (to address
dora-rs#748). "Already borrowed" seems to
happen when there are concurrent calls to read/write functions because
both currently require an exclusive `Node` ref.

1. Updates `DelayedCleanup.get_mut` and callers to use `&self` instead
of `&mut self`.
2. Adds a mutex to `Events.inner` so callers can use `&self` instead of
`&mut self`.
3. Adds an example dataflow that fails on main (with the "Already
borrowed" error) but works with this change

Author: haixuanTao

apis/python/node/src/lib.rs

@@ -4,6 +4,7 @@ use std::env::current_dir;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::Duration;
+use tokio::sync::Mutex;
 
 use arrow::pyarrow::{FromPyArrow, ToPyArrow};
 use dora_download::download_file;
@@ -135,7 +136,7 @@ impl Node {
 
         Ok(Node {
             events: Events {
-                inner: EventsInner::Dora(events),
+                inner: Arc::new(Mutex::new(EventsInner::Dora(events))),
                 _cleanup_handle: cleanup_handle,
             },
             dataflow_id,
@@ -167,7 +168,7 @@ impl Node {
     /// :rtype: dict
     #[pyo3(signature = (timeout=None))]
     #[allow(clippy::should_implement_trait)]
-    pub fn next(&mut self, py: Python, timeout: Option<f32>) -> PyResult<Option<Py<PyDict>>> {
+    pub fn next(&self, py: Python, timeout: Option<f32>) -> PyResult<Option<Py<PyDict>>> {
         let event = py.allow_threads(|| self.events.recv(timeout.map(Duration::from_secs_f32)));
         if let Some(event) = event {
             let dict = event
@@ -190,7 +191,7 @@ impl Node {
     ///
     /// :rtype: list[dict]
     #[allow(clippy::should_implement_trait)]
-    pub fn drain(&mut self, py: Python) -> PyResult<Vec<Py<PyDict>>> {
+    pub fn drain(&self, py: Python) -> PyResult<Vec<Py<PyDict>>> {
         let events = self
             .events
             .drain()
@@ -253,7 +254,7 @@ impl Node {
     /// :rtype: dict
     #[pyo3(signature = (timeout=None))]
     #[allow(clippy::should_implement_trait)]
-    pub async fn recv_async(&mut self, timeout: Option<f32>) -> PyResult<Option<Py<PyDict>>> {
+    pub async fn recv_async(&self, timeout: Option<f32>) -> PyResult<Option<Py<PyDict>>> {
         let event = self
             .events
             .recv_async_timeout(timeout.map(Duration::from_secs_f32))
@@ -264,6 +265,7 @@ impl Node {
                 let dict = event
                     .to_py_dict(py)
                     .context("Could not convert event into a dict")?;
+
                 Ok(Some(dict))
             })
         } else {
@@ -284,7 +286,7 @@ impl Node {
     /// Default behaviour is to timeout after 2 seconds.
     ///
     /// :rtype: dict
-    pub fn __next__(&mut self, py: Python) -> PyResult<Option<Py<PyDict>>> {
+    pub fn __next__(&self, py: Python) -> PyResult<Option<Py<PyDict>>> {
         self.next(py, None)
     }
 
@@ -324,7 +326,7 @@ impl Node {
     /// :rtype: None
     #[pyo3(signature = (output_id, data, metadata=None))]
     pub fn send_output(
-        &mut self,
+        &self,
         output_id: String,
         data: PyObject,
         metadata: Option<Bound<'_, PyDict>>,
@@ -356,7 +358,7 @@ impl Node {
     /// This method returns the parsed dataflow YAML file.
     ///
     /// :rtype: dict
-    pub fn dataflow_descriptor(&mut self, py: Python) -> eyre::Result<PyObject> {
+    pub fn dataflow_descriptor(&self, py: Python) -> eyre::Result<PyObject> {
         Ok(
             pythonize::pythonize(py, &self.node.get_mut().dataflow_descriptor()?)
                 .map(|x| x.unbind())?,
@@ -366,7 +368,7 @@ impl Node {
     /// Returns the node configuration.
     ///
     /// :rtype: dict
-    pub fn node_config(&mut self, py: Python) -> eyre::Result<PyObject> {
+    pub fn node_config(&self, py: Python) -> eyre::Result<PyObject> {
         Ok(pythonize::pythonize(py, &self.node.get_mut().node_config()).map(|x| x.unbind())?)
     }
 
@@ -382,10 +384,7 @@ impl Node {
     ///
     /// :type subscription: dora.Ros2Subscription
     /// :rtype: None
-    pub fn merge_external_events(
-        &mut self,
-        subscription: &mut Ros2Subscription,
-    ) -> eyre::Result<()> {
+    pub fn merge_external_events(&self, subscription: &mut Ros2Subscription) -> eyre::Result<()> {
         let subscription = subscription.into_stream()?;
         let stream = futures::stream::poll_fn(move |cx| {
             let s = subscription.as_stream().map(|item| {
@@ -404,12 +403,13 @@ impl Node {
         });
 
         // take out the event stream and temporarily replace it with a dummy
+        let mut inner = self.events.inner.blocking_lock();
         let events = std::mem::replace(
-            &mut self.events.inner,
+            &mut *inner,
             EventsInner::Merged(Box::new(futures::stream::empty())),
         );
         // update self.events with the merged stream
-        self.events.inner = EventsInner::Merged(events.merge_external_send(Box::pin(stream)));
+        *inner = EventsInner::Merged(events.merge_external_send(Box::pin(stream)));
 
         Ok(())
     }
@@ -424,13 +424,14 @@ fn err_to_pyany(err: eyre::Report, gil: Python<'_>) -> Py<PyAny> {
 }
 
 struct Events {
-    inner: EventsInner,
+    inner: Arc<Mutex<EventsInner>>,
     _cleanup_handle: NodeCleanupHandle,
 }
 
 impl Events {
-    fn recv(&mut self, timeout: Option<Duration>) -> Option<PyEvent> {
-        let event = match &mut self.inner {
+    fn recv(&self, timeout: Option<Duration>) -> Option<PyEvent> {
+        let mut inner = self.inner.blocking_lock();
+        let event = match &mut *inner {
             EventsInner::Dora(events) => match timeout {
                 Some(timeout) => events.recv_timeout(timeout).map(MergedEvent::Dora),
                 None => events.recv().map(MergedEvent::Dora),
@@ -440,8 +441,9 @@ impl Events {
         event.map(|event| PyEvent { event })
     }
 
-    fn try_recv(&mut self) -> Result<PyEvent, TryRecvError> {
-        let event = match &mut self.inner {
+    fn try_recv(&self) -> Result<PyEvent, TryRecvError> {
+        let mut inner = self.inner.blocking_lock();
+        let event = match &mut *inner {
             EventsInner::Dora(events) => events.try_recv().map(MergedEvent::Dora),
             EventsInner::Merged(_events) => {
                 todo!("try_recv on external event stream is not yet implemented!")
@@ -450,8 +452,9 @@ impl Events {
         event.map(|event| PyEvent { event })
     }
 
-    async fn recv_async_timeout(&mut self, timeout: Option<Duration>) -> Option<PyEvent> {
-        let event = match &mut self.inner {
+    async fn recv_async_timeout(&self, timeout: Option<Duration>) -> Option<PyEvent> {
+        let mut inner = self.inner.lock().await;
+        let event = match &mut *inner {
             EventsInner::Dora(events) => match timeout {
                 Some(timeout) => events
                     .recv_async_timeout(timeout)
@@ -464,8 +467,9 @@ impl Events {
         event.map(|event| PyEvent { event })
     }
 
-    fn drain(&mut self) -> Option<Vec<PyEvent>> {
-        match &mut self.inner {
+    fn drain(&self) -> Option<Vec<PyEvent>> {
+        let mut inner = self.inner.blocking_lock();
+        match &mut *inner {
             EventsInner::Dora(events) => match events.drain() {
                 Some(items) => {
                     return Some(
@@ -485,7 +489,8 @@ impl Events {
     }
 
     fn is_empty(&self) -> bool {
-        match &self.inner {
+        let inner = self.inner.blocking_lock();
+        match &*inner {
             EventsInner::Dora(events) => events.is_empty(),
             EventsInner::Merged(_events) => {
                 todo!("is_empty on external event stream is not yet implemented!")

apis/python/operator/src/lib.rs

@@ -40,7 +40,7 @@ impl<T> DelayedCleanup<T> {
         CleanupHandle(self.0.clone())
     }
 
-    pub fn get_mut(&mut self) -> std::sync::MutexGuard<T> {
+    pub fn get_mut(&self) -> std::sync::MutexGuard<T> {
         self.0.try_lock().expect("failed to lock DelayedCleanup")
     }
 }

examples/python-concurrent-rw/dataflow.yml

@@ -0,0 +1,14 @@
+nodes:
+  - id: node1
+    path: ./receive_data.py
+    inputs:
+      data: node2/data
+    outputs:
+      - data
+
+  - id: node2
+    path: ./receive_data.py
+    inputs:
+      data: node1/data
+    outputs:
+      - data

examples/python-concurrent-rw/pyproject.toml

@@ -0,0 +1,11 @@
+[project]
+name = "dora-test"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "asyncio>=4.0.0",
+    "dora-rs-cli>=0.3.13",
+    "numpy>=2.3.5",
+]

examples/python-concurrent-rw/receive_data.py

@@ -0,0 +1,47 @@
+from dora import Node
+
+import logging
+import threading
+import time
+
+import numpy as np
+import pyarrow as pa
+
+
+def read_data_task(node, log):
+    """Task that reads incoming events."""
+    while (event := node.next()) is not None:
+        if event["type"] == "INPUT":
+            print(f"info {event['value'].to_numpy()}")
+        del event
+    log.log(logging.INFO, "read_data_task done!")
+
+
+def publish_task(node, log):
+    """Task that publishes to a topic."""
+    while True:
+        time.sleep(1)  # Publish every 1s
+        now = time.perf_counter_ns()
+        node.send_output("data", pa.array([np.uint64(now)]))
+
+
+def main():
+    node = Node()
+    log = logging.getLogger(__name__)
+    
+    # Create thread for read task
+    read_thread = threading.Thread(target=read_data_task, args=(node, log))
+    read_thread.start()
+    
+    # Run publish task in a daemon thread (so it doesn't block main thread)
+    publish_thread = threading.Thread(target=publish_task, args=(node, log), daemon=True)
+    publish_thread.start()
+    
+    # Wait for read thread to complete
+    read_thread.join()
+    
+    log.log(logging.INFO, "done!")
+
+
+if __name__ == "__main__":
+    main()