Add task timeout (#52)

* Add timeout task param

* Add nice log for when timeouts happen

* Add tests for task timeout

* Implement review comments, better events data structure

Author: AlanCoding

dispatcher/main.py

@@ -71,6 +71,13 @@ async def alive(self, dispatcher, **data):
         return
 
 
+class DispatcherEvents:
+    "Benchmark tests have to re-create this because they use same object in different event loops"
+
+    def __init__(self) -> None:
+        self.exit_event: asyncio.Event = asyncio.Event()
+
+
 class DispatcherMain:
     def __init__(self, config: dict):
         self.delayed_messages: list[SimpleNamespace] = []
@@ -97,11 +104,7 @@ def __init__(self, config: dict):
             if 'scheduled' in producer_config:
                 self.producers.append(ScheduledProducer(producer_config['scheduled']))
 
-        self.events = self._create_events()
-
-    def _create_events(self):
-        "Benchmark tests have to re-create this because they use same object in different event loops"
-        return SimpleNamespace(exit_event=asyncio.Event())
+        self.events: DispatcherEvents = DispatcherEvents()
 
     def fatal_error_callback(self, *args) -> None:
         """Method to connect to error callbacks of other tasks, will kick out of main loop"""

dispatcher/pool.py

@@ -3,8 +3,8 @@
 import multiprocessing
 import os
 import signal
+import time
 from asyncio import Task
-from types import SimpleNamespace
 from typing import Iterator, Optional
 
 from dispatcher.utils import DuplicateBehavior, MessageAction
@@ -19,18 +19,28 @@ def __init__(self, worker_id: int, finished_queue: multiprocessing.Queue):
         # TODO: rename message_queue to call_queue, because this is what cpython ProcessPoolExecutor calls them
         self.message_queue: multiprocessing.Queue = multiprocessing.Queue()
         self.process = multiprocessing.Process(target=work_loop, args=(self.worker_id, self.message_queue, finished_queue))
+
+        # Info specific to the current task being ran
         self.current_task: Optional[dict] = None
+        self.started_at: Optional[int] = None
+        self.is_active_cancel: bool = False
+
+        # Tracking information for worker
         self.finished_count = 0
         self.status = 'initialized'
         self.exit_msg_event = asyncio.Event()
-        self.active_cancel = False
 
     async def start(self) -> None:
         self.status = 'spawned'
         self.process.start()
         logger.debug(f'Worker {self.worker_id} pid={self.process.pid} subprocess has spawned')
         self.status = 'starting'  # Not ready until it sends callback message
 
+    async def start_task(self, message: dict) -> None:
+        self.current_task = message  # NOTE: this marks this worker as busy
+        self.message_queue.put(message)
+        self.started_at = time.monotonic_ns()
+
     async def join(self) -> None:
         logger.debug(f'Joining worker {self.worker_id} pid={self.process.pid} subprocess')
         self.process.join()
@@ -65,12 +75,13 @@ async def stop(self) -> None:
         return
 
     def cancel(self) -> None:
-        self.active_cancel = True  # signal for result callback
+        self.is_active_cancel = True  # signal for result callback
         self.process.terminate()  # SIGTERM
 
     def mark_finished_task(self) -> None:
-        self.active_cancel = False
+        self.is_active_cancel = False
         self.current_task = None
+        self.started_at = None
         self.finished_count += 1
 
     @property
@@ -79,6 +90,16 @@ def inactive(self) -> bool:
         return self.status in ['exited', 'error', 'initialized']
 
 
+class PoolEvents:
+    "Benchmark tests have to re-create this because they use same object in different event loops"
+
+    def __init__(self) -> None:
+        self.queue_cleared: asyncio.Event = asyncio.Event()  # queue is now 0 length
+        self.work_cleared: asyncio.Event = asyncio.Event()  # Totally quiet, no blocked or queued messages, no busy workers
+        self.management_event: asyncio.Event = asyncio.Event()  # Process spawning is backgrounded, so this is the kicker
+        self.timeout_event: asyncio.Event = asyncio.Event()  # Anything that might affect the timeout watcher task
+
+
 class WorkerPool:
     def __init__(self, num_workers: int, fd_lock: Optional[asyncio.Lock] = None):
         self.num_workers = num_workers
@@ -97,7 +118,7 @@ def __init__(self, num_workers: int, fd_lock: Optional[asyncio.Lock] = None):
         self.management_lock = asyncio.Lock()
         self.fd_lock = fd_lock or asyncio.Lock()
 
-        self.events = self._create_events()
+        self.events: PoolEvents = PoolEvents()
 
     @property
     def processed_count(self):
@@ -107,19 +128,13 @@ def processed_count(self):
     def received_count(self):
         return self.processed_count + len(self.queued_messages) + sum(1 for w in self.workers.values() if w.current_task)
 
-    def _create_events(self):
-        "Benchmark tests have to re-create this because they use same object in different event loops"
-        return SimpleNamespace(
-            queue_cleared=asyncio.Event(),  # queue is now 0 length
-            work_cleared=asyncio.Event(),  # Totally quiet, no blocked or queued messages, no busy workers
-            management_event=asyncio.Event(),  # Process spawning is backgrounded, so this is the kicker
-        )
-
     async def start_working(self, dispatcher) -> None:
         self.read_results_task = asyncio.create_task(self.read_results_forever(), name='results_task')
         self.read_results_task.add_done_callback(dispatcher.fatal_error_callback)
         self.management_task = asyncio.create_task(self.manage_workers(), name='management_task')
         self.management_task.add_done_callback(dispatcher.fatal_error_callback)
+        self.timeout_task = asyncio.create_task(self.manage_timeout(), name='timeout_task')
+        self.timeout_task.add_done_callback(dispatcher.fatal_error_callback)
 
     async def manage_workers(self) -> None:
         """Enforces worker policy like min and max workers, and later, auto scale-down"""
@@ -140,6 +155,43 @@ async def manage_workers(self) -> None:
             self.events.management_event.clear()
         logger.debug('Pool worker management task exiting')
 
+    async def process_worker_timeouts(self, current_time: float) -> Optional[int]:
+        """
+        Cancels tasks that have exceeded their timeout.
+        Returns the system clock time of the next task timeout, for rescheduling.
+        """
+        next_deadline = None
+        for worker in self.workers.values():
+            if (not worker.is_active_cancel) and worker.current_task and worker.started_at and (worker.current_task.get('timeout')):
+                timeout: float = worker.current_task['timeout']
+                worker_deadline = worker.started_at + int(timeout * 1.0e9)
+
+                # Established that worker is running a task that has a timeout
+                if worker_deadline < current_time:
+                    uuid: str = worker.current_task.get('uuid', '<unknown>')
+                    delta: float = (current_time - worker.started_at) * 1.0e9
+                    logger.info(f'Worker {worker.worker_id} runtime {delta:.5f}(s) for task uuid={uuid} exceeded timeout {timeout}(s), canceling')
+                    worker.cancel()
+                elif next_deadline is None or worker_deadline < next_deadline:
+                    # worker timeout is closer than any yet seen
+                    next_deadline = worker_deadline
+
+        return next_deadline
+
+    async def manage_timeout(self) -> None:
+        while not self.shutting_down:
+            current_time = time.monotonic_ns()
+            pool_deadline = await self.process_worker_timeouts(current_time)
+            if pool_deadline:
+                time_until_deadline = (pool_deadline - current_time) * 1.0e-9
+                try:
+                    await asyncio.wait_for(self.events.timeout_event.wait(), timeout=time_until_deadline)
+                except asyncio.TimeoutError:
+                    pass  # will handle in next loop run
+            else:
+                await self.events.timeout_event.wait()
+            self.events.timeout_event.clear()
+
     async def up(self) -> None:
         worker = PoolWorker(worker_id=self.next_worker_id, finished_queue=self.finished_queue)
         self.workers[self.next_worker_id] = worker
@@ -166,6 +218,7 @@ async def force_shutdown(self) -> None:
     async def shutdown(self) -> None:
         self.shutting_down = True
         self.events.management_event.set()
+        self.events.timeout_event.set()
         await self.stop_workers()
         self.finished_queue.put('stop')
 
@@ -277,8 +330,9 @@ async def dispatch_task(self, message: dict) -> None:
 
             if worker := self.get_free_worker():
                 logger.debug(f"Dispatching task (uuid={uuid}) to worker (id={worker.worker_id})")
-                worker.current_task = message  # NOTE: this marks the worker as busy
-                worker.message_queue.put(message)
+                await worker.start_task(message)
+                if 'timeout' in message:
+                    self.events.timeout_event.set()  # kick timeout task to set wakeup
             else:
                 logger.warning(f'Queueing task (uuid={uuid}), ran out of workers, queued_ct={len(self.queued_messages)}')
                 self.queued_messages.append(message)
@@ -302,7 +356,7 @@ async def process_finished(self, worker, message) -> None:
         result = None
         if message.get("result"):
             result = message["result"]
-            if worker.active_cancel:
+            if worker.is_active_cancel:
                 msg += ', expected cancel'
             if result == '<cancel>':
                 msg += ', canceled'
@@ -312,7 +366,7 @@ async def process_finished(self, worker, message) -> None:
 
         # Mark the worker as no longer busy
         async with self.management_lock:
-            if worker.active_cancel and result == '<cancel>':
+            if worker.is_active_cancel and result == '<cancel>':
                 self.canceled_count += 1
             elif 'control' in worker.current_task:
                 self.control_count += 1
@@ -323,6 +377,9 @@ async def process_finished(self, worker, message) -> None:
         if not self.queued_messages and all(worker.current_task is None for worker in self.workers.values()):
             self.events.work_cleared.set()
 
+        if 'timeout' in message:
+            self.events.timeout_event.set()
+
     async def read_results_forever(self) -> None:
         """Perpetual task that continuously waits for task completions."""
         loop = asyncio.get_event_loop()

dispatcher/producers/base.py

@@ -1,7 +1,10 @@
 import asyncio
-from types import SimpleNamespace
+
+
+class ProducerEvents:
+    def __init__(self):
+        self.ready_event = asyncio.Event()
 
 
 class BaseProducer:
-    def _create_events(self):
-        return SimpleNamespace(ready_event=asyncio.Event())
+    pass

dispatcher/producers/brokered.py

@@ -3,14 +3,14 @@
 from typing import Optional
 
 from dispatcher.brokers.pg_notify import aget_connection, aprocess_notify, apublish_message
-from dispatcher.producers.base import BaseProducer
+from dispatcher.producers.base import BaseProducer, ProducerEvents
 
 logger = logging.getLogger(__name__)
 
 
 class BrokeredProducer(BaseProducer):
     def __init__(self, broker: str = 'pg_notify', config: Optional[dict] = None, channels: tuple = (), connection=None) -> None:
-        self.events = self._create_events()
+        self.events = ProducerEvents()
         self.production_task: Optional[asyncio.Task] = None
         self.broker = broker
         self.config = config

dispatcher/producers/scheduled.py

@@ -1,14 +1,14 @@
 import asyncio
 import logging
 
-from dispatcher.producers.base import BaseProducer
+from dispatcher.producers.base import BaseProducer, ProducerEvents
 
 logger = logging.getLogger(__name__)
 
 
 class ScheduledProducer(BaseProducer):
     def __init__(self, task_schedule: dict):
-        self.events = self._create_events()
+        self.events = ProducerEvents()
         self.task_schedule = task_schedule
         self.scheduled_tasks: list[asyncio.Task] = []
         self.produced_count = 0

dispatcher/publish.py

@@ -9,15 +9,18 @@
 
 
 class DispatcherDecorator:
-    def __init__(self, registry: DispatcherMethodRegistry, *, queue: Optional[str] = None, on_duplicate: Optional[str] = None) -> None:
+    def __init__(
+        self, registry: DispatcherMethodRegistry, *, queue: Optional[str] = None, on_duplicate: Optional[str] = None, timeout: Optional[float] = None
+    ) -> None:
         self.registry = registry
         self.queue = queue
         self.on_duplicate = on_duplicate
+        self.timeout = timeout
 
     def __call__(self, fn: DispatcherCallable, /) -> DispatcherCallable:
         "Concrete task decorator, registers method and glues on some methods from the registry"
 
-        dmethod = self.registry.register(fn, queue=self.queue, on_duplicate=self.on_duplicate)
+        dmethod = self.registry.register(fn, queue=self.queue, on_duplicate=self.on_duplicate, timeout=self.timeout)
 
         setattr(fn, 'apply_async', dmethod.apply_async)
         setattr(fn, 'delay', dmethod.delay)
@@ -29,6 +32,7 @@ def task(
     *,
     queue: Optional[str] = None,
     on_duplicate: Optional[str] = None,
+    timeout: Optional[float] = None,
     registry: DispatcherMethodRegistry = default_registry,
 ) -> DispatcherDecorator:
     """
@@ -68,4 +72,4 @@ def announce():
     # The on_duplicate kwarg controls behavior when multiple instances of the task running
     # options are documented in dispatcher.utils.DuplicateBehavior
     """
-    return DispatcherDecorator(registry, queue=queue, on_duplicate=on_duplicate)
+    return DispatcherDecorator(registry, queue=queue, on_duplicate=on_duplicate, timeout=timeout)

dispatcher/registry.py

@@ -43,15 +43,18 @@ def get_callable(self) -> Callable:
         return self.fn
 
     def publication_defaults(self) -> dict:
-        defaults = self.submission_defaults.copy()
+        defaults = {}
+        for k, v in self.submission_defaults.items():
+            if v:  # all None or falsy values have no effect
+                defaults[k] = v
         defaults['task'] = self.serialize_task()
         defaults['time_pub'] = time.time()
         return defaults
 
     def delay(self, *args, **kwargs) -> Tuple[dict, str]:
         return self.apply_async(args, kwargs)
 
-    def get_async_body(self, args=None, kwargs=None, uuid=None, on_duplicate: Optional[str] = None, delay: float = 0.0) -> dict:
+    def get_async_body(self, args=None, kwargs=None, uuid=None, on_duplicate: Optional[str] = None, timeout: Optional[float] = 0.0, delay: float = 0.0) -> dict:
         """
         Get the python dict to become JSON data in the pg_notify message
         This same message gets passed over the dispatcher IPC queue to workers
@@ -67,6 +70,8 @@ def get_async_body(self, args=None, kwargs=None, uuid=None, on_duplicate: Option
             body['on_duplicate'] = on_duplicate
         if delay:
             body['delay'] = delay
+        if timeout:
+            body['timeout'] = timeout
 
         return body