jqueue 0.1.0__py3-none-any.whl

jqueue/core/direct.py

@@ -0,0 +1,170 @@
+"""
+DirectQueue — one CAS write per operation.
+
+Every enqueue, dequeue, ack, nack, or heartbeat call does:
+  1. read current state + etag from storage
+  2. mutate state in memory
+  3. CAS write back with if_match=etag (retries on CASConflictError)
+
+Suitable for ~1-5 ops/sec workloads depending on storage backend latency.
+Use BrokerQueue for higher throughput.
+
+Retry policy
+------------
+Operations retry up to `max_retries` times (default 10) on CASConflictError
+with linear back-off (10ms × attempt). Raises CASConflictError if all retries
+are exhausted.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+from collections.abc import Callable
+from datetime import UTC, datetime, timedelta
+
+from jqueue.core import codec
+from jqueue.domain.errors import CASConflictError, JobNotFoundError
+from jqueue.domain.models import Job, JobStatus, QueueState
+from jqueue.ports.storage import ObjectStoragePort
+
+MutationFn = Callable[[QueueState], QueueState]
+
+
+@dataclasses.dataclass
+class DirectQueue:
+    """
+    Thin stateless wrapper around ObjectStoragePort.
+
+    All methods are async and safe to call from multiple coroutines;
+    each operation performs a full CAS cycle independently.
+    """
+
+    storage: ObjectStoragePort
+    max_retries: int = 10
+
+    # ------------------------------------------------------------------ #
+    # Write operations                                                     #
+    # ------------------------------------------------------------------ #
+
+    async def enqueue(
+        self,
+        entrypoint: str,
+        payload: bytes,
+        priority: int = 0,
+    ) -> Job:
+        """Add a new job to the queue. Returns the committed Job."""
+        job = Job.new(entrypoint, payload, priority)
+        await self._mutate(lambda state: state.with_job_added(job))
+        return job
+
+    async def dequeue(
+        self,
+        entrypoint: str | None = None,
+        *,
+        batch_size: int = 1,
+    ) -> list[Job]:
+        """
+        Claim up to `batch_size` QUEUED jobs and mark them IN_PROGRESS.
+
+        Optionally filter by entrypoint. Returns the list of claimed jobs.
+        Returns an empty list if no jobs are available.
+        """
+        claimed: list[Job] = []
+
+        def _fn(state: QueueState) -> QueueState:
+            nonlocal claimed
+            available = state.queued_jobs(entrypoint)[:batch_size]
+            claimed = []
+            new_state = state
+            for job in available:
+                updated = job.with_status(JobStatus.IN_PROGRESS).with_heartbeat(
+                    datetime.now(UTC)
+                )
+                new_state = new_state.with_job_replaced(updated)
+                claimed.append(updated)
+            return new_state
+
+        await self._mutate(_fn)
+        return claimed
+
+    async def ack(self, job_id: str) -> None:
+        """Mark a job as done and remove it from the queue."""
+        await self._mutate(lambda state: state.with_job_removed(job_id))
+
+    async def nack(self, job_id: str) -> None:
+        """Return a job to QUEUED status (worker failed or declined it)."""
+
+        def _fn(state: QueueState) -> QueueState:
+            job = state.find(job_id)
+            if job is None:
+                raise JobNotFoundError(job_id)
+            return state.with_job_replaced(
+                job.with_status(JobStatus.QUEUED).with_heartbeat(None)
+            )
+
+        await self._mutate(_fn)
+
+    async def heartbeat(self, job_id: str) -> None:
+        """Update the heartbeat timestamp of an IN_PROGRESS job."""
+
+        def _fn(state: QueueState) -> QueueState:
+            job = state.find(job_id)
+            if job is None:
+                raise JobNotFoundError(job_id)
+            return state.with_job_replaced(job.with_heartbeat(datetime.now(UTC)))
+
+        await self._mutate(_fn)
+
+    async def requeue_stale(self, timeout: timedelta) -> int:
+        """
+        Re-queue any IN_PROGRESS jobs whose heartbeat is older than `timeout`.
+
+        Returns the number of jobs re-queued.
+        """
+        cutoff = datetime.now(UTC) - timeout
+        requeued = 0
+
+        def _fn(state: QueueState) -> QueueState:
+            nonlocal requeued
+            old_in_progress = {j.id for j in state.in_progress_jobs()}
+            new_state = state.requeue_stale(cutoff)
+            new_in_progress = {j.id for j in new_state.in_progress_jobs()}
+            requeued = len(old_in_progress - new_in_progress)
+            return new_state
+
+        await self._mutate(_fn)
+        return requeued
+
+    # ------------------------------------------------------------------ #
+    # Read operations (no CAS needed)                                     #
+    # ------------------------------------------------------------------ #
+
+    async def read_state(self) -> QueueState:
+        """Read-only snapshot of the current queue state."""
+        content, _ = await self.storage.read()
+        return codec.decode(content)
+
+    # ------------------------------------------------------------------ #
+    # Internal CAS loop                                                   #
+    # ------------------------------------------------------------------ #
+
+    async def _mutate(self, fn: MutationFn) -> None:
+        """
+        Read-modify-write with CAS retry loop.
+
+        fn(state) -> new_state  (synchronous)
+        Retries up to self.max_retries on CASConflictError.
+        """
+        for attempt in range(self.max_retries):
+            content, etag = await self.storage.read()
+            state = codec.decode(content)
+            new_state = fn(state)
+            new_content = codec.encode(new_state)
+            try:
+                await self.storage.write(new_content, if_match=etag)
+                return
+            except CASConflictError:
+                if attempt == self.max_retries - 1:
+                    raise
+                await asyncio.sleep(0.01 * (attempt + 1))

jqueue/core/group_commit.py

@@ -0,0 +1,263 @@
+"""
+GroupCommitLoop — serialize all mutations through a single asyncio writer task.
+
+Algorithm (from the turbopuffer blog post)
+------------------------------------------
+When a write is in-flight, incoming operations accumulate in the pending buffer.
+As soon as the write finishes, the buffer is flushed as the next CAS write.
+This collapses N concurrent operations into O(1) storage writes.
+
+Concretely:
+
+  Caller 1: enqueue()    ──────────────────────────────────────────> [future]
+  Caller 2: enqueue()    ──────────────────────────────────────────> [future]
+  Caller 3: dequeue()    ──────────────────────────────────────────> [future]
+                           ↓ batch = [op1, op2, op3]
+  Writer:              read → apply op1, op2, op3 → CAS write → resolve futures
+
+If the CAS write fails (concurrent external writer), the whole batch is
+re-applied on a fresh state and retried.
+
+Per-operation error isolation
+------------------------------
+If one mutation in a batch raises (e.g., JobNotFoundError), that future gets
+the exception but the other mutations in the batch still commit normally.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+from collections.abc import Callable
+from datetime import UTC, datetime, timedelta
+
+from jqueue.core import codec
+from jqueue.domain.errors import CASConflictError, JQueueError
+from jqueue.domain.models import Job, JobStatus, QueueState
+from jqueue.ports.storage import ObjectStoragePort
+
+MutationFn = Callable[[QueueState], QueueState]
+
+_MAX_RETRIES: int = 20
+
+
+@dataclasses.dataclass
+class _PendingOp:
+    """A buffered mutation waiting to be committed to object storage."""
+
+    fn: MutationFn
+    future: asyncio.Future[None]
+
+
+@dataclasses.dataclass
+class GroupCommitLoop:
+    """
+    Serializes all storage mutations through a single asyncio writer task.
+
+    Usage
+    -----
+        loop = GroupCommitLoop(storage=my_storage)
+        await loop.start()
+        try:
+            job  = await loop.enqueue("send_email", b"payload")
+            jobs = await loop.dequeue("send_email", batch_size=5)
+        finally:
+            await loop.stop()   # drains pending ops before shutting down
+    """
+
+    storage: ObjectStoragePort
+    stale_timeout: timedelta = timedelta(minutes=5)
+
+    _pending: list[_PendingOp] = dataclasses.field(
+        default_factory=list, init=False, repr=False
+    )
+    _wakeup: asyncio.Event = dataclasses.field(
+        default_factory=asyncio.Event, init=False, repr=False
+    )
+    _task: asyncio.Task[None] | None = dataclasses.field(
+        default=None, init=False, repr=False
+    )
+    _stopped: bool = dataclasses.field(default=False, init=False, repr=False)
+
+    # ------------------------------------------------------------------ #
+    # Lifecycle                                                            #
+    # ------------------------------------------------------------------ #
+
+    async def start(self) -> None:
+        """Start the background writer task."""
+        if self._task is not None:
+            raise RuntimeError("GroupCommitLoop is already running")
+        self._task = asyncio.create_task(
+            self._writer_loop(), name="jqueue-group-commit-writer"
+        )
+
+    async def stop(self) -> None:
+        """Signal shutdown and wait for the writer to drain pending ops."""
+        self._stopped = True
+        self._wakeup.set()
+        if self._task is not None:
+            await self._task
+            self._task = None
+
+    # ------------------------------------------------------------------ #
+    # Public mutation API                                                  #
+    # ------------------------------------------------------------------ #
+
+    async def enqueue(
+        self,
+        entrypoint: str,
+        payload: bytes,
+        priority: int = 0,
+    ) -> Job:
+        """Add a new job. Returns the committed Job (UUID stable across retries)."""
+        job = Job.new(entrypoint, payload, priority)
+
+        def _fn(state: QueueState) -> QueueState:
+            return state.with_job_added(job)
+
+        await self._submit(_fn)
+        return job
+
+    async def dequeue(
+        self,
+        entrypoint: str | None = None,
+        *,
+        batch_size: int = 1,
+    ) -> list[Job]:
+        """Claim up to batch_size QUEUED jobs and mark them IN_PROGRESS."""
+        claimed: list[Job] = []
+
+        def _fn(state: QueueState) -> QueueState:
+            nonlocal claimed
+            available = state.queued_jobs(entrypoint)[:batch_size]
+            claimed = []
+            new_state = state
+            for job in available:
+                updated = job.with_status(JobStatus.IN_PROGRESS).with_heartbeat(
+                    datetime.now(UTC)
+                )
+                new_state = new_state.with_job_replaced(updated)
+                claimed.append(updated)
+            return new_state
+
+        await self._submit(_fn)
+        return claimed
+
+    async def ack(self, job_id: str) -> None:
+        """Remove a completed job from the queue."""
+        await self._submit(lambda state: state.with_job_removed(job_id))
+
+    async def nack(self, job_id: str) -> None:
+        """Return a job to QUEUED status."""
+        from jqueue.domain.errors import JobNotFoundError
+
+        def _fn(state: QueueState) -> QueueState:
+            job = state.find(job_id)
+            if job is None:
+                raise JobNotFoundError(job_id)
+            return state.with_job_replaced(
+                job.with_status(JobStatus.QUEUED).with_heartbeat(None)
+            )
+
+        await self._submit(_fn)
+
+    async def heartbeat(self, job_id: str) -> None:
+        """Refresh the heartbeat timestamp for an IN_PROGRESS job."""
+        from jqueue.domain.errors import JobNotFoundError
+
+        def _fn(state: QueueState) -> QueueState:
+            job = state.find(job_id)
+            if job is None:
+                raise JobNotFoundError(job_id)
+            return state.with_job_replaced(job.with_heartbeat(datetime.now(UTC)))
+
+        await self._submit(_fn)
+
+    async def read_state(self) -> QueueState:
+        """Read-only snapshot of current queue state (bypasses the write pipeline)."""
+        content, _ = await self.storage.read()
+        return codec.decode(content)
+
+    # ------------------------------------------------------------------ #
+    # Internal machinery                                                   #
+    # ------------------------------------------------------------------ #
+
+    async def _submit(self, fn: MutationFn) -> None:
+        """
+        Enqueue a mutation and block until it is committed.
+
+        Appends the op to _pending, wakes the writer, then awaits the future
+        that resolves when the batch containing this op successfully commits.
+        """
+        if self._stopped:
+            raise JQueueError("GroupCommitLoop is stopped")
+        future: asyncio.Future[None] = asyncio.get_running_loop().create_future()
+        self._pending.append(_PendingOp(fn=fn, future=future))
+        self._wakeup.set()
+        await future
+
+    async def _writer_loop(self) -> None:
+        """Background coroutine — runs until stopped and all pending ops drain."""
+        while not self._stopped or self._pending:
+            if not self._pending:
+                self._wakeup.clear()
+                await self._wakeup.wait()
+
+            if not self._pending:
+                continue
+
+            batch = list(self._pending)
+            self._pending.clear()
+            await self._commit_batch(batch)
+
+    async def _commit_batch(self, batch: list[_PendingOp]) -> None:
+        """
+        Apply all ops in `batch` to the current state and CAS write.
+
+        Retries on CASConflictError. Per-mutation exceptions only fail that
+        op's future; the rest of the batch still commits on the same write.
+        """
+        for attempt in range(_MAX_RETRIES):
+            try:
+                content, etag = await self.storage.read()
+                state = codec.decode(content)
+
+                # Sweep stale jobs on every write cycle (free — no extra I/O)
+                cutoff = datetime.now(UTC) - self.stale_timeout
+                state = state.requeue_stale(cutoff)
+
+                per_op_errors: dict[int, Exception] = {}
+                for i, op in enumerate(batch):
+                    try:
+                        state = op.fn(state)
+                    except Exception as exc:
+                        per_op_errors[i] = exc
+
+                await self.storage.write(codec.encode(state), if_match=etag)
+
+                for i, op in enumerate(batch):
+                    if op.future.done():
+                        continue
+                    if i in per_op_errors:
+                        op.future.set_exception(per_op_errors[i])
+                    else:
+                        op.future.set_result(None)
+                return
+
+            except CASConflictError:
+                if attempt == _MAX_RETRIES - 1:
+                    _fail_batch(batch, CASConflictError("Max CAS retries exceeded"))
+                    return
+                # Exponential back-off, capped at ~320 ms
+                await asyncio.sleep(0.005 * (2 ** min(attempt, 6)))
+
+            except Exception as exc:
+                _fail_batch(batch, exc)
+                return
+
+
+def _fail_batch(batch: list[_PendingOp], exc: Exception) -> None:
+    """Set exception on all unresolved futures in the batch."""
+    for op in batch:
+        if not op.future.done():
+            op.future.set_exception(exc)

