pydocket 0.0.2__py3-none-any.whl → 0.1.0__py3-none-any.whl

docket/worker.py

@@ -9,7 +9,6 @@ from typing import (
     Any,
     Protocol,
     Self,
-    Sequence,
     TypeVar,
     cast,
 )
@@ -18,10 +17,17 @@ from uuid import uuid4
 import redis.exceptions
 from opentelemetry import propagate, trace
 from opentelemetry.trace import Tracer
-from redis import RedisError
 
-from .docket import Docket, Execution
+from .docket import (
+    Docket,
+    Execution,
+    RedisMessage,
+    RedisMessageID,
+    RedisMessages,
+    RedisReadGroupResponse,
+)
 from .instrumentation import (
+    REDIS_DISRUPTIONS,
     TASK_DURATION,
     TASK_PUNCTUALITY,
     TASKS_COMPLETED,
@@ -29,6 +35,7 @@ from .instrumentation import (
     TASKS_RETRIED,
     TASKS_RUNNING,
     TASKS_STARTED,
+    TASKS_STRICKEN,
     TASKS_SUCCEEDED,
     message_getter,
 )
@@ -37,12 +44,6 @@ logger: logging.Logger = logging.getLogger(__name__)
 tracer: Tracer = trace.get_tracer(__name__)
 
 
-RedisStreamID = bytes
-RedisMessageID = bytes
-RedisMessage = dict[bytes, bytes]
-RedisStream = tuple[RedisStreamID, Sequence[tuple[RedisMessageID, RedisMessage]]]
-RedisReadGroupResponse = Sequence[RedisStream]
-
 if TYPE_CHECKING:  # pragma: no cover
     from .dependencies import Dependency
 
@@ -63,28 +64,20 @@ class Worker:
         self,
         docket: Docket,
         name: str | None = None,
-        prefetch_count: int = 10,
+        concurrency: int = 10,
         redelivery_timeout: timedelta = timedelta(minutes=5),
         reconnection_delay: timedelta = timedelta(seconds=5),
+        minimum_check_interval: timedelta = timedelta(milliseconds=10),
     ) -> None:
         self.docket = docket
         self.name = name or f"worker:{uuid4()}"
-        self.prefetch_count = prefetch_count
+        self.concurrency = concurrency
         self.redelivery_timeout = redelivery_timeout
         self.reconnection_delay = reconnection_delay
+        self.minimum_check_interval = minimum_check_interval
 
     async def __aenter__(self) -> Self:
-        async with self.docket.redis() as redis:
-            try:
-                await redis.xgroup_create(
-                    groupname=self.consumer_group_name,
-                    name=self.docket.stream_key,
-                    id="0-0",
-                    mkstream=True,
-                )
-            except RedisError as e:
-                if "BUSYGROUP" not in repr(e):
-                    raise
+        self._heartbeat_task = asyncio.create_task(self._heartbeat())
 
         return self
 
@@ -94,11 +87,12 @@ class Worker:
         exc_value: BaseException | None,
         traceback: TracebackType | None,
     ) -> None:
-        pass
-
-    @property
-    def consumer_group_name(self) -> str:
-        return "docket"
+        self._heartbeat_task.cancel()
+        try:
+            await self._heartbeat_task
+        except asyncio.CancelledError:
+            pass
+        del self._heartbeat_task
 
     @property
     def _log_context(self) -> dict[str, str]:
@@ -113,7 +107,7 @@ class Worker:
         docket_name: str = "docket",
         url: str = "redis://localhost:6379/0",
         name: str | None = None,
-        prefetch_count: int = 10,
+        concurrency: int = 10,
         redelivery_timeout: timedelta = timedelta(minutes=5),
         reconnection_delay: timedelta = timedelta(seconds=5),
         until_finished: bool = False,
@@ -126,7 +120,7 @@ class Worker:
             async with Worker(
                 docket=docket,
                 name=name,
-                prefetch_count=prefetch_count,
+                concurrency=concurrency,
                 redelivery_timeout=redelivery_timeout,
                 reconnection_delay=reconnection_delay,
             ) as worker:
