DM-45138: Add support for aborting arq jobs

changelog.d/20240711_120038_rra_DM_45138.md

@@ -0,0 +1,3 @@
+### New features
+
+- Add new `abort_job` method to instances of `safir.arq.ArqQueue`, which tells arq to abort a job that has been queued or in progress. To successfully abort a job that has already started, the arq worker must enable support for aborting jobs.

docs/user-guide/arq.rst

@@ -171,6 +171,9 @@ See `arq.worker.Worker` for details.
 The ``on_startup`` and ``on_shutdown`` handlers are ideal places to set up (and tear down) worker state, including network and database clients.
 The context variable, ``ctx``, passed to these functions are also passed to the worker functions.
 
+If you want to allow jobs to be aborted, add ``allow_abort_jobs = True`` to the ``WorkerSettings`` class.
+If a job is already running when it is aborted, it will be cancelled using asyncio task cancellation, which means that `asyncio.CancelledError` will be raised inside the job at the next opportunity.
+
 To run a worker, you run your application's Docker image with the ``arq`` command, followed by the fully-qualified namespace of the ``WorkerSettings`` class.
 
 Using the arq dependency in endpoint handlers
@@ -230,6 +233,17 @@ The `safir.dependencies.arq.arq_dependency` dependency provides your FastAPI end
 
         return response
 
+
+    @app.delete("/jobs/{job_id}", status_code=204)
+    async def delete_job(
+        job_id: str,
+        arq_queue: Annotated[ArqQueue, Depends(arq_dependency)],
+    ) -> None:
+        # This will only work if allow_abort_jobs is set to True in the worker
+        # configuration.
+        if not await arq_queue.abort_job(job_id):
+            raise HTTPException(status_code=404)
+
 For information on the metadata available from jobs, see `JobMetadata` and `JobResult`.
 
 Testing applications with an arq queue
@@ -238,7 +252,7 @@ Testing applications with an arq queue
 Unit testing an application with a running distributed queue is difficult since three components (two instances of the application and a redis database) must coordinate.
 A better unit testing approach is to test the front-end application separately from the worker functions.
 To help you do this, the arq dependency allows you to run a mocked version of an arq queue.
-With the mocked client, your front-end application can run the three basic client methods as normal: `ArqQueue.enqueue`, `ArqQueue.get_job_metadata`, and `ArqQueue.get_job_result`).
+With the mocked client, your front-end application can run the four basic client methods as normal: `ArqQueue.enqueue`, `ArqQueue.abort_job`, `ArqQueue.get_job_metadata`, and `ArqQueue.get_job_result`).
 This mocked client is a subclass of `ArqQueue` called `MockArqQueue`.
 
 Configuring the test mode

src/safir/arq.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import abc
+import asyncio
 import uuid
 from dataclasses import dataclass
 from datetime import datetime
