ltq 0.3.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

ltq/middleware.py

@@ -1,117 +1,115 @@
 from __future__ import annotations
 
-import asyncio
+from datetime import timedelta
+from functools import lru_cache
+import random
 import time
 from abc import ABC, abstractmethod
-from typing import Any, Awaitable, Callable
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, AsyncIterator
 
-from .errors import RetryMessage
-from .message import Message
-from .logger import get_logger
+from .errors import RejectError, RetryError
 
-logger = get_logger()
-Handler = Callable[[Message], Awaitable[Any]]
+if TYPE_CHECKING:
+    from .message import Message
+    from .task import Task
 
 
 class Middleware(ABC):
     @abstractmethod
-    async def handle(self, message: Message, next_handler: Handler) -> Any: ...
+    @asynccontextmanager
+    async def __call__(self, message: Message, task: Task) -> AsyncIterator[None]:
+        yield
 
 
-class Retry(Middleware):
-    def __init__(
-        self,
-        max_retries: int = 3,
-        min_delay: float = 1.0,
-        max_delay: float = 60.0,
-        backoff: float = 2.0,
-    ):
-        self.max_retries = max_retries
-        self.min_delay = min_delay
-        self.max_delay = max_delay
-        self.backoff = backoff
-
-    async def handle(self, message: Message, next_handler: Handler) -> Any:
-        retries = message.ctx.get("retries", 0)
+class MaxTries(Middleware):
+    @asynccontextmanager
+    async def __call__(self, message: Message, task: Task) -> AsyncIterator[None]:
+        max_tries = task.options.get("max_tries")
+        if max_tries is not None:
+            if message.ctx.get("tries", 0) >= max_tries:
+                raise RejectError(f"Message {message.id} exceeded max tries")
 
         try:
-            return await next_handler(message)
-        except Exception as e:
-            retries += 1
-            message.ctx["retries"] = retries
-            max_retries = max(self.max_retries - 1, 0)
-
-            if retries > max_retries:
-                raise
-
-            delay = min(
-                self.min_delay * (self.backoff ** (retries - 1)),
-                self.max_delay,
-            )
-            logger.warning(
-                f"Retry attempt {retries}/{max_retries} ({type(e).__name__})",
-                exc_info=True,
-            )
-            raise RetryMessage(delay, str(e))
-
-
-class RateLimit(Middleware):
-    def __init__(self, requests_per_second: float):
-        self.min_interval = 1.0 / requests_per_second
-        self._last_request: float = 0
-        self._lock = asyncio.Lock()
-
-    async def handle(self, message: Message, next_handler: Handler) -> Any:
-        async with self._lock:
-            now = time.monotonic()
-            elapsed = now - self._last_request
-            if elapsed < self.min_interval:
-                await asyncio.sleep(self.min_interval - elapsed)
-            self._last_request = time.monotonic()
-
-        return await next_handler(message)
+            yield
+        except Exception:
+            if not message.ctx.pop("rate_limited", False):
+                message.ctx["tries"] = message.ctx.get("tries", 0) + 1
+            raise
+
+
+class MaxAge(Middleware):
+    @asynccontextmanager
+    async def __call__(self, message: Message, task: Task) -> AsyncIterator[None]:
+        max_age: timedelta | None = task.options.get("max_age")
+        created_at = message.ctx.get("created_at")
+
+        if max_age is not None and created_at is not None:
+            age = time.time() - float(created_at)
+            if age > max_age.total_seconds():
+                raise RejectError(f"Message {message.id} too old")
+
+        yield
+
+
+class MaxRate(Middleware):
+    def __init__(self) -> None:
+        self.last_times: dict[str, float] = {}
+
+    @lru_cache(maxsize=128)
+    def _parse_rate(self, rate: str) -> float:
+        count, unit = rate.split("/")
+        count = float(count)
+        unit = unit.strip().lower()
+
+        if unit == "s":
+            return count
+        elif unit == "m":
+            return count / 60
+        elif unit == "h":
+            return count / 3600
+        else:
+            raise ValueError(f"Invalid rate unit: {unit}. Use 's', 'm', or 'h'")
+
+    @asynccontextmanager
+    async def __call__(self, message: Message, task: Task) -> AsyncIterator[None]:
+        max_rate = task.options.get("max_rate")
+        if max_rate:
+            now = time.time()
+            last = self.last_times.get(message.task_name, 0.0)
+            elapsed = now - last
+            rate_per_sec = self._parse_rate(max_rate)
+            interval = 1.0 / rate_per_sec
+
+            if elapsed < interval:
+                base_delay = interval - elapsed
+                delay = base_delay * 0.5 + random.uniform(0, base_delay * 0.5)
+                message.ctx["rate_limited"] = True
+                raise RetryError(delay=delay)
+
+            self.last_times[message.task_name] = now
+        yield
 
 
-class Timeout(Middleware):
-    def __init__(self, timeout: float):
-        self.timeout = timeout
-
-    async def handle(self, message: Message, next_handler: Handler) -> Any:
-        return await asyncio.wait_for(next_handler(message), timeout=self.timeout)
+class Sentry(Middleware):
+    def __init__(self, dsn: str) -> None:
+        self.sentry = None
+        try:
+            import sentry_sdk  # type: ignore
 
