fix: 修复 redis 连接导致的 pub sub 错误

apps/assets/signal_handlers/node_assets_mapping.py

@@ -1,7 +1,5 @@
 # -*- coding: utf-8 -*-
 #
-import os
-import threading
 
 from django.db.models.signals import (
     m2m_changed, post_save, post_delete
@@ -9,15 +7,15 @@ from django.db.models.signals import (
 from django.dispatch import receiver
 from django.utils.functional import LazyObject
 
-from common.signals import django_ready
-from common.utils.connection import RedisPubSub
-from common.utils import get_logger
 from assets.models import Asset, Node
+from common.signals import django_ready
+from common.utils import get_logger
+from common.utils.connection import RedisPubSub
 from orgs.models import Organization
 
-
 logger = get_logger(__file__)
 
+
 # clear node assets mapping for memory
 # ------------------------------------
 
@@ -78,9 +76,4 @@ def subscribe_node_assets_mapping_expire(sender, **kwargs):
         Node.expire_node_all_asset_ids_mapping_from_memory(org_id)
         Node.expire_node_all_asset_ids_mapping_from_memory(root_org_id)
 
-    def keep_subscribe_node_assets_relation():
-        node_assets_mapping_for_memory_pub_sub.subscribe(handle_node_relation_change)
-
-    t = threading.Thread(target=keep_subscribe_node_assets_relation)
-    t.daemon = True
-    t.start()
+    node_assets_mapping_for_memory_pub_sub.subscribe(handle_node_relation_change)

apps/common/utils/connection.py

@@ -1,8 +1,10 @@
 import json
 import threading
-import redis
+import time
 
+import redis
 from django.core.cache import cache
+from redis.client import PubSub
 
 from common.db.utils import safe_db_connection
 from common.utils import get_logger
@@ -16,15 +18,39 @@ def get_redis_client(db=0):
     return client
 
 
-class Subscription:
-    def __init__(self, ch, sub, ):
+class RedisPubSub:
+    def __init__(self, ch, db=10):
         self.ch = ch
+        self.db = db
+        self.redis = get_redis_client(db)
+
+    def subscribe(self, _next, error=None, complete=None):
+        ps = self.redis.pubsub()
+        ps.subscribe(self.ch)
+        sub = Subscription(self, ps)
+        sub.keep_handle_msg(_next, error, complete)
+        return sub
+
+    def resubscribe(self, _next, error=None, complete=None):
+        self.redis = get_redis_client(self.db)
+        self.subscribe(_next, error, complete)
+
+    def publish(self, data):
+        data_json = json.dumps(data)
+        self.redis.publish(self.ch, data_json)
+        return True
+
+
+class Subscription:
+    def __init__(self, pb: RedisPubSub, sub: PubSub):
+        self.pb = pb
+        self.ch = pb.ch
         self.sub = sub
+        self.unsubscribed = False
 
     def _handle_msg(self, _next, error, complete):
         """
         handle arg is the pub published
-
         :param _next: next msg handler
         :param error: error msg handler
         :param complete: complete msg handler
@@ -53,9 +79,12 @@ class Subscription:
                     error(msg, item)
                     logger.error('Subscribe handler handle msg error: {}'.format(e))
         except Exception as e:
-            # 正常的 websocket 断开时, redis 会断开连接，避免日志太多
-            # logger.error('Consume msg error: {}'.format(e))
-            pass
+            if self.unsubscribed:
+                logger.debug('Subscription unsubscribed')
+            else:
+                logger.error('Consume msg error: {}'.format(e))
+                self.retry(_next, error, complete)
+                return
 
         try:
             complete()
@@ -75,25 +104,22 @@ class Subscription:
         return t
 
     def unsubscribe(self):
+        self.unsubscribed = True
         try:
             self.sub.close()
         except Exception as e:
             logger.error('Unsubscribe msg error: {}'.format(e))
 
+    def retry(self, _next, error, complete):
+        logger.info('Retry subscribe channel: {}'.format(self.ch))
+        times = 0
 
-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()
-        ps.subscribe(self.ch)
-        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
+        while True:
+            try:
+                self.unsubscribe()
+                self.pb.resubscribe(_next, error, complete)
+                break
+            except Exception as e:
+                logger.error('Retry #{} {} subscribe channel error: {}'.format(times, self.ch, e))
+                times += 1
+                time.sleep(times * 2)

apps/orgs/signal_handlers/common.py

@@ -1,29 +1,27 @@
 # -*- coding: utf-8 -*-
 #
-import threading
 from collections import defaultdict
 from functools import partial
 
-from django.dispatch import receiver
-from django.utils.functional import LazyObject
 from django.db.models.signals import m2m_changed
 from django.db.models.signals import post_save, pre_delete
+from django.dispatch import receiver
+from django.utils.functional import LazyObject
 
-from orgs.utils import tmp_to_org
-from orgs.models import Organization
-from orgs.hands import set_current_org, Node, get_current_org
-from perms.models import (AssetPermission, ApplicationPermission)
-from users.models import UserGroup, User
+from assets.models import CommandFilterRule
 from assets.models import SystemUser
 from common.const.signals import PRE_REMOVE, POST_REMOVE
 from common.decorator import on_transaction_commit
 from common.signals import django_ready
 from common.utils import get_logger
 from common.utils.connection import RedisPubSub
-from assets.models import CommandFilterRule
+from orgs.hands import set_current_org, Node, get_current_org
+from orgs.models import Organization
+from orgs.utils import tmp_to_org
+from perms.models import (AssetPermission, ApplicationPermission)
+from users.models import UserGroup, User
 from users.signals import post_user_leave_org
 
-
 logger = get_logger(__file__)
 
 
@@ -47,14 +45,9 @@ def expire_orgs_mapping_for_memory(org_id):
 def subscribe_orgs_mapping_expire(sender, **kwargs):
     logger.debug("Start subscribe for expire orgs mapping from memory")
 
-    def keep_subscribe_org_mapping():
-        orgs_mapping_for_memory_pub_sub.subscribe(
-            lambda org_id: Organization.expire_orgs_mapping()
-        )
-
-    t = threading.Thread(target=keep_subscribe_org_mapping)
-    t.daemon = True
-    t.start()
+    orgs_mapping_for_memory_pub_sub.subscribe(
+        lambda org_id: Organization.expire_orgs_mapping()
+    )
 
 
 # 创建对应的root
@@ -85,7 +78,7 @@ def on_org_delete(sender, instance, **kwargs):
 def _remove_users(model, users, org, user_field_name='users'):
     with tmp_to_org(org):
         if not isinstance(users, (tuple, list, set)):
-            users = (users, )
+            users = (users,)
 
         user_field = getattr(model, user_field_name)
         m2m_model = user_field.through

apps/settings/signal_handlers.py

@@ -1,19 +1,18 @@
 # -*- coding: utf-8 -*-
 #
 import json
-import threading
 
 from django.conf import LazySettings
+from django.db.models.signals import post_save, pre_save
 from django.db.utils import ProgrammingError, OperationalError
 from django.dispatch import receiver
-from django.db.models.signals import post_save, pre_save
 from django.utils.functional import LazyObject
 
-from jumpserver.utils import current_request
 from common.decorator import on_transaction_commit
+from common.signals import django_ready
 from common.utils import get_logger, ssh_key_gen
 from common.utils.connection import RedisPubSub
-from common.signals import django_ready
+from jumpserver.utils import current_request
 from .models import Setting
 
 logger = get_logger(__file__)
@@ -81,12 +80,7 @@ def on_create_set_created_by(sender, instance=None, **kwargs):
 def subscribe_settings_change(sender, **kwargs):
     logger.debug("Start subscribe setting change")
 
-    def keep_subscribe_settings_change():
-        setting_pub_sub.subscribe(lambda name: Setting.refresh_item(name))
-
-    t = threading.Thread(target=keep_subscribe_settings_change)
-    t.daemon = True
-    t.start()
+    setting_pub_sub.subscribe(lambda name: Setting.refresh_item(name))
 
 
 @receiver(django_ready)