Fix task cancellation during graceful worker shutdown

CHANGES/+cancel-during-graceful-shutdown.bugfix

@@ -0,0 +1,2 @@
+Fixed canceling a task while its worker is shutting down gracefully. The cancellation never
+aborted the task, so the worker could not exit and kept holding the task's resource locks.

pulpcore/tasking/redis_worker.py

@@ -185,6 +185,13 @@ def __init__(self):
         # and set to None for fully graceful shutdown.
         self.task_grace_timeout = timezone.now()
 
+        # Deadline for force-killing an explicitly canceled (or shutdown-aborted) task
+        # process. Kept separate from task_grace_timeout: that field can be set to None
+        # to make the worker wait forever for a task to finish naturally on a graceful
+        # shutdown, but an explicit cancellation must still take effect on its own
+        # schedule regardless of that wait, or it can never actually abort the task.
+        self.cancel_kill_deadline = None
+
         self.worker_cleanup_countdown = random.randint(
             int(WORKER_CLEANUP_INTERVAL / 10), WORKER_CLEANUP_INTERVAL
         )
@@ -702,6 +709,7 @@ def supervise_task(self, task):
         self.task = task
         cancel_state = None
         cancel_reason = None
+        self.cancel_kill_deadline = None
         domain = task.pulp_domain
         with TemporaryDirectory(dir=".") as task_working_dir_rel_path:
             task_process = Process(target=perform_task, args=(task.pk, task_working_dir_rel_path))
@@ -710,10 +718,10 @@ def supervise_task(self, task):
             # Heartbeat while waiting for task to complete
             while task_process.is_alive():
                 if cancel_state:
-                    if self.task_grace_timeout is None or self.task_grace_timeout > timezone.now():
+                    if self.cancel_kill_deadline > timezone.now():
                         _logger.info("Wait for canceled task to abort.")
                     else:
