perf: db connection close if needt

apps/common/apps.py

@@ -4,6 +4,7 @@ import os
 import sys
 
 from django.apps import AppConfig
+from django.db import close_old_connections
 
 
 class CommonConfig(AppConfig):
@@ -21,3 +22,4 @@ class CommonConfig(AppConfig):
 
         if not os.environ.get('DJANGO_DEBUG_SHELL'):
             django_ready.send(CommonConfig)
+            close_old_connections()

apps/common/db/utils.py

@@ -56,7 +56,7 @@ def close_old_connections():
 
 
 @contextmanager
-def safe_db_connection():
+def safe_db_connection(auto_close=False):
     in_atomic_block = connection.in_atomic_block  # 当前是否处于事务中
     autocommit = transaction.get_autocommit()  # 是否启用了自动提交
     created = False
@@ -69,7 +69,7 @@ def safe_db_connection():
         yield
     finally:
         # 如果不是事务中（API 请求中可能需要提交事务），则关闭连接
-        if created and not in_atomic_block and autocommit:
+        if auto_close or (created and not in_atomic_block and autocommit):
             print("close connection in safe_db_connection")
             close_old_connections()
 

apps/notifications/ws.py

@@ -50,7 +50,7 @@ class SiteMsgWebsocket(JsonWebsocketConsumer):
         user_id = str(self.scope["user"].id)
 
         # 先发一个消息再说
-        with safe_db_connection():
+        with safe_db_connection(auto_close=True):
             self.send_unread_msg_count()
 
         def handle_new_site_msg_recv(msg):

apps/ops/signal_handlers.py

@@ -115,11 +115,10 @@ def on_celery_task_pre_run(task_id='', kwargs=None, **others):
 
 @signals.task_postrun.connect
 def on_celery_task_post_run(task_id='', state='', **kwargs):
-    close_old_connections()
-
     CeleryTaskExecution.objects.filter(id=task_id).update(
         state=state, date_finished=timezone.now(), is_finished=True
     )
+    close_old_connections()
 
 
 @signals.after_task_publish.connect

apps/terminal/ws.py

@@ -52,7 +52,7 @@ class TerminalTaskWebsocket(JsonWebsocketConsumer):
         self.send(bytes_data=content)
 
     def get_terminal_tasks(self, task_id=None):
-        with safe_db_connection():
+        with safe_db_connection(auto_close=True):
             critical_time = timezone.now() - datetime.timedelta(minutes=10)
             tasks = self.terminal.task_set.filter(is_finished=False, date_created__gte=critical_time)
             if task_id: