From 03cf356bab79b2401cadb07342dffeea5624b13e Mon Sep 17 00:00:00 2001
From: virusdefender <virusdefender@outlook.com>
Date: Sun, 24 Dec 2017 12:10:02 +0800
Subject: [PATCH] add judge server is_disabled field

---
 .../0003_judgeserver_is_disabled.py           | 20 +++++++++++++++++++
 conf/models.py                                |  1 +
 conf/serializers.py                           |  5 +++++
 conf/tests.py                                 | 11 +++++++---
 conf/views.py                                 |  8 +++++++-
 judge/dispatcher.py                           |  3 +--
 6 files changed, 42 insertions(+), 6 deletions(-)
 create mode 100644 conf/migrations/0003_judgeserver_is_disabled.py

diff --git a/conf/migrations/0003_judgeserver_is_disabled.py b/conf/migrations/0003_judgeserver_is_disabled.py
new file mode 100644
index 00000000..6a571c01
--- /dev/null
+++ b/conf/migrations/0003_judgeserver_is_disabled.py
@@ -0,0 +1,20 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.11.4 on 2017-12-24 03:44
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('conf', '0002_auto_20171011_1214'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='judgeserver',
+            name='is_disabled',
+            field=models.BooleanField(default=False),
+        ),
+    ]
diff --git a/conf/models.py b/conf/models.py
index 4c6348d5..2317f81c 100644
--- a/conf/models.py
+++ b/conf/models.py
@@ -13,6 +13,7 @@ class JudgeServer(models.Model):
     create_time = models.DateTimeField(auto_now_add=True)
     task_number = models.IntegerField(default=0)
     service_url = models.CharField(max_length=256, blank=True, null=True)
+    is_disabled = models.BooleanField(default=False)
 
     @property
     def status(self):
diff --git a/conf/serializers.py b/conf/serializers.py
index 3dd16650..50f259c1 100644
--- a/conf/serializers.py
+++ b/conf/serializers.py
@@ -44,3 +44,8 @@ class JudgeServerHeartbeatSerializer(serializers.Serializer):
     cpu = serializers.FloatField(min_value=0, max_value=100)
     action = serializers.ChoiceField(choices=("heartbeat", ))
     service_url = serializers.CharField(max_length=256)
+
+
+class EditJudgeServerSerializer(serializers.Serializer):
+    id = serializers.IntegerField()
+    is_disabled = serializers.BooleanField()
diff --git a/conf/tests.py b/conf/tests.py
index 4d5370df..f6debd62 100644
--- a/conf/tests.py
+++ b/conf/tests.py
@@ -105,9 +105,9 @@ class JudgeServerHeartbeatTest(APITestCase):
 
 class JudgeServerAPITest(APITestCase):
     def setUp(self):
-        JudgeServer.objects.create(**{"hostname": "testhostname", "judger_version": "1.0.4",
-                                      "cpu_core": 4, "cpu_usage": 90.5, "memory_usage": 80.3,
-                                      "last_heartbeat": timezone.now()})
+        self.server = JudgeServer.objects.create(**{"hostname": "testhostname", "judger_version": "1.0.4",
+                                                    "cpu_core": 4, "cpu_usage": 90.5, "memory_usage": 80.3,
+                                                    "last_heartbeat": timezone.now()})
         self.url = self.reverse("judge_server_api")
         self.create_super_admin()
 
@@ -121,6 +121,11 @@ class JudgeServerAPITest(APITestCase):
         self.assertSuccess(resp)
         self.assertFalse(JudgeServer.objects.filter(hostname="testhostname").exists())
 
+    def test_disabled_judge_server(self):
+        resp = self.client.put(self.url, data={"is_disabled": True, "id": self.server.id})
+        self.assertSuccess(resp)
+        self.assertTrue(JudgeServer.objects.get(id=self.server.id).is_disabled)
+
 
 class LanguageListAPITest(APITestCase):
     def test_get_languages(self):
diff --git a/conf/views.py b/conf/views.py
index 2bd059e5..a4f070c9 100644
--- a/conf/views.py
+++ b/conf/views.py
@@ -17,7 +17,7 @@ from .models import JudgeServer
 from .serializers import (CreateEditWebsiteConfigSerializer,
                           CreateSMTPConfigSerializer, EditSMTPConfigSerializer,
                           JudgeServerHeartbeatSerializer,
-                          JudgeServerSerializer, TestSMTPConfigSerializer)
+                          JudgeServerSerializer, TestSMTPConfigSerializer, EditJudgeServerSerializer)
 
 
 class SMTPAPI(APIView):
@@ -102,6 +102,12 @@ class JudgeServerAPI(APIView):
             JudgeServer.objects.filter(hostname=hostname).delete()
         return self.success()
 
+    @validate_serializer(EditJudgeServerSerializer)
+    @super_admin_required
+    def put(self, request):
+        JudgeServer.objects.filter(id=request.data["id"]).update(is_disabled=request.data["is_disabled"])
+        return self.success()
+
 
 class JudgeServerHeartbeatAPI(CSRFExemptAPIView):
     @validate_serializer(JudgeServerHeartbeatSerializer)
diff --git a/judge/dispatcher.py b/judge/dispatcher.py
index 682e054e..ecb5dbf6 100644
--- a/judge/dispatcher.py
+++ b/judge/dispatcher.py
@@ -6,7 +6,6 @@ from urllib.parse import urljoin
 import requests
 from django.db import transaction
 from django.db.models import F
-from django.conf import settings
 
 from account.models import User
 from conf.models import JudgeServer
@@ -47,7 +46,7 @@ class DispatcherBase(object):
     @staticmethod
     def choose_judge_server():
         with transaction.atomic():
-            servers = JudgeServer.objects.select_for_update().all().order_by("task_number")
+            servers = JudgeServer.objects.select_for_update().filter(is_disabled=False).order_by("task_number")
             servers = [s for s in servers if s.status == "normal"]
             if servers:
                 server = servers[0]