Fix org task locking: acquire non-blocking with default expiry, preserve last results on failure

dash/orgs/models.py

@@ -352,6 +352,7 @@ class TaskState(models.Model):
 
     last_successfully_started_on = models.DateTimeField(null=True)
 
+    # the results of the last successful run - preserved on failure so incremental tasks can resume from them
     last_results = models.TextField(null=True)
 
     is_failing = models.BooleanField(default=False)

dash/orgs/tasks.py

@@ -5,13 +5,14 @@
 
 from celery import shared_task, signature
 from django_valkey import get_valkey_connection
+from valkey.exceptions import LockError
 
 from django.apps import apps
 from django.utils import timezone
 
 from .models import Invitation, TaskState
 
-ORG_TASK_LOCK_KEY = "org-task-lock:%s:%s"
+DEFAULT_LOCK_TIMEOUT = 60 * 60 * 2  # 2 hours
 
 logger = logging.getLogger(__name__)
 
@@ -40,9 +41,15 @@ def trigger_org_task(task_name, queue="celery"):
     logger.info("Requested task '%s' for %d active orgs" % (task_name, len(active_orgs)))
 
 
-def org_task(task_key, lock_timeout=None):
+def org_task(task_key, lock_timeout=DEFAULT_LOCK_TIMEOUT):
     """
     Decorator to create an org task.
+
+    The task holds a lock while it runs so that it can't run concurrently for the same org. The lock expires after
+    lock_timeout seconds (2 hours by default) so that a dead worker can't hold it forever - which means a task that
+    runs longer than its lock timeout may be started concurrently. Set lock_timeout to comfortably exceed the task's
+    worst-case runtime.
+
     :param task_key: the task key used for state storage and locking, e.g. 'do-stuff'
     :param lock_timeout: the lock timeout in seconds
     """
@@ -57,64 +64,77 @@ def _decorator(org_id):
     return _org_task
 
 
-def maybe_run_for_org(org, task_func, task_key, lock_timeout):
+def maybe_run_for_org(org, task_func, task_key, lock_timeout=DEFAULT_LOCK_TIMEOUT):
     """
     Runs the given task function for the specified org provided it's not already running
     :param org: the org
     :param task_func: the task function
     :param task_key: the task key
-    :param lock_timeout: the lock timeout in seconds
+    :param lock_timeout: the lock timeout in seconds (defaults to 1 hour so dead workers can't hold the lock forever)
     """
     r = get_valkey_connection()
 
     key = TaskState.get_lock_key(org, task_key)
 
-    if r.get(key):
-        logger.warning("Skipping task %s for org #%d as it is still running" % (task_key, org.id))
-    else:
-        with r.lock(key, timeout=lock_timeout):
-            state = org.get_task_state(task_key)
-            if state.is_disabled:
-                logger.info("Skipping task %s for org #%d as is marked disabled" % (task_key, org.id))
-                return
-
-            logger.info("Started task %s for org #%d..." % (task_key, org.id))
-
-            prev_results = json.loads(state.last_results) if state.last_results else None
-            prev_started_on = state.last_successfully_started_on
-            this_started_on = timezone.now()
-
-            state.started_on = this_started_on
-            state.ended_on = None
-            state.save(update_fields=("started_on", "ended_on"))
-
-            num_task_args = len(inspect.getfullargspec(task_func).args)
+    lock = r.lock(key, timeout=lock_timeout)
 
