kickdown 0.4.0__py3-none-any.whl

kickdown/__init__.py

@@ -0,0 +1,14 @@
+from .client import Client
+from .models import Performable, Stats, Task, Worker
+from .queue import Queue
+from .server import Server
+
+__all__ = [
+    "Client",
+    "Performable",
+    "Queue",
+    "Server",
+    "Stats",
+    "Task",
+    "Worker",
+]

kickdown/client.py

@@ -0,0 +1,32 @@
+import asyncio
+
+from .log import default_logger
+from .models import Task
+from .queue import Queue
+from .store import Store
+
+
+class Client:
+    def __init__(self, redis_url: str):
+        self._store = Store(redis_url)
+        self.logger = default_logger()
+
+    def queue(self, name: str) -> Queue:
+        return Queue(name, self._store)
+
+    async def enqueue(self, task: Task) -> str:
+        await self.queue(task.queue).push(task)
+        self.logger.info(
+            f"jid={task.jid} accepted",
+            extra={"queue": task.queue, "operation": task.operation},
+        )
+        return task.jid
+
+    async def close(self):
+        await asyncio.to_thread(self._store.close)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *_exc_info):
+        await self.close()

kickdown/consumer.py

@@ -0,0 +1,206 @@
+import asyncio
+import logging
+import os
+import socket
+import time
+from collections import deque
+from collections.abc import Awaitable, Callable, Mapping
+
+from pydantic import ValidationError
+from uuid_extensions import uuid7str
+
+from .log import default_logger
+from .models import Task, Worker
+from .queue import Queue
+from .store import Store, StoreError
+
+_default_poll_interval = 0.1
+_retry_delay = 1
+_backoff_coefficient = 1.5
+_heartbeat_interval = 10
+_heartbeat_ttl = 30
+
+
+def _default_consumer_id() -> str:
+    return f"{socket.gethostname()}:{os.getpid()}:{uuid7str()[-8:]}"
+
+
+class Consumer:
+    def __init__(
+        self,
+        store: Store,
+        workers: Mapping[tuple[str, str], Worker],
+        concurrency: int = 1,
+        logger: logging.Logger | None = None,
+        consumer_id: str | None = None,
+    ):
+        self.id = consumer_id or _default_consumer_id()
+        self._store = store
+        self._workers = workers
+        self._queues: dict[str, Queue] = {
+            name: Queue(name, store)
+            for name in sorted({worker.queue for worker in workers.values()})
+        }
+        self._order: deque[str] = deque(self._queues)
+        self._semaphore = asyncio.Semaphore(concurrency)
+        self.logger = logger or default_logger()
+        self._tasks: set[asyncio.Task] = set()
+
+    async def start(self) -> None:
+        await asyncio.to_thread(self._store.register_consumer, self.id, _heartbeat_ttl)
+        self.logger.info(f"consumer {self.id} registered")
+
+    async def stop(self) -> None:
+        """Returns whatever is still in flight, then leaves the registry.
+
+        Without this a clean restart would leave its tasks sitting until the
+        heartbeat expired and another server reaped them.
+        """
+        try:
+            returned = await asyncio.to_thread(self._store.reap, self.id)
+            if returned:
+                self.logger.warning(
+                    f"consumer {self.id} returned {returned} unfinished task(s)"
+                )
+            await asyncio.to_thread(self._store.deregister_consumer, self.id)
+        except StoreError as err:
+            self.logger.error(
+                "redis error while deregistering consumer", extra={"error": str(err)}
+            )
+
+    async def heartbeat(self) -> None:
+        while True:
+            try:
+                await asyncio.to_thread(self._store.heartbeat, self.id, _heartbeat_ttl)
+            except StoreError as err:
+                self.logger.error(
+                    "redis error while refreshing heartbeat",
+                    extra={"error": str(err)},
+                )
+            await asyncio.sleep(_heartbeat_interval)
+
+    async def consume(self) -> None:
+        while True:
+            await self._semaphore.acquire()
+
+            try:
+                claimed = await asyncio.to_thread(
+                    self._store.claim, self._poll_order(), self.id
+                )
+            except StoreError as err:
+                self._semaphore.release()
+                self.logger.error(
+                    "redis error while claiming task", extra={"error": str(err)}
+                )
+                await asyncio.sleep(_default_poll_interval)
+                continue
+            except ValidationError as err:
+                self._semaphore.release()
+                self.logger.error(
+                    "failed to parse task payload", extra={"error": str(err)}
+                )
+                continue
+
+            if claimed is None:
+                self._semaphore.release()
+                await asyncio.sleep(_default_poll_interval)
+                continue
+
+            raw, task = claimed
+            handle = asyncio.create_task(self._run_task(task, raw))
+            self._tasks.add(handle)
+            handle.add_done_callback(self._tasks.discard)
+
+    async def drain(self) -> None:
+        if self._tasks:
+            await asyncio.gather(*self._tasks, return_exceptions=True)
+
+    def _poll_order(self) -> list[str]:
+        order = list(self._order)
+        self._order.rotate(-1)
+        return order
+
+    def _queue(self, name: str) -> Queue:
+        # a payload naming a queue we do not serve is malformed, but its stats
+        # should still land on the queue it claims
+        return self._queues.get(name) or Queue(name, self._store)
+
+    async def _run_task(self, task: Task, raw: bytes) -> None:
+        queue = self._queue(task.queue)
+        failure: Exception | None = None
+        try:
+            worker = self._workers.get((task.queue, task.operation))
+            if worker is None:
+                self.logger.error(
+                    f"no worker registered for task jid={task.jid}",
+                    extra={
+                        "jid": task.jid,
+                        "queue": task.queue,
+                        "operation": task.operation,
+                    },
+                )
+                await self._increment(queue.increment_failed)
+                await self._ack(raw)
+                return
+
+            self.logger.info(
+                f"jid={task.jid} started",
+                extra={"queue": task.queue, "operation": task.operation},
+            )
+            await worker.perform(task.params)
+            self.logger.info(f"jid={task.jid} done")
+            await self._increment(queue.increment_processed)
+        except Exception as err:  # noqa: BLE001 - worker code is arbitrary; retry boundary must catch anything
+            failure = err
+        finally:
+            self._semaphore.release()
+
+        if failure is None:
+            await self._ack(raw)
+            return
+
+        if task.retry_count > 0:
+            delay = _retry_delay * _backoff_coefficient**task.attempt
+            self.logger.warning(
+                f"jid={task.jid} failed, retrying in {delay:.1f}s "
+                f"({task.retry_count} attempt(s) left)",
+                extra={"error": str(failure)},
+            )
+            retry_task = task.model_copy(
+                update={
+                    "retry_count": task.retry_count - 1,
+                    "attempt": task.attempt + 1,
+                }
+            )
+            try:
+                await queue.schedule(retry_task, time.time() + delay)
+            except StoreError as schedule_err:
+                # leave it in flight: the reaper will put it back rather than
+                # drop it on the floor
+                self.logger.error(
+                    f"jid={task.jid} failed to schedule retry",
+                    extra={"error": str(schedule_err)},
+                )
+                return
+            await self._ack(raw)
+        else:
+            self.logger.error(
+                f"jid={task.jid} failed permanently", extra={"error": str(failure)}
+            )
+            await self._increment(queue.increment_failed)
+            await self._ack(raw)
+
+    async def _ack(self, raw: bytes) -> None:
+        try:
+            await asyncio.to_thread(self._store.ack, self.id, raw)
+        except StoreError as err:
+            self.logger.error(
+                "failed to remove task from the in-flight list",
+                extra={"error": str(err)},
+            )
+
+    async def _increment(self, increment: Callable[[], Awaitable[None]]) -> None:
+        try:
+            await increment()
+        except StoreError as err:
+            self.logger.error("failed to update stats", extra={"error": str(err)})

