Debounce and Cooldown admission control dependencies

loq.toml

@@ -10,7 +10,7 @@ max_lines = 750
 # Source files that still need exceptions above 750
 [[rules]]
 path = "src/docket/worker.py"
-max_lines = 1141
+max_lines = 1150
 
 [[rules]]
 path = "src/docket/cli/__init__.py"

src/docket/__init__.py

@@ -12,7 +12,9 @@
 from .annotations import Logged
 from .dependencies import (
     ConcurrencyLimit,
+    Cooldown,
     Cron,
+    Debounce,
     CurrentDocket,
     CurrentExecution,
     CurrentWorker,
@@ -37,7 +39,9 @@
     "__version__",
     "Agenda",
     "ConcurrencyLimit",
+    "Cooldown",
     "Cron",
+    "Debounce",
     "CurrentDocket",
     "CurrentExecution",
     "CurrentWorker",

src/docket/dependencies/__init__.py

@@ -19,6 +19,8 @@
     format_duration,
 )
 from ._concurrency import ConcurrencyBlocked, ConcurrencyLimit
+from ._cooldown import Cooldown, CooldownBlocked
+from ._debounce import Debounce, DebounceBlocked
 from ._cron import Cron
 from ._contextual import (
     CurrentDocket,
@@ -83,6 +85,10 @@
     "AdmissionBlocked",
     "ConcurrencyBlocked",
     "ConcurrencyLimit",
+    "Cooldown",
+    "CooldownBlocked",
+    "Debounce",
+    "DebounceBlocked",
     "Cron",
     "Perpetual",
     "Progress",

src/docket/dependencies/_base.py

@@ -55,8 +55,15 @@ class AdmissionBlocked(Exception):
 
     This is the base exception for admission control mechanisms like
     concurrency limits, rate limits, or health gates.
+
+    When ``reschedule`` is True (default), the worker re-queues the task
+    with a short delay.  When False, the task is silently acknowledged
+    and dropped (appropriate for debounce/cooldown where re-trying would
+    just hit the same window).
     """
 
+    reschedule: bool = True
+
     def __init__(self, execution: Execution, reason: str = "admission control"):
         self.execution = execution
         self.reason = reason

src/docket/dependencies/_cooldown.py

@@ -0,0 +1,94 @@
+"""Cooldown (trailing-edge) admission control dependency."""
+
+from __future__ import annotations
+
+from datetime import timedelta
+from types import TracebackType
+from typing import TYPE_CHECKING, Any
+
+from ._base import AdmissionBlocked, Dependency, current_docket, current_execution
+
+if TYPE_CHECKING:  # pragma: no cover
+    from ..execution import Execution
+
+
+class CooldownBlocked(AdmissionBlocked):
+    """Raised when a task is blocked by cooldown."""
+
+    reschedule = False
+
+    def __init__(self, execution: Execution, cooldown_key: str, window: timedelta):
+        self.cooldown_key = cooldown_key
+        self.window = window
+        reason = f"cooldown ({window}) on {cooldown_key}"
+        super().__init__(execution, reason=reason)
+
+
+class Cooldown(Dependency["Cooldown"]):
+    """Trailing-edge cooldown: blocks execution if one recently succeeded.
+
+    Checks for a Redis key on entry.  If present, the task is blocked.
+    The key is only set on *successful* exit, so failed tasks don't
+    trigger the cooldown — they can be retried immediately.
+
+    Works both as a default parameter and as ``Annotated`` metadata::
+
+        # Per-task: don't start if one succeeded in the last 60s
+        async def send_digest(
+            cooldown: Cooldown = Cooldown(timedelta(seconds=60)),
+        ) -> None: ...
+
+        # Per-parameter: don't start for this customer if one succeeded in the last 60s
+        async def send_notification(
+            customer_id: Annotated[int, Cooldown(timedelta(seconds=60))],
+        ) -> None: ...
+    """
+
+    single: bool = True
+
+    def __init__(self, window: timedelta, *, scope: str | None = None) -> None:
+        self.window = window
+        self.scope = scope
+        self._argument_name: str | None = None
+        self._argument_value: Any = None
+
+    def bind_to_parameter(self, name: str, value: Any) -> Cooldown:
+        bound = Cooldown(self.window, scope=self.scope)
+        bound._argument_name = name
+        bound._argument_value = value
+        return bound
+
+    def _cooldown_key(self, function_name: str) -> str:
+        scope = self.scope or current_docket.get().name
+        if self._argument_name is not None:
+            return f"{scope}:cooldown:{self._argument_name}:{self._argument_value}"
+        return f"{scope}:cooldown:{function_name}"
+
+    async def __aenter__(self) -> Cooldown:
+        execution = current_execution.get()
+        docket = current_docket.get()
+
+        self._key = self._cooldown_key(execution.function_name)
+
+        async with docket.redis() as redis:
+            exists = await redis.exists(self._key)
+
+        if exists:
+            raise CooldownBlocked(execution, self._key, self.window)
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        if exc_type is not None:
+            return
+
+        docket = current_docket.get()
+        window_ms = int(self.window.total_seconds() * 1000)
+
+        async with docket.redis() as redis:
+            await redis.set(self._key, 1, px=window_ms)

src/docket/dependencies/_debounce.py

@@ -0,0 +1,88 @@
+"""Debounce (leading-edge) admission control dependency."""
+
+from __future__ import annotations
+
+from datetime import timedelta
+from types import TracebackType
+from typing import TYPE_CHECKING, Any
+
+from ._base import AdmissionBlocked, Dependency, current_docket, current_execution
+
+if TYPE_CHECKING:  # pragma: no cover
+    from ..execution import Execution
+
+
+class DebounceBlocked(AdmissionBlocked):
+    """Raised when a task is blocked by debounce."""
+
+    reschedule = False
+
+    def __init__(self, execution: Execution, debounce_key: str, window: timedelta):
+        self.debounce_key = debounce_key
+        self.window = window
+        reason = f"debounce ({window}) on {debounce_key}"
+        super().__init__(execution, reason=reason)
+
+
+class Debounce(Dependency["Debounce"]):
+    """Leading-edge debounce: blocks execution if one was recently started.
+
+    Sets a Redis key on entry with a TTL equal to the window. If the key
+    already exists, the task is blocked via ``AdmissionBlocked``.
+
+    Works both as a default parameter and as ``Annotated`` metadata::
+
+        # Per-task: don't start if one started in the last 30s
+        async def process_webhooks(
+            debounce: Debounce = Debounce(timedelta(seconds=30)),
+        ) -> None: ...
+
+        # Per-parameter: don't start for this customer if one started in the last 30s
+        async def process_customer(
+            customer_id: Annotated[int, Debounce(timedelta(seconds=30))],
+        ) -> None: ...
+    """
+
+    single: bool = True
+
+    def __init__(self, window: timedelta, *, scope: str | None = None) -> None:
+        self.window = window
+        self.scope = scope
+        self._argument_name: str | None = None
+        self._argument_value: Any = None
+
+    def bind_to_parameter(self, name: str, value: Any) -> Debounce:
+        bound = Debounce(self.window, scope=self.scope)
+        bound._argument_name = name
+        bound._argument_value = value
+        return bound
+
+    async def __aenter__(self) -> Debounce:
+        execution = current_execution.get()
+        docket = current_docket.get()
+
+        scope = self.scope or docket.name
+        if self._argument_name is not None:
+            debounce_key = (
+                f"{scope}:debounce:{self._argument_name}:{self._argument_value}"
+            )
+        else:
+            debounce_key = f"{scope}:debounce:{execution.function_name}"
+
+        window_ms = int(self.window.total_seconds() * 1000)
+
+        async with docket.redis() as redis:
+            acquired = await redis.set(debounce_key, 1, nx=True, px=window_ms)
+
+        if not acquired:
+            raise DebounceBlocked(execution, debounce_key, self.window)
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        pass

src/docket/dependencies/_resolution.py

@@ -32,7 +32,7 @@ def get_single_dependency_parameter_of_type(
     for _, dependencies in get_annotation_dependencies(function).items():
         for dependency in dependencies:
             if isinstance(dependency, dependency_type):
-                return dependency  # type: ignore[return-value]
+                return dependency
     return None
 
 

src/docket/worker.py

@@ -528,15 +528,24 @@ async def process_completed_tasks() -> None:
                     await task
                     await ack_message(redis, message_id)
                 except AdmissionBlocked as e:
-                    logger.debug(
-                        "🔒 Task %s blocked by admission control, rescheduling",
-                        e.execution.key,
-                        extra=log_context,
-                    )
-                    e.execution.when = (
-                        datetime.now(timezone.utc) + ADMISSION_BLOCKED_RETRY_DELAY
-                    )
-                    await e.execution.schedule(reschedule_message=message_id)
+                    if e.reschedule:
+                        logger.debug(
+                            "⏳ Task %s blocked by admission control, rescheduling",
+                            e.execution.key,
+                            extra=log_context,
+                        )
+                        e.execution.when = (
+                            datetime.now(timezone.utc) + ADMISSION_BLOCKED_RETRY_DELAY
+                        )
+                        await e.execution.schedule(reschedule_message=message_id)
+                    else:
+                        logger.debug(
+                            "⏭ Task %s blocked by admission control, dropping",
+                            e.execution.key,
+                            extra=log_context,
+                        )
+                        await e.execution.mark_as_cancelled()
+                        await ack_message(redis, message_id)
 
         async def ack_message(redis: Redis, message_id: RedisMessageID) -> None:
             logger.debug("Acknowledging message", extra=log_context)