From f53cad423f0757d26749f7272d45abaa1ac8e7a0 Mon Sep 17 00:00:00 2001 From: Norbert Kwizera Date: Mon, 27 Jul 2026 13:34:27 +0200 Subject: [PATCH 1/3] Fix org task locking: acquire non-blocking with default expiry, preserve last results on failure - Acquire the org task lock non-blocking instead of a racy get + blocking lock, so a concurrent invocation skips immediately rather than blocking and re-running the task - Default the lock timeout to 1 hour so a killed worker can't leave the lock held forever - Release the lock on both success and failure paths - Stop clearing last_results on task failure so incremental tasks resume from their last successful results --- dash/orgs/tasks.py | 82 ++++++++++++++++++++++++-------------------- test_runner/tests.py | 54 ++++++++++++++++++++++++++--- 2 files changed, 95 insertions(+), 41 deletions(-) diff --git a/dash/orgs/tasks.py b/dash/orgs/tasks.py index feeb9678..9e9148ca 100644 --- a/dash/orgs/tasks.py +++ b/dash/orgs/tasks.py @@ -13,6 +13,8 @@ ORG_TASK_LOCK_KEY = "org-task-lock:%s:%s" +DEFAULT_LOCK_TIMEOUT = 60 * 60 # 1 hour + logger = logging.getLogger(__name__) @@ -60,58 +62,64 @@ def maybe_run_for_org(org, task_func, task_key, lock_timeout): :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): + lock = r.lock(key, timeout=lock_timeout if lock_timeout is not None else DEFAULT_LOCK_TIMEOUT) + + if not lock.acquire(blocking=False): 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 + 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)) + 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() + 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")) + 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) + 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)" + assert num_task_args >= 1, "task signature must be foo(org) or foo(org, since, until)" - task_args = [org] + 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) + 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) + 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")) + 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))) + 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")) + 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 + logger.exception("Task %s for org #%d failed" % (task_key, org.id)) + raise e # re-raise with original stack trace + finally: + lock.release() diff --git a/test_runner/tests.py b/test_runner/tests.py index 0dcfdf04..64999be1 100644 --- a/test_runner/tests.py +++ b/test_runner/tests.py @@ -1,8 +1,8 @@ import zoneinfo -from dash.tags.models import Tag -from unittest.mock import Mock, patch, call +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 @@ -25,8 +25,9 @@ from dash.orgs.tasks import org_task from dash.orgs.templatetags.dashorgs import display_time, national_phone from dash.stories.models import Story, StoryImage -from dash.utils import random_string +from dash.tags.models import Tag from dash.test import MockResponse +from dash.utils import random_string class UserTest(SmartminTest): @@ -1369,7 +1370,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]) @@ -1377,6 +1378,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) @@ -1392,6 +1403,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) @@ -1416,6 +1439,29 @@ 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=10) + 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()) + 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) + class TaskCRUDLTest(DashTest): def setUp(self): From 96caa5e1bc96e9746da3027659843605b27c360a Mon Sep 17 00:00:00 2001 From: Norbert Kwizera Date: Mon, 27 Jul 2026 16:29:05 +0200 Subject: [PATCH 2/3] Guard lock release against expiry, make default lock timeout part of the API - Catch LockError when releasing the org task lock so a task that outlives its lock timeout doesn't fail a successful run or mask its own exception, and warn instead - Move the default lock timeout into the org_task/maybe_run_for_org signatures and document that a task running longer than its timeout may be started concurrently - Remove the unused ORG_TASK_LOCK_KEY constant and document that last_results holds the results of the last successful run - Add tests for lock expiry during a run and for the lock TTL being set --- dash/orgs/models.py | 1 + dash/orgs/tasks.py | 24 ++++++++++++++++------ test_runner/tests.py | 49 +++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 67 insertions(+), 7 deletions(-) diff --git a/dash/orgs/models.py b/dash/orgs/models.py index 6e7ad295..52e5fdbf 100644 --- a/dash/orgs/models.py +++ b/dash/orgs/models.py @@ -358,6 +358,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 9e9148ca..6bd3c400 100644 --- a/dash/orgs/tasks.py +++ b/dash/orgs/tasks.py @@ -5,14 +5,13 @@ 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 # 1 hour logger = logging.getLogger(__name__) @@ -39,9 +38,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 (1 hour 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 """ @@ -56,7 +61,7 @@ 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 @@ -68,7 +73,7 @@ def maybe_run_for_org(org, task_func, task_key, lock_timeout): key = TaskState.get_lock_key(org, task_key) - lock = r.lock(key, timeout=lock_timeout if lock_timeout is not None else DEFAULT_LOCK_TIMEOUT) + lock = r.lock(key, timeout=lock_timeout) if not lock.acquire(blocking=False): logger.warning("Skipping task %s for org #%d as it is still running" % (task_key, org.id)) @@ -122,4 +127,11 @@ def maybe_run_for_org(org, task_func, task_key, lock_timeout): logger.exception("Task %s for org #%d failed" % (task_key, org.id)) raise e # re-raise with original stack trace finally: - lock.release() + 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 64999be1..f8c82116 100644 --- a/test_runner/tests.py +++ b/test_runner/tests.py @@ -1444,7 +1444,7 @@ 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=10) + lock = r.lock(TaskState.get_lock_key(self.org, "test-task-2"), timeout=60) self.assertTrue(lock.acquire(blocking=False)) try: @@ -1453,6 +1453,9 @@ def test_org_task_skipped_if_already_running(self, mock_over_time_window): 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() @@ -1462,6 +1465,50 @@ def test_org_task_skipped_if_already_running(self, mock_over_time_window): 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) # 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): From ddc1495024e5b2028c0b63ad5377c4992790774a Mon Sep 17 00:00:00 2001 From: Norbert Kwizera Date: Mon, 27 Jul 2026 23:09:19 +0200 Subject: [PATCH 3/3] Make default org task lock timeout 2 hours --- dash/orgs/tasks.py | 4 ++-- test_runner/tests.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dash/orgs/tasks.py b/dash/orgs/tasks.py index f929cefd..935255a8 100644 --- a/dash/orgs/tasks.py +++ b/dash/orgs/tasks.py @@ -12,7 +12,7 @@ from .models import Invitation, TaskState -DEFAULT_LOCK_TIMEOUT = 60 * 60 # 1 hour +DEFAULT_LOCK_TIMEOUT = 60 * 60 * 2 # 2 hours logger = logging.getLogger(__name__) @@ -46,7 +46,7 @@ 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 (1 hour by default) so that a dead worker can't hold it forever - which means a task that + 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. diff --git a/test_runner/tests.py b/test_runner/tests.py index 76e0df90..83430403 100644 --- a/test_runner/tests.py +++ b/test_runner/tests.py @@ -1632,7 +1632,7 @@ def test_org_task_locks_with_expiry(self, mock_over_time_window): def check_lock(org, prev_started_on, started_on): ttl = r.ttl(key) - self.assertTrue(0 < ttl <= 60 * 60) # lock is held with the default expiry whilst the task runs + 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