kickdown/log.py

@@ -0,0 +1,14 @@
+import logging
+import sys
+
+_logger_name = "kickdown"
+
+
+def default_logger() -> logging.Logger:
+    logger = logging.getLogger(_logger_name)
+    if not logger.handlers:
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
+        logger.addHandler(handler)
+        logger.setLevel(logging.INFO)
+    return logger

kickdown/models.py

@@ -0,0 +1,31 @@
+from typing import Protocol, runtime_checkable
+
+from pydantic import BaseModel, Field
+from uuid_extensions import uuid7str
+
+
+@runtime_checkable
+class Performable(Protocol):
+    queue: str
+    operation: str
+
+    async def perform(self, payload: dict) -> None: ...
+
+
+# Workers are registered either as classes (with `perform` as a classmethod) or
+# as instances; both are called the same way.
+type Worker = type[Performable] | Performable
+
+
+class Task(BaseModel):
+    queue: str
+    operation: str
+    params: dict
+    jid: str = Field(default_factory=lambda: uuid7str())
+    retry_count: int = 1
+    attempt: int = 0
+
+
+class Stats(BaseModel):
+    processed: int = 0
+    failed: int = 0

kickdown/queue.py

@@ -0,0 +1,60 @@
+import asyncio
+
+from .models import Stats, Task
+from .store import Store
+
+
+class Queue:
+    """
+    A named queue: every per-queue operation, bound to one name.
+
+    Wraps the synchronous `Store` calls in threads, so callers stay async.
+    """
+
+    def __init__(self, name: str, store: Store):
+        self._name = name
+        self._store = store
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def __repr__(self) -> str:
+        return f"Queue({self._name!r})"
+
+    async def push(self, task: Task) -> None:
+        await asyncio.to_thread(self._store.push, task)
+
+    async def schedule(self, task: Task, run_at: float) -> None:
+        await asyncio.to_thread(self._store.schedule, task, run_at)
+
+    async def pending(self) -> list[Task]:
+        return await asyncio.to_thread(self._store.pending, self._name)
+
+    async def length(self) -> int:
+        return await asyncio.to_thread(self._store.queue_length, self._name)
+
+    async def scheduled(self) -> list[Task]:
+        return await asyncio.to_thread(self._store.scheduled, self._name)
+
+    async def scheduled_length(self) -> int:
+        return await asyncio.to_thread(self._store.scheduled_length, self._name)
+
+    async def enqueue_due(self, now: float, limit: int) -> int:
+        return await asyncio.to_thread(self._store.enqueue_due, self._name, now, limit)
+
+    async def purge(self) -> int:
+        """Drops every pending task, returning how many were dropped.
+
+        Scheduled tasks and the processed/failed counters are left alone.
+        """
+        return await asyncio.to_thread(self._store.purge, self._name)
+
+    async def stats(self) -> Stats:
+        return await asyncio.to_thread(self._store.stats, self._name)
+
+    async def increment_processed(self) -> None:
+        await asyncio.to_thread(self._store.increment_processed, self._name)
+
+    async def increment_failed(self) -> None:
+        await asyncio.to_thread(self._store.increment_failed, self._name)

