baqueue 0.1.0__py3-none-any.whl

baqueue/serializer.py

@@ -0,0 +1,124 @@
+"""Job serialization and deserialization."""
+
+from __future__ import annotations
+
+import json
+import importlib
+from datetime import datetime, timezone
+from typing import Any
+from uuid import uuid4
+
+
+def _now_ts() -> float:
+    return datetime.now(timezone.utc).timestamp()
+
+
+class JobPayload:
+    """Serializable representation of a queued job."""
+
+    __slots__ = (
+        "id",
+        "job_class",
+        "data",
+        "queue",
+        "attempts",
+        "max_attempts",
+        "backoff",
+        "timeout",
+        "tags",
+        "batch_id",
+        "delay_until",
+        "created_at",
+        "updated_at",
+        "started_at",
+        "completed_at",
+        "failed_at",
+        "status",
+        "error",
+    )
+
+    def __init__(
+        self,
+        *,
+        id: str | None = None,
+        job_class: str = "",
+        data: dict[str, Any] | None = None,
+        queue: str = "default",
+        attempts: int = 0,
+        max_attempts: int = 3,
+        backoff: str | list[int] = "exponential",
+        timeout: int = 60,
+        tags: list[str] | None = None,
+        batch_id: str | None = None,
+        delay_until: float | None = None,
+        created_at: float | None = None,
+        updated_at: float | None = None,
+        started_at: float | None = None,
+        completed_at: float | None = None,
+        failed_at: float | None = None,
+        status: str = "pending",
+        error: str | None = None,
+    ):
+        self.id = id or uuid4().hex
+        self.job_class = job_class
+        self.data = data or {}
+        self.queue = queue
+        self.attempts = attempts
+        self.max_attempts = max_attempts
+        self.backoff = backoff
+        self.timeout = timeout
+        self.tags = tags or []
+        self.batch_id = batch_id
+        self.delay_until = delay_until
+        self.created_at = created_at or _now_ts()
+        self.updated_at = updated_at or self.created_at
+        self.started_at = started_at
+        self.completed_at = completed_at
+        self.failed_at = failed_at
+        self.status = status
+        self.error = error
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "job_class": self.job_class,
+            "data": self.data,
+            "queue": self.queue,
+            "attempts": self.attempts,
+            "max_attempts": self.max_attempts,
+            "backoff": self.backoff,
+            "timeout": self.timeout,
+            "tags": self.tags,
+            "batch_id": self.batch_id,
+            "delay_until": self.delay_until,
+            "created_at": self.created_at,
+            "updated_at": self.updated_at,
+            "started_at": self.started_at,
+            "completed_at": self.completed_at,
+            "failed_at": self.failed_at,
+            "status": self.status,
+            "error": self.error,
+        }
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> JobPayload:
+        return cls(**data)
+
+    @classmethod
+    def from_json(cls, raw: str) -> JobPayload:
+        return cls.from_dict(json.loads(raw))
+
+
+def resolve_job_class(class_path: str):
+    """Dynamically import and return a job class from its dotted path."""
+    module_path, class_name = class_path.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)
+
+
+def get_class_path(cls: type) -> str:
+    """Get the fully qualified dotted path for a class."""
+    return f"{cls.__module__}.{cls.__qualname__}"

baqueue/supervisor.py

