diff --git a/dash/orgs/models.py b/dash/orgs/models.py index cb8fabc..d1da7ee 100644 --- a/dash/orgs/models.py +++ b/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) diff --git a/dash/orgs/tasks.py b/dash/orgs/tasks.py index cfa1e0b..935255a 100644 --- a/dash/orgs/tasks.py +++ b/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) + ) diff --git a/test_runner/tests.py b/test_runner/tests.py index d9cc2f4..8343040 100644 --- a/test_runner/tests.py +++ b/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,7 +1530,7 @@ 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]) @@ -1537,6 +1538,16 @@ def test_org_task(self, mock_over_time_window): 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):