kickdown/reaper.py

@@ -0,0 +1,64 @@
+import asyncio
+import logging
+
+from .log import default_logger
+from .store import Store, StoreError
+
+_default_poll_interval = 15
+_reap_lock_ttl = 60
+
+
+class Reaper:
+    """Returns tasks stranded by consumers that died mid-task.
+
+    A consumer holds its claimed tasks in its own in-flight list and keeps a
+    heartbeat key alive. When the heartbeat expires the consumer is gone, and
+    whatever is left in its list is pushed back into the queues it came from.
+    """
+
+    def __init__(
+        self,
+        store: Store,
+        poll_interval: float = _default_poll_interval,
+        logger: logging.Logger | None = None,
+    ):
+        self._store = store
+        self._poll_interval = poll_interval
+        self.logger = logger or default_logger()
+
+    async def run(self) -> None:
+        while True:
+            await self._tick()
+            await asyncio.sleep(self._poll_interval)
+
+    async def _tick(self) -> None:
+        try:
+            consumers = await asyncio.to_thread(self._store.consumers)
+        except StoreError as err:
+            self.logger.error(
+                "redis error while listing consumers", extra={"error": str(err)}
+            )
+            return
+
+        for consumer_id in consumers:
+            try:
+                if await asyncio.to_thread(self._store.is_alive, consumer_id):
+                    continue
+                # one server reaps a given corpse; the rest skip it
+                if not await asyncio.to_thread(
+                    self._store.claim_reap, consumer_id, _reap_lock_ttl
+                ):
+                    continue
+                reaped = await asyncio.to_thread(self._store.reap, consumer_id)
+                await asyncio.to_thread(self._store.deregister_consumer, consumer_id)
+            except StoreError as err:
+                self.logger.error(
+                    "redis error while reaping consumer",
+                    extra={"consumer": consumer_id, "error": str(err)},
+                )
+                continue
+
+            self.logger.warning(
+                f"consumer {consumer_id} died, requeued {reaped} task(s)",
+                extra={"consumer": consumer_id},
+            )

