Dequeue atomic

mcserver/views.py

@@ -11,6 +11,7 @@
 
 from datetime import datetime, timedelta
 
+from django.db import transaction
 from django.shortcuts import get_object_or_404
 from django.contrib.auth import login
 from django.core.files.base import ContentFile
@@ -1654,7 +1655,7 @@ def get_permissions(self):
             return [permission() for permission in custom_permission_classes]
 
         return super().get_permissions()
-    
+
     @action(detail=False)
     def dequeue(self, request):
         try:
@@ -1663,86 +1664,88 @@ def dequeue(self, request):
             workerType = self.request.query_params.get('workerType', 'all')
             isMonoQuery = self.request.query_params.get('isMono', 'False')
 
+            with transaction.atomic():
+                # Trials are waiting-for-upload when a missing video was updated
+                # recently. Do not dequeue trials with non-uploaded videos
+                # or saved-local videos, within the 7 day updated-at window.
+                active_trial_cutoff = timezone.now() + timedelta(days=-4)
+                recent_video_cutoff = timezone.now() + timedelta(minutes=-15)
+                missing_video_q = Q(video='')
+                not_uploaded = Video.objects.filter(
+                    missing_video_q,
+                    trial__updated_at__gte=active_trial_cutoff,
+                ).filter(
+                    Q(updated_at__gte=recent_video_cutoff) |
+                    Q(saved_local=True)
+                ).values_list("trial__id", flat=True)
+
+                if isMonoQuery == 'False':
+                    uploaded_trials = Trial.objects.filter(updated_at__gte=active_trial_cutoff).exclude(
+                        id__in=not_uploaded).exclude(session__isMono=True)
+                else:
+                    uploaded_trials = Trial.objects.filter(updated_at__gte=active_trial_cutoff).exclude(
+                        id__in=not_uploaded).filter(session__isMono=True)
 
-
-            # Trials are waiting-for-upload when a missing video was updated
-            # recently. Do not dequeue trials with non-uploaded videos
-            # or saved-local videos, within the 7 day updated-at window.
-            active_trial_cutoff = timezone.now() + timedelta(days=-4)
-            recent_video_cutoff = timezone.now() + timedelta(minutes=-15)
-            missing_video_q = Q(video='')
-            not_uploaded = Video.objects.filter(
-                missing_video_q,
-                trial__updated_at__gte=active_trial_cutoff,
-            ).filter(
-                Q(updated_at__gte=recent_video_cutoff) |
-                Q(saved_local=True)
-            ).values_list("trial__id", flat=True)
-
-
-            if isMonoQuery == 'False':
-                uploaded_trials = Trial.objects.filter(updated_at__gte=active_trial_cutoff).exclude(
-                                                        id__in=not_uploaded).exclude(session__isMono=True)
-            else:
-                uploaded_trials = Trial.objects.filter(updated_at__gte=active_trial_cutoff).exclude(
-                                                        id__in=not_uploaded).filter(session__isMono=True)
-
-            if workerType != 'dynamic':
-                # Priority for 'calibration' and 'neutral'
-                trials = uploaded_trials.filter(status="stopped",
-                                          name__in=["calibration","neutral"],
-                                          result=None)
-
-                trialsReprocess = uploaded_trials.filter(status="reprocess",
-                                          name__in=["calibration","neutral"],
-                                          result=None)
-
-                if trials.count() == 0 and workerType != 'calibration':
+                if workerType != 'dynamic':
+                    # Priority for 'calibration' and 'neutral'
                     trials = uploaded_trials.filter(status="stopped",
-                                              result=None)
+                                                    name__in=["calibration", "neutral"],
+                                                    result=None)
 
-                if trials.count()==0 and trialsReprocess.count() == 0 and workerType != 'calibration':
                     trialsReprocess = uploaded_trials.filter(status="reprocess",
-                                              result=None)
-
-            else:
-                trials = uploaded_trials.filter(status="stopped",
-                                                result=None).exclude(name__in=["calibration", "neutral"])
-
-                trialsReprocess = uploaded_trials.filter(status="reprocess",
-                                                result=None).exclude(name__in=["calibration", "neutral"])
+                                                             name__in=["calibration", "neutral"],
+                                                             result=None)
 