@@ -326,6 +327,31 @@ async def enqueue(
         """
         raise NotImplementedError
 
+    @abc.abstractmethod
+    async def abort_job(
+        self,
+        job_id: str,
+        queue_name: str | None = None,
+        *,
+        timeout: float | None = None,
+    ) -> bool:
+        """Abort a queued or running job.
+
+        The worker must be configured to allow aborting jobs for this to
+        succeed.
+
+        Parameters
+        ----------
+        job_id
+            The job's identifier.
+        queue_name
+            Name of the queue.
+        timeout
+            How long to wait for the job result before raising `TimeoutError`.
+            If `None`, waits forever.
+        """
+        raise NotImplementedError
+
     @abc.abstractmethod
     async def get_job_metadata(
         self, job_id: str, queue_name: str | None = None
@@ -433,6 +459,16 @@ def _get_job(self, job_id: str, queue_name: str | None = None) -> Job:
             _queue_name=queue_name or self.default_queue_name,
         )
 
+    async def abort_job(
+        self,
+        job_id: str,
+        queue_name: str | None = None,
+        *,
+        timeout: float | None = None,
+    ) -> bool:
+        job = self._get_job(job_id, queue_name=queue_name)
+        return await job.abort(timeout=timeout)
+
     async def get_job_metadata(
         self, job_id: str, queue_name: str | None = None
     ) -> JobMetadata:
@@ -483,6 +519,46 @@ async def enqueue(
         self._job_metadata[queue_name][new_job.id] = new_job
         return new_job
 
+    async def abort_job(
+        self,
+        job_id: str,
+        queue_name: str | None = None,
+        *,
+        timeout: float | None = None,
+    ) -> bool:
+        queue_name = self._resolve_queue_name(queue_name)
+        try:
+            job_metadata = self._job_metadata[queue_name][job_id]
+        except KeyError:
+            return False
+
+        # If the job was started, simulate cancelling it.
+        if job_metadata.status == JobStatus.in_progress:
+            job_metadata.status = JobStatus.complete
+            result_info = JobResult(
+                id=job_metadata.id,
+                name=job_metadata.name,
+                args=job_metadata.args,
+                kwargs=job_metadata.kwargs,
+                status=job_metadata.status,
+                enqueue_time=job_metadata.enqueue_time,
+                start_time=current_datetime(microseconds=True),
+                finish_time=current_datetime(microseconds=True),
+                result=asyncio.CancelledError(),
+                success=False,
+                queue_name=queue_name,
+            )
+            self._job_results[queue_name][job_id] = result_info
+            return True
+
+        # If it was just queued, delete it.
+        if job_metadata.status in (JobStatus.deferred, JobStatus.queued):
+            del self._job_metadata[queue_name][job_id]
+            return True
+
+        # Otherwise, the job has already completed, so we can't abort it.
+        return False
+
     async def get_job_metadata(
         self, job_id: str, queue_name: str | None = None
     ) -> JobMetadata:

tests/dependencies/arq_test.py

@@ -80,6 +80,10 @@ async def get_result(
             )
         except (JobNotFound, JobResultUnavailable) as e:
             raise HTTPException(status_code=404, detail=str(e)) from e
+
+        # For testing purposes, turn exceptions into something serializable.
+        if isinstance(job_result.result, BaseException):
+            job_result.result = f"EXCEPTION {type(job_result.result).__name__}"
         return {
             "job_id": job_result.id,
             "job_status": job_result.status,
@@ -121,6 +125,19 @@ async def post_job_complete(
         except JobNotFound as e:
             raise HTTPException(status_code=404, detail=str(e)) from e
 
+    @app.post("/jobs/{job_id}/abort")
+    async def abort_job(
+        *,
+        job_id: str,
+        queue_name: str | None = None,
+        arq_queue: Annotated[MockArqQueue, Depends(arq_dependency)],
+    ) -> None:
+        """Abort a job, for testing."""
+        try:
+            await arq_queue.abort_job(job_id, queue_name=queue_name)
+        except JobNotFound as e:
+            raise HTTPException(status_code=404, detail=str(e)) from e
+
     transport = ASGITransport(app=app)  # type: ignore[arg-type]
     base_url = "http://example.com"
     async with LifespanManager(app):
@@ -165,3 +182,39 @@ async def post_job_complete(
             assert data["job_status"] == "complete"
             assert data["job_result"] == "done"
             assert data["job_success"] is True
+
+            # Aborting a completed job does nothing.
+            r = await c.post(f"/jobs/{job_id}/abort")
+            r = await c.get(f"/results/{job_id}")
+            assert r.status_code == 200
+            data = r.json()
+            assert data["job_status"] == "complete"
+
+            # Create a new job and abort it before starting it, which should
+            # delete the job.
+            r = await c.post("/")
+            assert r.status_code == 200
+            data = r.json()
+            job_id = data["job_id"]
+            r = await c.get(f"/jobs/{job_id}")
+            assert r.status_code == 200
+            r = await c.post(f"/jobs/{job_id}/abort")
+            r = await c.get(f"/results/{job_id}")
+            assert r.status_code == 404
+
+            # Create a new job, start it, and then abort it. This should keep
+            # the job but mark it complete and failed.
+            r = await c.post("/")
+            assert r.status_code == 200
+            data = r.json()
+            job_id = data["job_id"]
+            r = await c.get(f"/jobs/{job_id}")
+            assert r.status_code == 200
+            r = await c.post(f"/jobs/{job_id}/inprogress")
+            r = await c.post(f"/jobs/{job_id}/abort")
+            r = await c.get(f"/results/{job_id}")
+            assert r.status_code == 200
+            data = r.json()
+            assert data["job_status"] == "complete"
+            assert data["job_result"] == "EXCEPTION CancelledError"
+            assert data["job_success"] is False