From 7d5efa60c749674c9bb6c8a5613f9b1d2a586f4c Mon Sep 17 00:00:00 2001
From: cuom1999 <lephuocdinh99@gmail.com>
Date: Tue, 27 Dec 2022 13:19:02 -0600
Subject: [PATCH] Fix bridge (https://github.com/DMOJ/online-judge/pull/2080)

---
 judge/bridge/base_handler.py   | 5 +++++
 judge/bridge/django_handler.py | 6 ++++--
 judge/bridge/judge_handler.py  | 8 ++++----
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/judge/bridge/base_handler.py b/judge/bridge/base_handler.py
index 369a8e1..0e5c067 100644
--- a/judge/bridge/base_handler.py
+++ b/judge/bridge/base_handler.py
@@ -145,6 +145,9 @@ class ZlibPacketHandler(metaclass=RequestHandlerMeta):
     def on_timeout(self):
         pass
 
+    def on_cleanup(self):
+        pass
+
     def handle(self):
         try:
             tag = self.read_size()
@@ -204,6 +207,8 @@ class ZlibPacketHandler(metaclass=RequestHandlerMeta):
             if e.__class__.__name__ == "cancel_wait_ex":
                 return
             raise
+        finally:
+            self.on_cleanup()
 
     def send(self, data):
         compressed = zlib.compress(data.encode("utf-8"))
diff --git a/judge/bridge/django_handler.py b/judge/bridge/django_handler.py
index f72aba1..8a40c3f 100644
--- a/judge/bridge/django_handler.py
+++ b/judge/bridge/django_handler.py
@@ -2,6 +2,8 @@ import json
 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 @@ class DjangoHandler(ZlibPacketHandler):
     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()
diff --git a/judge/bridge/judge_handler.py b/judge/bridge/judge_handler.py
index 1eaeac2..cca520f 100644
--- a/judge/bridge/judge_handler.py
+++ b/judge/bridge/judge_handler.py
@@ -36,10 +36,7 @@ SubmissionData = namedtuple(
 
 
 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 @@ class JudgeHandler(ZlibPacketHandler):
                     "language": data["language__key"],
                 },
             )
+
+    def on_cleanup(self):
+        db.connection.close()