Add task metadata feat

Author: CCXLV

crates/fluxqueue-worker/src/task.rs

@@ -1,4 +1,5 @@
 use anyhow::{Context, Result, anyhow};
+use fluxqueue_common::Task;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
 use pyo3::types::{PyAnyMethods, PyDict, PyDictMethods, PyList, PyModule, PyTuple};
@@ -112,8 +113,7 @@ struct TaskRequest {
     executor_id: Arc<String>,
     task_data: Arc<TaskData>,
     task_name: Arc<String>,
-    raw_args: Arc<Vec<u8>>,
-    raw_kwargs: Arc<Vec<u8>>,
+    task: Arc<Task>,
     resp_tx: oneshot::Sender<Result<()>>,
 }
 
@@ -134,8 +134,7 @@ impl PythonDispatcher {
                     req.executor_id,
                     req.task_data,
                     req.task_name,
-                    req.raw_args,
-                    req.raw_kwargs,
+                    req.task,
                 )
                 .await
                 .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
@@ -159,8 +158,7 @@ impl PythonDispatcher {
         executor_id: Arc<String>,
         task_data: Arc<TaskData>,
         task_name: Arc<String>,
-        raw_args: Arc<Vec<u8>>,
-        raw_kwargs: Arc<Vec<u8>>,
+        task: Arc<Task>,
     ) -> Result<()> {
         let (resp_tx, resp_rx) = oneshot::channel();
 
@@ -170,8 +168,7 @@ impl PythonDispatcher {
                 executor_id,
                 task_data,
                 task_name,
-                raw_args,
-                raw_kwargs,
+                task,
                 resp_tx,
             })
             .await
@@ -182,29 +179,39 @@ impl PythonDispatcher {
     }
 }
 
