Support abort for deferred tasks (#307)

* Support abort for deferred tasks

* fix imports

* fix mypy

* fixes related to PR 307

* lint fixes

* lint fixes

* reversed changes

* fix linter

* fix avoiding race condition

* fix test

* fix lint

* fix lint

Author: antoine-collet

arq/jobs.py

@@ -156,7 +156,15 @@ async def abort(self, *, timeout: Optional[float] = None, poll_delay: float = 0.
         :param poll_delay: how often to poll redis for the job result
         :return: True if the job aborted properly, False otherwise
         """
+        job_info = await self.info()
+        if job_info and job_info.score and job_info.score > timestamp_ms():
+            async with self._redis.pipeline(transaction=True) as tr:
+                tr.zrem(self._queue_name, self.job_id)  # type: ignore[unused-coroutine]
+                tr.zadd(self._queue_name, {self.job_id: 1})  # type: ignore[unused-coroutine]
+                await tr.execute()
+
         await self._redis.zadd(abort_jobs_ss, {self.job_id: timestamp_ms()})
+
         try:
             await self.result(timeout=timeout, poll_delay=poll_delay)
         except asyncio.CancelledError:

tests/test_worker.py

@@ -4,6 +4,7 @@
 import re
 import signal
 import sys
+from datetime import datetime, timedelta
 from unittest.mock import MagicMock
 
 import msgpack
@@ -824,6 +825,34 @@ async def longfunc(ctx):
     assert worker.tasks == {}
 
 
+async def test_abort_deferred_job_before(arq_redis: ArqRedis, worker, caplog, loop):
+    async def longfunc(ctx):
+        await asyncio.sleep(3600)
+
+    caplog.set_level(logging.INFO)
+
+    job = await arq_redis.enqueue_job('longfunc', _job_id='testing', _defer_until=datetime.utcnow() + timedelta(days=1))
+
+    worker: Worker = worker(functions=[func(longfunc, name='longfunc')], allow_abort_jobs=True, poll_delay=0.1)
+    assert worker.jobs_complete == 0
+    assert worker.jobs_failed == 0
+    assert worker.jobs_retried == 0
+
+    with pytest.raises(asyncio.TimeoutError):
+        await job.abort(timeout=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 aborted before start' in log
+    await worker.main()
+    assert worker.aborting_tasks == set()
+    assert worker.job_tasks == {}
+    assert worker.tasks == {}
+
+
 async def test_not_abort_job(arq_redis: ArqRedis, worker, caplog, loop):
     async def shortfunc(ctx):
         await asyncio.sleep(0.2)