orchestrata_ultimate.core.threading.thread_pool

Reference › orchestrata_ultimate › core › threading › thread_pool

The thread-pool layer under the batched test runner (test_executor imports it) and the default pool of the package. It has an old and a new interface. The legacy ThreadPool (minimum and maximum threads, submit(Task) returning a CustomFuture with get, and shutdown) is kept for compatibility. The newer ThreadPoolManager(max_workers) wraps a ThreadPoolExecutor and tracks each submission as a ThreadPoolTask with a UUID task_id, a ThreadPoolTaskStatus (pending, running, completed, failed, cancelled), the result, the error and its traceback, and creation, start and completion times.

submit and submit_async return the task id. get_task returns the record, get_task_result returns the result (waiting, with an optional timeout, when asked), and wait_for_task is the awaitable form. stop(wait, cancel_futures) shuts the executor down, and task_count reports the number of tracked tasks. get_default_thread_pool(max_workers) returns a process-wide singleton manager. A task that raises is recorded as failed with its traceback instead of propagating at submission.

At a glance

Class diagram

Dependency diagram

API

Constants

Classes

ThreadPoolTaskStatus (enum)

Bases: Enum · Line 22

Lifecycle states of a thread-pool task.

ThreadPoolTask (dataclass)

Decorators: dataclass · Line 33

Tracked task submitted to ThreadPoolManager.

Task

Line 49

Legacy task wrapper (kept for backward compatibility).

CustomFuture

Line 57

Legacy future wrapper (kept for backward compatibility).

ThreadPool

Line 67

Legacy thread pool (kept for backward compatibility).

ThreadPoolManager

Line 81

Managed thread pool with task lifecycle tracking.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/core/threading/thread_pool.py

# mypy: ignore-errors
"""Thread pool management for orchestrata-ultimate.

Provides backward-compatible ThreadPool/Task/CustomFuture and a new
ThreadPoolManager with full task lifecycle tracking.
"""

from __future__ import annotations

import asyncio
import threading
import traceback
import uuid
from collections.abc import Callable
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any


