Sync TaskTiger worker heartbeat (#331)

Use a thread to heartbeat in the sync worker

tasktiger/executor.py

@@ -63,6 +63,25 @@ def __init__(self, worker: "Worker"):
         self.connection = worker.connection
         self.config = worker.config
 
+    def heartbeat(
+        self,
+        queue: str,
+        task_ids: Collection[str],
+        log: BoundLogger,
+        locks: Collection[Lock],
+        queue_lock: Optional[Semaphore],
+    ) -> None:
+        self.worker.heartbeat(queue, task_ids)
+        for lock in locks:
+            try:
+                lock.reacquire()
+            except LockError:
+                log.warning("could not reacquire lock", lock=lock.name)
+        if queue_lock:
+            acquired, current_locks = queue_lock.renew()
+            if not acquired:
+                log.debug("queue lock renew failure")
+
     def execute(
         self,
         queue: str,
@@ -351,18 +370,7 @@ def check_child_exit() -> Optional[int]:
                     break
 
                 try:
-                    self.worker.heartbeat(queue, all_task_ids)
-                    for lock in locks:
-                        try:
-                            lock.reacquire()
-                        except LockError:
-                            log.warning(
-                                "could not reacquire lock", lock=lock.name
-                            )
-                    if queue_lock:
-                        acquired, current_locks = queue_lock.renew()
-                        if not acquired:
-                            log.debug("queue lock renew failure")
+                    self.heartbeat(queue, all_task_ids, log, locks, queue_lock)
                 except OSError as e:
                     # EINTR happens if the task completed. Since we're just
                     # renewing locks/heartbeat it's okay if we get interrupted.
@@ -386,6 +394,19 @@ class SyncExecutor(Executor):
 
     exit_worker_on_job_timeout = True
 
+    def _periodic_heartbeat(
+        self,
+        queue: str,
+        task_ids: Collection[str],
+        log: BoundLogger,
+        locks: Collection[Lock],
+        queue_lock: Optional[Semaphore],
+        stop_event: threading.Event,
+    ) -> None:
+        while not stop_event.is_set():
+            stop_event.wait(self.config["ACTIVE_TASK_UPDATE_TIMER"])
+            self.heartbeat(queue, task_ids, log, locks, queue_lock)
+
     def execute(
         self,
         queue: str,
@@ -394,5 +415,27 @@ def execute(
         locks: Collection[Lock],
         queue_lock: Optional[Semaphore],
     ) -> bool:
+        # Run heartbeat thread.
+        all_task_ids = {task.id for task in tasks}
+        stop_event = threading.Event()
+        heartbeat_thread = threading.Thread(
+            target=self._periodic_heartbeat,
+            kwargs={
+                "queue": queue,
+                "task_ids": all_task_ids,
+                "log": log,
+                "locks": locks,
+                "queue_lock": queue_lock,
+                "stop_event": stop_event,
+            },
+        )
+        heartbeat_thread.start()
+
         # Run the tasks.
-        return self.execute_tasks(tasks, log)
+        result = self.execute_tasks(tasks, log)
+
+        # Stop the heartbeat thread.
+        stop_event.set()
+        heartbeat_thread.join()
+
+        return result

tests/test_workers.py

@@ -8,13 +8,16 @@
 from freezefrog import FreezeTime
 
 from tasktiger import Task, Worker
+from tasktiger._internal import ACTIVE
 from tasktiger.executor import SyncExecutor
 
+from .config import DELAY
 from .tasks import (
     exception_task,
     long_task_killed,
     long_task_ok,
     simple_task,
+    sleep_task,
     wait_for_long_task,
 )
 from .test_base import BaseTestCase
@@ -35,10 +38,12 @@ def test_max_workers(self):
 
         # Start two workers and wait until they start processing.
         worker1 = Process(
-            target=external_worker, kwargs={"max_workers_per_queue": 2}
+            target=external_worker,
+            kwargs={"worker_kwargs": {"max_workers_per_queue": 2}},
         )
         worker2 = Process(
-            target=external_worker, kwargs={"max_workers_per_queue": 2}
+            target=external_worker,
+            kwargs={"worker_kwargs": {"max_workers_per_queue": 2}},
         )
         worker1.start()
         worker2.start()
@@ -181,3 +186,28 @@ def test_handles_timeout(self, tiger, ensure_queues):
         with pytest.raises(SystemExit):
             worker.run(once=True, force_once=True)
         ensure_queues(error={"default": 1})
+
+    def test_heartbeat(self, tiger):
+        task = Task(tiger, sleep_task)
+        task.delay()
+
+        # Start a worker and wait until it starts processing.
+        worker = Process(
+            target=external_worker,
+            kwargs={
+                "patch_config": {"ACTIVE_TASK_UPDATE_TIMER": DELAY / 2},
+                "worker_kwargs": {"executor_class": SyncExecutor},
+            },
+        )
+        worker.start()
+
+        time.sleep(DELAY / 2)
+
+        key = tiger._key(ACTIVE, "default")
+        conn = tiger.connection
+        heartbeat_1 = conn.zscore(key, task.id)
+        time.sleep(DELAY / 2)
+        heartbeat_2 = conn.zscore(key, task.id)
+        assert heartbeat_2 > heartbeat_1 > 0
+
+        worker.kill()

tests/utils.py

@@ -70,7 +70,7 @@ def get_tiger():
     return tiger
 
 
-def external_worker(n=None, patch_config=None, max_workers_per_queue=None):
+def external_worker(n=None, patch_config=None, worker_kwargs=None):
     """
     Runs a worker. To be used with multiprocessing.Pool.map.
     """
@@ -79,12 +79,11 @@ def external_worker(n=None, patch_config=None, max_workers_per_queue=None):
     if patch_config:
         tiger.config.update(patch_config)
 
-    worker = Worker(tiger)
+    if worker_kwargs is None:
+        worker_kwargs = {}
 
-    if max_workers_per_queue is not None:
-        worker.max_workers_per_queue = max_workers_per_queue
+    Worker(tiger, **worker_kwargs).run(once=True, force_once=True)
 
-    worker.run(once=True, force_once=True)
     tiger.connection.close()