pydocket 0.7.1__py3-none-any.whl → 0.8.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/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()
@@ -284,7 +286,7 @@ class Worker:
                 count=available_slots,
             )
 
-        def start_task(message_id: RedisMessageID, message: RedisMessage) -> bool:
+        async def start_task(message_id: RedisMessageID, message: RedisMessage) -> bool:
             function_name = message[b"function"].decode()
             if not (function := self.docket.tasks.get(function_name)):
                 logger.warning(
@@ -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 = await 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)
 
@@ -534,6 +540,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 +592,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 +653,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 +674,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 +826,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.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.7.1
+Version: 0.8.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.8.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=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=RNfwvNC1zpGfTYq8-HOjicociuVvOJWMgj8w8DqmN3Y,34940
+pydocket-0.8.0.dist-info/METADATA,sha256=Y2tUVZmlDqPWmBCmmjRUTs56Rd6thT3_wSd9JMKb79s,5418
+pydocket-0.8.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydocket-0.8.0.dist-info/entry_points.txt,sha256=4WOk1nUlBsUT5O3RyMci2ImuC5XFswuopElYcLHtD5k,47
+pydocket-0.8.0.dist-info/licenses/LICENSE,sha256=YuVWU_ZXO0K_k2FG8xWKe5RGxV24AhJKTvQmKfqXuyk,1087
+pydocket-0.8.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,,