kickdown/scheduler.py

@@ -0,0 +1,47 @@
+import asyncio
+import logging
+import time
+
+from .log import default_logger
+from .queue import Queue
+from .store import StoreError
+
+_default_poll_interval = 1
+_default_batch_size = 100
+
+
+class Scheduler:
+    """Moves tasks whose scheduled time has come back into their queues."""
+
+    def __init__(
+        self,
+        queues: list[Queue],
+        poll_interval: float = _default_poll_interval,
+        logger: logging.Logger | None = None,
+    ):
+        self._queues = queues
+        self._poll_interval = poll_interval
+        self.logger = logger or default_logger()
+
+    async def run(self) -> None:
+        while True:
+            await self._tick()
+            await asyncio.sleep(self._poll_interval)
+
+    async def _tick(self) -> None:
+        now = time.time()
+        for queue in self._queues:
+            try:
+                moved = await queue.enqueue_due(now, _default_batch_size)
+            except StoreError as err:
+                self.logger.error(
+                    "redis error while enqueueing due tasks",
+                    extra={"queue": queue.name, "error": str(err)},
+                )
+                continue
+
+            if moved:
+                self.logger.info(
+                    f"scheduler enqueued {moved} due task(s)",
+                    extra={"queue": queue.name},
+                )

kickdown/server.py

