diff --git a/mcserver/views.py b/mcserver/views.py index cacadcd..137954e 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,7 +1664,19 @@ def dequeue(self, request): workerType = self.request.query_params.get('workerType', 'all') isMonoQuery = self.request.query_params.get('isMono', 'False') + # Use a transaction to ensure atomicity + with transaction.atomic(): + # find trials with some videos not uploaded + not_uploaded = Video.objects.filter(video='', + updated_at__gte=timezone.now() + timedelta( + minutes=-15)).values_list("trial__id", flat=True) + if isMonoQuery == 'False': + uploaded_trials = Trial.objects.filter(updated_at__gte=timezone.now() + timedelta(days=-7)).exclude( + id__in=not_uploaded).exclude(session__isMono=True).select_for_update(skip_locked=True) + else: + uploaded_trials = Trial.objects.filter(updated_at__gte=timezone.now() + timedelta(days=-7)).exclude( + id__in=not_uploaded).filter(session__isMono=True).select_for_update(skip_locked=True) # Trials are waiting-for-upload when a missing video was updated # recently. Do not dequeue trials with non-uploaded videos @@ -1699,50 +1712,64 @@ def dequeue(self, request): if trials.count() == 0 and workerType != 'calibration': 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) + name__in=["calibration", "neutral"], + result=None) - else: - trials = uploaded_trials.filter(status="stopped", - result=None).exclude(name__in=["calibration", "neutral"]) + if trials.count() == 0 and workerType != 'calibration': + trials = uploaded_trials.filter(status="stopped", + result=None) - trialsReprocess = uploaded_trials.filter(status="reprocess", - result=None).exclude(name__in=["calibration", "neutral"]) + 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"]) - 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() + trialsReprocess = uploaded_trials.filter(status="reprocess", + result=None).exclude(name__in=["calibration", "neutral"]) + + 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] + + # Double-check that the trial is still in a claimable state + # This prevents race conditions where another worker might have claimed it + if trial.status not in ["stopped", "reprocess"]: + raise Http404 + + 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() + 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) + 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())})