pydocket 0.9.1__py3-none-any.whl → 0.10.0__py3-none-any.whl

docket/docket.py

@@ -16,7 +16,6 @@ from typing import (
     Mapping,
     NoReturn,
     ParamSpec,
-    Protocol,
     Self,
     Sequence,
     TypedDict,
@@ -28,6 +27,7 @@ from typing import (
 import redis.exceptions
 from opentelemetry import propagate, trace
 from redis.asyncio import ConnectionPool, Redis
+from redis.asyncio.client import Pipeline
 from uuid_extensions import uuid7
 
 from .execution import (
@@ -55,18 +55,6 @@ logger: logging.Logger = logging.getLogger(__name__)
 tracer: trace.Tracer = trace.get_tracer(__name__)
 
 
-class _schedule_task(Protocol):
-    async def __call__(
-        self, keys: list[str], args: list[str | float | bytes]
-    ) -> str: ...  # pragma: no cover
-
-
-class _cancel_task(Protocol):
-    async def __call__(
-        self, keys: list[str], args: list[str]
-    ) -> str: ...  # pragma: no cover
-
-
 P = ParamSpec("P")
 R = TypeVar("R")
 
@@ -143,8 +131,6 @@ class Docket:
 
     _monitor_strikes_task: asyncio.Task[None]
     _connection_pool: ConnectionPool
-    _schedule_task_script: _schedule_task | None
-    _cancel_task_script: _cancel_task | None
 
     def __init__(
         self,
@@ -170,8 +156,6 @@ class Docket:
         self.url = url
         self.heartbeat_interval = heartbeat_interval
         self.missed_heartbeats = missed_heartbeats
-        self._schedule_task_script = None
-        self._cancel_task_script = None
 
     @property
     def worker_group_name(self) -> str:
@@ -316,7 +300,9 @@ class Docket:
             execution = Execution(function, args, kwargs, when, key, attempt=1)
 
             async with self.redis() as redis:
-                await self._schedule(redis, execution, replace=False)
+                async with redis.pipeline() as pipeline:
+                    await self._schedule(redis, pipeline, execution, replace=False)
+                    await pipeline.execute()
 
             TASKS_ADDED.add(1, {**self.labels(), **execution.general_labels()})
             TASKS_SCHEDULED.add(1, {**self.labels(), **execution.general_labels()})
@@ -375,7 +361,9 @@ class Docket:
             execution = Execution(function, args, kwargs, when, key, attempt=1)
 
             async with self.redis() as redis:
-                await self._schedule(redis, execution, replace=True)
+                async with redis.pipeline() as pipeline:
+                    await self._schedule(redis, pipeline, execution, replace=True)
+                    await pipeline.execute()
 
             TASKS_REPLACED.add(1, {**self.labels(), **execution.general_labels()})
             TASKS_CANCELLED.add(1, {**self.labels(), **execution.general_labels()})
@@ -395,7 +383,9 @@ class Docket:
             },
         ):
             async with self.redis() as redis:
-                await self._schedule(redis, execution, replace=False)
+                async with redis.pipeline() as pipeline:
+                    await self._schedule(redis, pipeline, execution, replace=False)
+                    await pipeline.execute()
 
         TASKS_SCHEDULED.add(1, {**self.labels(), **execution.general_labels()})
 
@@ -410,7 +400,9 @@ class Docket:
             attributes={**self.labels(), "docket.key": key},
         ):
             async with self.redis() as redis:
-                await self._cancel(redis, key)
+                async with redis.pipeline() as pipeline:
+                    await self._cancel(pipeline, key)
+                    await pipeline.execute()
 
         TASKS_CANCELLED.add(1, self.labels())
 
@@ -431,17 +423,10 @@ class Docket:
     async def _schedule(
         self,
         redis: Redis,
+        pipeline: Pipeline,
         execution: Execution,
         replace: bool = False,
     ) -> None:
-        """Schedule a task atomically.
-
-        Handles:
-        - Checking for task existence
-        - Cancelling existing tasks when replacing
-        - Adding tasks to stream (immediate) or queue (future)
-        - Tracking stream message IDs for later cancellation
-        """
         if self.strike_list.is_stricken(execution):
             logger.warning(
                 "%r is stricken, skipping schedule of %r",
@@ -464,133 +449,32 @@ class Docket:
         key = execution.key
         when = execution.when
         known_task_key = self.known_task_key(key)
-        is_immediate = when <= datetime.now(timezone.utc)
 
-        # Lock per task key to prevent race conditions between concurrent operations
         async with redis.lock(f"{known_task_key}:lock", timeout=10):
-            if self._schedule_task_script is None:
-                self._schedule_task_script = cast(
-                    _schedule_task,
-                    redis.register_script(
-                        # KEYS: stream_key, known_key, parked_key, queue_key
-                        # ARGV: task_key, when_timestamp, is_immediate, replace, ...message_fields
-                        """
-                        local stream_key = KEYS[1]
-                        local known_key = KEYS[2]
-                        local parked_key = KEYS[3]
-                        local queue_key = KEYS[4]
-
-                        local task_key = ARGV[1]
-                        local when_timestamp = ARGV[2]
-                        local is_immediate = ARGV[3] == '1'
-                        local replace = ARGV[4] == '1'
-
-                        -- Extract message fields from ARGV[5] onwards
-                        local message = {}
-                        for i = 5, #ARGV, 2 do
-                            message[#message + 1] = ARGV[i]     -- field name
-                            message[#message + 1] = ARGV[i + 1] -- field value
-                        end
-
-                        -- Handle replacement: cancel existing task if needed
-                        if replace then
-                            local existing_message_id = redis.call('HGET', known_key, 'stream_message_id')
-                            if existing_message_id then
-                                redis.call('XDEL', stream_key, existing_message_id)
-                            end
-                            redis.call('DEL', known_key, parked_key)
-                            redis.call('ZREM', queue_key, task_key)
-                        else
-                            -- Check if task already exists
-                            if redis.call('EXISTS', known_key) == 1 then
-                                return 'EXISTS'
-                            end
-                        end
-
-                        if is_immediate then
-                            -- Add to stream and store message ID for later cancellation
-                            local message_id = redis.call('XADD', stream_key, '*', unpack(message))
-                            redis.call('HSET', known_key, 'when', when_timestamp, 'stream_message_id', message_id)
-                            return message_id
-                        else
-                            -- Add to queue with task data in parked hash
-                            redis.call('HSET', known_key, 'when', when_timestamp)
-                            redis.call('HSET', parked_key, unpack(message))
-                            redis.call('ZADD', queue_key, when_timestamp, task_key)
-                            return 'QUEUED'
-                        end
-                        """
-                    ),
-                )
-            schedule_task = self._schedule_task_script
+            if replace:
+                await self._cancel(pipeline, key)
+            else:
+                # if the task is already in the queue or stream, retain it
+                if await redis.exists(known_task_key):
+                    logger.debug(
+                        "Task %r is already in the queue or stream, not scheduling",
+                        key,
+                        extra=self.labels(),
+                    )
+                    return
 
-            await schedule_task(
-                keys=[
-                    self.stream_key,
-                    known_task_key,
-                    self.parked_task_key(key),
-                    self.queue_key,
-                ],
-                args=[
-                    key,
-                    str(when.timestamp()),
-                    "1" if is_immediate else "0",
-                    "1" if replace else "0",
-                    *[
-                        item
-                        for field, value in message.items()
-                        for item in (field, value)
-                    ],
-                ],
-            )
+            pipeline.set(known_task_key, when.timestamp())
 
-    async def _cancel(self, redis: Redis, key: str) -> None:
-        """Cancel a task atomically.
+            if when <= datetime.now(timezone.utc):
+                pipeline.xadd(self.stream_key, message)  # type: ignore[arg-type]
+            else:
+                pipeline.hset(self.parked_task_key(key), mapping=message)  # type: ignore[arg-type]
+                pipeline.zadd(self.queue_key, {key: when.timestamp()})
 
-        Handles cancellation regardless of task location:
-        - From the stream (using stored message ID)
-        - From the queue (scheduled tasks)
-        - Cleans up all associated metadata keys
-        """
-        if self._cancel_task_script is None:
-            self._cancel_task_script = cast(
-                _cancel_task,
-                redis.register_script(
-                    # KEYS: stream_key, known_key, parked_key, queue_key
-                    # ARGV: task_key
-                    """
-                    local stream_key = KEYS[1]
-                    local known_key = KEYS[2]
-                    local parked_key = KEYS[3]
-                    local queue_key = KEYS[4]
-                    local task_key = ARGV[1]
-
-                    -- Delete from stream if message ID exists
-                    local message_id = redis.call('HGET', known_key, 'stream_message_id')
-                    if message_id then
-                        redis.call('XDEL', stream_key, message_id)
-                    end
-
-                    -- Clean up all task-related keys
-                    redis.call('DEL', known_key, parked_key)
-                    redis.call('ZREM', queue_key, task_key)
-
-                    return 'OK'
-                    """
-                ),
-            )
-        cancel_task = self._cancel_task_script
-
-        # Execute the cancellation script
-        await cancel_task(
-            keys=[
-                self.stream_key,
-                self.known_task_key(key),
-                self.parked_task_key(key),
-                self.queue_key,
-            ],
-            args=[key],
-        )
+    async def _cancel(self, pipeline: Pipeline, key: str) -> None:
+        pipeline.delete(self.known_task_key(key))
+        pipeline.delete(self.parked_task_key(key))
+        pipeline.zrem(self.queue_key, key)
 
     @property
     def strike_key(self) -> str:

{pydocket-0.9.1.dist-info → pydocket-0.10.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.9.1
+Version: 0.10.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.9.1.dist-info → pydocket-0.10.0.dist-info}/RECORD

@@ -3,14 +3,14 @@ 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=0nQCHDDHy7trv2a0eYygGgIKiA7fWq5GcOXye3_CPWM,30847
+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=pOBRoEbakUwAGVKAuCNPSMyHRBSalUxtMc93QZewX7M,34928
-pydocket-0.9.1.dist-info/METADATA,sha256=1MpHrin4Unc4IJ3c7R002wD3BE7Uz_SAAFqu7aX3SF0,5418
-pydocket-0.9.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydocket-0.9.1.dist-info/entry_points.txt,sha256=4WOk1nUlBsUT5O3RyMci2ImuC5XFswuopElYcLHtD5k,47
-pydocket-0.9.1.dist-info/licenses/LICENSE,sha256=YuVWU_ZXO0K_k2FG8xWKe5RGxV24AhJKTvQmKfqXuyk,1087
-pydocket-0.9.1.dist-info/RECORD,,
+pydocket-0.10.0.dist-info/METADATA,sha256=Tsm_S5NTj5yOPmt-q4KAKjdEDjH6ZRzz_ITVapnFk64,5419
+pydocket-0.10.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydocket-0.10.0.dist-info/entry_points.txt,sha256=4WOk1nUlBsUT5O3RyMci2ImuC5XFswuopElYcLHtD5k,47
+pydocket-0.10.0.dist-info/licenses/LICENSE,sha256=YuVWU_ZXO0K_k2FG8xWKe5RGxV24AhJKTvQmKfqXuyk,1087
+pydocket-0.10.0.dist-info/RECORD,,