+                    if not trials.exists() and workerType != 'calibration':
+                        trials = uploaded_trials.filter(status="stopped",
+                                                        result=None)
 
-            if trials.count() == 0 and trialsReprocess.count() == 0:
-                raise Http404
-
-            # prioritize admin and priority group trials (priority group doesn't exist yet, but should have same priv. as user)
-            trialsPrioritized = trials.filter(session__user__groups__name__in=["admin"])
-            # if no admin trials, go to priority group trials
-            if trialsPrioritized.count() == 0:
-                trialsPrioritized = trials.filter(session__user__groups__name__in=["priority"])
-            # if not priority trials, go to normal trials
-            if trialsPrioritized.count() == 0:
-                trialsPrioritized = trials
-            # if no normal trials, go to reprocess trials
-            if trials.count() == 0:
-                trialsPrioritized = trialsReprocess
-
-            trial = trialsPrioritized[0]
-            trial.status = "processing"
-            trial.server = ip
-            trial.processed_count += 1
-            trial.save()
+                    if not trials.exists() and not trialsReprocess.exists() and workerType != 'calibration':
+                        trialsReprocess = uploaded_trials.filter(status="reprocess",
+                                                                 result=None)
 
-            if (not trial.session.server) or len(trial.session.server) < 1:
-                session = Session.objects.get(id=trial.session.id)
-                session.server = ip
-                session.save()
+                else:
+                    trials = uploaded_trials.filter(status="stopped",
+                                                    result=None).exclude(name__in=["calibration", "neutral"])
 
-            serializer = TrialSerializer(trial, many=False)
+                    trialsReprocess = uploaded_trials.filter(status="reprocess",
+                                                             result=None).exclude(name__in=["calibration", "neutral"])
+
+                if not trials.exists() and not trialsReprocess.exists():
+                    raise Http404
+
+                # prioritize admin and priority group trials (priority group doesn't exist yet, but should have same priv. as user)
+                trialsPrioritized = trials.filter(session__user__groups__name__in=["admin"])
+                # if no admin trials, go to priority group trials
+                if not trialsPrioritized.exists():
+                    trialsPrioritized = trials.filter(session__user__groups__name__in=["priority"])
+                # if not priority trials, go to normal trials
+                if not trialsPrioritized.exists():
+                    trialsPrioritized = trials
+                # if no normal trials, go to reprocess trials
+                if not trials.exists():
+                    trialsPrioritized = trialsReprocess
+
+                trial = trialsPrioritized.select_for_update(
+                    skip_locked=True,
+                    of=("self",)
+                ).first()
+
+                if trial:
+                    trial.status = "processing"
+                    trial.server = ip
+                    trial.processed_count += 1
+                    trial.save()
+
+                    if (not trial.session.server) or len(trial.session.server) < 1:
+                        session = Session.objects.get(id=trial.session.id)
+                        session.server = ip
+                        session.save()
+
+                serializer = TrialSerializer(trial, many=False)
 
         except Http404:
-            raise Http404 # we use the 404 to tell app.py that there are no trials, so need to pass this thru
+            raise Http404  # we use the 404 to tell app.py that there are no trials, so need to pass this thru
         except Exception:
             if settings.DEBUG:
                 raise APIException(_("error") % {"error_message": str(traceback.format_exc())})

tests/test_atomic_dequeue.py

