[MUS-937] Porting fixes from upstream - fix jobs timeout

arq/worker.py

@@ -525,10 +525,9 @@ async def job_failed(exc: BaseException) -> None:
             logger.info('%6.2fs → %s(%s)%s', (start_ms - enqueue_time_ms) / 1000, ref, s, extra)
             self.job_tasks[job_id] = task = self.loop.create_task(function.coroutine(ctx, *args, **kwargs))
 
-            cancel_handler = self.loop.call_at(self.loop.time() + timeout_s, task.cancel)
             # run repr(result) and extra inside try/except as they can raise exceptions
             try:
-                result = await task
+                result = await asyncio.wait_for(task, timeout_s)
             except (Exception, asyncio.CancelledError) as e:
                 exc_extra = getattr(e, 'extra', None)
                 if callable(exc_extra):
@@ -538,7 +537,6 @@ async def job_failed(exc: BaseException) -> None:
                 result_str = '' if result is None else truncate(repr(result))
             finally:
                 del self.job_tasks[job_id]
-                cancel_handler.cancel()
 
         except (Exception, asyncio.CancelledError) as e:
             finished_ms = timestamp_ms()

tests/test_worker.py

@@ -837,3 +837,21 @@ async def wait_and_abort(job, delay=0.1):
     assert worker.aborting_tasks == set()
     assert worker.tasks == {}
     assert worker.job_tasks == {}
+
+
+async def test_job_timeout(arq_redis: ArqRedis, worker, caplog):
+    async def longfunc(ctx):
+        await asyncio.sleep(0.3)
+
+    caplog.set_level(logging.ERROR)
+    await arq_redis.enqueue_job('longfunc', _job_id='testing')
+    worker: Worker = worker(functions=[func(longfunc, name='longfunc')], job_timeout=0.2, poll_delay=0.1)
+    assert worker.jobs_complete == 0
+    assert worker.jobs_failed == 0
+    assert worker.jobs_retried == 0
+    await worker.main()
+    assert worker.jobs_complete == 0
+    assert worker.jobs_failed == 1
+    assert worker.jobs_retried == 0
+    log = re.sub(r'\d+.\d\ds', 'X.XXs', '\n'.join(r.message for r in caplog.records))
+    assert 'X.XXs ! testing:longfunc failed, TimeoutError:' in log