pydocket 0.15.3__py3-none-any.whl

docket/instrumentation.py

@@ -0,0 +1,225 @@
+from contextlib import contextmanager
+from threading import Thread
+from typing import Generator, cast
+
+from opentelemetry import metrics
+from opentelemetry.exporter.prometheus import PrometheusMetricReader
+from opentelemetry.metrics import set_meter_provider
+from opentelemetry.propagators.textmap import Getter, Setter
+from opentelemetry.sdk.metrics import MeterProvider
+
+meter: metrics.Meter = metrics.get_meter("docket")
+
+TASKS_ADDED = meter.create_counter(
+    "docket_tasks_added",
+    description="How many tasks added to the docket",
+    unit="1",
+)
+
+TASKS_REPLACED = meter.create_counter(
+    "docket_tasks_replaced",
+    description="How many tasks replaced on the docket",
+    unit="1",
+)
+
+TASKS_SCHEDULED = meter.create_counter(
+    "docket_tasks_scheduled",
+    description="How many tasks added or replaced on the docket",
+    unit="1",
+)
+
+TASKS_CANCELLED = meter.create_counter(
+    "docket_tasks_cancelled",
+    description="How many tasks cancelled from the docket",
+    unit="1",
+)
+
+TASKS_STARTED = meter.create_counter(
+    "docket_tasks_started",
+    description="How many tasks started",
+    unit="1",
+)
+
+TASKS_REDELIVERED = meter.create_counter(
+    "docket_tasks_redelivered",
+    description="How many tasks started that were redelivered from another worker",
+    unit="1",
+)
+
+TASKS_STRICKEN = meter.create_counter(
+    "docket_tasks_stricken",
+    description="How many tasks have been stricken from executing",
+    unit="1",
+)
+
+TASKS_COMPLETED = meter.create_counter(
+    "docket_tasks_completed",
+    description="How many tasks that have completed in any state",
+    unit="1",
+)
+
+TASKS_FAILED = meter.create_counter(
+    "docket_tasks_failed",
+    description="How many tasks that have failed",
+    unit="1",
+)
+
+TASKS_SUCCEEDED = meter.create_counter(
+    "docket_tasks_succeeded",
+    description="How many tasks that have succeeded",
+    unit="1",
+)
+
+TASKS_RETRIED = meter.create_counter(
+    "docket_tasks_retried",
+    description="How many tasks that have been retried",
+    unit="1",
+)
+
+TASKS_PERPETUATED = meter.create_counter(
+    "docket_tasks_perpetuated",
+    description="How many tasks that have been self-perpetuated",
+    unit="1",
+)
+
+TASK_DURATION = meter.create_histogram(
+    "docket_task_duration",
+    description="How long tasks take to complete",
+    unit="s",
+)
+
+TASK_PUNCTUALITY = meter.create_histogram(
+    "docket_task_punctuality",
+    description="How close a task was to its scheduled time",
+    unit="s",
+)
+
+TASKS_RUNNING = meter.create_up_down_counter(
+    "docket_tasks_running",
+    description="How many tasks that are currently running",
+    unit="1",
+)
+
+REDIS_DISRUPTIONS = meter.create_counter(
+    "docket_redis_disruptions",
+    description="How many times the Redis connection has been disrupted",
+    unit="1",
+)
+
+STRIKES_IN_EFFECT = meter.create_up_down_counter(
+    "docket_strikes_in_effect",
+    description="How many strikes are currently in effect",
+    unit="1",
+)
+
+QUEUE_DEPTH = meter.create_gauge(
+    "docket_queue_depth",
+    description="How many tasks are due to be executed now",
+    unit="1",
+)
+SCHEDULE_DEPTH = meter.create_gauge(
+    "docket_schedule_depth",
+    description="How many tasks are scheduled to be executed in the future",
+    unit="1",
+)
+
+CACHE_SIZE = meter.create_gauge(
+    "docket_cache_size",
+    description="Size of internal docket caches",
+    unit="1",
+)
+
+Message = dict[bytes, bytes]
+
+
+class MessageGetter(Getter[Message]):
+    def get(self, carrier: Message, key: str) -> list[str] | None:
+        val = carrier.get(key.encode(), None)
+        if val is None:
+            return None
+        return [val.decode()]
+
+    def keys(self, carrier: Message) -> list[str]:
+        return [key.decode() for key in carrier.keys()]
+
+
+class MessageSetter(Setter[Message]):
+    def set(
+        self,
+        carrier: Message,
+        key: str,
+        value: str,
+    ) -> None:
+        carrier[key.encode()] = value.encode()
+
+
+message_getter: MessageGetter = MessageGetter()
+message_setter: MessageSetter = MessageSetter()
+
+
+@contextmanager
+def healthcheck_server(
+    host: str = "0.0.0.0", port: int | None = None
+) -> Generator[None, None, None]:
+    if port is None:
+        yield
+        return
+
+    from http.server import BaseHTTPRequestHandler, HTTPServer
+
+    class HealthcheckHandler(BaseHTTPRequestHandler):
+        def do_GET(self):
+            self.send_response(200)
+            self.send_header("Content-type", "text/plain")
+            self.end_headers()
+            self.wfile.write(b"OK")
+
+        def log_message(self, format: str, *args: object) -> None:
+            # Suppress access logs from the webserver
+            pass
+
+    server = HTTPServer((host, port), HealthcheckHandler)
+    with server:
+        Thread(target=server.serve_forever, daemon=True).start()
+
+        yield
+
+
+@contextmanager
+def metrics_server(
+    host: str = "0.0.0.0", port: int | None = None
+) -> Generator[None, None, None]:
+    if port is None:
+        yield
+        return
+
+    import sys
+    from typing import Any
+
+    # wsgiref.types was added in Python 3.11
+    if sys.version_info >= (3, 11):  # pragma: no cover
+        from wsgiref.types import WSGIApplication
+    else:  # pragma: no cover
+        WSGIApplication = Any  # type: ignore[misc,assignment]
+
+    from prometheus_client import REGISTRY
+    from prometheus_client.exposition import (
+        ThreadingWSGIServer,
+        _SilentHandler,  # type: ignore[member-access]
+        make_server,  # type: ignore[import]
+        make_wsgi_app,  # type: ignore[import]
+    )
+
+    set_meter_provider(MeterProvider(metric_readers=[PrometheusMetricReader()]))
+
+    server = make_server(
+        host,
+        port,
+        cast(WSGIApplication, make_wsgi_app(registry=REGISTRY)),
+        ThreadingWSGIServer,
+        handler_class=_SilentHandler,
+    )
+    with server:
+        Thread(target=server.serve_forever, daemon=True).start()
+
+        yield