@@ -153,6 +147,9 @@ class Worker:
             try:
                 return await self._worker_loop(forever=forever)
             except redis.exceptions.ConnectionError:
+                REDIS_DISRUPTIONS.add(
+                    1, {"docket": self.docket.name, "worker": self.name}
+                )
                 logger.warning(
                     "Error connecting to redis, retrying in %s...",
                     self.reconnection_delay,
@@ -207,67 +204,121 @@ class Worker:
                 ),
             )
 
-            total_work, due_work = sys.maxsize, 0
-            while forever or total_work:
-                now = datetime.now(timezone.utc)
-                total_work, due_work = await stream_due_tasks(
-                    keys=[self.docket.queue_key, self.docket.stream_key],
-                    args=[now.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,
-                    )
+            active_tasks: dict[asyncio.Task[None], RedisMessageID] = {}
 
-                _, redeliveries, _ = await redis.xautoclaim(
-                    name=self.docket.stream_key,
-                    groupname=self.consumer_group_name,
-                    consumername=self.name,
-                    min_idle_time=int(self.redelivery_timeout.total_seconds() * 1000),
-                    start_id="0-0",
-                    count=self.prefetch_count,
-                )
+            async def process_completed_tasks() -> None:
+                completed_tasks = {task for task in active_tasks if task.done()}
+                for task in completed_tasks:
+                    message_id = active_tasks.pop(task)
 
-                new_deliveries: RedisReadGroupResponse = await redis.xreadgroup(
-                    groupname=self.consumer_group_name,
-                    consumername=self.name,
-                    streams={self.docket.stream_key: ">"},
-                    count=self.prefetch_count,
-                    block=10,
-                )
+                    await task
+
+                    async with redis.pipeline() as pipeline:
+                        pipeline.xack(
+                            self.docket.stream_key,
+                            self.docket.worker_group_name,
+                            message_id,
+                        )
+                        pipeline.xdel(
+                            self.docket.stream_key,
+                            message_id,
+                        )
+                        await pipeline.execute()
+
+            future_work, due_work = sys.maxsize, 0
+
+            try:
+                while forever or future_work or active_tasks:
+                    await process_completed_tasks()
+
+                    available_slots = self.concurrency - len(active_tasks)
+
+                    def start_task(
+                        message_id: RedisMessageID, message: RedisMessage
+                    ) -> None:
+                        task = asyncio.create_task(self._execute(message))
+                        active_tasks[task] = message_id
+
+                        nonlocal available_slots, future_work
+                        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,
+                        groupname=self.docket.worker_group_name,
+                        consumername=self.name,
+                        min_idle_time=int(
+                            self.redelivery_timeout.total_seconds() * 1000
+                        ),
+                        start_id="0-0",
+                        count=available_slots,
+                    )
 
-                for source in [[(b"redeliveries", redeliveries)], new_deliveries]:
-                    for _, messages in source:
+                    for message_id, message in redeliveries:
+                        start_task(message_id, message)
+                        if available_slots <= 0:
+                            break
+
+                    if available_slots <= 0:
+                        continue
+
+                    new_deliveries: RedisReadGroupResponse = await redis.xreadgroup(
+                        groupname=self.docket.worker_group_name,
+                        consumername=self.name,
+                        streams={self.docket.stream_key: ">"},
+                        block=(
+                            int(self.minimum_check_interval.total_seconds() * 1000)
+                            if forever or active_tasks
+                            else None
+                        ),
+                        count=available_slots,
+                    )
+                    for _, messages in new_deliveries:
                         for message_id, message in messages:
