diff --git a/crates/fluxqueue-worker/src/task.rs b/crates/fluxqueue-worker/src/task.rs index f2ddec1..bd2ae9c 100644 --- a/crates/fluxqueue-worker/src/task.rs +++ b/crates/fluxqueue-worker/src/task.rs @@ -1,8 +1,12 @@ -use anyhow::{Result, anyhow}; +use anyhow::{Context, Result, anyhow}; +use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; +use pyo3::types::{PyDict, PyList, PyTuple}; use pyo3_async_runtimes::tokio::into_future; +use pythonize::pythonize; +use rmp_serde::from_slice; +use rmpv::Value; use std::collections::HashMap; -use std::pin::Pin; use std::sync::{Arc, RwLock}; use tokio::sync::{mpsc, oneshot}; @@ -19,7 +23,7 @@ impl TaskRegistry { pub fn insert(&self, name: String, func: Py) -> Result<()> { let mut tasks = self.tasks.write().map_err(|_| { - anyhow::anyhow!("Internal Error: Task registry lock poisoned (a thread panicked)") + anyhow!("Internal Error: Task registry lock poisoned (a thread panicked)") })?; tasks.insert(name, Arc::new(func)); Ok(()) @@ -31,11 +35,12 @@ impl TaskRegistry { } } -type PyResponse = Pin>> + Send>>; - struct TaskRequest { - func: Py, - resp_tx: oneshot::Sender>, + func: Arc>, + task_name: Arc, + raw_args: Arc>, + raw_kwargs: Arc>, + resp_tx: oneshot::Sender>, } pub struct PythonDispatcher { @@ -43,39 +48,106 @@ pub struct PythonDispatcher { } impl PythonDispatcher { - pub fn new() -> Self { + pub fn new() -> Result { let logical_cores = num_cpus::get(); let (tx, mut rx) = mpsc::channel::(logical_cores * 2); - std::thread::spawn(move || { - Python::attach(|py| { - let dispatcher = async move { - while let Some(req) = rx.recv().await { - let res = Python::attach(|py| into_future(req.func.into_bound(py))); + let dispatcher = async move { + while let Some(req) = rx.recv().await { + run_task(req.func, req.task_name, req.raw_args, req.raw_kwargs) + .await + .map_err(|e| PyRuntimeError::new_err(e.to_string()))?; - let _ = req.resp_tx.send(res.map(|f| Box::pin(f) as PyResponse)); - } - Ok(()) - }; + let _ = req.resp_tx.send(Ok(())); + } + Ok(()) + }; + tokio::task::spawn_blocking(move || { + Python::attach(|py| { pyo3_async_runtimes::tokio::run(py, dispatcher).expect("Python loop failed"); }); }); - Self { tx } + Ok(Self { tx }) } - pub async fn execute(&self, func: Py) -> Result> { + pub async fn execute( + &self, + func: Arc>, + task_name: Arc, + raw_args: Arc>, + raw_kwargs: Arc>, + ) -> Result<()> { let (resp_tx, resp_rx) = oneshot::channel(); self.tx - .send(TaskRequest { func, resp_tx }) + .send(TaskRequest { + func, + task_name, + raw_args, + raw_kwargs, + resp_tx, + }) .await .map_err(|_| anyhow!("Dispatcher channel closed"))?; - let py_fut = resp_rx.await??; + resp_rx.await??; + Ok(()) + } +} + +async fn run_task( + task_function: Arc>, + task_name: Arc, + raw_args: Arc>, + raw_kwargs: Arc>, +) -> Result<()> { + let task_args: Value = from_slice(&raw_args).context(format!( + "Failed to deserialize task '{}' function args", + &task_name + ))?; + let task_kwargs: Value = from_slice(&raw_kwargs).context(format!( + "Failed to deserialize task '{}' function kwargs", + &task_name + ))?; + + let maybe_coro = Python::attach(|py| -> Result>> { + let py_args = pythonize(py, &task_args).context("Failed to pythonize args")?; + let py_kwargs = pythonize(py, &task_kwargs).context("Failed to pythonize kwargs")?; + + let args_tuple = if let Ok(list) = py_args.cast::() { + list.to_tuple() + } else if let Ok(tuple) = py_args.cast::() { + tuple.clone() + } else { + anyhow::bail!("Args must be an array/tuple, found {}", py_args.get_type()); + }; + + let kwargs_dict = py_kwargs + .cast_into::() + .map_err(|_| anyhow!("Kwargs must be a map/dict"))?; + + let result = task_function + .call(py, args_tuple, Some(&kwargs_dict)) + .map_err(|e| anyhow!("Failed to call Python function: {:?}", e))?; + + let bound_result = result.bind(py); + let is_coroutine = bound_result + .hasattr("__await__") + .map_err(|_| anyhow!("Failed to check if result is awaitable"))?; + + if is_coroutine { + Ok(Some(result)) + } else { + Ok(None) + } + })?; - let result = py_fut.await?; - Ok(result) + if let Some(coro) = maybe_coro { + let fut = Python::attach(|py| into_future(coro.into_bound(py)))?; + fut.await?; } + + Ok(()) } diff --git a/crates/fluxqueue-worker/src/worker.rs b/crates/fluxqueue-worker/src/worker.rs index b2a8b55..81d9c1f 100644 --- a/crates/fluxqueue-worker/src/worker.rs +++ b/crates/fluxqueue-worker/src/worker.rs @@ -1,8 +1,6 @@ -use anyhow::{Context, Result, anyhow}; -use pyo3::types::{PyAnyMethods, PyDict, PyDictMethods, PyList, PyListMethods, PyModule, PyTuple}; +use anyhow::{Result, anyhow}; +use pyo3::types::{PyAnyMethods, PyDict, PyDictMethods, PyModule}; use pyo3::{Bound, Py, PyAny, Python}; -use pythonize::pythonize; -use rmp_serde::from_slice; use std::ffi::CString; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -58,7 +56,7 @@ pub async fn run_worker( let executor_id = Arc::clone(&executor_ids[i]); let shutdown = shutdown.clone(); let task_registry = Arc::clone(&task_registry); - let python_dispatcher = Arc::new(PythonDispatcher::new()); + let python_dispatcher = Arc::new(PythonDispatcher::new()?); redis_client .register_executor(&queue_name, &executor_id) @@ -276,50 +274,13 @@ async fn run_task( task: &Task, task_function: Arc>, ) -> Result<()> { - let task_args: rmpv::Value = from_slice(&task.args).context(format!( - "Failed to deserialize task '{}' function args", - task.name - ))?; - let task_kwargs: rmpv::Value = from_slice(&task.kwargs).context(format!( - "Failed to deserialize task '{}' function kwargs", - task.name - ))?; - - let maybe_coro = Python::attach(|py| -> Result>> { - let py_args = pythonize(py, &task_args).context("Failed to pythonize args")?; - let py_kwargs = pythonize(py, &task_kwargs).context("Failed to pythonize kwargs")?; - - let args_tuple = if let Ok(list) = py_args.cast::() { - list.to_tuple() - } else if let Ok(tuple) = py_args.cast::() { - tuple.clone() - } else { - anyhow::bail!("Args must be an array/tuple, found {}", py_args.get_type()); - }; - - let kwargs_dict = py_kwargs - .cast_into::() - .map_err(|_| anyhow!("Kwargs must be a map/dict"))?; - - let result = task_function - .call(py, args_tuple, Some(&kwargs_dict)) - .map_err(|e| anyhow!("Failed to call Python function: {:?}", e))?; - - let bound_result = result.bind(py); - let is_coroutine = bound_result - .hasattr("__await__") - .map_err(|_| anyhow!("Failed to check if result is awaitable"))?; + let task_name = Arc::new(task.name.clone()); + let raw_args = Arc::new(task.args.clone()); + let raw_kwargs = Arc::new(task.kwargs.clone()); - if is_coroutine { - Ok(Some(result)) - } else { - Ok(None) - } - })?; - - if let Some(coro) = maybe_coro { - python_dispatcher.execute(coro).await?; - } + python_dispatcher + .execute(task_function, task_name, raw_args, raw_kwargs) + .await?; Ok(()) } @@ -356,7 +317,7 @@ fn get_task_functions(module_path: &str, queue_name: &str) -> Result = module .getattr("list_functions") @@ -505,7 +466,7 @@ mod tests { #[tokio::test] async fn test_run_task_with_sync_function() -> Result<()> { let task_registry = TaskRegistry::new(); - let python_dispatcher = Arc::new(PythonDispatcher::new()); + let python_dispatcher = Arc::new(PythonDispatcher::new()?); let module_path_str = get_test_module_path("test_tasks_module.py"); let task_functions = get_task_functions(&module_path_str, "default")?; @@ -537,7 +498,7 @@ mod tests { #[tokio::test] async fn test_run_task_with_async_function() -> Result<()> { let task_registry = TaskRegistry::new(); - let python_dispatcher = Arc::new(PythonDispatcher::new()); + let python_dispatcher = Arc::new(PythonDispatcher::new()?); let module_path_str = get_test_module_path("test_tasks_module.py"); let task_functions = get_task_functions(&module_path_str, "default")?;