jqueue/core/heartbeat.py

@@ -0,0 +1,90 @@
+"""
+HeartbeatManager — async context manager for background heartbeats.
+
+A worker holding a long-running job wraps its work in HeartbeatManager
+to send periodic heartbeats, preventing the broker from re-queueing the
+job as stale.
+
+Usage
+-----
+    async with BrokerQueue(storage) as q:
+        [job] = await q.dequeue("process_video")
+
+        async with HeartbeatManager(q, job.id, interval=timedelta(seconds=30)):
+            result = await do_long_work(job.payload)
+
+        await q.ack(job.id)
+
+If the worker raises, the heartbeat task is cancelled. Callers should
+nack() the job in an except/finally block.
+
+HeartbeatManager is typed against the structural Protocol _HasHeartbeat, so
+it works with BrokerQueue, DirectQueue, and GroupCommitLoop without any
+shared base class.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+from datetime import timedelta
+from types import TracebackType
+from typing import Protocol
+
+from jqueue.domain.errors import JQueueError
+
+
+class _HasHeartbeat(Protocol):
+    """Structural Protocol — any object with an async heartbeat(job_id) method."""
+
+    async def heartbeat(self, job_id: str) -> None: ...
+
+
+@dataclasses.dataclass
+class HeartbeatManager:
+    """
+    Sends periodic heartbeats for a single job.
+
+    Parameters
+    ----------
+    queue    : any object with async heartbeat(job_id: str) -> None
+    job_id   : the job to keep alive
+    interval : time between heartbeat sends (default 60 seconds)
+    """
+
+    queue: _HasHeartbeat
+    job_id: str
+    interval: timedelta = timedelta(seconds=60)
+
+    _task: asyncio.Task[None] | None = dataclasses.field(
+        default=None, init=False, repr=False
+    )
+
+    async def __aenter__(self) -> HeartbeatManager:
+        self._task = asyncio.create_task(
+            self._beat(), name=f"jqueue-heartbeat-{self.job_id}"
+        )
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        if self._task is not None:
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+            self._task = None
+
+    async def _beat(self) -> None:
+        while True:
+            await asyncio.sleep(self.interval.total_seconds())
+            try:
+                await self.queue.heartbeat(self.job_id)
+            except JQueueError:
+                # Job was acked or removed — stop silently.
+                return

jqueue/domain/errors.py

@@ -0,0 +1,46 @@
+"""
+Exception hierarchy for jqueue.
+
+JQueueError
+├── CASConflictError   — write rejected because etag did not match
+├── JobNotFoundError   — job_id not present in current QueueState
+└── StorageError       — underlying I/O failure (wraps original exception)
+"""
+
+from __future__ import annotations
+
+
+class JQueueError(Exception):
+    """Base class for all jqueue exceptions."""
+
+
+class CASConflictError(JQueueError):
+    """
+    Raised when a compare-and-set write is rejected by the storage backend.
+
+    The caller should re-read the current state and retry the operation.
+    This is the normal concurrency signal — not an error in the traditional sense.
+    """
+
+
+class JobNotFoundError(JQueueError):
+    """Raised when a job_id is not present in the current QueueState."""
+
+    def __init__(self, job_id: str) -> None:
+        self.job_id = job_id
+        super().__init__(f"Job {job_id!r} not found in queue state")
+
+
+class StorageError(JQueueError):
+    """
+    Wraps an underlying I/O failure from a storage adapter.
+
+    Attributes
+    ----------
+    cause : Exception
+        The original exception from the storage backend.
+    """
+
+    def __init__(self, message: str, cause: Exception) -> None:
+        self.cause = cause
+        super().__init__(f"{message}: {cause}")