+            sentry_sdk.init(dsn=dsn)
+            self.sentry = sentry_sdk
+        except ImportError:
+            pass
 
-class Sentry(Middleware):
-    def __init__(self, dsn: str, **kwargs: Any) -> None:
+    @asynccontextmanager
+    async def __call__(self, message: Message, task: Task) -> AsyncIterator[None]:
         try:
-            import sentry_sdk  # type: ignore[import-not-found]
-        except ModuleNotFoundError as exc:
-            raise ModuleNotFoundError(
-                "Sentry middleware requires optional dependency 'sentry-sdk'. "
-                "Install with 'ltq[sentry]'."
-            ) from exc
-
-        self.sentry = sentry_sdk
-        self.sentry.init(dsn=dsn, send_default_pii=True, **kwargs)
-
-    async def handle(self, message: Message, next_handler: Handler) -> Any:
-        with self.sentry.push_scope() as scope:
-            scope.set_tag("task", message.task_name)
-            scope.set_tag("message_id", message.id)
-            scope.set_context(
-                "message",
-                {
-                    "id": message.id,
-                    "task": message.task_name,
-                    "args": message.args,
-                    "kwargs": message.kwargs,
-                    "ctx": message.ctx,
-                },
-            )
-
-            try:
-                return await next_handler(message)
-            except RetryMessage:
-                raise
-            except Exception as e:
+            yield
+        except Exception as e:
+            if self.sentry:
                 self.sentry.capture_exception(e)
-                raise
+            raise
+
+
+DEFAULT: list[Middleware] = [MaxTries(), MaxAge(), MaxRate()]

ltq/scheduler.py

@@ -1,27 +1,22 @@
 from __future__ import annotations
 
 import asyncio
-import time
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
+from .broker import Broker
 from .message import Message
-from .utils import dispatch
 from .logger import get_logger
 
 try:
-    from croniter import croniter
+    from croniter import croniter  # type: ignore
 except ImportError:
     croniter = None
 
-if TYPE_CHECKING:
-    from .task import Task
-
 
 @dataclass
 class ScheduledJob:
-    task: Task
     msg: Message
     expr: str
     _cron: Any = field(init=False, repr=False)  # croniter instance
@@ -36,11 +31,16 @@ class ScheduledJob:
 
 
 class Scheduler:
-    def __init__(self, poll_interval: float = 10.0) -> None:
+    def __init__(
+        self,
+        broker_url: str = "redis://localhost:6379",
+        poll_interval: float = 10.0,
+    ) -> None:
+        self.broker = Broker.from_url(broker_url)
         self.poll_interval = poll_interval
         self.jobs: list[ScheduledJob] = []
-        self.logger = get_logger("ltq.scheduler")
-        self._running = False
+        self.logger = get_logger("scheduler")
+        self.task: asyncio.Task[None] | None = None
 
     def cron(self, expr: str, msg: Message) -> None:
         if croniter is None:
@@ -48,35 +48,54 @@ class Scheduler:
                 "Scheduler requires optional dependency 'croniter'. "
                 "Install with 'ltq[scheduler]'."
             )
-        if msg.task is None:
-            raise ValueError("Message must have a task assigned to use with scheduler")
-        self.jobs.append(ScheduledJob(msg.task, msg, expr))
+        self.jobs.append(ScheduledJob(msg, expr))
 
-    def run(self) -> None:
-        self._running = True
+    async def run(self) -> None:
         self.logger.info("Starting scheduler")
         for job in self.jobs:
             self.logger.info(
-                f"{job.task.name} [{job.expr}] next={job.next_run:%H:%M:%S}"
+                f"{job.msg.task_name} [{job.expr}] next={job.next_run:%H:%M:%S}"
             )
 