@@ -0,0 +1,140 @@
+import time
+import threading
+from unittest.mock import patch
+
+from django.test import TransactionTestCase
+from django.db import connection
+from django.urls import reverse, NoReverseMatch
+from rest_framework.test import APIClient
+
+from mcserver.models import Trial, Session
+from django.contrib.auth import get_user_model
+from django.contrib.auth.models import Group
+
+
+class DequeueConcurrencyTest(TransactionTestCase):
+    def setUp(self):
+        User = get_user_model()
+
+        # 1. Setup mock data and permissions
+        self.user = User.objects.create(
+            username="worker_test",
+            is_staff=True,
+            is_superuser=True
+        )
+
+        backend_group, _ = Group.objects.get_or_create(name="backend")
+        admin_group, _ = Group.objects.get_or_create(name="admin")
+        self.user.groups.add(backend_group, admin_group)
+
+        self.session = Session.objects.create(user=self.user, isMono=False)
+
+        # 2. Create valid trials
+        self.trial_1 = Trial.objects.create(
+            session=self.session,
+            name="calibration",
+            status="stopped",
+            result=None
+        )
+        self.trial_2 = Trial.objects.create(
+            session=self.session,
+            name="calibration",
+            status="stopped",
+            result=None
+        )
+
+        # 3. Dynamically resolve the URL to guarantee we hit the right endpoint
+        try:
+            self.dequeue_url = reverse('trial-dequeue')
+        except NoReverseMatch:
+            try:
+                self.dequeue_url = reverse('trials-dequeue')
+            except NoReverseMatch:
+                self.dequeue_url = '/api/trials/dequeue/'
+
+        print(f"\n[DEBUG] Using dequeue URL: {self.dequeue_url}")
+
+    def test_concurrent_dequeue_skips_locked_rows(self):
+        client1 = APIClient()
+        client2 = APIClient()
+
+        client1.force_authenticate(user=self.user)
+        client2.force_authenticate(user=self.user)
+
+        results = {}
+        thread1_locked = threading.Event()
+
+        def worker_1():
+            original_save = Trial.save
+
+            def delayed_save(self_instance, *args, **kwargs):
+                thread1_locked.set()
+                time.sleep(1.5)  # Hold the lock slightly longer
+                return original_save(self_instance, *args, **kwargs)
+
+            try:
+                with patch('mcserver.models.Trial.save', new=delayed_save):
+                    response = client1.get(self.dequeue_url, REMOTE_ADDR='127.0.0.1')
+                    if response.status_code == 200:
+                        results['worker1'] = response.data.get('id')
+                    else:
+                        results[
+                            'worker1_error'] = f"HTTP {response.status_code}: {response.content.decode('utf-8')[:200]}"
+            except Exception as e:
+                results['worker1_error'] = f"Exception: {str(e)}"
+            finally:
+                connection.close()
+
+        def worker_2():
+            # Wait for thread 1 to start its save() and grab the row lock
+            thread1_locked.wait(timeout=3.0)
+            time.sleep(0.2)  # Ensure select_for_update is fully engaged
+
+            try:
+                response = client2.get(self.dequeue_url, REMOTE_ADDR='127.0.0.1')
+                if response.status_code == 200:
+                    results['worker2'] = response.data.get('id')
+                else:
+                    results['worker2_error'] = f"HTTP {response.status_code}: {response.content.decode('utf-8')[:200]}"
+            except Exception as e:
+                results['worker2_error'] = f"Exception: {str(e)}"
+            finally:
+                connection.close()
+
+        t1 = threading.Thread(target=worker_1)
+        t2 = threading.Thread(target=worker_2)
+
+        t1.start()
+        t2.start()
+
+        t1.join()
+        t2.join()
+
+        # Check results
+        self.assertIsNotNone(
+            results.get('worker1'),
+            f"Worker 1 failed! Reason: {results.get('worker1_error')}"
+        )
+        self.assertIsNotNone(
+            results.get('worker2'),
+            f"Worker 2 failed! Reason: {results.get('worker2_error')}"
+        )
+
+        # Confirm they grabbed different trials
+        self.assertEqual(
+            results['worker1'],
+            str(self.trial_1.id),
+            "Worker 1 did not get Trial 1"
+        )
+
+        self.assertEqual(
+            results['worker2'],
+            str(self.trial_2.id),
+            "Worker 2 did not skip locked Trial 1 to get Trial 2"
+        )
+
+        # Confirm they updated correctly in the database
+        self.trial_1.refresh_from_db()
+        self.trial_2.refresh_from_db()
+        self.assertEqual(self.trial_1.status, "processing")
+        self.assertEqual(self.trial_2.status, "processing")