@@ -0,0 +1,206 @@
+"""Supervisor - manages a pool of workers with auto-scaling and graceful shutdown."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import signal
+import os
+from typing import Any
+
+from baqueue.config import SupervisorConfig
+from baqueue.drivers.base import BaseDriver
+from baqueue.events import EventBus
+from baqueue.pruner import Pruner
+from baqueue.worker import Worker
+
+logger = logging.getLogger("baqueue.supervisor")
+
+
+class Supervisor:
+    """Manages a pool of workers for one or more queues.
+
+    Supports auto-balancing, scaling, and graceful shutdown.
+    """
+
+    def __init__(
+        self,
+        driver: BaseDriver,
+        config: SupervisorConfig | None = None,
+        events: EventBus | None = None,
+        balancer: Any | None = None,
+        pruner: Pruner | None = None,
+    ):
+        self.driver = driver
+        self.config = config or SupervisorConfig()
+        self.events = events or EventBus.default()
+        self.balancer = balancer
+        self.pruner = pruner
+        self._workers: list[Worker] = []
+        self._tasks: list[asyncio.Task] = []
+        self._running = False
+        self._delayed_task: asyncio.Task | None = None
+        self._heartbeat_task: asyncio.Task | None = None
+        self._balance_task: asyncio.Task | None = None
+        self._pruner_task: asyncio.Task | None = None
+
+    @property
+    def is_running(self) -> bool:
+        return self._running
+
+    @property
+    def worker_count(self) -> int:
+        return len(self._workers)
+
+    @property
+    def stats(self) -> dict[str, Any]:
+        return {
+            "name": self.config.name,
+            "queues": self.config.queues,
+            "balance": self.config.balance,
+            "workers": len(self._workers),
+            "running": self._running,
+            "worker_stats": [w.stats for w in self._workers],
+        }
+
+    async def start(self) -> None:
+        """Start the supervisor and its worker pool."""
+        self._running = True
+        logger.info(
+            "Supervisor '%s' starting with %d workers on queues %s",
+            self.config.name, self.config.min_workers, self.config.queues,
+        )
+        await self.events.emit("supervisor.started", supervisor=self.config.name)
+
+        self._setup_signal_handlers()
+
+        for i in range(self.config.min_workers):
+            self._spawn_worker(i)
+
+        await self._report_stats()
+        self._delayed_task = asyncio.create_task(self._poll_delayed())
+        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
+
+        if self.balancer:
+            self._balance_task = asyncio.create_task(self._balance_loop())
+
+        if self.pruner:
+            self._pruner_task = asyncio.create_task(self.pruner.start())
+
+        try:
+            await asyncio.gather(*self._tasks, return_exceptions=True)
+        except asyncio.CancelledError:
+            pass
+
+    async def stop(self) -> None:
+        """Gracefully stop all workers."""
+        logger.info("Supervisor '%s' shutting down...", self.config.name)
+        self._running = False
+
+        await self._report_stats()
+        for w in self._workers:
+            w.stop()
+
+        if self.pruner:
+            self.pruner.stop()
+
+        if self._delayed_task:
+            self._delayed_task.cancel()
+        if self._heartbeat_task:
+            self._heartbeat_task.cancel()
+        if self._balance_task:
+            self._balance_task.cancel()
+        if self._pruner_task:
+            self._pruner_task.cancel()
+
+        aux_tasks = [
+            t for t in (self._delayed_task, self._heartbeat_task, self._balance_task, self._pruner_task)
+            if t is not None
+        ]
+        if aux_tasks:
+            await asyncio.gather(*aux_tasks, return_exceptions=True)
+        self._delayed_task = None
+        self._heartbeat_task = None
+        self._balance_task = None
+        self._pruner_task = None
+
+        for task in self._tasks:
+            task.cancel()
+
+        await asyncio.gather(*self._tasks, return_exceptions=True)
+        self._tasks.clear()
+        self._workers.clear()
+        await self._report_stats()
+
+        await self.events.emit("supervisor.stopped", supervisor=self.config.name)
+        logger.info("Supervisor '%s' stopped", self.config.name)
+
+    def _spawn_worker(self, index: int) -> Worker:
+        worker = Worker(
+            driver=self.driver,
+            queues=list(self.config.queues),
+            events=self.events,
+            sleep_interval=self.config.sleep,
+            timeout=self.config.timeout,
+            name=f"{self.config.name}-worker-{index}",
+        )
+        self._workers.append(worker)
+        task = asyncio.create_task(worker.start())
+        self._tasks.append(task)
+        return worker
+
+    async def scale(self, count: int) -> None:
+        """Scale worker pool to the specified count."""
+        count = max(self.config.min_workers, min(count, self.config.max_workers))
+        current = len(self._workers)
+
+        if count > current:
+            for i in range(current, count):
+                self._spawn_worker(i)
+            logger.info("Scaled up to %d workers", count)
+        elif count < current:
+            for _ in range(current - count):
+                worker = self._workers.pop()
+                worker.stop()
+            logger.info("Scaled down to %d workers", count)
+
+    async def _poll_delayed(self) -> None:
+        """Periodically move delayed jobs into their queues."""
+        while self._running:
+            try:
+                await self.driver.pop_delayed()
+            except Exception:
+                logger.exception("Error polling delayed jobs")
+            await asyncio.sleep(1)
+
+    async def _balance_loop(self) -> None:
+        """Periodically rebalance workers across queues."""
+        while self._running:
+            try:
+                if self.balancer:
+                    new_count = await self.balancer.recommend(
+                        self.driver, self.config.queues, len(self._workers)
+                    )
+                    if new_count != len(self._workers):
+                        await self.scale(new_count)
+            except Exception:
+                logger.exception("Error in balance loop")
+            await asyncio.sleep(5)
+
+    async def _heartbeat_loop(self) -> None:
+        while self._running:
+            await self._report_stats()
+            await asyncio.sleep(1)
+
+    async def _report_stats(self) -> None:
+        try:
+            await self.driver.report_supervisor(self.stats)
+        except Exception:
+            logger.exception("Failed to report supervisor stats")
+
+    def _setup_signal_handlers(self) -> None:
+        if os.name == "nt":
+            return
+        loop = asyncio.get_running_loop()
+        for sig in (signal.SIGTERM, signal.SIGINT):
+            loop.add_signal_handler(sig, lambda: asyncio.create_task(self.stop()))

baqueue/worker.py

@@ -0,0 +1,165 @@
+"""Worker - processes jobs from queues."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import traceback
+from typing import Any
+
+from baqueue.drivers.base import BaseDriver
+from baqueue.events import EventBus
+from baqueue.job import Job, FunctionJob
+from baqueue.retry import compute_delay, should_retry
+from baqueue.serializer import JobPayload, resolve_job_class
+
+logger = logging.getLogger("baqueue.worker")
+
+
+class Worker:
+    """Pulls and executes jobs from one or more queues."""
+
+    def __init__(
+        self,
+        driver: BaseDriver,
+        queues: list[str],
+        events: EventBus | None = None,
+        sleep_interval: float = 1.0,
+        timeout: int = 60,
+        name: str = "worker-0",
+    ):
+        self.driver = driver
+        self.queues = queues
+        self.events = events or EventBus.default()
+        self.sleep_interval = sleep_interval
+        self.timeout = timeout
+        self.name = name
+        self._running = False
+        self._current_job: JobPayload | None = None
+        self._jobs_processed = 0
+        self._jobs_failed = 0
+
+    @property
+    def is_running(self) -> bool:
+        return self._running
+
+    @property
+    def stats(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "queues": self.queues,
+            "running": self._running,
+            "current_job": self._current_job.id if self._current_job else None,
+            "jobs_processed": self._jobs_processed,
+            "jobs_failed": self._jobs_failed,
+        }
+
+    async def start(self) -> None:
+        """Start the worker loop."""
+        self._running = True
+        await self.events.emit("worker.started", worker=self.name)
+        logger.info("Worker %s started on queues %s", self.name, self.queues)
+
+        try:
+            while self._running:
+                job = await self._fetch_next()
+                if job:
+                    await self._process(job)
+                else:
+                    await asyncio.sleep(self.sleep_interval)
+        except asyncio.CancelledError:
+            pass
+        finally:
+            self._running = False
+            await self.events.emit("worker.stopped", worker=self.name)
+            logger.info("Worker %s stopped", self.name)
+
+    def stop(self) -> None:
+        self._running = False
+
+    async def _fetch_next(self) -> JobPayload | None:
+        for queue in self.queues:
+            job = await self.driver.pop(queue)
+            if job:
+                return job
+        return None
+
+    async def _process(self, payload: JobPayload) -> None:
+        self._current_job = payload
+        job_timeout = payload.timeout or self.timeout
+
+        try:
+            await self.events.emit("job.started", payload=payload, worker=self.name)
+            await self.driver.record_metric(payload.queue, "processing", 1)
+            logger.debug("Processing job %s (%s)", payload.id, payload.job_class)
+
+            job_instance = self._instantiate(payload)
+            result = await asyncio.wait_for(
+                job_instance.handle(**payload.data),
+                timeout=job_timeout,
+            )
+
+            await self.driver.complete(payload)
+            await self.driver.record_metric(payload.queue, "completed", 1)
+            await self.events.emit("job.completed", payload=payload, result=result, worker=self.name)
+            self._jobs_processed += 1
+
+            if hasattr(job_instance, "on_success"):
+                try:
+                    await job_instance.on_success(result, payload)
+                except Exception:
+                    logger.exception("Error in on_success for job %s", payload.id)
+
+            await self._check_batch_completion(payload)
+
+        except Exception as exc:
+            error_msg = f"{type(exc).__name__}: {exc}\n{traceback.format_exc()}"
+            logger.warning("Job %s failed (attempt %d): %s", payload.id, payload.attempts, exc)
+
+            if should_retry(payload.attempts, payload.max_attempts):
+                delay = compute_delay(payload.backoff, payload.attempts)
+                await self.driver.release(payload, delay=delay)
+                await self.events.emit("job.retrying", payload=payload, error=error_msg, delay=delay)
+            else:
+                await self.driver.fail(payload, error_msg)
+                await self.driver.record_metric(payload.queue, "failed", 1)
+                await self.events.emit("job.failed", payload=payload, error=error_msg, worker=self.name)
+                self._jobs_failed += 1
+
+                job_instance = self._instantiate(payload)
+                try:
+                    await job_instance.on_failure(exc, payload)
+                except Exception:
+                    logger.exception("Error in on_failure for job %s", payload.id)
+
+                await self._check_batch_failure(payload)
+        finally:
+            self._current_job = None
+
+    def _instantiate(self, payload: JobPayload) -> Job:
+        cls = resolve_job_class(payload.job_class)
+        if isinstance(cls, FunctionJob):
+            return cls
+        return cls()
+
+    async def _check_batch_completion(self, payload: JobPayload) -> None:
+        if not payload.batch_id:
+            return
+        batch = await self.driver.increment_batch_counter(payload.batch_id, "completed_count", 1)
+        if not batch:
+            return
+        done = batch.get("completed_count", 0) + batch.get("failed_count", 0)
+        # Equality, not >=, so only the worker that pushed `done` to total fires the event.
+        if done == batch.get("total", 0):
+            await self.events.emit("batch.completed", batch_id=payload.batch_id, batch=batch)
+
+    async def _check_batch_failure(self, payload: JobPayload) -> None:
+        if not payload.batch_id:
+            return
+        batch = await self.driver.increment_batch_counter(payload.batch_id, "failed_count", 1)
+        if not batch:
+            return
+        # Fire batch.failed exactly once on the first failure (the increment that
+        # transitioned failed_count from 0 to 1).
+        if batch.get("allow_failures", False) is False and batch.get("failed_count", 0) == 1:
+            await self.events.emit("batch.failed", batch_id=payload.batch_id, batch=batch)