pydocket 0.3.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

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())

docket/docket.py

@@ -132,6 +132,9 @@ class Docket:
                 - "redis://user:password@localhost:6379/0?ssl=true"
                 - "rediss://localhost:6379/0"
                 - "unix:///path/to/redis.sock"
+            heartbeat_interval: How often workers send heartbeat messages to the docket.
+            missed_heartbeats: How many heartbeats a worker can miss before it is
+                considered dead.
         """
         self.name = name
         self.url = url
@@ -305,6 +308,9 @@ class Docket:
     def stream_key(self) -> str:
         return f"{self.name}:stream"
 
+    def known_task_key(self, key: str) -> str:
+        return f"{self.name}:known:{key}"
+
     def parked_task_key(self, key: str) -> str:
         return f"{self.name}:{key}"
 
@@ -340,30 +346,36 @@ class Docket:
             when = execution.when
 
             async with self.redis() as redis:
-                # if the task is already in the queue, retain it
-                if await redis.zscore(self.queue_key, key) is not None:
+                # if the task is already in the queue or stream, retain it
+                if await redis.exists(self.known_task_key(key)):
+                    logger.debug(
+                        "Task %r is already in the queue or stream, skipping schedule",
+                        key,
+                        extra=self.labels(),
+                    )
                     return
 
-                if when <= datetime.now(timezone.utc):
-                    await redis.xadd(self.stream_key, message)  # type: ignore[arg-type]
-                else:
-                    async with redis.pipeline() as pipe:
+                async with redis.pipeline() as pipe:
+                    pipe.set(self.known_task_key(key), when.timestamp())
+
+                    if when <= datetime.now(timezone.utc):
+                        pipe.xadd(self.stream_key, message)  # type: ignore[arg-type]
+                    else:
                         pipe.hset(self.parked_task_key(key), mapping=message)  # type: ignore[arg-type]
                         pipe.zadd(self.queue_key, {key: when.timestamp()})
-                        await pipe.execute()
+
+                    await pipe.execute()
 
         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={
-                **self.labels(),
-                "docket.key": key,
-            },
+            attributes={**self.labels(), "docket.key": key},
         ):
             async with self.redis() as redis:
                 async with redis.pipeline() as pipe:
+                    pipe.delete(self.known_task_key(key))
                     pipe.delete(self.parked_task_key(key))
                     pipe.zrem(self.queue_key, key)
                     await pipe.execute()

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()
@@ -569,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.3.2.dist-info → pydocket-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.3.2
+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.3.2.dist-info → pydocket-0.5.0.dist-info}/RECORD

@@ -1,16 +1,16 @@
 docket/__init__.py,sha256=7oruGALDoU6W_ntF-mMxxv3FFtO970DVzj3lUgoVIiM,775
 docket/__main__.py,sha256=Vkuh7aJ-Bl7QVpVbbkUksAd_hn05FiLmWbc-8kbhZQ4,34
 docket/annotations.py,sha256=GZwOPtPXyeIhnsLh3TQMBnXrjtTtSmF4Ratv4vjPx8U,950
-docket/cli.py,sha256=EseF0Sj7IEgd9QDC-FSbHSffvF7DNsrmDGYGgZBdJc8,19413
+docket/cli.py,sha256=OWql6QFthSbvRCGkIg-ufo26F48z0eCmzRXJYOdyAEc,20309
 docket/dependencies.py,sha256=S3KqXxEF0Q2t_jO3R-kI5IIA3M-tqybtiSod2xnRO4o,4991
-docket/docket.py,sha256=zva6ofTm7i5hRwAaAnNtlgIqoMPaNLqCTs2PXGka_8s,19723
+docket/docket.py,sha256=xjmXZtE2QY6Pet5Gv5hEJNr8d5wp3zL1GJqRbJYwnKs,20321
 docket/execution.py,sha256=PDrlAr8VzmB6JvqKO71YhXUcTcGQW7eyXrSKiTcAexE,12508
 docket/instrumentation.py,sha256=bZlGA02JoJcY0J1WGm5_qXDfY0AXKr0ZLAYu67wkeKY,4611
 docket/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 docket/tasks.py,sha256=RIlSM2omh-YDwVnCz6M5MtmK8T_m_s1w2OlRRxDUs6A,1437
-docket/worker.py,sha256=MO-j5NQWvhlZh0vCVeUCwLas3i6AVsY672MYgyoU620,21882
-pydocket-0.3.2.dist-info/METADATA,sha256=F0cKpadZ6gQ-mJYXde7XTogoy_uaaTpprpJ5uum9z-w,13092
-pydocket-0.3.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydocket-0.3.2.dist-info/entry_points.txt,sha256=4WOk1nUlBsUT5O3RyMci2ImuC5XFswuopElYcLHtD5k,47
-pydocket-0.3.2.dist-info/licenses/LICENSE,sha256=YuVWU_ZXO0K_k2FG8xWKe5RGxV24AhJKTvQmKfqXuyk,1087
-pydocket-0.3.2.dist-info/RECORD,,
+docket/worker.py,sha256=lPrx2G87tQDHgbIKZLsjsRQn56u-L8_8gGz4UxWQK68,23710
+pydocket-0.5.0.dist-info/METADATA,sha256=Jh28apfadHchCF5Xr_YV-57COunujrhtdDEF0UbjzmU,13092
+pydocket-0.5.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydocket-0.5.0.dist-info/entry_points.txt,sha256=4WOk1nUlBsUT5O3RyMci2ImuC5XFswuopElYcLHtD5k,47
+pydocket-0.5.0.dist-info/licenses/LICENSE,sha256=YuVWU_ZXO0K_k2FG8xWKe5RGxV24AhJKTvQmKfqXuyk,1087
+pydocket-0.5.0.dist-info/RECORD,,