perf: 优化订阅处理，形成框架

apps/notifications/ws.py

@@ -1,10 +1,9 @@
 import threading
 import json
-from redis.exceptions import ConnectionError
 from channels.generic.websocket import JsonWebsocketConsumer
 
-from common.db.utils import close_old_connections
 from common.utils import get_logger
+from common.db.utils import safe_db_connection
 from .site_msg import SiteMessageUtil
 from .signals_handler import new_site_msg_chan
 
@@ -13,15 +12,13 @@ logger = get_logger(__name__)
 
 class SiteMsgWebsocket(JsonWebsocketConsumer):
     refresh_every_seconds = 10
-    chan = None
 
     def connect(self):
         user = self.scope["user"]
         if user.is_authenticated:
             self.accept()
-            self.chan = new_site_msg_chan.subscribe()
 
-            thread = threading.Thread(target=self.unread_site_msg_count)
+            thread = threading.Thread(target=self.watch_recv_new_site_msg)
             thread.start()
         else:
             self.close()
@@ -45,45 +42,18 @@ class SiteMsgWebsocket(JsonWebsocketConsumer):
         logger.debug('Send unread count to user: {} {}'.format(user_id, unread_count))
         self.send_json({'type': 'unread_count', 'unread_count': unread_count})
 
-    def unread_site_msg_count(self):
+    def watch_recv_new_site_msg(self):
+        ws = self
         user_id = str(self.scope["user"].id)
-        self.send_unread_msg_count()
-
-        try:
-            msgs = self.chan.listen()
-            # 开始之前关闭连接，因为server端可能关闭了连接，而 client 还在 CONN_MAX_AGE 中
-            for message in msgs:
-                if message['type'] != 'message':
-                    continue
-                close_old_connections()
-                try:
-                    msg = json.loads(message['data'].decode())
-                except json.JSONDecoder as e:
-                    logger.debug('Decode json error: ', e)
-                    continue
-                if not msg:
-                    continue
-
-                logger.debug('New site msg recv, may be mine: {}'.format(msg))
-                users = msg.get('users', [])
-                logger.debug('Message users: {}'.format(users))
-                if user_id in users:
-                    self.send_unread_msg_count()
-                close_old_connections()
-        except ConnectionError:
-            logger.error('Redis chan closed')
-        finally:
-            logger.info('Notification ws thread end')
-            close_old_connections()
-
-    def disconnect(self, close_code):
-        try:
-            if self.chan is not None:
-                self.chan.close()
-            self.close()
-        finally:
-            close_old_connections()
-            logger.info('Notification websocket disconnect')
 
+        # 先发一个消息再说
+        with safe_db_connection():
+            self.send_unread_msg_count()
 
+        def handle_new_site_msg_recv(msg):
+            users = msg.get('users', [])
+            logger.debug('New site msg recv, message users: {}'.format(users))
+            if user_id in users:
+                ws.send_unread_msg_count()
 
+        new_site_msg_chan.keep_handle_msg(handle_new_site_msg_recv)