docket/tasks.py

@@ -0,0 +1,59 @@
+import asyncio
+import logging
+from datetime import datetime, timezone
+
+from .dependencies import (
+    CurrentDocket,
+    CurrentExecution,
+    CurrentWorker,
+    Retry,
+    TaskLogger,
+)
+from .docket import Docket, TaskCollection
+from .execution import Execution
+from .worker import Worker
+
+
+async def trace(
+    message: str,
+    logger: "logging.LoggerAdapter[logging.Logger]" = TaskLogger(),
+    docket: Docket = CurrentDocket(),
+    worker: Worker = CurrentWorker(),
+    execution: Execution = CurrentExecution(),
+) -> None:
+    logger.info(
+        "%s: %r added to docket %r %s ago now running on worker %r",
+        message,
+        execution.key,
+        docket.name,
+        (datetime.now(timezone.utc) - execution.when),
+        worker.name,
+    )
+
+
+async def fail(
+    message: str,
+    docket: Docket = CurrentDocket(),
+    worker: Worker = CurrentWorker(),
+    execution: Execution = CurrentExecution(),
+    retry: Retry = Retry(attempts=2),
+) -> None:
+    raise Exception(
+        f"{message}: {execution.key} added to docket "
+        f"{docket.name} {datetime.now(timezone.utc) - execution.when} "
+        f"ago now running on worker {worker.name}"
+    )
+
+
+async def sleep(
+    seconds: float, logger: "logging.LoggerAdapter[logging.Logger]" = TaskLogger()
+) -> None:
+    logger.info("Sleeping for %s seconds", seconds)
+    await asyncio.sleep(seconds)
+
+
+standard_tasks: TaskCollection = [
+    trace,
+    fail,
+    sleep,
+]