-            assert num_task_args >= 1, "task signature must be foo(org) or foo(org, since, until)"
-
-            task_args = [org]
-
-            try:
-                if num_task_args >= 3:
-                    task_args += [prev_started_on, this_started_on]
-                if num_task_args >= 4:
-                    task_args.append(prev_results)
-
-                results = task_func(*task_args)
-
-                state.ended_on = timezone.now()
-                state.last_successfully_started_on = this_started_on
-                state.last_results = json.dumps(results)
-                state.is_failing = False
-                state.save(update_fields=("ended_on", "last_successfully_started_on", "last_results", "is_failing"))
-
-                logger.info("Finished task %s for org #%d with result: %s" % (task_key, org.id, json.dumps(results)))
-
-            except Exception as e:
-                state.ended_on = timezone.now()
-                state.last_results = None
-                state.is_failing = True
-                state.save(update_fields=("ended_on", "last_results", "is_failing"))
-
-                logger.exception("Task %s for org #%d failed" % (task_key, org.id))
-                raise e  # re-raise with original stack trace
+    if not lock.acquire(blocking=False):
+        logger.warning("Skipping task %s for org #%d as it is still running" % (task_key, org.id))
+        return
+
+    try:
+        state = org.get_task_state(task_key)
+        if state.is_disabled:
+            logger.info("Skipping task %s for org #%d as is marked disabled" % (task_key, org.id))
+            return
+
+        logger.info("Started task %s for org #%d..." % (task_key, org.id))
+
+        prev_results = json.loads(state.last_results) if state.last_results else None
+        prev_started_on = state.last_successfully_started_on
+        this_started_on = timezone.now()
+
+        state.started_on = this_started_on
+        state.ended_on = None
+        state.save(update_fields=("started_on", "ended_on"))
+
+        num_task_args = len(inspect.getfullargspec(task_func).args)
+
+        assert num_task_args >= 1, "task signature must be foo(org) or foo(org, since, until)"
+
+        task_args = [org]
+
+        try:
+            if num_task_args >= 3:
+                task_args += [prev_started_on, this_started_on]
+            if num_task_args >= 4:
+                task_args.append(prev_results)
+
+            results = task_func(*task_args)
+
+            state.ended_on = timezone.now()
+            state.last_successfully_started_on = this_started_on
+            state.last_results = json.dumps(results)
+            state.is_failing = False
+            state.save(update_fields=("ended_on", "last_successfully_started_on", "last_results", "is_failing"))
+
+            logger.info("Finished task %s for org #%d with result: %s" % (task_key, org.id, json.dumps(results)))
+
+        except Exception as e:
+            # note we don't clear last_results here so that incremental tasks can resume from their last
+            # successful results after a transient failure
+            state.ended_on = timezone.now()
+            state.is_failing = True
+            state.save(update_fields=("ended_on", "is_failing"))
+
+            logger.exception("Task %s for org #%d failed" % (task_key, org.id))
+            raise e  # re-raise with original stack trace
+    finally:
+        try:
+            lock.release()
+        except LockError:
+            # the lock expired before we finished (i.e. the task ran longer than its lock timeout) - don't let that
+            # fail an otherwise successful run or mask an in-flight exception
+            logger.warning(
+                "Unable to release lock for task %s for org #%d as it is no longer owned" % (task_key, org.id)
+            )

test_runner/tests.py

@@ -2,6 +2,7 @@
 from unittest.mock import Mock, call, patch
 
 import valkey
+from django_valkey import get_valkey_connection
 from smartmin.tests import SmartminTest
 from temba_client.v2 import TembaClient
 
@@ -1529,14 +1530,24 @@ def test_org_task(self, mock_over_time_window):
         self.assertGreater(task2_state3.ended_on, task2_state2.ended_on)
         self.assertEqual(task2_state3.last_successfully_started_on, task2_state2.started_on)  # hasn't changed
         self.assertFalse(task2_state3.is_running())
-        self.assertEqual(task2_state3.get_last_results(), None)
+        self.assertEqual(task2_state3.get_last_results(), {"foo": "bar", "zed": 123})  # previous results preserved
         self.assertTrue(task2_state3.is_failing)
 
         self.assertEqual(list(TaskState.get_failing()), [task2_state3])
 
         mock_over_time_window.assert_called_once_with(self.org, task2_state2.started_on, task2_state3.started_on)
         mock_over_time_window.reset_mock()
 
+        # a failed run of a task that takes prev_results doesn't clear its previous successful results
+        self.assertRaises(ValueError, test_org_task_3, self.org.pk)
+
+        task3_state3 = TaskState.objects.get(org=self.org, task_key="test-task-3")
+
+        self.assertTrue(task3_state3.is_failing)
+        self.assertEqual(task3_state3.get_last_results(), {"foo": "bar", "zed": 123})
+
+        mock_over_time_window.reset_mock()
+
         # test when called, again, start time is from last successful run
         self.assertRaises(ValueError, test_org_task_2, self.org.pk)
 
