Bridge: only update necessary problems when judge sends update signal

judge/bridge/judge_handler.py

@@ -25,6 +25,7 @@ from judge.models import (
     SubmissionTestCase,
 )
 from judge.bridge.utils import VanishedSubmission
+from judge.caching import cache_wrapper
 
 logger = logging.getLogger("judge.bridge")
 json_log = logging.getLogger("judge.json.bridge")
@@ -144,16 +145,41 @@ class JudgeHandler(ZlibPacketHandler):
         self.problems = set(p for p, _ in problem_packet)
 
     def _update_judge_problems(self):
-        problem_codes = list(self.problems)
         chunk_size = 500
-        self.judge.problems.clear()
 
-        for i in range(0, len(problem_codes), chunk_size):
-            chunk = problem_codes[i : i + chunk_size]
+        target_problem_codes = self.problems
+        current_problems = _get_judge_problems(self.judge)
+
+        updated = False
+        problems_to_add = list(target_problem_codes - current_problems)
+        problems_to_remove = list(current_problems - target_problem_codes)
+
+        if problems_to_add:
+            for i in range(0, len(problems_to_add), chunk_size):
+                chunk = problems_to_add[i : i + chunk_size]
                 problem_ids = Problem.objects.filter(code__in=chunk).values_list(
                     "id", flat=True
                 )
+                if not problem_ids:
+                    continue
+                logger.info("%s: Add %d problems", self.name, len(problem_ids))
                 self.judge.problems.add(*problem_ids)
+                updated = True
+
+        if problems_to_remove:
+            for i in range(0, len(problems_to_remove), chunk_size):
+                chunk = problems_to_remove[i : i + chunk_size]
+                problem_ids = Problem.objects.filter(code__in=chunk).values_list(
+                    "id", flat=True
+                )
+                if not problem_ids:
+                    continue
+                logger.info("%s: Remove %d problems", self.name, len(problem_ids))
+                self.judge.problems.remove(*problem_ids)
+                updated = True
+
+        if updated:
+            _get_judge_problems.dirty(self.judge)
 
     def _connected(self):
         judge = self.judge = Judge.objects.get(name=self.name)
@@ -194,6 +220,8 @@ class JudgeHandler(ZlibPacketHandler):
     def _disconnected(self):
         Judge.objects.filter(id=self.judge.id).update(online=False)
         RuntimeVersion.objects.filter(judge=self.judge).delete()
+        self.judge.problems.clear()
+        _get_judge_problems.dirty(self.judge)
 
     def _update_ping(self):
         try:
@@ -931,3 +959,8 @@ class JudgeHandler(ZlibPacketHandler):
 
     def on_cleanup(self):
         db.connection.close()
+
+
+@cache_wrapper(prefix="gjp", timeout=3600)
+def _get_judge_problems(judge):
+    return set(judge.problems.values_list("code", flat=True))

judge/models/problem.py

@@ -561,24 +561,37 @@ class Problem(models.Model, PageVotable, Bookmarkable):
         cache.set(key, result)
         return result
 
-    def save(self, should_move_data=True, *args, **kwargs):
-        code_changed = self.__original_code and self.code != self.__original_code
-        super(Problem, self).save(*args, **kwargs)
-        if code_changed and should_move_data:
-            if hasattr(self, "data_files") or self.pdf_description:
+    def handle_code_change(self):
+        has_data = hasattr(self, "data_files")
+        has_pdf = bool(self.pdf_description)
+        if not has_data and not has_pdf:
+            return
+
         try:
             problem_data_storage.rename(self.__original_code, self.code)
         except OSError as e:
             if e.errno != errno.ENOENT:
                 raise
-                if self.pdf_description:
+
+        if has_pdf:
             self.pdf_description.name = problem_directory_file_helper(
                 self.code, self.pdf_description.name
             )
             super().save(update_fields=["pdf_description"])
-                if hasattr(self, "data_files"):
+
+        if has_data:
             self.data_files._update_code(self.__original_code, self.code)
 
+    def save(self, should_move_data=True, *args, **kwargs):
+        code_changed = self.__original_code and self.code != self.__original_code
+        super(Problem, self).save(*args, **kwargs)
+        if code_changed and should_move_data:
+            self.handle_code_change()
+
+    def delete(self, *args, **kwargs):
+        super().delete(*args, **kwargs)
+        problem_data_storage.delete_directory(self.code)
+
     save.alters_data = True
 
     class Meta:

judge/utils/problem_data.py

@@ -4,6 +4,7 @@ import os
 import re
 import yaml
 import zipfile
+import shutil
 
 from django.conf import settings
 from django.core.files.base import ContentFile
@@ -48,6 +49,10 @@ class ProblemDataStorage(FileSystemStorage):
     def rename(self, old, new):
         return os.rename(self.path(old), self.path(new))
 
+    def delete_directory(self, name):
+        directory_path = self.path(name)
+        shutil.rmtree(directory_path)
+
 
 class ProblemDataError(Exception):
     def __init__(self, message):