Start working on context in the worker

Author: CCXLV

crates/fluxqueue-worker/scripts/get_functions.py

@@ -0,0 +1,37 @@
+import importlib
+import inspect
+import sys
+from pathlib import Path
+
+
+def get_registry(module_path: str, queue: str, module_dir: str | None = None):
+    if module_dir:
+        module_dir_path = Path(module_dir).resolve()
+        if str(module_dir_path) not in sys.path:
+            sys.path.insert(0, str(module_dir_path))
+
+    module = importlib.import_module(module_path)
+    registry = {"tasks": {}, "contexts": {}}
+    for _name, obj in inspect.getmembers(module):
+        if inspect.isfunction(obj):
+            task_name = getattr(obj, "task_name", None)
+            task_queue = getattr(obj, "queue", None)
+            if not task_queue or task_queue != queue:
+                continue
+
+            if registry["tasks"].get(task_name):
+                raise ValueError(f"Task '{task_name}' is duplicated")
+
+            original_func = getattr(obj, "__wrapped__", obj)
+            registry["tasks"][task_name] = original_func
+        elif inspect.isclass(obj):
+            context_name = getattr(obj, "__fluxqueue_context__", None)
+            if not context_name:
+                continue
+
+            if registry["contexts"].get(context_name):
+                raise ValueError(f"Context '{context_name}' is duplicated")
+
+            registry["contexts"][context_name] = obj
+
+    return registry

crates/fluxqueue-worker/scripts/get_registry.py

@@ -34,13 +34,15 @@ pub fn initial_logs(
     concurrency: usize,
     redis_url: &str,
     tasks_module_path: &str,
-    tasks: &Vec<&String>,
+    tasks: Vec<String>,
+    contexts: Vec<String>,
 ) {
     info!("Queue: {}", queue_name);
     info!("Concurrency: {}", concurrency);
     info!("Redis: {}", redis_url);
     info!("Tasks module: {}", tasks_module_path);
     info!("Tasks found: {:?}", tasks);
+    info!("Contexts found: {:?}", contexts);
     info!("Starting up the executors...");
 }
 

crates/fluxqueue-worker/src/logger.rs

@@ -33,7 +33,7 @@ impl RedisClient {
         Ok(())
     }
 