docket/testing.py

@@ -0,0 +1,235 @@
+"""Testing utilities for making assertions about scheduled tasks.
+
+Example usage:
+    from docket import Docket, testing
+
+    docket = Docket("redis://localhost:6379/0")
+
+    # Schedule a task
+    await docket.add(my_task)("arg1", kwarg1="value1")
+
+    # Assert it's scheduled
+    await testing.assert_task_scheduled(docket, my_task, args=("arg1",))
+
+    # After completion
+    await worker.run_until_finished()
+    await testing.assert_no_tasks(docket)
+"""
+
+from collections.abc import Callable
+from typing import Any
+
+from docket.docket import Docket
+from docket.execution import Execution
+
+
+def _matches_criteria(
+    execution: Execution,
+    function: str | Callable[..., Any],
+    args: tuple[Any, ...] | None,
+    kwargs: dict[str, Any] | None,
+    key: str | None,
+) -> bool:
+    """Check if an execution matches the given criteria."""
+    # Check function name
+    function_name = function if isinstance(function, str) else function.__name__
+    if execution.function.__name__ != function_name:
+        return False
+
+    # Check key if specified
+    if key is not None and execution.key != key:
+        return False
+
+    # Check args if specified
+    if args is not None and execution.args != args:
+        return False
+
+    # Check kwargs if specified (subset matching)
+    if kwargs is not None:
+        for k, v in kwargs.items():
+            if k not in execution.kwargs or execution.kwargs[k] != v:
+                return False
+
+    return True
+
+
+def _format_criteria(
+    function: str | Callable[..., Any],
+    args: tuple[Any, ...] | None,
+    kwargs: dict[str, Any] | None,
+    key: str | None,
+) -> str:
+    """Format criteria for error messages."""
+    parts: list[str] = []
+
+    function_name = function if isinstance(function, str) else function.__name__
+    parts.append(f"function={function_name}")
+
+    if key is not None:
+        parts.append(f"key={key!r}")
+    if args is not None:
+        parts.append(f"args={args!r}")
+    if kwargs is not None:
+        parts.append(f"kwargs={kwargs!r}")
+
+    return ", ".join(parts)
+
+
+async def assert_task_scheduled(
+    docket: Docket,
+    function: str | Callable[..., Any],
+    *,
+    args: tuple[Any, ...] | None = None,
+    kwargs: dict[str, Any] | None = None,
+    key: str | None = None,
+) -> None:
+    """Assert that a task matching the criteria is scheduled.
+
+    Args:
+        docket: The Docket instance to check
+        function: The task function or function name (string)
+        args: Optional tuple of positional arguments to match
+        kwargs: Optional dict of keyword arguments to match (subset matching)
+        key: Optional task key to match
+
+    Raises:
+        AssertionError: If no matching task is found
+
+    Example:
+        await assert_task_scheduled(docket, my_task)
+        await assert_task_scheduled(docket, my_task, args=("foo",))
+        await assert_task_scheduled(docket, "my_task", key="task-123")
+    """
+    snapshot = await docket.snapshot()
+
+    # Check all scheduled tasks (both immediate and future)
+    all_tasks = list(snapshot.future)
+
+    for execution in all_tasks:
+        if _matches_criteria(execution, function, args, kwargs, key):
+            return
+
+    # Build error message
+    criteria = _format_criteria(function, args, kwargs, key)
+    function_name = function if isinstance(function, str) else function.__name__
+
+    if not all_tasks:
+        raise AssertionError(
+            f"Task {function_name} not found: no tasks scheduled on docket"
+        )
+
+    # Show what we found instead
+    found_tasks = [
+        f"  - {e.function.__name__}(args={e.args!r}, kwargs={e.kwargs!r}, key={e.key!r})"
+        for e in all_tasks
+    ]
+    found_str = "\n".join(found_tasks)
+
+    raise AssertionError(
+        f"Task {function_name} not found with {criteria}\n\n"
+        f"Scheduled tasks:\n{found_str}"
+    )
+
+
+async def assert_task_not_scheduled(
+    docket: Docket,
+    function: str | Callable[..., Any],
+    *,
+    args: tuple[Any, ...] | None = None,
+    kwargs: dict[str, Any] | None = None,
+    key: str | None = None,
+) -> None:
+    """Assert that no task matching the criteria is scheduled.
+
+    Args:
+        docket: The Docket instance to check
+        function: The task function or function name (string)
+        args: Optional tuple of positional arguments to match
+        kwargs: Optional dict of keyword arguments to match (subset matching)
+        key: Optional task key to match
+
+    Raises:
+        AssertionError: If a matching task is found
+
+    Example:
+        await assert_task_not_scheduled(docket, my_task)
+        await assert_task_not_scheduled(docket, my_task, args=("foo",))
+    """
+    snapshot = await docket.snapshot()
+
+    # Check all scheduled tasks (both immediate and future)
+    all_tasks = list(snapshot.future)
+
+    for execution in all_tasks:
+        if _matches_criteria(execution, function, args, kwargs, key):
+            function_name = function if isinstance(function, str) else function.__name__
+            raise AssertionError(
+                f"Task {function_name} found but should not be scheduled\n"
+                f"Found: {execution.function.__name__}(args={execution.args!r}, "
+                f"kwargs={execution.kwargs!r}, key={execution.key!r})"
+            )
+
+
+async def assert_task_count(
+    docket: Docket,
+    function: str | Callable[..., Any] | None = None,
+    *,
+    count: int,
+) -> None:
+    """Assert the number of scheduled tasks matches the expected count.
+
+    Args:
+        docket: The Docket instance to check
+        function: Optional task function or name to count (if None, counts all tasks)
+        count: Expected number of tasks
+
+    Raises:
+        AssertionError: If the count doesn't match
+
+    Example:
+        await assert_task_count(docket, count=5)  # All tasks
+        await assert_task_count(docket, my_task, count=2)  # Specific function
+    """
+    snapshot = await docket.snapshot()
+
+    # Check all scheduled tasks (both immediate and future)
+    all_tasks = list(snapshot.future)
+
+    if function is None:
+        actual_count = len(all_tasks)
+        function_desc = "all tasks"
+    else:
+        function_name = function if isinstance(function, str) else function.__name__
+        actual_count = sum(1 for e in all_tasks if e.function.__name__ == function_name)
+        function_desc = f"tasks for {function_name}"
+
+    if actual_count != count:
+        raise AssertionError(f"Expected {count} {function_desc}, found {actual_count}")
+
+
+async def assert_no_tasks(docket: Docket) -> None:
+    """Assert that no tasks are scheduled on the docket.
+
+    Args:
+        docket: The Docket instance to check
+
+    Raises:
+        AssertionError: If any tasks are scheduled
+
+    Example:
+        await assert_no_tasks(docket)
+    """
+    snapshot = await docket.snapshot()
+
+    # Check all scheduled tasks (both immediate and future)
+    all_tasks = list(snapshot.future)
+
+    if all_tasks:
+        found_tasks = [
+            f"  - {e.function.__name__}(args={e.args!r}, kwargs={e.kwargs!r}, key={e.key!r})"
+            for e in all_tasks
+        ]
+        found_str = "\n".join(found_tasks)
+        raise AssertionError(
+            f"Expected no tasks, found {len(all_tasks)} task(s) scheduled:\n{found_str}"
+        )