fix(get_recurringtasks_to_process): Add last_picked_at (#37)

* fix(get_recurringtasks_to_process): Add last_picked_at

Closes: #35

* make last_picked_at nullable

* remove task_runs var

Author: gagantrivedi

src/task_processor/migrations/0013_add_last_picked_at.py

@@ -0,0 +1,34 @@
+# Generated by Django 4.2.18 on 2025-04-03 07:34
+
+from django.db import migrations, models
+import django.utils.timezone
+import os
+from common.migrations.helpers import PostgresOnlyRunSQL
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("task_processor", "0012_add_locked_at_and_timeout"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="recurringtask",
+            name="last_picked_at",
+            field=models.DateTimeField(blank=True, null=True),
+            preserve_default=False,
+        ),
+        PostgresOnlyRunSQL.from_sql_file(
+            os.path.join(
+                os.path.dirname(__file__),
+                "sql",
+                "0013_get_recurringtasks_to_process.sql",
+            ),
+            reverse_sql=os.path.join(
+                os.path.dirname(__file__),
+                "sql",
+                "0012_get_recurringtasks_to_process.sql",
+            ),
+        ),
+    ]

src/task_processor/migrations/sql/0013_get_recurringtasks_to_process.sql

@@ -0,0 +1,33 @@
+CREATE OR REPLACE FUNCTION get_recurringtasks_to_process()
+RETURNS SETOF task_processor_recurringtask AS $$
+DECLARE
+    row_to_return task_processor_recurringtask;
+BEGIN
+    -- Select the tasks that needs to be processed
+    FOR row_to_return IN
+        SELECT *
+        FROM task_processor_recurringtask
+        -- Add one minute to the timeout as a grace period for overhead
+        WHERE is_locked = FALSE OR (locked_at IS NOT NULL AND locked_at < NOW() - timeout + INTERVAL '1 minute')
+        ORDER BY last_picked_at NULLS FIRST
+        LIMIT 1
+        -- Select for update to ensure that no other workers can select these tasks while in this transaction block
+        FOR UPDATE SKIP LOCKED
+    LOOP
+        -- Lock every selected task(by updating `is_locked` to true)
+        UPDATE task_processor_recurringtask
+        -- Lock this row by setting is_locked True, so that no other workers can select these tasks after this
+        -- transaction is complete (but the tasks are still being executed by the current worker)
+        SET is_locked = TRUE, locked_at = NOW(), last_picked_at = NOW()
+        WHERE id = row_to_return.id;
+        -- If we don't explicitly update the columns here, the client will receive a row
+        -- that is locked but still shows `is_locked` as `False` and `locked_at` as `None`.
+        row_to_return.is_locked := TRUE;
+        row_to_return.locked_at := NOW();
+        RETURN NEXT row_to_return;
+    END LOOP;
+
+    RETURN;
+END;
+$$ LANGUAGE plpgsql
+

src/task_processor/models.py

@@ -163,6 +163,7 @@ class RecurringTask(AbstractBaseTask):
     locked_at = models.DateTimeField(blank=True, null=True)
     timeout = models.DurationField(default=timedelta(minutes=30))
 
+    last_picked_at = models.DateTimeField(blank=True, null=True)
     objects: RecurringTaskManager = RecurringTaskManager()
 
     class Meta:

tests/unit/task_processor/test_unit_task_processor_processor.py

@@ -303,6 +303,39 @@ def _dummy_recurring_task() -> None:
         assert task_run.error_details is None
 
 
+@pytest.mark.django_db(transaction=True)
+def test_run_recurring_tasks_loops_over_all_tasks(
+    db: None,
+    run_by_processor: None,
+) -> None:
+    # Given, Three recurring tasks
+    @register_recurring_task(run_every=timedelta(milliseconds=200))
+    def _dummy_recurring_task_1() -> None:
+        pass
+
+    @register_recurring_task(run_every=timedelta(milliseconds=200))
+    def _dummy_recurring_task_2() -> None:
+        pass
+
+    @register_recurring_task(run_every=timedelta(milliseconds=200))
+    def _dummy_recurring_task_3() -> None:
+        pass
+
+    initialise()
+
+    # When, we call run_recurring_tasks in a loop few times
+    for _ in range(4):
+        run_recurring_tasks()
+
+    # Then - we should have exactly one RecurringTaskRun for each task
+    for i in range(1, 4):
+        task = RecurringTask.objects.get(
+            task_identifier=f"test_unit_task_processor_processor._dummy_recurring_task_{i}",
+        )
+
+        assert RecurringTaskRun.objects.filter(task=task).count() == 1
+
+
 def test_run_recurring_tasks_only_executes_tasks_after_interval_set_by_run_every(
     db: None,
     run_by_processor: None,