class ThreadPoolTaskStatus(Enum):
    """Lifecycle states of a thread-pool task."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class ThreadPoolTask:
    """Tracked task submitted to ThreadPoolManager."""

    task_id: str
    status: ThreadPoolTaskStatus
    func: Callable
    args: tuple
    kwargs: dict
    result: Any = None
    error: Exception | None = None
    traceback: str | None = None
    created_at: datetime = field(default_factory=datetime.now)
    started_at: datetime | None = None
    completed_at: datetime | None = None


class Task:
    """Legacy task wrapper (kept for backward compatibility)."""

    def __init__(self, func: Callable, args: tuple = ()):
        self.func = func
        self.args = args


class CustomFuture:
    """Legacy future wrapper (kept for backward compatibility)."""

    def __init__(self, future: Future):
        self.future = future

    def get(self) -> Any:
        return self.future.result()


class ThreadPool:
    """Legacy thread pool (kept for backward compatibility)."""

    def __init__(self, min_threads: int = 1, max_threads: int = 5):
        self.executor = ThreadPoolExecutor(max_workers=max_threads)

    def submit(self, task: Task) -> CustomFuture:
        future = self.executor.submit(task.func, *task.args)
        return CustomFuture(future)

    def shutdown(self, wait: bool = True) -> None:
        self.executor.shutdown(wait=wait)


class ThreadPoolManager:
    """Managed thread pool with task lifecycle tracking."""

    def __init__(self, max_workers: int = 4):
        self._executor = ThreadPoolExecutor(max_workers=max_workers)
        self._tasks: dict[str, ThreadPoolTask] = {}
        self._futures: dict[str, Future] = {}
        self._lock = threading.RLock()
        self._shutdown = False

    def submit(self, fn: Callable, *args: Any, **kwargs: Any) -> str:
        """Submit a callable and return a UUID task_id."""
        task_id = str(uuid.uuid4())
        task = ThreadPoolTask(
            task_id=task_id,
            status=ThreadPoolTaskStatus.PENDING,
            func=fn,
            args=args,
            kwargs=kwargs,
        )
        with self._lock:
            if self._shutdown:
                raise RuntimeError("ThreadPoolManager is shut down")
            self._tasks[task_id] = task

        def _wrapper() -> Any:
            with self._lock:
                task.status = ThreadPoolTaskStatus.RUNNING
                task.started_at = datetime.now()
            try:
                result = fn(*args, **kwargs)
                with self._lock:
                    task.result = result
                    task.status = ThreadPoolTaskStatus.COMPLETED
                    task.completed_at = datetime.now()
                return result
            except Exception as exc:
                with self._lock:
                    task.error = exc
                    task.traceback = traceback.format_exc()
                    task.status = ThreadPoolTaskStatus.FAILED
                    task.completed_at = datetime.now()
                raise

        future = self._executor.submit(_wrapper)
        with self._lock:
            self._futures[task_id] = future
        return task_id

    async def submit_async(self, fn: Callable, *args: Any, **kwargs: Any) -> str:
        """Async version of submit."""
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, self.submit, fn, *args, **kwargs)

    def get_task(self, task_id: str) -> ThreadPoolTask | None:
        """Return the task record or None if unknown."""
        with self._lock:
            return self._tasks.get(task_id)

    def _wait_future_blocking(
        self,
        task_id: str,
        timeout: float | None,
    ) -> None:
        """Block until the underlying future resolves (no return value)."""
        future: Future | None = None
        with self._lock:
            future = self._futures.get(task_id)
        if future is not None:
            try:
                future.result(timeout=timeout)
            except Exception:
                pass

    def get_task_result(
        self,
        task_id: str,
        wait: bool = True,
        timeout: float | None = None,
    ) -> Any:
        """Return the result of a completed task, optionally waiting.

        Raises:
            KeyError: if the task_id is unknown.
            TimeoutError: if wait is True and the task does not finish in time.
            Exception: the original exception if the task failed.
        """
        task = self.get_task(task_id)
        if task is None:
            raise KeyError(f"Task {task_id} not found")

        if not wait:
            if task.status == ThreadPoolTaskStatus.COMPLETED:
                return task.result
            if task.status == ThreadPoolTaskStatus.FAILED:
                # `error` is set alongside the FAILED transition (see the
                # worker wrapper above); the `is not None` fallback is a
                # narrowing guard for pyright, not an expected runtime path.
                raise (
                    task.error
                    if task.error is not None
                    else RuntimeError(f"Task {task_id} failed with no recorded error")
                )
            raise TimeoutError("Task not yet complete")

        self._wait_future_blocking(task_id, timeout)

        task = self.get_task(task_id)
        if task is None:
            raise KeyError(f"Task {task_id} not found")
        if task.status == ThreadPoolTaskStatus.COMPLETED:
            return task.result
        if task.status == ThreadPoolTaskStatus.FAILED:
            # `error` is set alongside the FAILED transition (see the worker
            # wrapper above); the `is not None` fallback is a narrowing guard
            # for pyright, not an expected runtime path.
            raise (
                task.error
                if task.error is not None
                else RuntimeError(f"Task {task_id} failed with no recorded error")
            )
        if task.status == ThreadPoolTaskStatus.CANCELLED:
            raise RuntimeError("Task was cancelled")
        raise TimeoutError("Timeout waiting for task")

    async def _wait_future_async(
        self,
        task_id: str,
        timeout: float | None,
    ) -> None:
        """Async wait until the underlying future resolves."""
        future: Future | None = None
        with self._lock:
            future = self._futures.get(task_id)
        if future is not None:
            loop = asyncio.get_running_loop()
            try:
                await asyncio.wait_for(
                    loop.run_in_executor(None, future.result),
                    timeout=timeout,
                )
            except TimeoutError:
                raise TimeoutError(f"Timeout waiting for task {task_id}")
            except Exception:
                pass

    async def wait_for_task(self, task_id: str, timeout: float | None = None) -> Any:
        """Async wait for a task to finish and return its result.

        Raises:
            KeyError: if the task_id is unknown.
            TimeoutError: if the timeout expires.
            Exception: the original exception if the task failed.
        """
        task = self.get_task(task_id)
        if task is None:
            raise KeyError(f"Task {task_id} not found")

        if task.status == ThreadPoolTaskStatus.COMPLETED:
            return task.result
        if task.status == ThreadPoolTaskStatus.FAILED:
            # `error` is set alongside the FAILED transition (see the worker
            # wrapper above); the `is not None` fallback is a narrowing guard
            # for pyright, not an expected runtime path.
            raise (
                task.error
                if task.error is not None
                else RuntimeError(f"Task {task_id} failed with no recorded error")
            )
        if task.status == ThreadPoolTaskStatus.CANCELLED:
            raise RuntimeError("Task was cancelled")

        await self._wait_future_async(task_id, timeout)

        task = self.get_task(task_id)
        if task is None:
            raise KeyError(f"Task {task_id} not found")
        if task.status == ThreadPoolTaskStatus.COMPLETED:
            return task.result
        if task.status == ThreadPoolTaskStatus.FAILED:
            # `error` is set alongside the FAILED transition (see the worker
            # wrapper above); the `is not None` fallback is a narrowing guard
            # for pyright, not an expected runtime path.
            raise (
                task.error
                if task.error is not None
                else RuntimeError(f"Task {task_id} failed with no recorded error")
            )
        if task.status == ThreadPoolTaskStatus.CANCELLED:
            raise RuntimeError("Task was cancelled")
        raise TimeoutError("Timeout waiting for task")

    def stop(self, wait: bool = True, cancel_futures: bool = False) -> None:
        """Shut down the executor.

        Args:
            wait: Block until running tasks finish.
            cancel_futures: Cancel pending futures (Python 3.9+).
        """
        with self._lock:
            if self._shutdown:
                return
            self._shutdown = True
            if cancel_futures:
                for t in self._tasks.values():
                    if t.status in (
                        ThreadPoolTaskStatus.PENDING,
                        ThreadPoolTaskStatus.RUNNING,
                    ):
                        t.status = ThreadPoolTaskStatus.CANCELLED
        self._executor.shutdown(wait=wait, cancel_futures=cancel_futures)

    @property
    def task_count(self) -> int:
        """Number of tasks tracked by this manager."""
        with self._lock:
            return len(self._tasks)


_default_pool: ThreadPoolManager | None = None
_default_pool_lock = threading.Lock()


def get_default_thread_pool(max_workers: int = 4) -> ThreadPoolManager:
    """Return the process-wide singleton ThreadPoolManager."""
    global _default_pool
    with _default_pool_lock:
        if _default_pool is None:
            _default_pool = ThreadPoolManager(max_workers=max_workers)
        return _default_pool