Race condition has a chance of multi task creation off of a schedule

Author: vegu

src/fullctl/django/admin.py

@@ -21,6 +21,7 @@
     TaskClaim,
     TaskHeartbeat,
     TaskSchedule,
+    TaskScheduleClaim,
     UserSettings,
 )
 from fullctl.django.models.concrete.service_bridge import (
@@ -188,6 +189,9 @@ def recent_tasks(self, obj):
         tasks = obj.tasks.all()[:5]
         return f"{tasks}"
 
+@admin.register(TaskScheduleClaim)
+class TaskScheduleClaimAdmin(BaseAdmin):
+    list_display = ("task_schedule", "worker_id", "schedule_date", "created", "updated")
 
 @admin.register(AuditLog)
 class AuditLogAdmin(admin.ModelAdmin):

src/fullctl/django/migrations/0039_taskscheduleclaim.py

@@ -0,0 +1,67 @@
+# Generated by Django 4.2.13 on 2025-05-14 18:15
+
+from django.db import migrations, models
+import django.db.models.deletion
+import django.db.models.manager
+import django_handleref.models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("django_fullctl", "0038_taskheartbeat_fullctl_tas_timesta_1a9e0d_idx"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="TaskScheduleClaim",
+            fields=[
+                ("id", models.AutoField(primary_key=True, serialize=False)),
+                (
+                    "created",
+                    django_handleref.models.CreatedDateTimeField(
+                        auto_now_add=True, verbose_name="Created"
+                    ),
+                ),
+                (
+                    "updated",
+                    django_handleref.models.UpdatedDateTimeField(
+                        auto_now=True, verbose_name="Updated"
+                    ),
+                ),
+                ("version", models.IntegerField(default=0)),
+                (
+                    "status",
+                    models.CharField(
+                        choices=[
+                            ("ok", "Ok"),
+                            ("pending", "Pending"),
+                            ("deactivated", "Deactivated"),
+                            ("failed", "Failed"),
+                            ("expired", "Expired"),
+                        ],
+                        default="ok",
+                        max_length=12,
+                    ),
+                ),
+                ("worker_id", models.CharField(max_length=255)),
+                ("schedule_date", models.DateTimeField()),
+                (
+                    "task_schedule",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        to="django_fullctl.taskschedule",
+                    ),
+                ),
+            ],
+            options={
+                "verbose_name": "Task Schedule Claim",
+                "verbose_name_plural": "Task Schedule Claims",
+                "db_table": "fullctl_task_schedule_claim",
+                "unique_together": {("task_schedule", "schedule_date")},
+            },
+            managers=[
+                ("handleref", django.db.models.manager.Manager()),
+            ],
+        ),
+    ]

src/fullctl/django/models/concrete/tasks.py

@@ -13,7 +13,7 @@
 from django.contrib.auth import get_user_model
 from django.core.exceptions import ValidationError
 from django.core.management import call_command
-from django.db import models
+from django.db import IntegrityError, models
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 
@@ -36,6 +36,7 @@
     "TaskClaim",
     "TaskHeartbeat",
     "TaskSchedule",
+    "TaskScheduleClaim",
     "CallCommand",
 ]
 
@@ -51,6 +52,9 @@ class TaskClaimed(IOError):
     def __init__(self, task):
         super().__init__(f"Task already claimed by another worker: {task}")
 
+class TaskScheduleClaimed(IOError):
+    def __init__(self, task_schedule):
+        super().__init__(f"Task schedule already claimed by another worker: {task_schedule}")
 
 class WorkerUnqualified(IOError):
     def __init__(self, task, qualifier):
@@ -759,6 +763,17 @@ def spawn_tasks(self):
 
         if self.are_limited_tasks_pending():
             return []
+        
+        # try to create a claim for the schedule
+        try:
+            TaskScheduleClaim.objects.create(
+                task_schedule=self,
+                worker_id=worker_id(),
+                schedule_date=self.schedule,
+            )
+        except IntegrityError:
+            raise TaskScheduleClaimed(self)
+
 
         for task in self.tasks.all():
             if task.status in ["pending", "running"]:
@@ -781,6 +796,19 @@ def spawn_tasks(self):
 
         return tasks
 
+class TaskScheduleClaim(HandleRefModel):
+    task_schedule = models.ForeignKey(TaskSchedule, on_delete=models.CASCADE)
+    worker_id = models.CharField(max_length=255)
+    schedule_date = models.DateTimeField()
+
+    class Meta:
+        db_table = "fullctl_task_schedule_claim"
+        verbose_name = _("Task Schedule Claim")
+        verbose_name_plural = _("Task Schedule Claims")
+        unique_together = (("task_schedule", "schedule_date"))
+
+    class HandleRef:
+        tag = "task_schedule_claim"
 
 class Monitor(HandleRefModel):
     email = models.EmailField(

src/fullctl/django/tasks/orm.py

@@ -17,7 +17,7 @@
     TaskSchedule,
     WorkerUnqualified,
 )