-    pub async fn set_executors_heartbeat(&self, executor_ids: Arc<Vec<Arc<str>>>) -> Result<()> {
+    pub async fn set_executors_heartbeat(&self, executor_ids: Arc<Vec<Arc<String>>>) -> Result<()> {
         for id in executor_ids.iter() {
             self.set_executor_heartbeat(id).await?;
         }
@@ -60,7 +60,7 @@ impl RedisClient {
     pub async fn cleanup_executors_registry(
         &self,
         queue_name: &str,
-        ids: Arc<Vec<Arc<str>>>,
+        ids: Arc<Vec<Arc<String>>>,
     ) -> Result<()> {
         let mut conn = self.redis_pool.get().await?;
         let executors_key = keys::get_executors_key(queue_name);

crates/fluxqueue-worker/src/redis_client.rs

@@ -1,32 +1,120 @@
 use anyhow::{Context, Result, anyhow};
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyList, PyTuple};
+use pyo3::types::{PyAnyMethods, PyDict, PyDictMethods, PyList, PyModule, 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::ffi::CString;
+use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, RwLock};
+use std::time::Instant;
 use tokio::sync::{mpsc, oneshot};
 
+use crate::logger::Logger;
+
+#[derive(Debug)]
 pub struct TaskRegistry {
     tasks: Arc<RwLock<HashMap<String, Arc<Py<PyAny>>>>>,
+    contexts: Arc<RwLock<HashMap<String, Arc<Py<PyAny>>>>>,
 }
 
 impl TaskRegistry {
-    pub fn new() -> Self {
-        Self {
-            tasks: Arc::new(RwLock::new(HashMap::new())),
+    pub fn new(module_path: &str, queue_name: &str) -> Result<Self> {
+        let script = include_str!("../scripts/get_registry.py");
+        let script_cstr = CString::new(script)?;
+        let filename = CString::new("get_registry.py")?;
+        let module_name = CString::new("get_registry")?;
+
+        let full_current_dir = std::env::current_dir().unwrap();
+        let full_module_path = full_current_dir.join(module_path);
+        let clean_module_path = normalize_path(&full_module_path);
+        let project_root = full_current_dir
+            .ancestors()
+            .find(|p| p.join("tests").exists())
+            .unwrap_or(&full_current_dir);
+        let real_module_path = path_to_module_path(project_root, &clean_module_path);
+
+        if !clean_module_path.exists() || real_module_path.is_none() {
+            return Err(anyhow!(
+                "Tasks module path {:?} doesn't exist.",
+                clean_module_path
+            ));
         }
+
+        let real_module_path = real_module_path.unwrap();
+        let module_dir = project_root.to_string_lossy().to_string();
+
+        let (tasks, contexts) = Python::attach(
+            |py| -> Result<(
+                HashMap<String, Arc<Py<PyAny>>>,
+                HashMap<String, Arc<Py<PyAny>>>,
+            )> {
+                let module = PyModule::from_code(
+                    py,
+                    script_cstr.as_c_str(),
+                    filename.as_c_str(),
+                    module_name.as_c_str(),
+                )
+                .map_err(|e| anyhow!("Failed to import python module: {}", e))?;
+
+                let registry: Bound<'_, PyDict> = module
+                    .getattr("get_registry")
+                    .map_err(|e| anyhow!("Failed to get 'get_registry' script: {}", e))?
+                    .call1((real_module_path, queue_name, module_dir))
+                    .map_err(|e| anyhow!("Failed to get tasks: {}", e))?
+                    .cast_into::<PyDict>()
+                    .map_err(|_| anyhow!("Failed to cast result to a Python Dictionary"))?;
+
+                let tasks: HashMap<String, Arc<Py<PyAny>>> = registry
+                    .get_item("tasks")?
+                    .expect("tasks missing")
+                    .cast::<PyDict>()
+                    .map_err(|e| anyhow!("tasks is not a dict: {}", e))?
+                    .iter()
+                    .filter_map(|(key, value)| {
+                        let name: String = key.extract().ok()?;
+                        let func: Py<PyAny> = value.unbind();
+                        Some((name, Arc::new(func)))
+                    })
+                    .collect();
+
+                let contexts: HashMap<String, Arc<Py<PyAny>>> = registry
+                    .get_item("contexts")?
+                    .expect("contexts missing")
+                    .cast::<PyDict>()
+                    .map_err(|e| anyhow!("contexts is not a dict: {}", e))?
+                    .iter()
+                    .filter_map(|(key, value): (Bound<PyAny>, Bound<PyAny>)| {
+                        let name: String = key.extract().ok()?;
+                        let func: Py<PyAny> = value.unbind();
+                        Some((name, Arc::new(func)))
+                    })
+                    .collect();
+
+                Ok((tasks, contexts))
+            },
+        )?;
+
+        Ok(Self {
+            tasks: Arc::new(RwLock::new(tasks)),
+            contexts: Arc::new(RwLock::new(contexts)),
+        })
     }
 
-    pub fn insert(&self, name: String, func: Py<PyAny>) -> Result<()> {
-        let mut tasks = self.tasks.write().map_err(|_| {
-            anyhow!("Internal Error: Task registry lock poisoned (a thread panicked)")
-        })?;
-        tasks.insert(name, Arc::new(func));
-        Ok(())
+    pub fn get_registered_tasks(&self) -> Result<Vec<String>> {
+        let tasks = self.tasks.read().map_err(|e| anyhow!(e.to_string()))?;
+        let task_names: Vec<_> = tasks.iter().map(|t| t.0.to_string()).collect();
+        Ok(task_names)
+    }
+
+    pub fn get_registered_contexts(&self) -> Result<Vec<String>> {
+        let contexts = self.contexts.read().map_err(|e| anyhow!(e.to_string()))?;
+        let context_names: Vec<_> = contexts.iter().map(|t| t.0.to_string()).collect();
+        Ok(context_names)
     }
 
     pub fn get(&self, name: &str) -> Option<Arc<Py<PyAny>>> {
@@ -36,6 +124,7 @@ impl TaskRegistry {
 }
 
 struct TaskRequest {
+    executor_id: Arc<String>,
     func: Arc<Py<PyAny>>,
     task_name: Arc<String>,
     raw_args: Arc<Vec<u8>>,
@@ -54,9 +143,15 @@ impl PythonDispatcher {
 
         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()))?;
+                run_task(
+                    req.executor_id,
+                    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(Ok(()));
             }
@@ -74,6 +169,7 @@ impl PythonDispatcher {
 
     pub async fn execute(
         &self,
+        executor_id: Arc<String>,
         func: Arc<Py<PyAny>>,
         task_name: Arc<String>,
         raw_args: Arc<Vec<u8>>,
@@ -83,6 +179,7 @@ impl PythonDispatcher {
 
         self.tx
             .send(TaskRequest {
+                executor_id,
                 func,
                 task_name,
                 raw_args,
@@ -98,11 +195,15 @@ impl PythonDispatcher {
 }
 
 async fn run_task(
+    executor_id: Arc<String>,
     task_function: Arc<Py<PyAny>>,
     task_name: Arc<String>,
     raw_args: Arc<Vec<u8>>,
     raw_kwargs: Arc<Vec<u8>>,
 ) -> Result<()> {
+    let logger = Logger::new(format!("EXECUTOR {}", &executor_id));
+    let duration_start = Instant::now();
+
     let task_args: Value = from_slice(&raw_args).context(format!(
         "Failed to deserialize task '{}' function args",
         &task_name
@@ -140,14 +241,99 @@ async fn run_task(
         if is_coroutine {
             Ok(Some(result))
         } else {
+            let duration_end = duration_start.elapsed();
+            logger.info(format_args!(
+                "Task '{}' successfully finished in {}ms",
+                &task_name,
+                duration_end.as_millis()
+            ));
+
             Ok(None)
         }
+    })
+    .map_err(|e| {
+        let duration_end = duration_start.elapsed();
+        logger.error(format_args!(
+            "Task '{}' failed in {}ms: {}",
+            &task_name,
+            duration_end.as_millis(),
+            e
+        ));
+        anyhow!(e.to_string())
     })?;
 
     if let Some(coro) = maybe_coro {
         let fut = Python::attach(|py| into_future(coro.into_bound(py)))?;
         fut.await?;
+
+        let duration_end = duration_start.elapsed();
+        logger.info(format_args!(
+            "Task '{}' successfully finished in {}ms",
+            &task_name,
+            duration_end.as_millis()
+        ));
     }
 
     Ok(())
 }
+
+pub struct DispatcherPool {
+    dispatchers: Vec<Arc<PythonDispatcher>>,
+    index: AtomicUsize,
+}
+
+impl DispatcherPool {
+    pub fn new(concurrency: usize) -> Result<Self> {
+        // let logical_cores = num_cpus::get();
+        // let pool_size = (concurrency / 4)
+        //     .max(1)
+        //     .min(logical_cores * 2)
+        //     .max(logical_cores);
+
+        let mut dispatchers = Vec::with_capacity(concurrency);
+        for _ in 0..concurrency {
+            dispatchers.push(Arc::new(PythonDispatcher::new()?));
+        }
+
+        Ok(Self {
+            dispatchers,
+            index: AtomicUsize::new(0),
+        })
+    }
+
+    pub fn get(&self) -> Arc<PythonDispatcher> {
+        let idx = self.index.fetch_add(1, Ordering::Relaxed) % self.dispatchers.len();
+        self.dispatchers[idx].clone()
+    }
+}
+
+fn normalize_path(path: &Path) -> PathBuf {
+    let mut components = Vec::new();
+    for comp in path.components() {
+        match comp {
+            std::path::Component::ParentDir => {
+                components.pop();
+            }
+            std::path::Component::CurDir => {}
+            other => components.push(other),
+        }
+    }
+    components.iter().collect()
+}
+
+fn path_to_module_path(current_dir: &Path, target_path: &Path) -> Option<String> {
+    let rel_path = target_path.strip_prefix(current_dir).ok()?;
+
+    let mut components: Vec<String> = rel_path
+        .components()
+        .map(|c| c.as_os_str().to_string_lossy().to_string())
+        .collect();
+
+    if let Some(last) = components.last_mut()
+        && let Some(pos) = last.rfind('.')
+    {
+        last.truncate(pos);
+    }
+
+    Some(components.join("."))
+}