Add renew_org_task_lock so org tasks can hold short, renewable locks

dash/orgs/tasks.py

@@ -1,6 +1,7 @@
 import inspect
 import json
 import logging
+import threading
 from functools import wraps
 
 from celery import shared_task, signature
@@ -16,6 +17,32 @@
 
 logger = logging.getLogger(__name__)
 
+# the lock held by the org task currently running in this worker process, so that task code
+# at any depth can renew it via renew_org_task_lock()
+_current_org_task = threading.local()
+
+
+def renew_org_task_lock():
+    """
+    Renews the lease of the lock held by the org task currently running in this worker.
+
+    Tasks that checkpoint their progress can call this after each unit of work and pass a
+    lock_timeout that only needs to outlive one unit rather than a worst-case full run - a
+    hard-killed worker then frees the org after one short lease instead of hours.
+
+    :return: whether the lock is still owned - False means the lease already expired and the
+             same task may have been started concurrently, so the caller should stop cleanly
+    """
+    lock = getattr(_current_org_task, "lock", None)
+    if not lock:
+        return True
+
+    try:
+        lock.extend(lock.timeout, replace_ttl=True)
+        return True
+    except LockError:
+        return False
+
 
 @shared_task(track_started=True, name="send_invitation_email_task")
 def send_invitation_email_task(invitation_id):
@@ -48,7 +75,8 @@ def org_task(task_key, lock_timeout=DEFAULT_LOCK_TIMEOUT):
     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.
+    worst-case runtime, or have the task call renew_org_task_lock() after each unit of work so that lock_timeout
+    only needs to exceed one unit.
 
     :param task_key: the task key used for state storage and locking, e.g. 'do-stuff'
     :param lock_timeout: the lock timeout in seconds
@@ -82,6 +110,8 @@ def maybe_run_for_org(org, task_func, task_key, lock_timeout=DEFAULT_LOCK_TIMEOU
         logger.warning("Skipping task %s for org #%d as it is still running" % (task_key, org.id))
         return
 
+    _current_org_task.lock = lock
+
     try:
         state = org.get_task_state(task_key)
         if state.is_disabled:
@@ -130,6 +160,8 @@ def maybe_run_for_org(org, task_func, task_key, lock_timeout=DEFAULT_LOCK_TIMEOU
             logger.exception("Task %s for org #%d failed" % (task_key, org.id))
             raise e  # re-raise with original stack trace
     finally:
+        _current_org_task.lock = None
+
         try:
             lock.release()
         except LockError:

test_runner/tests.py

@@ -26,7 +26,7 @@
 from dash.orgs.context_processors import GroupPermWrapper
 from dash.orgs.middleware import SetOrgMiddleware
 from dash.orgs.models import Invitation, Org, OrgBackend, OrgBackground, TaskState
-from dash.orgs.tasks import org_task
+from dash.orgs.tasks import org_task, renew_org_task_lock
 from dash.orgs.templatetags.dashorgs import display_time, national_phone
 from dash.orgs.views import OrgBackendForm, OrgCRUDL
 from dash.stories.models import Story, StoryImage
@@ -1885,6 +1885,40 @@ def delete_lock_and_fail(org, prev_started_on, started_on):
 
         self.assertTrue(TaskState.objects.get(org=self.org, task_key="test-task-2").is_failing)
 
+    def test_renew_org_task_lock(self):
+        r = get_valkey_connection()
+        key = TaskState.get_lock_key(self.org, "test-renew-task")
+        renewals = []
+
+        @org_task("test-renew-task", lock_timeout=10)
+        def renewing_task(org):
+            renewals.append(renew_org_task_lock())
+            renewals.append(0 < r.ttl(key) <= 10)  # the lease was refreshed to the lock timeout
+
+        renewing_task(self.org.id)
+
+        self.assertEqual([True, True], renewals)
+        self.assertEqual(r.ttl(key), -2)  # lock released when the task finished
+
+        # outside of a running org task there is nothing to renew
+        self.assertTrue(renew_org_task_lock())
+
+    def test_renew_org_task_lock_lost_lease(self):
+        r = get_valkey_connection()
+        key = TaskState.get_lock_key(self.org, "test-renew-lost")
+        renewals = []
+
+        @org_task("test-renew-lost", lock_timeout=10)
+        def losing_task(org):
+            r.delete(key)  # simulate the lease expiring mid-run
+            renewals.append(renew_org_task_lock())
+
+        # a lost lease reports False so the task can stop cleanly, and doesn't fail the run
+        losing_task(self.org.id)
+
+        self.assertEqual([False], renewals)
+        self.assertFalse(TaskState.objects.get(org=self.org, task_key="test-renew-lost").is_failing)
+
 
 class TaskCRUDLTest(DashTest):
     def setUp(self):