Do not handle modes if job is to be cancelled.

The db-maintenance and pause modes have special handling at the top of
the handle_job method that eagerly handle those options.

This introduced a bug that job that are to be be cancelled do not get
cancelled if those modes affect them. This means that user's db jobs are
not cancelled when in db-maintenance mode, for example.

This quick fix simply protects the eager mode handling to apply only to
jobs that are not about to be cancelled, which will let the cancellation
logic proceed as normal.

This is probably not the ideal solution, but a) this is affecting users
right now with the extended db-maintenance mode and b) we are about to
completely refactor all the things, so we can live with a less than
ideal solution

Author: bloodearnest

jobrunner/run.py

@@ -215,32 +215,34 @@ def handle_job(job, api, mode=None, paused=None):
     synchronous_transitions = getattr(api, "synchronous_transitions", [])
     is_synchronous = False
 
-    # handle special modes before considering executor state, as they ignore it
-    if paused:
-        if job.state == State.PENDING:
-            # do not start the job, keep it pending
+    # only consider these modes if we are not about to cancel the job
+    if not job_definition.cancelled:
+        # handle special modes before considering executor state, as they ignore it
+        if paused:
+            if job.state == State.PENDING:
+                # do not start the job, keep it pending
+                set_code(
+                    job,
+                    StatusCode.WAITING_PAUSED,
+                    "Backend is currently paused for maintenance, job will start once this is completed",
+                )
+                return
+
+        if mode == "db-maintenance" and job_definition.allow_database_access:
+            if job.state == State.RUNNING:
+                log.warning(f"DB maintenance mode active, killing db job {job.id}")
+                # we ignore the JobStatus returned from these API calls, as this is not a hard error
+                api.terminate(job_definition)
+                api.cleanup(job_definition)
+
+            # reset state to pending and exit
             set_code(
                 job,
-                StatusCode.WAITING_PAUSED,
-                "Backend is currently paused for maintenance, job will start once this is completed",
+                StatusCode.WAITING_DB_MAINTENANCE,
+                "Waiting for database to finish maintenance",
             )
             return
 
-    if mode == "db-maintenance" and job_definition.allow_database_access:
-        if job.state == State.RUNNING:
-            log.warning(f"DB maintenance mode active, killing db job {job.id}")
-            # we ignore the JobStatus returned from these API calls, as this is not a hard error
-            api.terminate(job_definition)
-            api.cleanup(job_definition)
-
-        # reset state to pending and exit
-        set_code(
-            job,
-            StatusCode.WAITING_DB_MAINTENANCE,
-            "Waiting for database to finish maintenance",
-        )
-        return
-
     try:
         initial_status = api.get_status(job_definition)
     except ExecutorRetry as retry:

tests/test_run.py

@@ -731,15 +731,36 @@ def test_handle_pending_db_maintenance_mode(db, backend_db_config):
 
     # executor state
     assert api.get_status(job).state == ExecutorState.UNKNOWN
-    # our state
     assert job.state == State.PENDING
+    assert job.status_code == StatusCode.WAITING_DB_MAINTENANCE
     assert job.status_message == "Waiting for database to finish maintenance"
     assert job.started_at is None
 
     spans = get_trace("jobs")
     assert spans[-1].name == "CREATED"
 
 
+def test_handle_pending_cancelled_db_maintenance_mode(db, backend_db_config):
+    api = StubExecutorAPI()
+    job = api.add_test_job(
+        ExecutorState.UNKNOWN,
+        State.PENDING,
+        run_command="cohortextractor:latest generate_cohort",
+        requires_db=True,
+        cancelled=True,
+    )
+
+    run.handle_job(job, api, mode="db-maintenance")
+
+    # executor state
+    assert api.get_status(job).state == ExecutorState.UNKNOWN
+    # our state
+    assert job.state == State.FAILED
+    assert job.status_code == StatusCode.CANCELLED_BY_USER
+    assert job.status_message == "Cancelled by user"
+    assert job.started_at is None
+
+
 def test_handle_running_db_maintenance_mode(db, backend_db_config):
     api = StubExecutorAPI()
     job = api.add_test_job(
@@ -759,13 +780,41 @@ def test_handle_running_db_maintenance_mode(db, backend_db_config):
 
     # our state
     assert job.state == State.PENDING
+    assert job.status_code == StatusCode.WAITING_DB_MAINTENANCE
     assert job.status_message == "Waiting for database to finish maintenance"
     assert job.started_at is None
 
     spans = get_trace("jobs")
     assert spans[-1].name == "EXECUTING"
 
 
+def test_handle_running_cancelled_db_maintenance_mode(db, backend_db_config):
+    api = StubExecutorAPI()
+    job = api.add_test_job(
+        ExecutorState.EXECUTING,
+        State.RUNNING,
+        StatusCode.EXECUTING,
+        run_command="cohortextractor:latest generate_cohort",
+        requires_db=True,
+        cancelled=True,
+    )
+
+    run.handle_job(job, api, mode="db-maintenance")
+
+    # cancellation of running jobs puts it into EXECUTED for later finalization
+    # executor state
+    assert job.id in api.tracker["terminate"]
+    assert api.get_status(job).state == ExecutorState.EXECUTED
+
+    # our state
+    assert job.state == State.RUNNING
+    assert job.status_code == StatusCode.EXECUTED
+    assert job.status_message == "Cancelled whilst executing"
+
+    spans = get_trace("jobs")
+    assert spans[-1].name == "EXECUTING"
+
+
 def test_handle_pending_pause_mode(db, backend_db_config):
     api = StubExecutorAPI()
     job = api.add_test_job(