-                        self.task_grace_timeout = timezone.now() + timezone.timedelta(
+                        self.cancel_kill_deadline = timezone.now() + timezone.timedelta(
                             seconds=TASK_KILL_INTERVAL
                         )
                         _logger.info(
@@ -741,6 +749,7 @@ def supervise_task(self, task):
                         domain.name,
                     )
                     cancel_state = TASK_STATES.CANCELED
+                    self.cancel_kill_deadline = timezone.now()
                     clear_cancel_signal(task.pk)
 
                 if self.sentinel in r:
@@ -751,7 +760,9 @@ def supervise_task(self, task):
                         break
 
                 # If shutdown was requested, handle gracefully or abort
-                if self.shutdown_requested:
+                # Once the task is being aborted, keep its state (e.g. CANCELED) and let the
+                # cancellation branch above drive the kill.
+                if self.shutdown_requested and not cancel_state:
                     if self.task_grace_timeout is None or self.task_grace_timeout > timezone.now():
                         msg = (
                             "Worker shutdown requested, waiting for task {pk} in domain: {name} "
@@ -766,6 +777,7 @@ def supervise_task(self, task):
                         )
                         cancel_state = TASK_STATES.FAILED
                         cancel_reason = "Aborted during worker shutdown."
+                        self.cancel_kill_deadline = timezone.now()
 
             task_process.join()
             if not cancel_state and task_process.exitcode != 0:

pulpcore/tasking/worker.py

@@ -107,6 +107,14 @@ def __init__(self, auxiliary=False):
         # It will be set into the future on moderately graceful worker shutdown,
         # and set to None for fully graceful shutdown.
         self.task_grace_timeout = timezone.now()
+
+        # Deadline for force-killing an explicitly canceled (or shutdown-aborted) task
+        # process. Kept separate from task_grace_timeout: that field can be set to None
+        # to make the worker wait forever for a task to finish naturally on a graceful
+        # shutdown, but an explicit cancellation must still take effect on its own
+        # schedule regardless of that wait, or it can never actually abort the task.
+        self.cancel_kill_deadline = None
+
         self.worker_cleanup_countdown = random.randint(
             int(WORKER_CLEANUP_INTERVAL / 10), WORKER_CLEANUP_INTERVAL
         )
@@ -460,16 +468,17 @@ def supervise_task(self, task):
         self.task = task
         cancel_state = None
         cancel_reason = None
+        self.cancel_kill_deadline = None
         domain = task.pulp_domain
         with TemporaryDirectory(dir=".") as task_working_dir_rel_path:
             task_process = Process(target=perform_task, args=(task.pk, task_working_dir_rel_path))
             task_process.start()
             while True:
                 if cancel_state:
-                    if self.task_grace_timeout is None or self.task_grace_timeout > timezone.now():
+                    if self.cancel_kill_deadline > timezone.now():
                         _logger.info("Wait for canceled task to abort.")
                     else:
-                        self.task_grace_timeout = timezone.now() + timezone.timedelta(
+                        self.cancel_kill_deadline = timezone.now() + timezone.timedelta(
                             seconds=TASK_KILL_INTERVAL
                         )
                         _logger.info(
@@ -496,14 +505,17 @@ def supervise_task(self, task):
                     )
                     cancel_state = TASK_STATES.CANCELED
                     self.cancel_task = False
+                    self.cancel_kill_deadline = timezone.now()
                 if self.wakeup_unblock:
                     self.unblock_tasks()
                 if task_process.sentinel in r:
                     if not task_process.is_alive():
                         break
                 if self.sentinel in r:
                     os.read(self.sentinel, 256)
-                if self.shutdown_requested:
+                # Once the task is being aborted, keep its state (e.g. CANCELED) and let the
+                # cancellation branch above drive the kill.
+                if self.shutdown_requested and not cancel_state:
                     if self.task_grace_timeout is None or self.task_grace_timeout > timezone.now():
                         msg = (
                             "Worker shutdown requested, waiting for task {pk} in domain: {name} "
@@ -519,6 +531,7 @@ def supervise_task(self, task):
                         )
                         cancel_state = TASK_STATES.FAILED
                         cancel_reason = "Aborted during worker shutdown."
+                        self.cancel_kill_deadline = timezone.now()
 
             task_process.join()
             if not cancel_state and task_process.exitcode != 0:

pulpcore/tests/unit/tasking/test_cancel_during_graceful_shutdown.py

@@ -0,0 +1,189 @@
+"""
+Unit tests for task cancellation while a worker is shutting down gracefully.
+
+A graceful shutdown (SIGTERM/SIGHUP) sets ``task_grace_timeout`` to ``None`` so the worker waits
+indefinitely for the in-flight task to finish. An explicit cancellation of that task must still
+abort it (SIGUSR1 to the task process); previously the cancellation reused
+``task_grace_timeout`` as its kill deadline, so it waited forever too and the worker deadlocked,
+heartbeating and holding the task's resource locks.
+"""
+
+import signal
+from datetime import timedelta
+from unittest import mock
+
+import pytest
+
+from pulpcore.constants import TASK_STATES
+from pulpcore.tasking import redis_worker as redis_worker_module
+from pulpcore.tasking import worker as worker_module
+
+MAX_LOOP_ITERATIONS = 50
+
+
+class FakeProcess:
+    """Stand-in for the task ``multiprocessing.Process``; it only exits once it gets SIGUSR1."""
+
+    def __init__(self, *args, **kwargs):
+        self.pid = 4242
+        self.sentinel = object()
+        self.alive = True
+        self.exitcode = None
+
+    def start(self):
+        pass
+
+    def is_alive(self):
+        return self.alive
+
+    def join(self):
+        pass
+
+    def finish(self, exitcode=0):
+        self.alive = False
+        self.exitcode = exitcode
+
+
+class Harness:
+    """Patches the process/select/kill machinery used by ``supervise_task``."""
+
+    def __init__(self, monkeypatch, module, on_first_select=None, finish_after=None):
+        self.process = None
+        self.kills = []
+        self.select_calls = 0
+        self.on_first_select = on_first_select
+        self.finish_after = finish_after
+
+        def make_process(*args, **kwargs):
+            self.process = FakeProcess()
+            return self.process
+
+        def fake_kill(pid, sig):
+            assert pid == self.process.pid
+            self.kills.append(sig)
+            if sig == signal.SIGUSR1:
+                self.process.finish(exitcode=0)
+
+        def fake_select(rlist, wlist, xlist, timeout=None):
+            self.select_calls += 1
+            if self.select_calls > MAX_LOOP_ITERATIONS:
+                raise AssertionError("supervise_task never let the task process exit (deadlock)")
+            if self.select_calls == 1 and self.on_first_select:
+                self.on_first_select()
+            if self.finish_after and self.select_calls >= self.finish_after:
+                self.process.finish(exitcode=0)
+            return ([self.process.sentinel] if not self.process.alive else [], [], [])
+
+        monkeypatch.setattr(module, "Process", make_process)
+        monkeypatch.setattr(module.os, "kill", fake_kill)
+        monkeypatch.setattr(module.select, "select", fake_select)
+        monkeypatch.setattr(module, "connection", mock.MagicMock())
+
+
+def _make_task():
+    task = mock.MagicMock()
+    task.pk = "01a1841c-b360-7f5d-adb6-4ee443942c8e"
+    task.pulp_domain.name = "default"
+    task.reserved_resources_record = None
+    return task
+
+
+def _pulpcore_worker(graceful_shutdown):
+    worker = object.__new__(worker_module.PulpcoreWorker)
+    worker.sentinel = -1
+    worker.heartbeat_period = timedelta(seconds=1)
+    worker.wakeup_unblock = False
+    worker.shutdown_requested = graceful_shutdown
+    worker.task_grace_timeout = None if graceful_shutdown else worker_module.timezone.now()
+    worker.cancel_kill_deadline = None
+    worker.beat = mock.Mock()
+    worker.cancel_abandoned_task = mock.Mock()
+    worker.notify_workers = mock.Mock()
+    return worker
+
+
+def _redis_worker(graceful_shutdown):
+    worker = object.__new__(redis_worker_module.RedisWorker)
+    worker.sentinel = -1
+    worker.heartbeat_period = timedelta(seconds=1)
+    worker.shutdown_requested = graceful_shutdown
+    worker.task_grace_timeout = None if graceful_shutdown else redis_worker_module.timezone.now()
+    worker.cancel_kill_deadline = None
+    worker.beat = mock.Mock()
+    worker._maybe_release_locks = mock.Mock()
+    return worker
+
+
+@pytest.fixture
+def in_tmp_dir(tmp_path, monkeypatch):
+    # supervise_task creates its TemporaryDirectory in the current directory.
+    monkeypatch.chdir(tmp_path)
+
+
+def _patch_redis_cancel_signal(monkeypatch, pending):
+    from pulpcore.tasking import _util, redis_tasks
+
+    monkeypatch.setattr(redis_tasks, "check_cancel_signal", lambda pk: pending["value"])
+    monkeypatch.setattr(redis_tasks, "clear_cancel_signal", lambda pk: pending.update(value=False))
+    monkeypatch.setattr(_util, "delete_incomplete_resources", mock.Mock())
+
+
+# --- PulpcoreWorker ------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("graceful_shutdown", [True, False], ids=["graceful-shutdown", "running"])
+def test_pulpcore_worker_cancel_aborts_task(monkeypatch, in_tmp_dir, graceful_shutdown):
+    worker = _pulpcore_worker(graceful_shutdown)
+    harness = Harness(
+        monkeypatch, worker_module, on_first_select=lambda: setattr(worker, "cancel_task", True)
+    )
+    task = _make_task()
+
+    worker.supervise_task(task)
+
+    assert harness.kills == [signal.SIGUSR1]
+    worker.cancel_abandoned_task.assert_called_once_with(task, TASK_STATES.CANCELED, None)
+
+
+def test_pulpcore_worker_graceful_shutdown_waits_for_task(monkeypatch, in_tmp_dir):
+    worker = _pulpcore_worker(graceful_shutdown=True)
+    harness = Harness(monkeypatch, worker_module, finish_after=5)
+
+    worker.supervise_task(_make_task())
+
+    assert harness.kills == []
+    assert harness.select_calls == 5
+    worker.cancel_abandoned_task.assert_not_called()
+
+
+# --- RedisWorker ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("graceful_shutdown", [True, False], ids=["graceful-shutdown", "running"])
+def test_redis_worker_cancel_aborts_task(monkeypatch, in_tmp_dir, graceful_shutdown):
+    pending = {"value": False}
+    _patch_redis_cancel_signal(monkeypatch, pending)
+    worker = _redis_worker(graceful_shutdown)
+    harness = Harness(
+        monkeypatch, redis_worker_module, on_first_select=lambda: pending.update(value=True)
+    )
+    task = _make_task()
+
+    worker.supervise_task(task)
+
+    assert harness.kills == [signal.SIGUSR1]
+    task.set_canceled.assert_called_once_with(final_state=TASK_STATES.CANCELED, reason=None)
+
+
+def test_redis_worker_graceful_shutdown_waits_for_task(monkeypatch, in_tmp_dir):
+    pending = {"value": False}
+    _patch_redis_cancel_signal(monkeypatch, pending)
+    worker = _redis_worker(graceful_shutdown=True)
+    harness = Harness(monkeypatch, redis_worker_module, finish_after=5)
+    task = _make_task()
+
+    worker.supervise_task(task)
+
+    assert harness.kills == []
+    assert harness.select_calls == 5
+    task.set_canceled.assert_not_called()