Fix bridge (DMOJ/online-judge#2080)

judge/bridge/base_handler.py

@@ -145,6 +145,9 @@ def on_disconnect(self):
     def on_timeout(self):
         pass
 
+    def on_cleanup(self):
+        pass
+
     def handle(self):
         try:
             tag = self.read_size()
@@ -204,6 +207,8 @@ def handle(self):
             if e.__class__.__name__ == "cancel_wait_ex":
                 return
             raise
+        finally:
+            self.on_cleanup()
 
     def send(self, data):
         compressed = zlib.compress(data.encode("utf-8"))

judge/bridge/django_handler.py

@@ -2,6 +2,8 @@
 import logging
 import struct
 
+from django import db
+
 from judge.bridge.base_handler import Disconnect, ZlibPacketHandler
 
 logger = logging.getLogger("judge.bridge")
@@ -60,5 +62,5 @@ def on_disconnect_request(self, data):
     def on_malformed(self, packet):
         logger.error("Malformed packet: %s", packet)
 
-    def on_close(self):
-        self._to_kill = False
+    def on_cleanup(self):
+        db.connection.close()

judge/bridge/judge_handler.py

@@ -36,10 +36,7 @@
 
 
 def _ensure_connection():
-    try:
-        db.connection.cursor().execute("SELECT 1").fetchall()
-    except Exception:
-        db.connection.close()
+    db.connection.close_if_unusable_or_obsolete()
 
 
 class JudgeHandler(ZlibPacketHandler):
@@ -903,3 +900,6 @@ def _post_update_submission(self, id, state, done=False):
                     "language": data["language__key"],
                 },
             )
+
+    def on_cleanup(self):
+        db.connection.close()