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]
-            problem_ids = Problem.objects.filter(code__in=chunk).values_list(
-                "id", flat=True
-            )
-            self.judge.problems.add(*problem_ids)
+        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,23 +561,36 @@ class Problem(models.Model, PageVotable, Bookmarkable):
         cache.set(key, result)
         return result
 
+    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 has_pdf:
+            self.pdf_description.name = problem_directory_file_helper(
+                self.code, self.pdf_description.name
+            )
+            super().save(update_fields=["pdf_description"])
+
+        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:
-            if hasattr(self, "data_files") or self.pdf_description:
-                try:
-                    problem_data_storage.rename(self.__original_code, self.code)
-                except OSError as e:
-                    if e.errno != errno.ENOENT:
-                        raise
-                if self.pdf_description:
-                    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"):
-                    self.data_files._update_code(self.__original_code, self.code)
+            self.handle_code_change()
+
+    def delete(self, *args, **kwargs):
+        super().delete(*args, **kwargs)
+        problem_data_storage.delete_directory(self.code)
 
     save.alters_data = True
 

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):