pydocket 0.10.0__tar.gz → 0.11.0__tar.gz

pydocket-0.11.0/.github/workflows/claude-code-review.yml

@@ -0,0 +1,40 @@
+name: Claude Code Review
+
+on:
+  pull_request:
+    types: [opened, synchronize]
+
+jobs:
+  claude-review:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      pull-requests: read
+      issues: read
+      id-token: write
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 1
+
+      - name: Run Claude Code Review
+        id: claude-review
+        uses: anthropics/claude-code-action@beta
+        with:
+          claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
+          model: "claude-opus-4-1-20250805"
+
+          # Direct prompt for automated review (no @claude mention needed)
+          direct_prompt: |
+            Please review this pull request and provide feedback on:
+            - Code quality and best practices
+            - Potential bugs or issues
+            - Performance considerations
+            - Security concerns
+            - Test coverage, which must be maintained at 100% for this project
+
+            Be constructive and helpful in your feedback.
+
+          use_sticky_comment: true

pydocket-0.11.0/.github/workflows/claude.yml

@@ -0,0 +1,42 @@
+name: Claude Code
+
+on:
+  issue_comment:
+    types: [created]
+  pull_request_review_comment:
+    types: [created]
+  issues:
+    types: [opened, assigned]
+  pull_request_review:
+    types: [submitted]
+
+jobs:
+  claude:
+    if: |
+      (github.event_name == 'issue_comment' && contains(github.event.comment.body, '@claude')) ||
+      (github.event_name == 'pull_request_review_comment' && contains(github.event.comment.body, '@claude')) ||
+      (github.event_name == 'pull_request_review' && contains(github.event.review.body, '@claude')) ||
+      (github.event_name == 'issues' && (contains(github.event.issue.body, '@claude') || contains(github.event.issue.title, '@claude')))
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      pull-requests: read
+      issues: read
+      id-token: write
+      actions: read # Required for Claude to read CI results on PRs
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 1
+
+      - name: Run Claude Code
+        id: claude
+        uses: anthropics/claude-code-action@beta
+        with:
+          claude_code_oauth_token: ${{ secrets.CLAUDE_CODE_OAUTH_TOKEN }}
+
+          additional_permissions: |
+            actions: read
+
+          model: "claude-opus-4-1-20250805"

{pydocket-0.10.0 → pydocket-0.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydocket
-Version: 0.10.0
+Version: 0.11.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.10.0 → pydocket-0.11.0}/src/docket/docket.py

@@ -16,6 +16,7 @@ from typing import (
     Mapping,
     NoReturn,
     ParamSpec,
+    Protocol,
     Self,
     Sequence,
     TypedDict,
@@ -27,7 +28,6 @@ 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,6 +55,18 @@ 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")
 
@@ -131,6 +143,8 @@ 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,
@@ -156,6 +170,8 @@ 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:
@@ -300,9 +316,7 @@ class Docket:
             execution = Execution(function, args, kwargs, when, key, attempt=1)
 
             async with self.redis() as redis:
-                async with redis.pipeline() as pipeline:
-                    await self._schedule(redis, pipeline, execution, replace=False)
-                    await pipeline.execute()
+                await self._schedule(redis, execution, replace=False)
 
             TASKS_ADDED.add(1, {**self.labels(), **execution.general_labels()})
             TASKS_SCHEDULED.add(1, {**self.labels(), **execution.general_labels()})
@@ -361,9 +375,7 @@ class Docket:
             execution = Execution(function, args, kwargs, when, key, attempt=1)
 
             async with self.redis() as redis:
-                async with redis.pipeline() as pipeline:
-                    await self._schedule(redis, pipeline, execution, replace=True)
-                    await pipeline.execute()
+                await self._schedule(redis, execution, replace=True)
 
             TASKS_REPLACED.add(1, {**self.labels(), **execution.general_labels()})
             TASKS_CANCELLED.add(1, {**self.labels(), **execution.general_labels()})
@@ -383,9 +395,7 @@ class Docket:
             },
         ):
             async with self.redis() as redis:
-                async with redis.pipeline() as pipeline:
-                    await self._schedule(redis, pipeline, execution, replace=False)
-                    await pipeline.execute()
+                await self._schedule(redis, execution, replace=False)
 
         TASKS_SCHEDULED.add(1, {**self.labels(), **execution.general_labels()})
 
