pydocket 0.4.0__tar.gz → 0.5.0__tar.gz

{pydocket-0.4.0 → pydocket-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.4.0
+Version: 0.5.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.4.0 → pydocket-0.5.0}/chaos/driver.py

@@ -76,9 +76,9 @@ async def run_redis(version: str) -> AsyncGenerator[tuple[str, Container], None]
 
 async def main(
     mode: Literal["performance", "chaos"] = "chaos",
-    tasks: int = 5000,
-    producers: int = 4,
-    workers: int = 7,
+    tasks: int = 20000,
+    producers: int = 5,
+    workers: int = 10,
 ):
     async with (
         run_redis("7.4.2") as (redis_url, redis_container),
@@ -97,9 +97,7 @@ async def main(
         # Add in some random strikes to performance test
         for _ in range(100):
             parameter = f"param_{random.randint(1, 100)}"
-            operator: Operator = random.choice(
-                ["==", "!=", ">", ">=", "<", "<=", "between"]
-            )
+            operator = random.choice(list(Operator))
             value = f"val_{random.randint(1, 1000)}"
             await docket.strike("rando", parameter, operator, value)
 
@@ -141,11 +139,9 @@ async def main(
                 redis_url,
                 "--tasks",
                 "chaos.tasks:chaos_tasks",
-                env=environment
-                | {
-                    "OTEL_SERVICE_NAME": "chaos-worker",
-                    "DOCKET_WORKER_REDELIVERY_TIMEOUT": "5s",
-                },
+                "--redelivery-timeout",
+                "5s",
+                env=environment | {"OTEL_SERVICE_NAME": "chaos-worker"},
                 stdout=subprocess.DEVNULL,
                 stderr=subprocess.DEVNULL,
             )

{pydocket-0.4.0 → pydocket-0.5.0}/chaos/producer.py

@@ -1,8 +1,11 @@
 import asyncio
+import datetime
 import logging
 import os
+import random
 import sys
 import time
+from datetime import timedelta
 
 import redis.exceptions
 
@@ -14,6 +17,10 @@ logging.getLogger().setLevel(logging.INFO)
 logger = logging.getLogger("chaos.producer")
 
 
+def now() -> datetime.datetime:
+    return datetime.datetime.now(datetime.timezone.utc)
+
+
 async def main(tasks_to_produce: int):
     docket = Docket(
         name=os.environ["DOCKET_NAME"],
@@ -25,7 +32,9 @@ async def main(tasks_to_produce: int):
             async with docket:
                 async with docket.redis() as r:
                     for _ in range(tasks_sent, tasks_to_produce):
-                        execution = await docket.add(hello)()
+                        jitter = 5 * ((random.random() * 2) - 1)
+                        when = now() + timedelta(seconds=jitter)
+                        execution = await docket.add(hello, when=when)()
                         await r.zadd("hello:sent", {execution.key: time.time()})
                         logger.info("Added task %s", execution.key)
                         tasks_sent += 1

{pydocket-0.4.0 → pydocket-0.5.0}/src/docket/cli.py

@@ -162,6 +162,7 @@ def worker(
                 "This can be specified multiple times.  A task collection is any "
                 "iterable of async functions."
             ),
+            envvar="DOCKET_TASKS",
         ),
     ] = ["docket.tasks:standard_tasks"],
     docket_: Annotated[
@@ -236,6 +237,14 @@ def worker(
             envvar="DOCKET_WORKER_MINIMUM_CHECK_INTERVAL",
         ),
     ] = timedelta(milliseconds=100),
+    scheduling_resolution: Annotated[
+        timedelta,
+        typer.Option(
+            parser=duration,
+            help="How frequently to check for future tasks to be scheduled",
+            envvar="DOCKET_WORKER_SCHEDULING_RESOLUTION",
+        ),
+    ] = timedelta(milliseconds=250),
     until_finished: Annotated[
         bool,
         typer.Option(
@@ -260,6 +269,7 @@ def worker(
             redelivery_timeout=redelivery_timeout,
             reconnection_delay=reconnection_delay,
             minimum_check_interval=minimum_check_interval,
+            scheduling_resolution=scheduling_resolution,
             until_finished=until_finished,
             metrics_port=metrics_port,
             tasks=tasks,
@@ -542,6 +552,18 @@ def relative_time(now: datetime, when: datetime) -> str:
 
 @app.command(help="Shows a snapshot of what's on the docket right now")
 def snapshot(
+    tasks: Annotated[
+        list[str],
+        typer.Option(
+            "--tasks",
+            help=(
+                "The dotted path of a task collection to register with the docket. "
+                "This can be specified multiple times.  A task collection is any "
+                "iterable of async functions."
+            ),
+            envvar="DOCKET_TASKS",
+        ),
+    ] = ["docket.tasks:standard_tasks"],
     docket_: Annotated[
         str,
         typer.Option(
@@ -560,6 +582,9 @@ def snapshot(
 ) -> None:
     async def run() -> DocketSnapshot:
         async with Docket(name=docket_, url=url) as docket:
+            for task_path in tasks:
+                docket.register_collection(task_path)
+
             return await docket.snapshot()
 
     snapshot = asyncio.run(run())

{pydocket-0.4.0 → pydocket-0.5.0}/src/docket/worker.py

@@ -18,6 +18,7 @@ from uuid import uuid4
 import redis.exceptions
 from opentelemetry import propagate, trace
 from opentelemetry.trace import Tracer
+from redis.asyncio import Redis
 
 from .docket import (
     Docket,
@@ -68,6 +69,7 @@ class Worker:
     redelivery_timeout: timedelta
     reconnection_delay: timedelta
     minimum_check_interval: timedelta
+    scheduling_resolution: timedelta
 
     def __init__(
         self,
@@ -77,6 +79,7 @@ class Worker:
         redelivery_timeout: timedelta = timedelta(minutes=5),
         reconnection_delay: timedelta = timedelta(seconds=5),
         minimum_check_interval: timedelta = timedelta(milliseconds=100),
+        scheduling_resolution: timedelta = timedelta(milliseconds=250),
     ) -> None:
         self.docket = docket
         self.name = name or f"worker:{uuid4()}"
@@ -84,6 +87,7 @@ class Worker:
         self.redelivery_timeout = redelivery_timeout
         self.reconnection_delay = reconnection_delay
         self.minimum_check_interval = minimum_check_interval
+        self.scheduling_resolution = scheduling_resolution
 
     async def __aenter__(self) -> Self:
         self._heartbeat_task = asyncio.create_task(self._heartbeat())
@@ -128,6 +132,7 @@ class Worker:
         redelivery_timeout: timedelta = timedelta(minutes=5),
         reconnection_delay: timedelta = timedelta(seconds=5),
         minimum_check_interval: timedelta = timedelta(milliseconds=100),
+        scheduling_resolution: timedelta = timedelta(milliseconds=250),
         until_finished: bool = False,
         metrics_port: int | None = None,
         tasks: list[str] = ["docket.tasks:standard_tasks"],
@@ -144,6 +149,7 @@ class Worker:
                     redelivery_timeout=redelivery_timeout,
                     reconnection_delay=reconnection_delay,
                     minimum_check_interval=minimum_check_interval,
+                    scheduling_resolution=scheduling_resolution,
                 ) as worker:
                     if until_finished:
                         await worker.run_until_finished()
@@ -210,57 +216,24 @@ class Worker:
                 await asyncio.sleep(self.reconnection_delay.total_seconds())
 
     async def _worker_loop(self, forever: bool = False):
-        async with self.docket.redis() as redis:
-            stream_due_tasks: _stream_due_tasks = cast(
-                _stream_due_tasks,
-                redis.register_script(
-                    # Lua script to atomically move scheduled tasks to the stream
-                    # KEYS[1]: queue key (sorted set)
-                    # KEYS[2]: stream key
-                    # ARGV[1]: current timestamp
-                    # ARGV[2]: docket name prefix
-                    """
-                local total_work = redis.call('ZCARD', KEYS[1])
-                local due_work = 0
-
-                if total_work > 0 then
-                    local tasks = redis.call('ZRANGEBYSCORE', KEYS[1], 0, ARGV[1])
-
-                    for i, key in ipairs(tasks) do
-                        local hash_key = ARGV[2] .. ":" .. key
-                        local task_data = redis.call('HGETALL', hash_key)
-
-                        if #task_data > 0 then
-                            local task = {}
-                            for j = 1, #task_data, 2 do
-                                task[task_data[j]] = task_data[j+1]
-                            end
-
-                            redis.call('XADD', KEYS[2], '*',
-                                'key', task['key'],
-                                'when', task['when'],
-                                'function', task['function'],
-                                'args', task['args'],
-                                'kwargs', task['kwargs'],
-                                'attempt', task['attempt']
-                            )
-                            redis.call('DEL', hash_key)
-                            due_work = due_work + 1
-                        end
-                    end
-                end
-
-                if due_work > 0 then
-                    redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, ARGV[1])
-                end
+        should_stop = asyncio.Event()
 
-                return {total_work, due_work}
-                """
-                ),
+        async with self.docket.redis() as redis:
+            scheduler_task = asyncio.create_task(
+                self._scheduler_loop(redis, should_stop)
             )
 
             active_tasks: dict[asyncio.Task[None], RedisMessageID] = {}
 
+            async def check_for_work() -> bool:
+                async with redis.pipeline() as pipeline:
+                    pipeline.xlen(self.docket.stream_key)
+                    pipeline.zcard(self.docket.queue_key)
+                    results: list[int] = await pipeline.execute()
+                    stream_len = results[0]
+                    queue_len = results[1]
+                    return stream_len > 0 or queue_len > 0
+
             async def process_completed_tasks() -> None:
                 completed_tasks = {task for task in active_tasks if task.done()}
                 for task in completed_tasks:
@@ -280,10 +253,13 @@ class Worker:
                         )
                         await pipeline.execute()
 
-            future_work, due_work = sys.maxsize, 0
+            has_work: bool = True
+
+            if not forever:  # pragma: no branch
+                has_work = await check_for_work()
 
             try:
-                while forever or future_work or active_tasks:
+                while forever or has_work or active_tasks:
                     await process_completed_tasks()
 
                     available_slots = self.concurrency - len(active_tasks)
@@ -297,28 +273,13 @@ class Worker:
                         task = asyncio.create_task(self._execute(message))
                         active_tasks[task] = message_id
 
-                        nonlocal available_slots, future_work
+                        nonlocal available_slots
                         available_slots -= 1
-                        future_work += 1
 
                     if available_slots <= 0:
                         await asyncio.sleep(self.minimum_check_interval.total_seconds())
                         continue
 
-                    future_work, due_work = await stream_due_tasks(
-                        keys=[self.docket.queue_key, self.docket.stream_key],
-                        args=[datetime.now(timezone.utc).timestamp(), self.docket.name],
-                    )
-                    if due_work > 0:
-                        logger.debug(
-                            "Moved %d/%d due tasks from %s to %s",
-                            due_work,
-                            future_work,
-                            self.docket.queue_key,
-                            self.docket.stream_key,
-                            extra=self._log_context(),
-                        )
-
                     redeliveries: RedisMessages
                     _, redeliveries, *_ = await redis.xautoclaim(
                         name=self.docket.stream_key,
@@ -348,10 +309,14 @@ class Worker:
                         ),
                         count=available_slots,
                     )
+
                     for _, messages in new_deliveries:
                         for message_id, message in messages:
                             start_task(message_id, message)
 
+                    if not forever and not active_tasks and not new_deliveries:
+                        has_work = await check_for_work()
+
             except asyncio.CancelledError:
                 if active_tasks:  # pragma: no cover
                     logger.info(
@@ -364,7 +329,98 @@ class Worker:
                     await asyncio.gather(*active_tasks, return_exceptions=True)
                     await process_completed_tasks()
 
+                should_stop.set()
+                await scheduler_task
+
+    async def _scheduler_loop(
+        self,
+        redis: Redis,
+        should_stop: asyncio.Event,
+    ) -> None:
+        """Loop that moves due tasks from the queue to the stream."""
+
+        stream_due_tasks: _stream_due_tasks = cast(
+            _stream_due_tasks,
+            redis.register_script(
+                # Lua script to atomically move scheduled tasks to the stream
+                # KEYS[1]: queue key (sorted set)
+                # KEYS[2]: stream key
+                # ARGV[1]: current timestamp
+                # ARGV[2]: docket name prefix
+                """
+            local total_work = redis.call('ZCARD', KEYS[1])
+            local due_work = 0
+
+            if total_work > 0 then
+                local tasks = redis.call('ZRANGEBYSCORE', KEYS[1], 0, ARGV[1])
+
+                for i, key in ipairs(tasks) do
+                    local hash_key = ARGV[2] .. ":" .. key
+                    local task_data = redis.call('HGETALL', hash_key)
+
+                    if #task_data > 0 then
+                        local task = {}
+                        for j = 1, #task_data, 2 do
+                            task[task_data[j]] = task_data[j+1]
+                        end
+
+                        redis.call('XADD', KEYS[2], '*',
+                            'key', task['key'],
+                            'when', task['when'],
+                            'function', task['function'],
+                            'args', task['args'],
+                            'kwargs', task['kwargs'],
+                            'attempt', task['attempt']
+                        )
+                        redis.call('DEL', hash_key)
+                        due_work = due_work + 1
+                    end
+                end
+            end
+
+            if due_work > 0 then
+                redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, ARGV[1])
+            end
+
+            return {total_work, due_work}
+            """
+            ),
+        )
+
+        total_work: int = sys.maxsize
+
+        while not should_stop.is_set() or total_work:
+            try:
+                total_work, due_work = await stream_due_tasks(
+                    keys=[self.docket.queue_key, self.docket.stream_key],
+                    args=[datetime.now(timezone.utc).timestamp(), self.docket.name],
+                )
+
+                if due_work > 0:
+                    logger.debug(
+                        "Moved %d/%d due tasks from %s to %s",
+                        due_work,
+                        total_work,
+                        self.docket.queue_key,
+                        self.docket.stream_key,
+                        extra=self._log_context(),
+                    )
+            except Exception:  # pragma: no cover
+                logger.exception(
+                    "Error in scheduler loop",
+                    exc_info=True,
+                    extra=self._log_context(),
+                )
+            finally:
+                await asyncio.sleep(self.scheduling_resolution.total_seconds())
+
+        logger.debug("Scheduler loop finished", extra=self._log_context())
+
     async def _execute(self, message: RedisMessage) -> None:
+        key = message[b"key"].decode()
+        async with self.docket.redis() as redis:
+            await redis.delete(self.docket.known_task_key(key))
+
         log_context: Mapping[str, str | float] = self._log_context()
 
         function_name = message[b"function"].decode()
@@ -377,9 +433,6 @@ class Worker:
 
         execution = Execution.from_message(function, message)
 
-        async with self.docket.redis() as redis:
-            await redis.delete(self.docket.known_task_key(execution.key))
-
         log_context = {**log_context, **execution.specific_labels()}
         counter_labels = {**self.labels(), **execution.general_labels()}
 
@@ -572,11 +625,10 @@ class Worker:
                         pipeline.zcount(self.docket.queue_key, 0, now)
                         pipeline.zcount(self.docket.queue_key, now, "+inf")
 
-                        (
-                            stream_depth,
-                            overdue_depth,
-                            schedule_depth,
-                        ) = await pipeline.execute()
+                        results: list[int] = await pipeline.execute()
+                        stream_depth = results[0]
+                        overdue_depth = results[1]
+                        schedule_depth = results[2]
 
                         QUEUE_DEPTH.set(
                             stream_depth + overdue_depth, self.docket.labels()

{pydocket-0.4.0 → pydocket-0.5.0}/tests/conftest.py

@@ -155,7 +155,9 @@ async def docket(redis_url: str, aiolib: str) -> AsyncGenerator[Docket, None]:
 @pytest.fixture
 async def worker(docket: Docket) -> AsyncGenerator[Worker, None]:
     async with Worker(
-        docket, minimum_check_interval=timedelta(milliseconds=10)
+        docket,
+        minimum_check_interval=timedelta(milliseconds=10),
+        scheduling_resolution=timedelta(milliseconds=10),
     ) as worker:
         yield worker