@@ -0,0 +1,155 @@
+import asyncio
+import signal
+from collections.abc import Awaitable, Callable
+
+from .consumer import Consumer
+from .log import default_logger
+from .models import Task, Worker
+from .queue import Queue
+from .reaper import Reaper
+from .scheduler import Scheduler
+from .store import Store, StoreError
+from .web import Web
+
+Hook = Callable[[], Awaitable[None]]
+
+
+class Server:
+    def __init__(
+        self,
+        redis_url: str,
+        concurrency: int = 1,
+        web_port: int = 3030,
+        admin_username: str | None = None,
+        admin_password: str | None = None,
+    ):
+        self._store = Store(redis_url)
+        self._concurrency = concurrency
+        self._web_port = web_port
+        self._admin_username = admin_username
+        self._admin_password = admin_password
+        self._startup_hooks: list[Hook] = []
+        self._shutdown_hooks: list[Hook] = []
+        self._worker: dict[tuple[str, str], Worker] = {}
+        self.logger = default_logger()
+
+    @property
+    def store(self) -> Store:
+        return self._store
+
+    @property
+    def queues(self) -> list[Queue]:
+        return [
+            Queue(name, self._store)
+            for name in sorted({worker.queue for worker in self._worker.values()})
+        ]
+
+    def queue(self, name: str) -> Queue:
+        return Queue(name, self._store)
+
+    def add_workers(self, *args: Worker):
+        for worker in args:
+            self._worker[(worker.queue, worker.operation)] = worker
+
+    async def enqueue(self, task: Task) -> str:
+        await self.queue(task.queue).push(task)
+        self.logger.info(
+            f"jid={task.jid} accepted",
+            extra={"queue": task.queue, "operation": task.operation},
+        )
+        return task.jid
+
+    def on_startup(self, fn: Hook) -> Hook:
+        self._startup_hooks.append(fn)
+        return fn
+
+    def on_shutdown(self, fn: Hook) -> Hook:
+        self._shutdown_hooks.append(fn)
+        return fn
+
+    async def run(self):
+        if not self._worker:
+            raise RuntimeError(
+                "No workers registered. Register them with add_workers() before running the server."
+            )
+
+        try:
+            await asyncio.to_thread(self._store.ping)
+        except StoreError as e:
+            raise RuntimeError(f"Redis connection failed: {e}") from e
+
+        consumer = Consumer(
+            store=self._store,
+            workers=self._worker,
+            concurrency=self._concurrency,
+            logger=self.logger,
+        )
+        scheduler = Scheduler(queues=self.queues, logger=self.logger)
+        reaper = Reaper(store=self._store, logger=self.logger)
+        web = Web(
+            port=self._web_port,
+            server=self,
+            admin_username=self._admin_username,
+            admin_password=self._admin_password,
+        )
+
+        await self._run_hooks(self._startup_hooks, "startup")
+        await consumer.start()
+
+        stop_event = asyncio.Event()
+        loop = asyncio.get_running_loop()
+        for sig in (signal.SIGTERM, signal.SIGINT):
+            loop.add_signal_handler(sig, stop_event.set)
+
+        self.logger.info(
+            f"server starting (queues={[q.name for q in self.queues]}, "
+            f"concurrency={self._concurrency}, web_port={self._web_port})"
+        )
+
+        consume_task = asyncio.create_task(consumer.consume(), name="consumer")
+        beat_task = asyncio.create_task(consumer.heartbeat(), name="heartbeat")
+        schedule_task = asyncio.create_task(scheduler.run(), name="scheduler")
+        reap_task = asyncio.create_task(reaper.run(), name="reaper")
+        web_task = asyncio.create_task(web.run(), name="web")
+        stop_task = asyncio.create_task(stop_event.wait(), name="stop")
+
+        background = (consume_task, beat_task, schedule_task, reap_task, web_task)
+
+        try:
+            await asyncio.wait(
+                [*background, stop_task],
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+
+            for task in background:
+                if task.done() and not task.cancelled():
+                    exc = task.exception()
+                    if exc is not None:
+                        self.logger.error(
+                            f"{task.get_name()} task failed unexpectedly",
+                            extra={"error": str(exc)},
+                        )
+        finally:
+            self.logger.info("server shutting down")
+            for task in (*background, stop_task):
+                task.cancel()
+            await asyncio.gather(*background, stop_task, return_exceptions=True)
+            await consumer.drain()
+            await consumer.stop()
+
+            await self._run_hooks(self._shutdown_hooks, "shutdown")
+
+            await asyncio.to_thread(self._store.close)
+
+    async def _run_hooks(self, hooks: list[Hook], phase: str) -> None:
+        for hook in hooks:
+            try:
+                await hook()
+            except Exception as err:  # noqa: BLE001 - hook code is arbitrary; must not abort the server
+                self.logger.error(
+                    f"{phase} hook failed",
+                    extra={
+                        "hook": getattr(hook, "__name__", repr(hook)),
+                        "error": str(err),
+                    },
+                )