fixing "stop" and adding test to verify multithreading

Author: marcomq

src/lib.rs

@@ -37,24 +37,6 @@ pub(crate) struct PyCommand {
     responder: oneshot::Sender<Result<Value, String>>,
 }
 
-/// A boxed, send-able future that resolves to a PyRunnerResult.
-type Task = Box<dyn FnOnce(&Runtime) -> Result<Value, PyRunnerError> + Send>;
-
-/// A lazily-initialized worker thread for handling synchronous function calls.
-/// This thread has its own private Tokio runtime to safely block on async operations
-/// without interfering with any existing runtime the user might be in.
-static SYNC_WORKER: Lazy<std_mpsc::Sender<Task>> = Lazy::new(|| {
-    let (tx, rx) = std_mpsc::channel::<Task>();
-
-    thread::spawn(move || {
-        let rt = Runtime::new().expect("Failed to create Tokio runtime for sync worker");
-        // When the sender (tx) is dropped, rx.recv() will return an Err, ending the loop.
-        while let Ok(task) = rx.recv() {
-            let _ = task(&rt); // The result is sent back via a channel inside the task.
-        }
-    });
-    tx
-});
 /// Custom error types for the `PyRunner`.
 #[derive(Error, Debug, Clone)]
 pub enum PyRunnerError {
@@ -86,6 +68,27 @@ pub fn print_path_for_python(path: &PathBuf) -> String {
     }
 }
 
+/// A boxed, send-able future that resolves to a PyRunnerResult.
+type Task = Box<dyn FnOnce(&Runtime) -> Result<Value, PyRunnerError> + Send>;
+
+/// A lazily-initialized worker thread for handling synchronous function calls
+/// to functions that otherwise return a future. It will only be engaged when 
+/// calling `..._sync()` functions of PyRunner.
+/// This thread has its own private Tokio runtime to safely block on async operations
+/// without interfering with any existing runtime the user might be in.
+static SYNC_WORKER: Lazy<std_mpsc::Sender<Task>> = Lazy::new(|| {
+    let (tx, rx) = std_mpsc::channel::<Task>();
+
+    thread::spawn(move || {
+        let rt = Runtime::new().expect("Failed to create Tokio runtime for sync worker");
+        // When the sender (tx) is dropped, rx.recv() will return an Err, ending the loop.
+        while let Ok(task) = rx.recv() {
+            let _ = task(&rt); // The result is sent back via a channel inside the task.
+        }
+    });
+    tx
+});
+
 /// Manages a dedicated thread for executing Python code asynchronously.
 #[derive(Clone)]
 pub struct PyRunner {
@@ -734,4 +737,59 @@ result = mymodule.my_func()
 
         assert_eq!(result, Value::String("hello from venv".to_string()));
     }
-}
+
+    #[tokio::test]
+    async fn test_pyrunner_thread_safety() {
+        let runner = PyRunner::new();
+        runner.run("x = 0").await.unwrap();
+
+        let mut handles = vec![];
+
+        for i in 0..5 {
+            let runner_clone = runner.clone();
+            let handle = tokio::spawn(async move {
+                // Each task increments 'x' in the shared Python interpreter
+                runner_clone.run(&format!("x += {}", i)).await.unwrap();
+            });
+            handles.push(handle);
+        }
+
+        // Wait for all tasks to complete
+        for handle in handles {
+            handle.await.unwrap();
+        }
+
+        let final_x = runner.read_variable("x").await.unwrap();
+        // Expected: 0 + 1 + 2 + 3 + 4 = 10
+        assert_eq!(final_x, Value::Number(10.into()));
+    }
+
+    #[test]
+    fn test_pyrunner_std_thread_safety() {
+        let runner = PyRunner::new();
+        runner.run_sync("x = 0").unwrap();
+
+        let mut handles = vec![];
+
+        for i in 0..5 {
+            let runner_clone = runner.clone();
+            // Spawn a new OS thread
+            let handle = thread::spawn(move || {
+                // Use the _sync version for non-async contexts
+                runner_clone.run_sync(&format!("x += {}", i)).unwrap();
+            });
+            handles.push(handle);
+        }
+
+        // Wait for all threads to complete
+        for handle in handles {
+            handle.join().unwrap();
+        }
+
+        let final_x = runner.read_variable_sync("x").unwrap();
+        // Expected: 0 + 1 + 2 + 3 + 4 = 10
+        assert_eq!(final_x, Value::Number(10.into()));
+
+        runner.stop_sync().unwrap();
+    }
+}

src/pyo3_runner.rs

@@ -54,7 +54,10 @@ pub(crate) async fn python_thread_main(mut receiver: mpsc::Receiver<PyCommand>)
                         Err(e) => Err(e),
                     }
                 }
-                CmdType::Stop => return receiver.close(),
+                CmdType::Stop => {
+                    receiver.close();
+                    Ok(Value::Null)
+                }
             };
 
             // Convert PyErr to a string representation to avoid exposing it outside this module.

src/rustpython_runner.rs

@@ -51,7 +51,10 @@ pub(crate) fn python_thread_main(mut receiver: mpsc::Receiver<PyCommand>) {
                     dbg!(name, args);
                     unimplemented!("Async functions are not supported yet in RustPython")
                 }
-                CmdType::Stop => break,
+                CmdType::Stop => {
+                    receiver.close();
+                    Ok(Value::Null)
+                }
             };
             let response = result.map_err(|err| {
                 format!(