pydocket 0.1.4__tar.gz → 0.2.0__tar.gz

{pydocket-0.1.4 → pydocket-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.1.4
+Version: 0.2.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

{pydocket-0.1.4 → pydocket-0.2.0}/chaos/driver.py

@@ -6,9 +6,9 @@ import socket
 import sys
 from asyncio import subprocess
 from asyncio.subprocess import Process
-from contextlib import contextmanager
+from contextlib import asynccontextmanager
 from datetime import timedelta
-from typing import Any, Generator, Literal, Sequence
+from typing import Any, AsyncGenerator, Literal, Sequence
 from uuid import uuid4
 
 import redis.exceptions
@@ -46,8 +46,8 @@ def python_entrypoint() -> list[str]:
     return [sys.executable]
 
 
-@contextmanager
-def run_redis(version: str) -> Generator[tuple[str, Container], None, None]:
+@asynccontextmanager
+async def run_redis(version: str) -> AsyncGenerator[tuple[str, Container], None]:
     def get_free_port() -> int:
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
             s.bind(("", 0))
@@ -80,12 +80,14 @@ async def main(
     producers: int = 4,
     workers: int = 7,
 ):
-    with run_redis("7.4.2") as (redis_url, redis_container):
-        logger.info("Redis running at %s", redis_url)
-        docket = Docket(
+    async with (
+        run_redis("7.4.2") as (redis_url, redis_container),
+        Docket(
             name=f"test-docket-{uuid4()}",
             url=redis_url,
-        )
+        ) as docket,
+    ):
+        logger.info("Redis running at %s", redis_url)
         environment = {
             **os.environ,
             "DOCKET_NAME": docket.name,
@@ -93,14 +95,13 @@ async def main(
         }
 
         # Add in some random strikes to performance test
-        async with docket:
-            for _ in range(100):
-                parameter = f"param_{random.randint(1, 100)}"
-                operator: Operator = random.choice(
-                    ["==", "!=", ">", ">=", "<", "<=", "between"]
-                )
-                value = f"val_{random.randint(1, 1000)}"
-                await docket.strike("rando", parameter, operator, value)
+        for _ in range(100):
+            parameter = f"param_{random.randint(1, 100)}"
+            operator: Operator = random.choice(
+                ["==", "!=", ">", ">=", "<", "<=", "between"]
+            )
+            value = f"val_{random.randint(1, 1000)}"
+            await docket.strike("rando", parameter, operator, value)
 
         if tasks % producers != 0:
             raise ValueError("total_tasks must be divisible by total_producers")
@@ -199,8 +200,7 @@ async def main(
                 elif chaos_chance < 0.15:
                     logger.warning("CHAOS: Queuing a toxic task...")
                     try:
-                        async with docket:
-                            await docket.add(toxic)()
+                        await docket.add(toxic)()
                     except redis.exceptions.ConnectionError:
                         pass
 

{pydocket-0.1.4 → pydocket-0.2.0}/pyproject.toml

@@ -45,7 +45,7 @@ dev = [
     "pre-commit>=4.1.0",
     "pyright>=1.1.396",
     "pytest>=8.3.4",
-    "pytest-asyncio>=0.25.3",
+    "pytest-aio>=1.9.0",
     "pytest-cov>=6.0.0",
     "pytest-xdist>=3.6.1",
     "ruff>=0.9.7",
@@ -66,8 +66,6 @@ packages = ["src/docket"]
 
 [tool.pytest.ini_options]
 addopts = "--cov=src/docket --cov=tests --cov-report=term-missing --cov-branch"
-asyncio_mode = "auto"
-asyncio_default_fixture_loop_scope = "session"
 filterwarnings = ["error"]
 
 [tool.pyright]

{pydocket-0.1.4 → pydocket-0.2.0}/src/docket/cli.py

@@ -243,6 +243,13 @@ def worker(
             help="Exit after the current docket is finished",
         ),
     ] = False,
+    metrics_port: Annotated[
+        int | None,
+        typer.Option(
+            "--metrics-port",
+            help="The port to serve Prometheus metrics on",
+        ),
+    ] = None,
 ) -> None:
     asyncio.run(
         Worker.run(
@@ -254,6 +261,7 @@ def worker(
             reconnection_delay=reconnection_delay,
             minimum_check_interval=minimum_check_interval,
             until_finished=until_finished,
+            metrics_port=metrics_port,
             tasks=tasks,
         )
     )

{pydocket-0.1.4 → pydocket-0.2.0}/src/docket/dependencies.py

@@ -61,15 +61,14 @@ class _TaskLogger(Dependency):
         self, docket: Docket, worker: Worker, execution: Execution
     ) -> logging.LoggerAdapter[logging.Logger]:
         logger = logging.getLogger(f"docket.task.{execution.function.__name__}")
-
-        extra = {
-            "execution.key": execution.key,
-            "execution.attempt": execution.attempt,
-            "worker.name": worker.name,
-            "docket.name": docket.name,
-        }
-
-        return logging.LoggerAdapter(logger, extra)
+        return logging.LoggerAdapter(
+            logger,
+            {
+                **docket.labels(),
+                **worker.labels(),
+                **execution.specific_labels(),
+            },
+        )
 
 
 def TaskLogger() -> logging.LoggerAdapter[logging.Logger]:

{pydocket-0.1.4 → pydocket-0.2.0}/src/docket/docket.py

@@ -13,6 +13,7 @@ from typing import (
     Collection,
     Hashable,
     Iterable,
+    Mapping,
     NoReturn,
     ParamSpec,
     Self,
@@ -26,7 +27,7 @@ from uuid import uuid4
 
 import redis.exceptions
 from opentelemetry import propagate, trace
-from redis.asyncio import Redis
+from redis.asyncio import ConnectionPool, Redis
 
 from .execution import (
     Execution,
@@ -112,11 +113,14 @@ class Docket:
     tasks: dict[str, Callable[..., Awaitable[Any]]]
     strike_list: StrikeList
 
+    _monitor_strikes_task: asyncio.Task[None]
+    _connection_pool: ConnectionPool
+
     def __init__(
         self,
         name: str = "docket",
         url: str = "redis://localhost:6379/0",
-        heartbeat_interval: timedelta = timedelta(seconds=1),
+        heartbeat_interval: timedelta = timedelta(seconds=2),
         missed_heartbeats: int = 5,
     ) -> None:
         """
@@ -144,6 +148,7 @@ class Docket:
         self.tasks = {fn.__name__: fn for fn in standard_tasks}
         self.strike_list = StrikeList()
 
+        self._connection_pool = ConnectionPool.from_url(self.url)  # type: ignore
         self._monitor_strikes_task = asyncio.create_task(self._monitor_strikes())
 
         # Ensure that the stream and worker group exist
@@ -176,23 +181,17 @@ class Docket:
         except asyncio.CancelledError:
             pass
 
+        await asyncio.shield(self._connection_pool.disconnect())
+        del self._connection_pool
+
     @asynccontextmanager
     async def redis(self) -> AsyncGenerator[Redis, None]:
-        redis: Redis | None = None
+        r = Redis(connection_pool=self._connection_pool)
+        await r.__aenter__()
         try:
-            redis = await Redis.from_url(
-                self.url,
-                single_connection_client=True,
-            )
-            await redis.__aenter__()
-            try:
-                yield redis
-            finally:
-                await asyncio.shield(redis.__aexit__(None, None, None))
+            yield r
         finally:
-            # redis 4.6.0 doesn't automatically disconnect and leaves connections open
-            if redis:
-                await asyncio.shield(redis.connection_pool.disconnect())
+            await asyncio.shield(r.__aexit__(None, None, None))
 
     def register(self, function: Callable[..., Awaitable[Any]]) -> None:
         from .dependencies import validate_dependencies
@@ -214,6 +213,11 @@ class Docket:
         for function in collection:
             self.register(function)
 
+    def labels(self) -> Mapping[str, str]:
+        return {
+            "docket.name": self.name,
+        }
+
     @overload
     def add(
         self,
@@ -251,7 +255,7 @@ class Docket:
             execution = Execution(function, args, kwargs, when, key, attempt=1)
             await self.schedule(execution)
 
-            TASKS_ADDED.add(1, {"docket": self.name, "task": function.__name__})
+            TASKS_ADDED.add(1, {**self.labels(), **execution.general_labels()})
 
             return execution
 
@@ -287,7 +291,7 @@ class Docket:
             await self.cancel(key)
             await self.schedule(execution)
 
-            TASKS_REPLACED.add(1, {"docket": self.name, "task": function.__name__})
+            TASKS_REPLACED.add(1, {**self.labels(), **execution.general_labels()})
 
             return execution
 
@@ -314,9 +318,9 @@ class Docket:
             TASKS_STRICKEN.add(
                 1,
                 {
-                    "docket": self.name,
-                    "task": execution.function.__name__,
-                    "where": "docket",
+                    **self.labels(),
+                    **execution.specific_labels(),
+                    "docket.where": "docket",
                 },
             )
             return
@@ -327,10 +331,8 @@ class Docket:
         with tracer.start_as_current_span(
             "docket.schedule",
             attributes={
-                "docket.name": self.name,
-                "docket.execution.when": execution.when.isoformat(),
-                "docket.execution.key": execution.key,
-                "docket.execution.attempt": execution.attempt,
+                **self.labels(),
+                **execution.specific_labels(),
                 "code.function.name": execution.function.__name__,
             },
         ):
@@ -350,16 +352,14 @@ class Docket:
                         pipe.zadd(self.queue_key, {key: when.timestamp()})
                         await pipe.execute()
 
-        TASKS_SCHEDULED.add(
-            1, {"docket": self.name, "task": execution.function.__name__}
-        )
+        TASKS_SCHEDULED.add(1, {**self.labels(), **execution.general_labels()})
 
     async def cancel(self, key: str) -> None:
         with tracer.start_as_current_span(
             "docket.cancel",
             attributes={
-                "docket.name": self.name,
-                "docket.execution.key": key,
+                **self.labels(),
+                "docket.key": key,
             },
         ):
             async with self.redis() as redis:
@@ -368,7 +368,7 @@ class Docket:
                     pipe.zrem(self.queue_key, key)
                     await pipe.execute()
 
-        TASKS_CANCELLED.add(1, {"docket": self.name})
+        TASKS_CANCELLED.add(1, self.labels())
 
     @property
     def strike_key(self) -> str:
@@ -408,8 +408,8 @@ class Docket:
         with tracer.start_as_current_span(
             f"docket.{instruction.direction}",
             attributes={
-                "docket.name": self.name,
-                **instruction.as_span_attributes(),
+                **self.labels(),
+                **instruction.labels(),
             },
         ):
             async with self.redis() as redis:
@@ -441,18 +441,15 @@ class Docket:
                                         else "Restoring"
                                     ),
                                     instruction.call_repr(),
-                                    extra={"docket": self.name},
+                                    extra=self.labels(),
                                 )
 
-                                counter_labels = {"docket": self.name}
-                                if instruction.function:
-                                    counter_labels["task"] = instruction.function
-                                if instruction.parameter:
-                                    counter_labels["parameter"] = instruction.parameter
-
                                 STRIKES_IN_EFFECT.add(
                                     1 if instruction.direction == "strike" else -1,
-                                    counter_labels,
+                                    {
+                                        **self.labels(),
+                                        **instruction.labels(),
+                                    },
                                 )
 
             except redis.exceptions.ConnectionError:  # pragma: no cover

{pydocket-0.1.4 → pydocket-0.2.0}/src/docket/execution.py

@@ -3,7 +3,7 @@ import enum
 import inspect
 import logging
 from datetime import datetime
-from typing import Any, Awaitable, Callable, Hashable, Literal, Self, cast
+from typing import Any, Awaitable, Callable, Hashable, Literal, Mapping, Self, cast
 
 import cloudpickle  # type: ignore[import]
 
@@ -55,6 +55,17 @@ class Execution:
             attempt=int(message[b"attempt"].decode()),
         )
 
+    def general_labels(self) -> Mapping[str, str]:
+        return {"docket.task": self.function.__name__}
+
+    def specific_labels(self) -> Mapping[str, str | int]:
+        return {
+            "docket.task": self.function.__name__,
+            "docket.key": self.key,
+            "docket.when": self.when.isoformat(),
+            "docket.attempt": self.attempt,
+        }
+
     def call_repr(self) -> str:
         arguments: list[str] = []
         signature = inspect.signature(self.function)
@@ -131,17 +142,17 @@ class StrikeInstruction(abc.ABC):
         else:
             return Restore(function, parameter, operator, value)
 
-    def as_span_attributes(self) -> dict[str, str]:
-        span_attributes: dict[str, str] = {}
+    def labels(self) -> Mapping[str, str]:
+        labels: dict[str, str] = {}
         if self.function:
-            span_attributes["docket.function"] = self.function
+            labels["docket.task"] = self.function
 
         if self.parameter:
-            span_attributes["docket.parameter"] = self.parameter
-            span_attributes["docket.operator"] = self.operator
-            span_attributes["docket.value"] = repr(self.value)
+            labels["docket.parameter"] = self.parameter
+            labels["docket.operator"] = self.operator
+            labels["docket.value"] = repr(self.value)
 
-        return span_attributes
+        return labels
 
     def call_repr(self) -> str:
         return (

{pydocket-0.1.4 → pydocket-0.2.0}/src/docket/instrumentation.py

@@ -1,5 +1,12 @@
+import threading
+from contextlib import contextmanager
+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")
 
@@ -93,6 +100,17 @@ STRIKES_IN_EFFECT = meter.create_up_down_counter(
     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",
+)
+
 Message = dict[bytes, bytes]
 
 
@@ -119,3 +137,38 @@ class MessageSetter(Setter[Message]):
 
 message_getter: MessageGetter = MessageGetter()
 message_setter: MessageSetter = MessageSetter()
+
+
+@contextmanager
+def metrics_server(
+    host: str = "0.0.0.0", port: int | None = None
+) -> Generator[None, None, None]:
+    if port is None:
+        yield
+        return
+
+    from wsgiref.types import WSGIApplication
+
+    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:
+        t = threading.Thread(target=server.serve_forever)
+        t.daemon = True
+        t.start()
+
+        yield

{pydocket-0.1.4 → pydocket-0.2.0}/src/docket/tasks.py

@@ -2,16 +2,21 @@ import asyncio
 import logging
 from datetime import datetime, timezone
 
-from .dependencies import CurrentDocket, CurrentExecution, CurrentWorker, Retry
+from .dependencies import (
+    CurrentDocket,
+    CurrentExecution,
+    CurrentWorker,
+    Retry,
+    TaskLogger,
+)
 from .docket import Docket, TaskCollection
 from .execution import Execution
 from .worker import Worker
 
-logger: logging.Logger = logging.getLogger(__name__)
-
 
 async def trace(
     message: str,
+    logger: logging.LoggerAdapter[logging.Logger] = TaskLogger(),
     docket: Docket = CurrentDocket(),
     worker: Worker = CurrentWorker(),
     execution: Execution = CurrentExecution(),
@@ -23,11 +28,6 @@ async def trace(
         docket.name,
         (datetime.now(timezone.utc) - execution.when),
         worker.name,
-        extra={
-            "docket.name": docket.name,
-            "worker.name": worker.name,
-            "execution.key": execution.key,
-        },
     )
 
 
@@ -45,7 +45,9 @@ async def fail(
     )
 
 
-async def sleep(seconds: float) -> None:
+async def sleep(
+    seconds: float, logger: logging.LoggerAdapter[logging.Logger] = TaskLogger()
+) -> None:
     logger.info("Sleeping for %s seconds", seconds)
     await asyncio.sleep(seconds)