-        loop = asyncio.new_event_loop()
-        while self._running:
-            now = datetime.now()
-            due = [job for job in self.jobs if now >= job.next_run]
-            if due:
-                try:
-                    loop.run_until_complete(dispatch([job.msg for job in due]))
-                    for job in due:
-                        self.logger.info(
-                            f"Enqueued {job.task.name} scheduled={job.next_run:%H:%M:%S}"
-                        )
-                except Exception:
-                    self.logger.exception("Failed to dispatch scheduled jobs")
-                for job in due:
-                    job.advance()
-            time.sleep(self.poll_interval)
-        loop.close()
-
-    def stop(self) -> None:
-        self._running = False
+        try:
+            while True:
+                now = datetime.now()
+                due = [job for job in self.jobs if now >= job.next_run]
+
+                if due:
+                    try:
+                        for job in due:
+                            await self.broker.publish(job.msg)
+                            self.logger.info(
+                                f"Enqueued {job.msg.task_name} scheduled={job.next_run:%H:%M:%S}"
+                            )
+                            job.advance()
+                    except Exception:
+                        self.logger.exception("Failed to send scheduled jobs")
+                        # Don't advance jobs on failure - they'll retry next poll
+
+                await asyncio.sleep(self.poll_interval)
+        finally:
+            await self.broker.close()
+
+    def start(self) -> None:
+        try:
+            asyncio.run(self.run())
+        except KeyboardInterrupt:
+            self.logger.info("Scheduler stopped")
+
+    def start_background(self) -> None:
+        if self.task is not None:
+            raise RuntimeError("Scheduler is already running")
+        self.task = asyncio.create_task(self.run())
+
+    async def stop(self) -> None:
+        if self.task is None:
+            return
+        self.task.cancel()
+        try:
+            await self.task
+        except asyncio.CancelledError:
+            pass
+        self.task = None
+        self.logger.info("Scheduler stopped")

ltq/task.py

@@ -2,8 +2,8 @@ from __future__ import annotations
 
 from typing import Awaitable, Callable, Generic, ParamSpec, TypeVar
 
+from .broker import Broker
 from .message import Message
-from .q import Queue
 
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -12,28 +12,25 @@ R = TypeVar("R")
 class Task(Generic[P, R]):
     def __init__(
         self,
+        broker: Broker,
         name: str,
         fn: Callable[P, Awaitable[R]],
-        queue: Queue,
-        ttl: int | None = None,
+        options: dict | None = None,
     ) -> None:
         self.name = name
         self.fn = fn
-        self.queue = queue
-        self.ttl = ttl
+        self.options = options or {}
+        self.broker = broker
 
     def message(self, *args: P.args, **kwargs: P.kwargs) -> Message:
         return Message(
             args=args,
             kwargs=kwargs,
-            task=self,
             task_name=self.name,
         )
 
-    async def send(self, *args: P.args, **kwargs: P.kwargs) -> str:
-        message = self.message(*args, **kwargs)
-        await self.queue.put([message], ttl=self.ttl)
-        return message.id
+    async def send(self, *args: P.args, **kwargs: P.kwargs) -> None:
+        await self.broker.publish(self.message(*args, **kwargs))
 
     async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
         return await self.fn(*args, **kwargs)

ltq/utils.py

@@ -1,19 +1 @@
-from __future__ import annotations
 
-from collections import defaultdict
-
-from .message import Message
-from .q import Queue
-
-
-async def dispatch(messages: list[Message]) -> list[str]:
-    by_queue: defaultdict[Queue, list[Message]] = defaultdict(list)
-    for msg in messages:
-        if msg.task is None:
-            raise ValueError(f"Message {msg.id} has no task assigned")
-        by_queue[msg.task.queue].append(msg)
-
-    for queue, batch in by_queue.items():
-        await queue.put(batch)
-
-    return [msg.id for msg in messages]

ltq/worker.py

@@ -1,23 +1,15 @@
 from __future__ import annotations
 
 import asyncio
-from functools import partial
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, ParamSpec, TypeVar
+from contextlib import AsyncExitStack
+from typing import Awaitable, Callable, ParamSpec, TypeVar
 
-import redis.asyncio as redis
-
-from .errors import RetryMessage
-from .task import Task
-from .message import Message
-from .middleware import Handler, Middleware
-from .q import Queue
+from .broker import Broker
+from .errors import RejectError, RetryError
 from .logger import get_logger
-
-if TYPE_CHECKING:
-    from redis.asyncio import Redis as AsyncRedis
-
-logger = get_logger()
-
+from .message import Message
+from .middleware import DEFAULT, Middleware
+from .task import Task
 
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -26,80 +18,98 @@ R = TypeVar("R")
 class Worker:
     def __init__(
         self,
-        url: str = "redis://localhost:6379",
+        name: str,
+        broker_url: str = "redis://localhost:6379",
+        concurrency: int = 100,
         middlewares: list[Middleware] | None = None,
-        concurrency: int = 250,
-        poll_sleep: float = 0.1,
     ) -> None:
