diff --git a/mcserver/views.py b/mcserver/views.py index cacadcd..04bf9fc 100644 --- a/mcserver/views.py +++ b/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())}) diff --git a/tests/test_atomic_dequeue.py b/tests/test_atomic_dequeue.py new file mode 100644 index 0000000..1e79da7 --- /dev/null +++ b/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") \ No newline at end of file