perf: 优化消息中心未读数量

apps/notifications/ws.py

@@ -0,0 +1,70 @@
+import threading
+import json
+
+from channels.generic.websocket import JsonWebsocketConsumer
+
+from common.utils import get_logger
+from .models import SiteMessage
+from .site_msg import SiteMessageUtil
+from .signals_handler import new_site_msg_chan
+
+logger = get_logger(__name__)
+
+
+class SiteMsgWebsocket(JsonWebsocketConsumer):
+    disconnected = False
+    refresh_every_seconds = 10
+
+    def connect(self):
+        user = self.scope["user"]
+        if user.is_authenticated:
+            self.accept()
+
+            thread = threading.Thread(target=self.unread_site_msg_count)
+            thread.start()
+        else:
+            self.close()
+
+    def receive(self, text_data=None, bytes_data=None, **kwargs):
+        data = json.loads(text_data)
+        refresh_every_seconds = data.get('refresh_every_seconds')
+
+        try:
+            refresh_every_seconds = int(refresh_every_seconds)
+        except Exception as e:
+            logger.error(e)
+            return
+
+        if refresh_every_seconds > 0:
+            self.refresh_every_seconds = refresh_every_seconds
+
+    def send_unread_msg_count(self):
+        user_id = self.scope["user"].id
+        unread_count = SiteMessageUtil.get_user_unread_msgs_count(user_id)
+        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):
+        user_id = str(self.scope["user"].id)
+        self.send_unread_msg_count()
+
+        while not self.disconnected:
+            subscribe = new_site_msg_chan.subscribe()
+            for message in subscribe.listen():
+                if message['type'] != 'message':
+                    continue
+                try:
+                    msg = json.loads(message['data'].decode())
+                    logger.debug('New site msg recv, may be mine: {}'.format(msg))
+                    if not msg:
+                        continue
+                    users = msg.get('users', [])
+                    logger.debug('Message users: {}'.format(users))
+                    if user_id in users:
+                        self.send_unread_msg_count()
+                except json.JSONDecoder as e:
+                    logger.debug('Decode json error: ', e)
+
+    def disconnect(self, close_code):
+        self.disconnected = True
+        self.close()