@@ -1552,6 +1563,18 @@ def test_org_task(self, mock_over_time_window):
         mock_over_time_window.side_effect = None
         mock_over_time_window.return_value = {"foo": "bar", "zed": 123}
 
+        # when the task next succeeds it receives the results from the last successful run, despite the failure
+        test_org_task_3(self.org.pk)
+
+        task3_state4 = TaskState.objects.get(org=self.org, task_key="test-task-3")
+
+        self.assertFalse(task3_state4.is_failing)
+
+        mock_over_time_window.assert_called_once_with(
+            self.org, task3_state2.started_on, task3_state4.started_on, {"foo": "bar", "zed": 123}
+        )
+        mock_over_time_window.reset_mock()
+
         # disable the task for our org
         TaskState.objects.filter(org=self.org, task_key="test-task-2").update(is_disabled=True)
 
@@ -1576,6 +1599,76 @@ def test_org_task(self, mock_over_time_window):
 
         mock_over_time_window.assert_called_once_with(self.org, task2_state2.started_on, state6.started_on)
 
+    @patch("test_runner.tests.test_over_time_window")
+    def test_org_task_skipped_if_already_running(self, mock_over_time_window):
+        mock_over_time_window.return_value = {"foo": "bar"}
+
+        r = get_valkey_connection()
+        lock = r.lock(TaskState.get_lock_key(self.org, "test-task-2"), timeout=60)
+        self.assertTrue(lock.acquire(blocking=False))
+
+        try:
+            # while another worker holds the lock, invoking the task returns immediately without running it
+            test_org_task_2(self.org.id)
+
+            mock_over_time_window.assert_not_called()
+            self.assertFalse(TaskState.objects.filter(org=self.org, task_key="test-task-2").exists())
+
+            # and the skipped invocation didn't steal or release the other worker's lock
+            self.assertTrue(lock.owned())
+        finally:
+            lock.release()
+
+        # once the lock is released the task can run again
+        test_org_task_2(self.org.id)
+
+        mock_over_time_window.assert_called_once()
+        self.assertFalse(TaskState.objects.get(org=self.org, task_key="test-task-2").is_failing)
+
+    @patch("test_runner.tests.test_over_time_window")
+    def test_org_task_locks_with_expiry(self, mock_over_time_window):
+        r = get_valkey_connection()
+        key = TaskState.get_lock_key(self.org, "test-task-2")
+
+        def check_lock(org, prev_started_on, started_on):
+            ttl = r.ttl(key)
+            self.assertTrue(0 < ttl <= 60 * 60 * 2)  # lock is held with the default expiry whilst the task runs
+            return {}
+
+        mock_over_time_window.side_effect = check_lock
+
+        test_org_task_2(self.org.id)
+
+        mock_over_time_window.assert_called_once()
+        self.assertEqual(r.ttl(key), -2)  # lock released when the task finished
+
+    @patch("test_runner.tests.test_over_time_window")
+    def test_org_task_lock_expires_mid_run(self, mock_over_time_window):
+        r = get_valkey_connection()
+        key = TaskState.get_lock_key(self.org, "test-task-2")
+
+        def delete_lock_and_succeed(org, prev_started_on, started_on):
+            r.delete(key)  # simulate the lock expiring before the task finishes
+            return {"foo": "bar"}
+
+        mock_over_time_window.side_effect = delete_lock_and_succeed
+
+        # a run that outlives its lock still completes successfully
+        test_org_task_2(self.org.id)
+
+        self.assertFalse(TaskState.objects.get(org=self.org, task_key="test-task-2").is_failing)
+
+        def delete_lock_and_fail(org, prev_started_on, started_on):
+            r.delete(key)  # simulate the lock expiring before the task fails
+            raise ValueError("DOH!")
+
+        mock_over_time_window.side_effect = delete_lock_and_fail
+
+        # and a failing run's own exception isn't masked by the failure to release the expired lock
+        self.assertRaises(ValueError, test_org_task_2, self.org.id)
+
+        self.assertTrue(TaskState.objects.get(org=self.org, task_key="test-task-2").is_failing)
+
 
 class TaskCRUDLTest(DashTest):
     def setUp(self):