fix: 修复 redis 连接过多的问题

apps/common/utils/connection.py

@@ -1,4 +1,5 @@
 import json
+import threading
 
 import redis
 from django.conf import settings
@@ -19,49 +20,83 @@ def get_redis_client(db):
     return rc
 
 
-class RedisPubSub:
-    def __init__(self, ch, db=10):
+class Subscription:
+    def __init__(self, ch, sub, ):
         self.ch = ch
-        self.redis = get_redis_client(db)
-
-    def subscribe(self):
-        ps = self.redis.pubsub()
-        ps.subscribe(self.ch)
-        return ps
+        self.sub = sub
 
-    def publish(self, data):
-        data_json = json.dumps(data)
-        self.redis.publish(self.ch, data_json)
-        return True
-
-    def keep_handle_msg(self, handle):
+    def _handle_msg(self, _next, error, complete):
         """
         handle arg is the pub published
 
-        :param handle: lambda item: do_something
+        :param _next: next msg handler
+        :param error: error msg handler
+        :param complete: complete msg handler
         :return:
         """
-        sub = self.subscribe()
-        msgs = sub.listen()
+        msgs = self.sub.listen()
+
+        if error is None:
+            error = lambda m, i: None
+
+        if complete is None:
+            complete = lambda: None
 
         try:
             for msg in msgs:
                 if msg["type"] != "message":
                     continue
+                item = None
                 try:
                     item_json = msg['data'].decode()
                     item = json.loads(item_json)
 
                     with safe_db_connection():
-                        handle(item)
+                        _next(item)
                 except Exception as e:
+                    error(msg, item)
                     logger.error('Subscribe handler handle msg error: ', e)
-
         except Exception as e:
             logger.error('Consume msg error: ', e)
 
         try:
-            sub.close()
+            complete()
         except Exception as e:
-            logger.error("Redis observer close error: ", e)
+            logger.error('Complete subscribe error: {}'.format(e))
+            pass
+
+        try:
+            self.unsubscribe()
+        except Exception as e:
+            logger.error("Redis observer close error: {}".format(e))
+
+    def keep_handle_msg(self, _next, error, complete):
+        t = threading.Thread(target=self._handle_msg, args=(_next, error, complete))
+        t.daemon = True
+        t.start()
+        return t
+
+    def unsubscribe(self):
+        try:
+            self.sub.close()
+        except Exception as e:
+            logger.error('Unsubscribe msg error: {}'.format(e))
+
+
+class RedisPubSub:
+    def __init__(self, ch, db=10):
+        self.ch = ch
+        self.redis = get_redis_client(db)
+
+    def subscribe(self, _next, error=None, complete=None):
+        ps = self.redis.pubsub()
+        sub = Subscription(self.ch, ps)
+        sub.keep_handle_msg(_next, error, complete)
+        return sub
+
+    def publish(self, data):
+        data_json = json.dumps(data)
+        self.redis.publish(self.ch, data_json)
+        return True
 
+    

apps/notifications/ws.py

@@ -12,14 +12,13 @@ logger = get_logger(__name__)
 
 class SiteMsgWebsocket(JsonWebsocketConsumer):
     refresh_every_seconds = 10
+    sub = None
 
     def connect(self):
         user = self.scope["user"]
         if user.is_authenticated:
             self.accept()
-
-            thread = threading.Thread(target=self.watch_recv_new_site_msg)
-            thread.start()
+            self.sub = self.watch_recv_new_site_msg()
         else:
             self.close()
 
@@ -56,4 +55,9 @@ class SiteMsgWebsocket(JsonWebsocketConsumer):
             if user_id in users:
                 ws.send_unread_msg_count()
 
-        new_site_msg_chan.keep_handle_msg(handle_new_site_msg_recv)
+        return new_site_msg_chan.subscribe(handle_new_site_msg_recv)
+
+    def disconnect(self, code):
+        if self.sub:
+            self.sub.unsubscribe()
+