-from fullctl.django.models.concrete.tasks import TaskClaim, TaskLimitError
+from fullctl.django.models.concrete.tasks import TaskClaim, TaskLimitError, TaskScheduleClaimed
 from fullctl.django.tasks import requeue as requeue_task
 from fullctl.django.tasks import specify_task
 from fullctl.django.tasks.util import worker_id
@@ -231,5 +231,7 @@ def progress_schedules(**filters):
 
     try:
         schedule.spawn_tasks()
+    except TaskScheduleClaimed as exc:
+        log.info(f"Task schedule already claimed by another worker: {exc}")
     except (TaskAlreadyStarted, TaskLimitError):
         schedule.reschedule()

tests/django_tests/test_tasks.py

@@ -1,6 +1,8 @@
 import pytest
 from django.conf import settings
+from django.db import IntegrityError, transaction
 from django.utils import timezone
+from unittest.mock import patch
 
 import fullctl.django.tasks.orm as orm
 import tests.django_tests.testapp.models as models
@@ -9,12 +11,17 @@
     health_check_task_stack_queue,
 )
 from fullctl.django.models.concrete.tasks import (
+    TaskClaim,
     TaskClaimed,
     TaskHeartbeatError,
     TaskLimitError,
     TaskMaxAgeError,
     TaskSchedule,
+    TaskScheduleClaim,
+    TaskScheduleClaimed,
+    TaskAlreadyStarted,
 )
+from fullctl.django.tasks.util import worker_id
 
 
 @pytest.mark.django_db
@@ -194,6 +201,41 @@ def test_schedule_limited_task_manually(dj_account_objects):
     task_schedule.spawn_tasks()
 
 
+@pytest.mark.django_db
+def test_task_schedule_cannot_be_claimed_twice(dj_account_objects):
+    """Test that a task schedule cannot be claimed twice."""
+    org = dj_account_objects.org
+    
+    # Create a task schedule
+    task_schedule = TaskSchedule.objects.create(
+        org=org,
+        task_config={
+            # NO tasks so it is forced to raise a TaskScheduleClaimed and not a TaskAlreadyStarted
+            "tasks": [
+            ],
+        },
+        description="test claiming",
+        repeat=True,
+        interval=3600,
+        schedule=timezone.now(),
+    )
+
+    schedule_date = task_schedule.schedule
+    
+    # First claim should succeed
+    task_schedule.spawn_tasks()
+    
+    # reset the schedule date to the original
+    task_schedule.schedule = schedule_date
+    task_schedule.save()
+
+    # Second claim should fail with TaskScheduleClaimed exception
+    with pytest.raises(TaskScheduleClaimed) as exc_info:
+        task_schedule.spawn_tasks()
+    
+    assert "Task schedule already claimed by another worker:" in str(exc_info.value)
+
+
 @pytest.mark.django_db
 def test_task_stack_queue_for_maximum_pending_tasks():
     settings.MAX_PENDING_TASKS = 10
@@ -244,3 +286,56 @@ def test_task_heartbeat_timeout():
 
     with pytest.raises(TaskHeartbeatError):
         health_check_task_heartbeat()
+
+
+@pytest.mark.django_db
+def test_task_schedule_unique_claim_constraint(dj_account_objects):
+    """Test that the TaskScheduleClaim enforces the unique constraint."""
+    
+    org = dj_account_objects.org
+    
+    # Create a task schedule
+    task_schedule = TaskSchedule.objects.create(
+        org=org,
+        task_config={
+            "tasks": [],
+        },
+        description="test claiming constraint",
+        repeat=True,
+        interval=3600,
+        schedule=timezone.now(),
+    )
+    
+    # Create first claim
+    schedule_date = timezone.now()
+    claim1 = TaskScheduleClaim.objects.create(
+        task_schedule=task_schedule,
+        worker_id=worker_id(),
+        schedule_date=schedule_date,
+    )
+    assert claim1.id is not None
+    
+    # Use an atomic block to contain the IntegrityError
+    try:
+        with transaction.atomic():
+            # Attempting to create a second claim with the same task_schedule and schedule_date
+            # should raise an IntegrityError
+            TaskScheduleClaim.objects.create(
+                task_schedule=task_schedule,
+                worker_id=worker_id(),
+                schedule_date=schedule_date,
+            )
+            # If we get here, the test should fail
+            assert False, "Expected IntegrityError was not raised"
+    except IntegrityError:
+        # This is expected, continue with the test
+        pass
+    
+    # Creating a claim with a different schedule_date should succeed
+    different_date = timezone.now() + timezone.timedelta(hours=1)
+    claim2 = TaskScheduleClaim.objects.create(
+        task_schedule=task_schedule,
+        worker_id=worker_id(),
+        schedule_date=different_date,
+    )
+    assert claim2.id is not None