-        self.client: AsyncRedis = redis.from_url(url)
+        self.name = name
+        self.broker = Broker.from_url(broker_url)
         self.tasks: list[Task] = []
-        self.middlewares: list[Middleware] = middlewares or []
+        self.middlewares: list[Middleware] = middlewares or list(DEFAULT)
         self.concurrency: int = concurrency
-        self.poll_sleep: float = poll_sleep
+        self.logger = get_logger(name)
+
+    def register_middleware(self, middleware: Middleware, pos: int = -1) -> None:
+        if pos == -1:
+            self.middlewares.append(middleware)
+        else:
+            self.middlewares.insert(pos, middleware)
 
     def task(
         self,
-        queue_name: str | None = None,
-        ttl: int | None = None,
+        **options,
     ) -> Callable[[Callable[P, Awaitable[R]]], Task[P, R]]:
         def decorator(fn: Callable[P, Awaitable[R]]) -> Task[P, R]:
-            task_name = f"{fn.__module__}:{fn.__qualname__}"
-            queue = Queue(self.client, queue_name or task_name)
+            task_name = f"{self.name}:{fn.__qualname__}"
             task = Task(
                 name=task_name,
                 fn=fn,
-                queue=queue,
-                ttl=ttl,
+                options=options,
+                broker=self.broker,
             )
             self.tasks.append(task)
             return task
 
         return decorator
 
-    async def processor(self, task: Task):
-        async def base(message: Message) -> Any:
-            return await task.fn(*message.args, **message.kwargs)
-
-        handler: Handler = base
-        for middleware in reversed(self.middlewares):
-            handler = partial(middleware.handle, next_handler=handler)
-
+    async def _poll(self, task: Task, broker: Broker) -> None:
         sem = asyncio.Semaphore(self.concurrency)
-        pending: dict[asyncio.Task, Message] = {}
+        self.logger.info(f"Polling for Task {task.name}")
 
-        async def process(msg: Message) -> None:
-            async with sem:
-                try:
-                    await handler(msg)
-                except RetryMessage as e:
-                    logger.warning(f"Retrying in {e.delay}s: {e}")
-                    await task.queue.put([msg], delay=e.delay)
-                except Exception as e:
-                    logger.error(
-                        f"Rejected after error in {task.name}: {e}",
-                        exc_info=True,
-                    )
-
-        while True:
-            messages = await task.queue.get(self.concurrency)
-            if not messages:
-                await asyncio.sleep(self.poll_sleep)
-                continue
-
-            logger.debug(f"Processing {len(messages)} messages for {task.name}")
+        try:
+            while True:
+                message = await broker.consume(task.name)
+                # concurrency limiter, without, queue would be drained in one go.
+                await sem.acquire()
+                asyncio.create_task(self._process(task, broker, sem, message))
+        except asyncio.CancelledError:
+            self.logger.info(f"Worker {task.name} cancelled...")
+            raise
 
-            for msg in messages:
-                t = asyncio.create_task(process(msg))
-                pending[t] = msg
+    async def _process(
+        self,
+        task: Task,
+        broker: Broker,
+        sem: asyncio.Semaphore,
+        message: Message,
+    ) -> None:
+        try:
+            self.logger.debug(f"Processing message {message.id}")
+            try:
+                if message.task_name != task.name:
+                    # This should never happen.
+                    raise RejectError(
+                        f"Message {message.id} for unknown task '{message.task_name}' (expected '{task.name}')"
+                    )
 
-            done, _ = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
-            await task.queue.ack([pending.pop(t) for t in done])
+                async with AsyncExitStack() as stack:
+                    for middleware in self.middlewares:
+                        await stack.enter_async_context(middleware(message, task))
+                    await task.fn(*message.args, **message.kwargs)
+
+                await broker.ack(message)
+            except RejectError as e:
+                self.logger.warning(f"Message {message.id} rejected: {e}")
+                await broker.nack(message, drop=True)
+            except RetryError as e:
+                self.logger.debug(f"Retrying in {e.delay}s: {e}")
+                await broker.nack(message, delay=e.delay or 0)
+            except Exception as e:
+                self.logger.error(
+                    f"Rejected after error in {task.name}: {e}",
+                    exc_info=True,
+                )
+                await broker.nack(message, drop=True)
+        finally:
+            sem.release()
 
     async def run(self) -> None:
         try:
-            processors = (self.processor(task) for task in self.tasks)
-            await asyncio.gather(*processors)
+            await asyncio.gather(
+                *[self._poll(task, self.broker) for task in self.tasks]
+            )
         except asyncio.CancelledError:
-            logger.info("Worker shutting down...")
+            self.logger.info("Worker shutting down...")
         finally:
-            await self.client.aclose()
+            await self.broker.close()