@@ -400,9 +410,7 @@ class Docket:
             attributes={**self.labels(), "docket.key": key},
         ):
             async with self.redis() as redis:
-                async with redis.pipeline() as pipeline:
-                    await self._cancel(pipeline, key)
-                    await pipeline.execute()
+                await self._cancel(redis, key)
 
         TASKS_CANCELLED.add(1, self.labels())
 
@@ -420,13 +428,23 @@ class Docket:
     def parked_task_key(self, key: str) -> str:
         return f"{self.name}:{key}"
 
+    def stream_id_key(self, key: str) -> str:
+        return f"{self.name}:stream-id:{key}"
+
     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",
@@ -449,32 +467,138 @@ 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 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
+            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, stream_id_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 stream_id_key = KEYS[5]
+
+                        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('GET', stream_id_key)
+                            if existing_message_id then
+                                redis.call('XDEL', stream_key, existing_message_id)
+                            end
+                            redis.call('DEL', known_key, parked_key, stream_id_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('SET', known_key, when_timestamp)
+                            redis.call('SET', stream_id_key, message_id)
+                            return message_id
+                        else
+                            -- Add to queue with task data in parked hash
+                            redis.call('SET', known_key, 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
 
-            pipeline.set(known_task_key, when.timestamp())
+            await schedule_task(
+                keys=[
+                    self.stream_key,
+                    known_task_key,
+                    self.parked_task_key(key),
+                    self.queue_key,
+                    self.stream_id_key(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)
+                    ],
+                ],
+            )
 
-            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()})
+    async def _cancel(self, redis: Redis, key: str) -> None:
+        """Cancel a task atomically.
 
-    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)
+        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, stream_id_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 stream_id_key = KEYS[5]
+                    local task_key = ARGV[1]
+
+                    -- Delete from stream if message ID exists
+                    local message_id = redis.call('GET', stream_id_key)
+                    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, stream_id_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,
+                self.stream_id_key(key),
+            ],
+            args=[key],
+        )
 
     @property
     def strike_key(self) -> str:
@@ -781,6 +905,7 @@ class Docket:
                         key = key_bytes.decode()
                         pipeline.delete(self.parked_task_key(key))
                         pipeline.delete(self.known_task_key(key))
+                        pipeline.delete(self.stream_id_key(key))
 
                     await pipeline.execute()
 

{pydocket-0.10.0 → pydocket-0.11.0}/src/docket/execution.py

@@ -51,6 +51,7 @@ class Execution:
         key: str,
         attempt: int,
         trace_context: opentelemetry.context.Context | None = None,
+        redelivered: bool = False,
     ) -> None:
         self.function = function
         self.args = args
@@ -59,6 +60,7 @@ class Execution:
         self.key = key
         self.attempt = attempt
         self.trace_context = trace_context
+        self.redelivered = redelivered
 
     def as_message(self) -> Message:
         return {
@@ -80,6 +82,7 @@ class Execution:
             key=message[b"key"].decode(),
             attempt=int(message[b"attempt"].decode()),
             trace_context=propagate.extract(message, getter=message_getter),
+            redelivered=False,  # Default to False, will be set to True in worker if it's a redelivery
         )
 
     def general_labels(self) -> Mapping[str, str]:

{pydocket-0.10.0 → pydocket-0.11.0}/src/docket/instrumentation.py

@@ -40,6 +40,12 @@ TASKS_STARTED = meter.create_counter(
     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",

{pydocket-0.10.0 → pydocket-0.11.0}/src/docket/worker.py

@@ -47,6 +47,7 @@ from .instrumentation import (
     TASKS_COMPLETED,
     TASKS_FAILED,
     TASKS_PERPETUATED,
+    TASKS_REDELIVERED,
     TASKS_RETRIED,
     TASKS_RUNNING,
     TASKS_STARTED,
@@ -286,7 +287,11 @@ class Worker:
                 count=available_slots,
             )
 
-        def start_task(message_id: RedisMessageID, message: RedisMessage) -> bool:
+        def start_task(
+            message_id: RedisMessageID,
+            message: RedisMessage,
+            is_redelivery: bool = False,
+        ) -> bool:
             function_name = message[b"function"].decode()
             if not (function := self.docket.tasks.get(function_name)):
                 logger.warning(
@@ -297,6 +302,7 @@ class Worker:
                 return False
 
             execution = Execution.from_message(function, message)
+            execution.redelivered = is_redelivery
 
             task = asyncio.create_task(self._execute(execution), name=execution.key)
             active_tasks[task] = message_id
@@ -342,12 +348,15 @@ class Worker:
                     continue
 
                 for source in [get_redeliveries, get_new_deliveries]:
-                    for _, messages in await source(redis):
+                    for stream_key, messages in await source(redis):
+                        is_redelivery = stream_key == b"__redelivery__"
                         for message_id, message in messages:
                             if not message:  # pragma: no cover
                                 continue
 
-                            task_started = start_task(message_id, message)
+                            task_started = start_task(
+                                message_id, message, is_redelivery
+                            )
                             if not task_started:
                                 # Other errors - delete and ack
                                 await self._delete_known_task(redis, message)
@@ -495,7 +504,8 @@ class Worker:
 
         logger.debug("Deleting known task", extra=self._log_context())
         known_task_key = self.docket.known_task_key(key)
-        await redis.delete(known_task_key)
+        stream_id_key = self.docket.stream_id_key(key)
+        await redis.delete(known_task_key, stream_id_key)
 
     async def _execute(self, execution: Execution) -> None:
         log_context = {**self._log_context(), **execution.specific_labels()}
@@ -520,6 +530,8 @@ class Worker:
         duration = 0.0
 
         TASKS_STARTED.add(1, counter_labels)
+        if execution.redelivered:
+            TASKS_REDELIVERED.add(1, counter_labels)
         TASKS_RUNNING.add(1, counter_labels)
         TASK_PUNCTUALITY.record(punctuality, counter_labels)
 

{pydocket-0.10.0 → pydocket-0.11.0}/tests/test_fundamentals.py

@@ -104,13 +104,9 @@ async def test_adding_is_idempotent(
     assert soon <= now() < later
 
 
-@pytest.mark.skip(
-    "Temporarily skipping due to test flake for task rescheduling. "
-    "See https://github.com/chrisguidry/docket/issues/149"
-)
 async def test_rescheduling_later(
     docket: Docket, worker: Worker, the_task: AsyncMock, now: Callable[[], datetime]
-):  # pragma: no cover
+):
     """docket should allow for rescheduling a task for later"""
 
     key = f"my-cool-task:{uuid4()}"
@@ -254,10 +250,10 @@ async def test_cancelling_future_task(
     the_task.assert_not_called()
 
 
-async def test_cancelling_current_task_not_supported(
+async def test_cancelling_immediate_task(
     docket: Docket, worker: Worker, the_task: AsyncMock, now: Callable[[], datetime]
 ):
-    """docket does not allow cancelling a task that is schedule now"""
+    """docket can cancel a task that is scheduled immediately"""
 
     execution = await docket.add(the_task, now())("a", "b", c="c")
 
@@ -265,7 +261,28 @@ async def test_cancelling_current_task_not_supported(
 
     await worker.run_until_finished()
 
-    the_task.assert_awaited_once_with("a", "b", c="c")
+    the_task.assert_not_called()
+
+
+async def test_cancellation_is_idempotent(
+    docket: Docket, worker: Worker, the_task: AsyncMock, now: Callable[[], datetime]
+):
+    """Test that canceling the same task twice doesn't error."""
+    key = f"test-task:{uuid4()}"
+
+    # Schedule a task
+    later = now() + timedelta(seconds=1)
+    await docket.add(the_task, later, key=key)("test")
+
+    # Cancel it twice - both should succeed without error
+    await docket.cancel(key)
+    await docket.cancel(key)  # Should be idempotent
+
+    # Run worker to ensure the task was actually cancelled
+    await worker.run_until_finished()
+
+    # Task should not have been executed since it was cancelled
+    the_task.assert_not_called()
 
 
 async def test_errors_are_logged(

{pydocket-0.10.0 → pydocket-0.11.0}/tests/test_instrumentation.py

@@ -376,6 +376,14 @@ def TASKS_RETRIED(monkeypatch: pytest.MonkeyPatch) -> Mock:
     return mock
 
 
+@pytest.fixture
+def TASKS_REDELIVERED(monkeypatch: pytest.MonkeyPatch) -> Mock:
+    """Mock for the TASKS_REDELIVERED counter."""
+    mock = Mock(spec=Counter.add)
+    monkeypatch.setattr("docket.instrumentation.TASKS_REDELIVERED.add", mock)
+    return mock
+
+
 async def test_worker_execution_increments_task_counters(
     docket: Docket,
     worker: Worker,
@@ -386,6 +394,7 @@ async def test_worker_execution_increments_task_counters(
     TASKS_SUCCEEDED: Mock,
     TASKS_FAILED: Mock,
     TASKS_RETRIED: Mock,
+    TASKS_REDELIVERED: Mock,
 ):
     """Should increment the appropriate task counters when a worker executes a task."""
     await docket.add(the_task)()
@@ -397,6 +406,7 @@ async def test_worker_execution_increments_task_counters(
     TASKS_SUCCEEDED.assert_called_once_with(1, worker_labels)
     TASKS_FAILED.assert_not_called()
     TASKS_RETRIED.assert_not_called()
+    TASKS_REDELIVERED.assert_not_called()
 
 
 async def test_failed_task_increments_failure_counter(
@@ -409,6 +419,7 @@ async def test_failed_task_increments_failure_counter(
     TASKS_SUCCEEDED: Mock,
     TASKS_FAILED: Mock,
     TASKS_RETRIED: Mock,
+    TASKS_REDELIVERED: Mock,
 ):
     """Should increment the TASKS_FAILED counter when a task fails."""
     the_task.side_effect = ValueError("Womp")
@@ -422,6 +433,7 @@ async def test_failed_task_increments_failure_counter(
     TASKS_FAILED.assert_called_once_with(1, worker_labels)
     TASKS_SUCCEEDED.assert_not_called()
     TASKS_RETRIED.assert_not_called()
+    TASKS_REDELIVERED.assert_not_called()
 
 
 async def test_retried_task_increments_retry_counter(
@@ -433,6 +445,7 @@ async def test_retried_task_increments_retry_counter(
     TASKS_SUCCEEDED: Mock,
     TASKS_FAILED: Mock,
     TASKS_RETRIED: Mock,
+    TASKS_REDELIVERED: Mock,
 ):
     """Should increment the TASKS_RETRIED counter when a task is retried."""
 
@@ -448,6 +461,7 @@ async def test_retried_task_increments_retry_counter(
     assert TASKS_FAILED.call_count == 2
     assert TASKS_RETRIED.call_count == 1
     TASKS_SUCCEEDED.assert_not_called()
+    TASKS_REDELIVERED.assert_not_called()
 
 
 async def test_exhausted_retried_task_increments_retry_counter(
@@ -459,6 +473,7 @@ async def test_exhausted_retried_task_increments_retry_counter(
     TASKS_SUCCEEDED: Mock,
     TASKS_FAILED: Mock,
     TASKS_RETRIED: Mock,
+    TASKS_REDELIVERED: Mock,
 ):
     """Should increment the appropriate counters when retries are exhausted."""
 
@@ -474,6 +489,41 @@ async def test_exhausted_retried_task_increments_retry_counter(
     TASKS_FAILED.assert_called_once_with(1, worker_labels)
     TASKS_RETRIED.assert_not_called()
     TASKS_SUCCEEDED.assert_not_called()
+    TASKS_REDELIVERED.assert_not_called()
+
+
+async def test_redelivered_tasks_increment_redelivered_counter(
+    docket: Docket,
+    worker_labels: dict[str, str],
+    TASKS_STARTED: Mock,
+    TASKS_COMPLETED: Mock,
+    TASKS_SUCCEEDED: Mock,
+    TASKS_FAILED: Mock,
+    TASKS_RETRIED: Mock,
+    TASKS_REDELIVERED: Mock,
+):
+    """Should increment the TASKS_REDELIVERED counter for redelivered tasks."""
+
+    async def test_task():
+        await asyncio.sleep(0.01)
+
+    await docket.add(test_task)()
+
+    worker = Worker(docket, redelivery_timeout=timedelta(milliseconds=50))
+
+    async with worker:
+        worker._execute = AsyncMock(side_effect=Exception("Simulated worker failure"))  # type: ignore[assignment]
+
+        with pytest.raises(Exception, match="Simulated worker failure"):
+            await worker.run_until_finished()
+
+    await asyncio.sleep(0.075)
+
+    worker2 = Worker(docket, redelivery_timeout=timedelta(milliseconds=100))
+    async with worker2:
+        await worker2.run_until_finished()
+
+    assert TASKS_REDELIVERED.call_count >= 1
 
 
 @pytest.fixture