-                            await self._execute(message)
-
-                            async with redis.pipeline() as pipeline:
-                                pipeline.xack(
-                                    self.docket.stream_key,
-                                    self.consumer_group_name,
-                                    message_id,
-                                )
-                                pipeline.xdel(
-                                    self.docket.stream_key,
-                                    message_id,
-                                )
-                                await pipeline.execute()
-
-                            # When executing a task, there's always a chance that it was
-                            # either retried or it scheduled another task, so let's give
-                            # ourselves one more iteration of the loop to handle that.
-                            total_work += 1
+                            start_task(message_id, message)
+                            if available_slots <= 0:
+                                break
+            except asyncio.CancelledError:
+                if active_tasks:  # pragma: no cover
+                    logger.info(
+                        "Shutdown requested, finishing %d active tasks...",
+                        len(active_tasks),
+                        extra=self._log_context,
+                    )
+            finally:
+                if active_tasks:
+                    await asyncio.gather(*active_tasks, return_exceptions=True)
+                    await process_completed_tasks()
 
     async def _execute(self, message: RedisMessage) -> None:
-        execution = Execution.from_message(
-            self.docket.tasks[message[b"function"].decode()],
-            message,
-        )
+        function_name = message[b"function"].decode()
+        function = self.docket.tasks.get(function_name)
+        if function is None:
+            logger.warning(
+                "Task function %r not found", function_name, extra=self._log_context
+            )
+            return
+
+        execution = Execution.from_message(function, message)
         name = execution.function.__name__
         key = execution.key
 
@@ -282,6 +333,15 @@ class Worker:
             "task": name,
         }
 
+        arrow = "↬" if execution.attempt > 1 else "↪"
+        call = execution.call_repr()
+
+        if self.docket.strike_list.is_stricken(execution):
+            arrow = "🗙"
+            logger.warning("%s %s", arrow, call, extra=log_context)
+            TASKS_STRICKEN.add(1, counter_labels | {"where": "worker"})
+            return
+
         dependencies = self._get_dependencies(execution)
 
         context = propagate.extract(message, getter=message_getter)
@@ -297,8 +357,6 @@ class Worker:
         TASKS_RUNNING.add(1, counter_labels)
         TASK_PUNCTUALITY.record(punctuality.total_seconds(), counter_labels)
 
-        arrow = "↬" if execution.attempt > 1 else "↪"
-        call = execution.call_repr()
         logger.info("%s [%s] %s", arrow, punctuality, call, extra=log_context)
 
         try:
@@ -387,3 +445,52 @@ class Worker:
             return True
 
         return False
+
+    @property
+    def workers_set(self) -> str:
+        return self.docket.workers_set
+
+    def worker_tasks_set(self, worker_name: str) -> str:
+        return self.docket.worker_tasks_set(worker_name)
+
+    def task_workers_set(self, task_name: str) -> str:
+        return self.docket.task_workers_set(task_name)
+
+    async def _heartbeat(self) -> None:
+        while True:
+            await asyncio.sleep(self.docket.heartbeat_interval.total_seconds())
+            try:
+                now = datetime.now(timezone.utc).timestamp()
+                maximum_age = (
+                    self.docket.heartbeat_interval * self.docket.missed_heartbeats
+                )
+                oldest = now - maximum_age.total_seconds()
+
+                task_names = list(self.docket.tasks)
+
+                async with self.docket.redis() as r:
+                    async with r.pipeline() as pipeline:
+                        pipeline.zremrangebyscore(self.workers_set, 0, oldest)
+                        pipeline.zadd(self.workers_set, {self.name: now})
+
+                        for task_name in task_names:
+                            task_workers_set = self.task_workers_set(task_name)
+                            pipeline.zremrangebyscore(task_workers_set, 0, oldest)
+                            pipeline.zadd(task_workers_set, {self.name: now})
+
+                        pipeline.sadd(self.worker_tasks_set(self.name), *task_names)
+                        pipeline.expire(
+                            self.worker_tasks_set(self.name),
+                            max(maximum_age, timedelta(seconds=1)),
+                        )
+
+                        await pipeline.execute()
+            except asyncio.CancelledError:  # pragma: no cover
+                return
+            except redis.exceptions.ConnectionError:
+                REDIS_DISRUPTIONS.add(
+                    1, {"docket": self.docket.name, "worker": self.name}
+                )
+                logger.exception("Error sending worker heartbeat", exc_info=True)
+            except Exception:
+                logger.exception("Error sending worker heartbeat", exc_info=True)