+struct CoroWithContext {
+    args: Py<PyTuple>,
+    kwargs: Py<PyDict>,
+    context: Arc<Py<PyAny>>,
+}
+
+struct MaybeCoro {
+    func: Arc<Py<PyAny>>,
+    with_context: Option<CoroWithContext>,
+}
+
 async fn run_task(
     task_registry: Arc<TaskRegistry>,
     executor_id: Arc<String>,
     task_data: Arc<TaskData>,
     task_name: Arc<String>,
-    raw_args: Arc<Vec<u8>>,
-    raw_kwargs: Arc<Vec<u8>>,
+    task: Arc<Task>,
 ) -> Result<()> {
     let logger = Logger::new(format!("EXECUTOR {}", &executor_id));
     let duration_start = Instant::now();
 
-    let task_args: Value = from_slice(&raw_args).context(format!(
+    let task_args: Value = from_slice(&task.args).context(format!(
         "Failed to deserialize task '{}' function args",
         &task_name
     ))?;
-    let task_kwargs: Value = from_slice(&raw_kwargs).context(format!(
+    let task_kwargs: Value = from_slice(&task.kwargs).context(format!(
         "Failed to deserialize task '{}' function kwargs",
         &task_name
     ))?;
 
     let context = task_registry.get_task_context(task_data.clone())?;
 
-    let maybe_coro = Python::attach(|py| -> Result<Option<Py<PyAny>>> {
+    let maybe_coro = Python::attach(|py| -> Result<Option<MaybeCoro>> {
         let py_args = pythonize(py, &task_args).context("Failed to pythonize args")?;
         let py_kwargs = pythonize(py, &task_kwargs).context("Failed to pythonize kwargs")?;
 
@@ -216,7 +223,11 @@ async fn run_task(
             anyhow::bail!("Args must be an array/tuple, found {}", py_args.get_type());
         };
 
-        if let Some(context) = context {
+        if let Some(context) = context.as_ref() {
+            let task_metadata = get_task_metadata(py, task.clone())?;
+            let metadata_var = context.getattr(py, "_metadata_var")?;
+            metadata_var.call_method1(py, "set", (task_metadata,))?;
+
             let context = context.as_any();
             let prefix = PyTuple::new(py, [context])?;
 
@@ -230,19 +241,25 @@ async fn run_task(
             .cast_into::<PyDict>()
             .map_err(|_| anyhow!("Kwargs must be a map/dict"))?;
 
-        let result = task_data
-            .func
-            .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 is_coroutine = is_coroutine(py, task_data.func.clone())?;
 
         if is_coroutine {
-            Ok(Some(result))
+            let with_context = context.map(|context| CoroWithContext {
+                args: args_tuple.unbind(),
+                kwargs: kwargs_dict.unbind(),
+                context,
+            });
+
+            Ok(Some(MaybeCoro {
+                func: task_data.func.clone(),
+                with_context,
+            }))
         } else {
+            task_data
+                .func
+                .call(py, args_tuple.clone(), Some(&kwargs_dict))
+                .map_err(|e| anyhow!("Failed to call Python function: {:?}", e))?;
+
             let duration_end = duration_start.elapsed();
             logger.info(format_args!(
                 "Task '{}' successfully finished in {}ms",
@@ -264,8 +281,27 @@ async fn run_task(
         anyhow!(e.to_string())
     })?;
 
-    if let Some(coro) = maybe_coro {
-        let fut = Python::attach(|py| into_future(coro.into_bound(py)))?;
+    if let Some(maybe_coro) = maybe_coro {
+        let fut = Python::attach(|py| {
+            if let Some(with_context) = maybe_coro.with_context {
+                let task_metadata = get_task_metadata(py, task.clone())
+                    .map_err(|e| PyRuntimeError::new_err(e.to_string()))?;
+                let result = with_context.context.call_method1(
+                    py,
+                    "_run_async_task",
+                    (
+                        maybe_coro.func.as_any(),
+                        task_metadata,
+                        with_context.args,
+                        Some(with_context.kwargs),
+                    ),
+                )?;
+                into_future(result.into_bound(py))
+            } else {
+                let func = maybe_coro.func.clone_ref(py);
+                into_future(func.into_bound(py))
+            }
+        })?;
         fut.await?;
 
         let duration_end = duration_start.elapsed();
@@ -370,6 +406,30 @@ fn get_registry(module_path: &str, queue_name: &str) -> Result<TasksAndContexts>
     Ok(result)
 }
 
+fn get_task_metadata(py: Python<'_>, task: Arc<Task>) -> Result<Py<PyAny>> {
+    let module = py.import("fluxqueue.models")?.unbind();
+    let task_metadata = module.call_method1(
+        py,
+        "TaskMetadata",
+        (
+            task.id.clone(),
+            task.retries,
+            task.max_retries,
+            task.created_at,
+        ),
+    )?;
+
+    Ok(task_metadata)
+}
+
+fn is_coroutine(py: Python<'_>, func: Arc<Py<PyAny>>) -> Result<bool> {
+    let inspect = py.import("inspect")?;
+    let is_coro: bool = inspect
+        .call_method1("iscoroutinefunction", (func.as_any(),))?
+        .extract()?;
+    Ok(is_coro)
+}
+
 fn normalize_path(path: &Path) -> PathBuf {
     let mut components = Vec::new();
     for comp in path.components() {

crates/fluxqueue-worker/src/worker.rs

@@ -24,10 +24,12 @@ pub async fn run_worker(
     })?;
     let redis_client = Arc::new(redis_client);
 
-    let task_registry = Arc::new(TaskRegistry::new(&tasks_module_path, &queue_name).map_err(|e| {
-        tracing::error!("{}", e);
-        std::process::exit(1);
-    })?);
+    let task_registry = Arc::new(TaskRegistry::new(&tasks_module_path, &queue_name).map_err(
+        |e| {
+            tracing::error!("{}", e);
+            std::process::exit(1);
+        },
+    )?);
     let registered_tasks = task_registry.get_registered_tasks()?;
     let registered_contexts = task_registry.get_registered_contexts()?;
 
@@ -164,6 +166,7 @@ async fn executor_loop(
                     Ok(Some(raw_data)) => {
                         let task = deserialize_raw_task_data(&raw_data)?;
                         let task_name = format!("{}:{}", &task.name, &task.id);
+                        let actual_task_name = task.name.clone();
 
                         logger.info(format_args!(
                             "Received a task '{}' with a total of {} Bytes",
@@ -181,7 +184,7 @@ async fn executor_loop(
                             return Ok(());
                         };
 
-                        let task_result = run_task(ctx.executor_id.clone(), ctx.python_dispatcher.clone(), &task, task_data.clone()).await;
+                        let task_result = run_task(ctx.executor_id.clone(), ctx.python_dispatcher.clone(), Arc::new(task), task_data.clone()).await;
 
                         match task_result {
                             Ok(_) => {
@@ -195,7 +198,7 @@ async fn executor_loop(
                                 if let Err(err) = ctx.redis_client
                                     .mark_as_failed(&ctx.queue_name, &ctx.executor_id, &raw_data)
                                     .await {
-                                        logger.error(format_args!("Failed to mark the task '{}' as failed: {}", &task.name, err));
+                                        logger.error(format_args!("Failed to mark the task '{}' as failed: {}", actual_task_name, err));
                                 }
                             }
                         }
@@ -294,15 +297,13 @@ async fn janitor_loop(
 async fn run_task(
     executor_id: Arc<String>,
     python_dispatcher: Arc<PythonDispatcher>,
-    task: &Task,
+    task: Arc<Task>,
     task_data: Arc<TaskData>,
 ) -> Result<()> {
     let task_name = Arc::new(format!("{}:{}", &task.name, &task.id));
-    let raw_args = Arc::new(task.args.clone());
-    let raw_kwargs = Arc::new(task.kwargs.clone());
 
     python_dispatcher
-        .execute(executor_id, task_data, task_name, raw_args, raw_kwargs)
+        .execute(executor_id, task_data, task_name, task)
         .await?;
 
     Ok(())
@@ -365,7 +366,7 @@ mod tests {
             let result = run_task(
                 Arc::new("test".to_string()),
                 dispatcher_pool.clone(),
-                &task,
+                Arc::new(task),
                 task_func,
             )
             .await;
@@ -398,7 +399,7 @@ mod tests {
             let result = run_task(
                 Arc::new("test".to_string()),
                 dispatcher_pool.clone(),
-                &task,
+                Arc::new(task),
                 task_func,
             )
             .await;

python/fluxqueue/__init__.py

@@ -1,4 +1,3 @@
-__all__ = ["Context", "FluxQueue"]
-
-from .client import FluxQueue
-from .context import Context
+from .client import FluxQueue as FluxQueue
+from .context import Context as Context
+from .models import TaskMetadata as TaskMetadata

python/fluxqueue/_task.py

@@ -45,14 +45,13 @@ def _task_decorator(
     if return_type and return_type is not type(None):
         raise TypeError(f"Task function must return None, got {return_type}")
 
-    is_async = inspect.iscoroutinefunction(func)
     task_name = get_task_name(func, name)
 
     # TODO: Add unique identifier 'fluxqueue' just to be 100% sure
     cast(Any, func).task_name = task_name
     cast(Any, func).queue = queue
 
-    if is_async:
+    if inspect.iscoroutinefunction(func):
 
         @wraps(func)
         async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> None:

python/fluxqueue/context.py

@@ -1,10 +1,12 @@
 import inspect
+import threading
 from collections.abc import Callable, Coroutine
 from contextvars import ContextVar
 from typing import Any, Concatenate, ParamSpec, TypeVar, cast, get_type_hints, overload
 
 from ._core import FluxQueueCore
 from ._task import _task_decorator
+from .models import TaskMetadata
 
 P = ParamSpec("P")
 
@@ -13,28 +15,47 @@ class Context:
     __fluxqueue_context__: str | None = None
 
     def __init__(self) -> None:
-        self._thread_storage: ContextVar[dict[str, Any]] = ContextVar(
-            "thread_storage", default=None
+        self._thread_local = threading.local()
+        self._metadata_var: ContextVar[TaskMetadata] = ContextVar(
+            "task_metadata", default=None
         )
 
     @property
     def thread_storage(self) -> dict[str, Any]:
         """
-        Context-scoped storage dictionary.
+        Retrieves the thread-persistent storage dictionary.
 
-        Returns a mutable dictionary associated with the current
-        execution context. The storage is isolated per thread or
-        async task using ContextVar, ensuring safe concurrent usage.
+        Returns a dictionary that persists across all tasks executed by the current worker.
+        Used for storing long-lived resources like database engines and connection
+        pools to avoid re-initialization overhead.
+        """
+        if not hasattr(self._thread_local, "storage"):
+            self._thread_local.storage = {}
+
+        return self._thread_local.storage
 
-        The dictionary is initialized lazily on first access.
+    @property
+    def metadata(self) -> TaskMetadata:
         """
-        storage = self._thread_storage.get(None)
+        Returns metadata isolated to the current task.
 
-        if storage is None:
-            storage = {}
-            self._thread_storage.set(storage)
+        Returns a TaskMetadata instance containing execution details like
+        retry counts and task IDs. This property uses ContextVars to ensure
+        data isolation during concurrent task execution on the same thread.
+        """
+        return self._metadata_var.get()
 
-        return storage
+    async def _run_async_task(
+        self, func: Callable, metadata: TaskMetadata, args, kwargs
+    ):
+        """
+        This function is for internal use only.
+        """
+        token = self._metadata_var.set(metadata)
+        try:
+            await func(*args, **kwargs)
+        finally:
+            self._metadata_var.reset(token)
 
     def __init_subclass__(cls) -> None:
         if not cls.__fluxqueue_context__:

python/fluxqueue/models.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+
+
+@dataclass(slots=True)
+class TaskMetadata:
+    """
+    Read-only metadata for a FluxQueue task.
+    """
+
+    task_id: str
+    """Unique identifier for the current task execution."""
+    retry_count: int
+    """Number of times this task has been retried."""
+    max_retries: int
+    """Maximum number of retry attempts allowed before failure."""
+    enqueued_at: int
+    """ISO 8601 timestamp of when the task was originally enqueued."""