pydocket 0.7.1__py3-none-any.whl → 0.9.0__py3-none-any.whl

docket/__init__.py

@@ -10,6 +10,7 @@ __version__ = version("pydocket")
 
 from .annotations import Logged
 from .dependencies import (
+    ConcurrencyLimit,
     CurrentDocket,
     CurrentExecution,
     CurrentWorker,
@@ -28,6 +29,7 @@ from .worker import Worker
 
 __all__ = [
     "__version__",
+    "ConcurrencyLimit",
     "CurrentDocket",
     "CurrentExecution",
     "CurrentWorker",

docket/cli.py

@@ -594,6 +594,73 @@ def relative_time(now: datetime, when: datetime) -> str:
         return f"at {local_time(when)}"
 
 
+def get_task_stats(
+    snapshot: DocketSnapshot,
+) -> dict[str, dict[str, int | datetime | None]]:
+    """Get task count statistics by function name with timestamp data."""
+    stats: dict[str, dict[str, int | datetime | None]] = {}
+
+    # Count running tasks by function
+    for execution in snapshot.running:
+        func_name = execution.function.__name__
+        if func_name not in stats:
+            stats[func_name] = {
+                "running": 0,
+                "queued": 0,
+                "total": 0,
+                "oldest_queued": None,
+                "latest_queued": None,
+                "oldest_started": None,
+                "latest_started": None,
+            }
+        stats[func_name]["running"] += 1
+        stats[func_name]["total"] += 1
+
+        # Track oldest/latest started times for running tasks
+        started = execution.started
+        if (
+            stats[func_name]["oldest_started"] is None
+            or started < stats[func_name]["oldest_started"]
+        ):
+            stats[func_name]["oldest_started"] = started
+        if (
+            stats[func_name]["latest_started"] is None
+            or started > stats[func_name]["latest_started"]
+        ):
+            stats[func_name]["latest_started"] = started
+
+    # Count future tasks by function
+    for execution in snapshot.future:
+        func_name = execution.function.__name__
+        if func_name not in stats:
+            stats[func_name] = {
+                "running": 0,
+                "queued": 0,
+                "total": 0,
+                "oldest_queued": None,
+                "latest_queued": None,
+                "oldest_started": None,
+                "latest_started": None,
+            }
+        stats[func_name]["queued"] += 1
+        stats[func_name]["total"] += 1
+
+        # Track oldest/latest queued times for future tasks
+        when = execution.when
+        if (
+            stats[func_name]["oldest_queued"] is None
+            or when < stats[func_name]["oldest_queued"]
+        ):
+            stats[func_name]["oldest_queued"] = when
+        if (
+            stats[func_name]["latest_queued"] is None
+            or when > stats[func_name]["latest_queued"]
+        ):
+            stats[func_name]["latest_queued"] = when
+
+    return stats
+
+
 @app.command(help="Shows a snapshot of what's on the docket right now")
 def snapshot(
     tasks: Annotated[
@@ -623,6 +690,13 @@ def snapshot(
             envvar="DOCKET_URL",
         ),
     ] = "redis://localhost:6379/0",
+    stats: Annotated[
+        bool,
+        typer.Option(
+            "--stats",
+            help="Show task count statistics by function name",
+        ),
+    ] = False,
 ) -> None:
     async def run() -> DocketSnapshot:
         async with Docket(name=docket_, url=url) as docket:
@@ -672,6 +746,44 @@ def snapshot(
 
     console.print(table)
 
+    # Display task statistics if requested
+    if stats:
+        task_stats = get_task_stats(snapshot)
+        if task_stats:
+            console.print()  # Add spacing between tables
+            stats_table = Table(title="Task Count Statistics by Function")
+            stats_table.add_column("Function", style="cyan")
+            stats_table.add_column("Total", style="bold magenta", justify="right")
+            stats_table.add_column("Running", style="green", justify="right")
+            stats_table.add_column("Queued", style="yellow", justify="right")
+            stats_table.add_column("Oldest Queued", style="dim yellow", justify="right")
+            stats_table.add_column("Latest Queued", style="dim yellow", justify="right")
+
+            # Sort by total count descending to highlight potential runaway tasks
+            for func_name in sorted(
+                task_stats.keys(), key=lambda x: task_stats[x]["total"], reverse=True
+            ):
+                counts = task_stats[func_name]
+
+                # Format timestamp columns
+                oldest_queued = ""
+                latest_queued = ""
+                if counts["oldest_queued"] is not None:
+                    oldest_queued = relative(counts["oldest_queued"])
+                if counts["latest_queued"] is not None:
+                    latest_queued = relative(counts["latest_queued"])
+
+                stats_table.add_row(
+                    func_name,
+                    str(counts["total"]),
+                    str(counts["running"]),
+                    str(counts["queued"]),
+                    oldest_queued,
+                    latest_queued,
+                )
+
+            console.print(stats_table)
+
 
 workers_app: typer.Typer = typer.Typer(
     help="Look at the workers on a docket", no_args_is_help=True

docket/dependencies.py

@@ -39,9 +39,9 @@ class Dependency(abc.ABC):
 
     async def __aexit__(
         self,
-        exc_type: type[BaseException] | None,
-        exc_value: BaseException | None,
-        traceback: TracebackType | None,
+        _exc_type: type[BaseException] | None,
+        _exc_value: BaseException | None,
+        _traceback: TracebackType | None,
     ) -> bool: ...  # pragma: no cover
 
 
@@ -505,6 +505,92 @@ def Depends(dependency: DependencyFunction[R]) -> R:
     return cast(R, _Depends(dependency))
 
 
+class ConcurrencyLimit(Dependency):
+    """Configures concurrency limits for a task based on specific argument values.
+
+    This allows fine-grained control over task execution by limiting concurrent
+    tasks based on the value of specific arguments.
+
+    Example:
+
+    ```python
+    async def process_customer(
+        customer_id: int,
+        concurrency: ConcurrencyLimit = ConcurrencyLimit("customer_id", max_concurrent=1)
+    ) -> None:
+        # Only one task per customer_id will run at a time
+        ...
+
+    async def backup_db(
+        db_name: str,
+        concurrency: ConcurrencyLimit = ConcurrencyLimit("db_name", max_concurrent=3)
+    ) -> None:
+        # Only 3 backup tasks per database name will run at a time
+        ...
+    ```
+    """
+
+    single: bool = True
+
+    def __init__(
+        self, argument_name: str, max_concurrent: int = 1, scope: str | None = None
+    ) -> None:
+        """
+        Args:
+            argument_name: The name of the task argument to use for concurrency grouping
+            max_concurrent: Maximum number of concurrent tasks per unique argument value
+            scope: Optional scope prefix for Redis keys (defaults to docket name)
+        """
+        self.argument_name = argument_name
+        self.max_concurrent = max_concurrent
+        self.scope = scope
+        self._concurrency_key: str | None = None
+        self._initialized: bool = False
+
+    async def __aenter__(self) -> "ConcurrencyLimit":
+        execution = self.execution.get()
+        docket = self.docket.get()
+
+        # Get the argument value to group by
+        try:
+            argument_value = execution.get_argument(self.argument_name)
+        except KeyError:
+            # If argument not found, create a bypass limit that doesn't apply concurrency control
+            limit = ConcurrencyLimit(
+                self.argument_name, self.max_concurrent, self.scope
+            )
+            limit._concurrency_key = None  # Special marker for bypassed concurrency
+            limit._initialized = True  # Mark as initialized but bypassed
+            return limit
+
+        # Create a concurrency key for this specific argument value
+        scope = self.scope or docket.name
+        self._concurrency_key = (
+            f"{scope}:concurrency:{self.argument_name}:{argument_value}"
+        )
+
+        limit = ConcurrencyLimit(self.argument_name, self.max_concurrent, self.scope)
+        limit._concurrency_key = self._concurrency_key
+        limit._initialized = True  # Mark as initialized
+        return limit
+
+    @property
+    def concurrency_key(self) -> str | None:
+        """Redis key used for tracking concurrency for this specific argument value.
+        Returns None when concurrency control is bypassed due to missing arguments.
+        Raises RuntimeError if accessed before initialization."""
+        if not self._initialized:
+            raise RuntimeError(
+                "ConcurrencyLimit not initialized - use within task context"
+            )
+        return self._concurrency_key
+
+    @property
+    def is_bypassed(self) -> bool:
+        """Returns True if concurrency control is bypassed due to missing arguments."""
+        return self._initialized and self._concurrency_key is None
+
+
 D = TypeVar("D", bound=Dependency)
 
 

docket/docket.py

@@ -16,6 +16,7 @@ from typing import (
     Mapping,
     NoReturn,
     ParamSpec,
+    Protocol,
     Self,
     Sequence,
     TypedDict,
@@ -27,7 +28,6 @@ from typing import (
 import redis.exceptions
 from opentelemetry import propagate, trace
 from redis.asyncio import ConnectionPool, Redis
-from redis.asyncio.client import Pipeline
 from uuid_extensions import uuid7
 
 from .execution import (
@@ -55,6 +55,18 @@ logger: logging.Logger = logging.getLogger(__name__)
 tracer: trace.Tracer = trace.get_tracer(__name__)
 
 
+class _schedule_task(Protocol):
+    async def __call__(
+        self, keys: list[str], args: list[str | float | bytes]
+    ) -> str: ...  # pragma: no cover
+
+
+class _cancel_task(Protocol):
+    async def __call__(
+        self, keys: list[str], args: list[str]
+    ) -> str: ...  # pragma: no cover
+
+
 P = ParamSpec("P")
 R = TypeVar("R")
 
@@ -131,6 +143,8 @@ class Docket:
 
     _monitor_strikes_task: asyncio.Task[None]
     _connection_pool: ConnectionPool
+    _schedule_task_script: _schedule_task | None
+    _cancel_task_script: _cancel_task | None
 
     def __init__(
         self,
@@ -156,6 +170,8 @@ class Docket:
         self.url = url
         self.heartbeat_interval = heartbeat_interval
         self.missed_heartbeats = missed_heartbeats
+        self._schedule_task_script = None
+        self._cancel_task_script = None
 
     @property
     def worker_group_name(self) -> str:
@@ -300,9 +316,7 @@ class Docket:
             execution = Execution(function, args, kwargs, when, key, attempt=1)
 
             async with self.redis() as redis:
-                async with redis.pipeline() as pipeline:
-                    await self._schedule(redis, pipeline, execution, replace=False)
-                    await pipeline.execute()
+                await self._schedule(redis, execution, replace=False)
 
             TASKS_ADDED.add(1, {**self.labels(), **execution.general_labels()})
             TASKS_SCHEDULED.add(1, {**self.labels(), **execution.general_labels()})
@@ -361,9 +375,7 @@ class Docket:
             execution = Execution(function, args, kwargs, when, key, attempt=1)
 
             async with self.redis() as redis:
-                async with redis.pipeline() as pipeline:
-                    await self._schedule(redis, pipeline, execution, replace=True)
-                    await pipeline.execute()
+                await self._schedule(redis, execution, replace=True)
 
             TASKS_REPLACED.add(1, {**self.labels(), **execution.general_labels()})
             TASKS_CANCELLED.add(1, {**self.labels(), **execution.general_labels()})
@@ -383,9 +395,7 @@ class Docket:
             },
         ):
             async with self.redis() as redis:
-                async with redis.pipeline() as pipeline:
-                    await self._schedule(redis, pipeline, execution, replace=False)
-                    await pipeline.execute()
+                await self._schedule(redis, execution, replace=False)
 
         TASKS_SCHEDULED.add(1, {**self.labels(), **execution.general_labels()})
 
@@ -400,9 +410,7 @@ class Docket:
             attributes={**self.labels(), "docket.key": key},
         ):
             async with self.redis() as redis:
-                async with redis.pipeline() as pipeline:
-                    await self._cancel(pipeline, key)
-                    await pipeline.execute()
+                await self._cancel(redis, key)
 
         TASKS_CANCELLED.add(1, self.labels())
 
@@ -423,10 +431,17 @@ class Docket:
     async def _schedule(
         self,
         redis: Redis,
-        pipeline: Pipeline,
         execution: Execution,
         replace: bool = False,
     ) -> None:
+        """Schedule a task atomically.
+
+        Handles:
+        - Checking for task existence
+        - Cancelling existing tasks when replacing
+        - Adding tasks to stream (immediate) or queue (future)
+        - Tracking stream message IDs for later cancellation
+        """
         if self.strike_list.is_stricken(execution):
             logger.warning(
                 "%r is stricken, skipping schedule of %r",
@@ -449,32 +464,133 @@ class Docket:
         key = execution.key
         when = execution.when
         known_task_key = self.known_task_key(key)
+        is_immediate = when <= datetime.now(timezone.utc)
 
+        # Lock per task key to prevent race conditions between concurrent operations
         async with redis.lock(f"{known_task_key}:lock", timeout=10):
-            if replace:
-                await self._cancel(pipeline, key)
-            else:
-                # if the task is already in the queue or stream, retain it
-                if await redis.exists(known_task_key):
-                    logger.debug(
-                        "Task %r is already in the queue or stream, not scheduling",
-                        key,
-                        extra=self.labels(),
-                    )
-                    return
+            if self._schedule_task_script is None:
+                self._schedule_task_script = cast(
+                    _schedule_task,
+                    redis.register_script(
+                        # KEYS: stream_key, known_key, parked_key, queue_key
+                        # ARGV: task_key, when_timestamp, is_immediate, replace, ...message_fields
+                        """
+                        local stream_key = KEYS[1]
+                        local known_key = KEYS[2]
+                        local parked_key = KEYS[3]
+                        local queue_key = KEYS[4]
+
+                        local task_key = ARGV[1]
+                        local when_timestamp = ARGV[2]
+                        local is_immediate = ARGV[3] == '1'
+                        local replace = ARGV[4] == '1'
+
+                        -- Extract message fields from ARGV[5] onwards
+                        local message = {}
+                        for i = 5, #ARGV, 2 do
+                            message[#message + 1] = ARGV[i]     -- field name
+                            message[#message + 1] = ARGV[i + 1] -- field value
+                        end
+
+                        -- Handle replacement: cancel existing task if needed
+                        if replace then
+                            local existing_message_id = redis.call('HGET', known_key, 'stream_message_id')
+                            if existing_message_id then
+                                redis.call('XDEL', stream_key, existing_message_id)
+                            end
+                            redis.call('DEL', known_key, parked_key)
+                            redis.call('ZREM', queue_key, task_key)
+                        else
+                            -- Check if task already exists
+                            if redis.call('EXISTS', known_key) == 1 then
+                                return 'EXISTS'
+                            end
+                        end
+
+                        if is_immediate then
+                            -- Add to stream and store message ID for later cancellation
+                            local message_id = redis.call('XADD', stream_key, '*', unpack(message))
+                            redis.call('HSET', known_key, 'when', when_timestamp, 'stream_message_id', message_id)
+                            return message_id
+                        else
+                            -- Add to queue with task data in parked hash
+                            redis.call('HSET', known_key, 'when', when_timestamp)
+                            redis.call('HSET', parked_key, unpack(message))
+                            redis.call('ZADD', queue_key, when_timestamp, task_key)
+                            return 'QUEUED'
+                        end
+                        """
+                    ),
+                )
+            schedule_task = self._schedule_task_script
 
-            pipeline.set(known_task_key, when.timestamp())
+            await schedule_task(
+                keys=[
+                    self.stream_key,
+                    known_task_key,
+                    self.parked_task_key(key),
+                    self.queue_key,
+                ],
+                args=[
+                    key,
+                    str(when.timestamp()),
+                    "1" if is_immediate else "0",
+                    "1" if replace else "0",
+                    *[
+                        item
+                        for field, value in message.items()
+                        for item in (field, value)
+                    ],
+                ],
+            )
 
-            if when <= datetime.now(timezone.utc):
-                pipeline.xadd(self.stream_key, message)  # type: ignore[arg-type]
-            else:
-                pipeline.hset(self.parked_task_key(key), mapping=message)  # type: ignore[arg-type]
-                pipeline.zadd(self.queue_key, {key: when.timestamp()})
+    async def _cancel(self, redis: Redis, key: str) -> None:
+        """Cancel a task atomically.
 
-    async def _cancel(self, pipeline: Pipeline, key: str) -> None:
-        pipeline.delete(self.known_task_key(key))
-        pipeline.delete(self.parked_task_key(key))
-        pipeline.zrem(self.queue_key, key)
+        Handles cancellation regardless of task location:
+        - From the stream (using stored message ID)
+        - From the queue (scheduled tasks)
+        - Cleans up all associated metadata keys
+        """
+        if self._cancel_task_script is None:
+            self._cancel_task_script = cast(
+                _cancel_task,
+                redis.register_script(
+                    # KEYS: stream_key, known_key, parked_key, queue_key
+                    # ARGV: task_key
+                    """
+                    local stream_key = KEYS[1]
+                    local known_key = KEYS[2]
+                    local parked_key = KEYS[3]
+                    local queue_key = KEYS[4]
+                    local task_key = ARGV[1]
+
+                    -- Delete from stream if message ID exists
+                    local message_id = redis.call('HGET', known_key, 'stream_message_id')
+                    if message_id then
+                        redis.call('XDEL', stream_key, message_id)
+                    end
+
+                    -- Clean up all task-related keys
+                    redis.call('DEL', known_key, parked_key)
+                    redis.call('ZREM', queue_key, task_key)
+
+                    return 'OK'
+                    """
+                ),
+            )
+        cancel_task = self._cancel_task_script
+
+        # Execute the cancellation script
+        await cancel_task(
+            keys=[
+                self.stream_key,
+                self.known_task_key(key),
+                self.parked_task_key(key),
+                self.queue_key,
+            ],
+            args=[key],
+        )
 
     @property
     def strike_key(self) -> str:

docket/worker.py

@@ -20,6 +20,7 @@ from redis.asyncio import Redis
 from redis.exceptions import ConnectionError, LockError
 
 from .dependencies import (
+    ConcurrencyLimit,
     Dependency,
     FailedDependency,
     Perpetual,
@@ -248,6 +249,7 @@ class Worker:
         )
 
         active_tasks: dict[asyncio.Task[None], RedisMessageID] = {}
+        task_executions: dict[asyncio.Task[None], Execution] = {}
         available_slots = self.concurrency
 
         log_context = self._log_context()
@@ -298,6 +300,7 @@ class Worker:
 
             task = asyncio.create_task(self._execute(execution), name=execution.key)
             active_tasks[task] = message_id
+            task_executions[task] = execution
 
             nonlocal available_slots
             available_slots -= 1
@@ -308,6 +311,7 @@ class Worker:
             completed_tasks = {task for task in active_tasks if task.done()}
             for task in completed_tasks:
                 message_id = active_tasks.pop(task)
+                task_executions.pop(task)
                 await task
                 await ack_message(redis, message_id)
 
@@ -343,7 +347,9 @@ class Worker:
                             if not message:  # pragma: no cover
                                 continue
 
-                            if not start_task(message_id, message):
+                            task_started = start_task(message_id, message)
+                            if not task_started:
+                                # Other errors - delete and ack
                                 await self._delete_known_task(redis, message)
                                 await ack_message(redis, message_id)
 
@@ -400,7 +406,7 @@ class Worker:
                             task[task_data[j]] = task_data[j+1]
                         end
 
-                        redis.call('XADD', KEYS[2], '*',
+                        local message_id = redis.call('XADD', KEYS[2], '*',
                             'key', task['key'],
                             'when', task['when'],
                             'function', task['function'],
@@ -408,6 +414,9 @@ class Worker:
                             'kwargs', task['kwargs'],
                             'attempt', task['attempt']
                         )
+                        -- Store the message ID in the known task key
+                        local known_key = ARGV[2] .. ":known:" .. key
+                        redis.call('HSET', known_key, 'stream_message_id', message_id)
                         redis.call('DEL', hash_key)
                         due_work = due_work + 1
                     end
@@ -534,6 +543,32 @@ class Worker:
         ) as span:
             try:
                 async with resolved_dependencies(self, execution) as dependencies:
+                    # Check concurrency limits after dependency resolution
+                    concurrency_limit = get_single_dependency_of_type(
+                        dependencies, ConcurrencyLimit
+                    )
+                    if concurrency_limit and not concurrency_limit.is_bypassed:
+                        async with self.docket.redis() as redis:
+                            # Check if we can acquire a concurrency slot
+                            if not await self._can_start_task(redis, execution):
+                                # Task cannot start due to concurrency limits - reschedule
+                                logger.debug(
+                                    "🔒 Task %s blocked by concurrency limit, rescheduling",
+                                    execution.key,
+                                    extra=log_context,
+                                )
+                                # Reschedule for a few milliseconds in the future
+                                when = datetime.now(timezone.utc) + timedelta(
+                                    milliseconds=50
+                                )
+                                await self.docket.add(execution.function, when=when)(
+                                    *execution.args, **execution.kwargs
+                                )
+                                return
+                            else:
+                                # Successfully acquired slot
+                                pass
+
                     # Preemptively reschedule the perpetual task for the future, or clear
                     # the known task key for this task
                     rescheduled = await self._perpetuate_if_requested(
@@ -560,17 +595,34 @@ class Worker:
                             ],
                         )
 
-                    if timeout := get_single_dependency_of_type(dependencies, Timeout):
-                        await self._run_function_with_timeout(
-                            execution, dependencies, timeout
-                        )
+                    # Apply timeout logic - either user's timeout or redelivery timeout
+                    user_timeout = get_single_dependency_of_type(dependencies, Timeout)
+                    if user_timeout:
+                        # If user timeout is longer than redelivery timeout, limit it
+                        if user_timeout.base > self.redelivery_timeout:
+                            # Create a new timeout limited by redelivery timeout
+                            # Remove the user timeout from dependencies to avoid conflicts
+                            limited_dependencies = {
+                                k: v
+                                for k, v in dependencies.items()
+                                if not isinstance(v, Timeout)
+                            }
+                            limited_timeout = Timeout(self.redelivery_timeout)
+                            limited_timeout.start()
+                            await self._run_function_with_timeout(
+                                execution, limited_dependencies, limited_timeout
+                            )
+                        else:
+                            # User timeout is within redelivery timeout, use as-is
+                            await self._run_function_with_timeout(
+                                execution, dependencies, user_timeout
+                            )
                     else:
-                        await execution.function(
-                            *execution.args,
-                            **{
-                                **execution.kwargs,
-                                **dependencies,
-                            },
+                        # No user timeout - apply redelivery timeout as hard limit
+                        redelivery_timeout = Timeout(self.redelivery_timeout)
+                        redelivery_timeout.start()
+                        await self._run_function_with_timeout(
+                            execution, dependencies, redelivery_timeout
                         )
 
                     duration = log_context["duration"] = time.time() - start
@@ -604,6 +656,15 @@ class Worker:
                     "%s [%s] %s", arrow, ms(duration), call, extra=log_context
                 )
             finally:
+                # Release concurrency slot if we acquired one
+                if dependencies:
+                    concurrency_limit = get_single_dependency_of_type(
+                        dependencies, ConcurrencyLimit
+                    )
+                    if concurrency_limit and not concurrency_limit.is_bypassed:
+                        async with self.docket.redis() as redis:
+                            await self._release_concurrency_slot(redis, execution)
+
                 TASKS_RUNNING.add(-1, counter_labels)
                 TASKS_COMPLETED.add(1, counter_labels)
                 TASK_DURATION.record(duration, counter_labels)
@@ -616,7 +677,13 @@ class Worker:
     ) -> None:
         task_coro = cast(
             Coroutine[None, None, None],
-            execution.function(*execution.args, **execution.kwargs, **dependencies),
+            execution.function(
+                *execution.args,
+                **{
+                    **execution.kwargs,
+                    **dependencies,
+                },
+            ),
         )
         task = asyncio.create_task(task_coro)
         try:
@@ -762,6 +829,95 @@ class Worker:
                     extra=self._log_context(),
                 )
 
+    async def _can_start_task(self, redis: Redis, execution: Execution) -> bool:
+        """Check if a task can start based on concurrency limits."""
+        # Check if task has a concurrency limit dependency
+        concurrency_limit = get_single_dependency_parameter_of_type(
+            execution.function, ConcurrencyLimit
+        )
+
+        if not concurrency_limit:
+            return True  # No concurrency limit, can always start
+
+        # Get the concurrency key for this task
+        try:
+            argument_value = execution.get_argument(concurrency_limit.argument_name)
+        except KeyError:
+            # If argument not found, let the task fail naturally in execution
+            return True
+
+        scope = concurrency_limit.scope or self.docket.name
+        concurrency_key = (
+            f"{scope}:concurrency:{concurrency_limit.argument_name}:{argument_value}"
+        )
+
+        # Use Redis sorted set with timestamps to track concurrency and handle expiration
+        lua_script = """
+        local key = KEYS[1]
+        local max_concurrent = tonumber(ARGV[1])
+        local worker_id = ARGV[2]
+        local task_key = ARGV[3]
+        local current_time = tonumber(ARGV[4])
+        local expiration_time = tonumber(ARGV[5])
+
+        -- Remove expired entries
+        local expired_cutoff = current_time - expiration_time
+        redis.call('ZREMRANGEBYSCORE', key, 0, expired_cutoff)
+
+        -- Get current count
+        local current = redis.call('ZCARD', key)
+
+        if current < max_concurrent then
+            -- Add this worker's task to the sorted set with current timestamp
+            redis.call('ZADD', key, current_time, worker_id .. ':' .. task_key)
+            return 1
+        else
+            return 0
+        end
+        """
+
+        current_time = datetime.now(timezone.utc).timestamp()
+        expiration_seconds = self.redelivery_timeout.total_seconds()
+
+        result = await redis.eval(  # type: ignore
+            lua_script,
+            1,
+            concurrency_key,
+            str(concurrency_limit.max_concurrent),
+            self.name,
+            execution.key,
+            current_time,
+            expiration_seconds,
+        )
+
+        return bool(result)
+
+    async def _release_concurrency_slot(
+        self, redis: Redis, execution: Execution
+    ) -> None:
+        """Release a concurrency slot when task completes."""
+        # Check if task has a concurrency limit dependency
+        concurrency_limit = get_single_dependency_parameter_of_type(
+            execution.function, ConcurrencyLimit
+        )
+
+        if not concurrency_limit:
+            return  # No concurrency limit to release
+
+        # Get the concurrency key for this task
+        try:
+            argument_value = execution.get_argument(concurrency_limit.argument_name)
+        except KeyError:
+            return  # If argument not found, nothing to release
+
+        scope = concurrency_limit.scope or self.docket.name
+        concurrency_key = (
+            f"{scope}:concurrency:{concurrency_limit.argument_name}:{argument_value}"
+        )
+
+        # Remove this worker's task from the sorted set
+        await redis.zrem(concurrency_key, f"{self.name}:{execution.key}")  # type: ignore
+
 
 def ms(seconds: float) -> str:
     if seconds < 100:

{pydocket-0.7.1.dist-info → pydocket-0.9.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.7.1
+Version: 0.9.0
 Summary: A distributed background task system for Python functions
 Project-URL: Homepage, https://github.com/chrisguidry/docket
 Project-URL: Bug Tracker, https://github.com/chrisguidry/docket/issues
@@ -93,6 +93,8 @@ reference](https://chrisguidry.github.io/docket/api-reference/).
 
 🧩 Fully type-complete and type-aware for your background task functions
 
+💉 Dependency injection like FastAPI, Typer, and FastMCP for reusable resources
+
 ## Installing `docket`
 
 Docket is [available on PyPI](https://pypi.org/project/pydocket/) under the package name

pydocket-0.9.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+docket/__init__.py,sha256=onwZzh73tESWoFBukbcW-7gjxoXb-yI7dutRD7tPN6g,915
+docket/__main__.py,sha256=wcCrL4PjG51r5wVKqJhcoJPTLfHW0wNbD31DrUN0MWI,28
+docket/annotations.py,sha256=wttix9UOeMFMAWXAIJUfUw5GjESJZsACb4YXJCozP7Q,2348
+docket/cli.py,sha256=rTfri2--u4Q5PlXyh7Ub_F5uh3-TtZOWLUp9WY_TvAE,25750
+docket/dependencies.py,sha256=BC0bnt10cr9_S1p5JAP_bnC9RwZkTr9ulPBrxC7eZnA,20247
+docket/docket.py,sha256=0nQCHDDHy7trv2a0eYygGgIKiA7fWq5GcOXye3_CPWM,30847
+docket/execution.py,sha256=r_2RGC1qhtAcBUg7E6wewLEgftrf3hIxNbH0HnYPbek,14961
+docket/instrumentation.py,sha256=ogvzrfKbWsdPGfdg4hByH3_r5d3b5AwwQkSrmXw0hRg,5492
+docket/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+docket/tasks.py,sha256=RIlSM2omh-YDwVnCz6M5MtmK8T_m_s1w2OlRRxDUs6A,1437
+docket/worker.py,sha256=jqVYqtQyxbk-BIy3shY8haX-amVT9Np97VhJuaQTfpM,35174
+pydocket-0.9.0.dist-info/METADATA,sha256=kymp9PKG7UwMj0i0qGSSCHKu-g-tS__qydr6mYuMLtg,5418
+pydocket-0.9.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydocket-0.9.0.dist-info/entry_points.txt,sha256=4WOk1nUlBsUT5O3RyMci2ImuC5XFswuopElYcLHtD5k,47
+pydocket-0.9.0.dist-info/licenses/LICENSE,sha256=YuVWU_ZXO0K_k2FG8xWKe5RGxV24AhJKTvQmKfqXuyk,1087
+pydocket-0.9.0.dist-info/RECORD,,

pydocket-0.7.1.dist-info/RECORD

@@ -1,16 +0,0 @@
-docket/__init__.py,sha256=sY1T_NVsXQNOmOhOnfYmZ95dcE_52Ov6DSIVIMZp-1w,869
-docket/__main__.py,sha256=wcCrL4PjG51r5wVKqJhcoJPTLfHW0wNbD31DrUN0MWI,28
-docket/annotations.py,sha256=wttix9UOeMFMAWXAIJUfUw5GjESJZsACb4YXJCozP7Q,2348
-docket/cli.py,sha256=XG_mbjcqNRO0F0hh6l3AwH9bIZv9xJofZaeaAj9nChc,21608
-docket/dependencies.py,sha256=GBwyEY198JFrfm7z5GkLbd84hv7sJktKBMJXv4veWig,17007
-docket/docket.py,sha256=Cw7QB1d0eDwSgwn0Rj26WjFsXSe7MJtfsUBBHGalL7A,26262
-docket/execution.py,sha256=r_2RGC1qhtAcBUg7E6wewLEgftrf3hIxNbH0HnYPbek,14961
-docket/instrumentation.py,sha256=ogvzrfKbWsdPGfdg4hByH3_r5d3b5AwwQkSrmXw0hRg,5492
-docket/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-docket/tasks.py,sha256=RIlSM2omh-YDwVnCz6M5MtmK8T_m_s1w2OlRRxDUs6A,1437
-docket/worker.py,sha256=CY5Z9p8FZw-6WUwp7Ws4A0V7IFTmonSnBmYP-Cp8Fdw,28079
-pydocket-0.7.1.dist-info/METADATA,sha256=00KHm5Er2R6dmjHLTYBUF13kKAeCRPHmDTdAcv5oRcQ,5335
-pydocket-0.7.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydocket-0.7.1.dist-info/entry_points.txt,sha256=4WOk1nUlBsUT5O3RyMci2ImuC5XFswuopElYcLHtD5k,47
-pydocket-0.7.1.dist-info/licenses/LICENSE,sha256=YuVWU_ZXO0K_k2FG8xWKe5RGxV24AhJKTvQmKfqXuyk,1087
-pydocket-0.7.1.dist-info/RECORD,,