fix: 修复ws引起的redis连接增加

apps/assets/serializers/asset.py

@@ -71,7 +71,7 @@ class AssetSerializer(BulkOrgResourceModelSerializer):
     """
     class Meta:
         model = Asset
-        fields_mini = ['id', 'hostname', 'ip']
+        fields_mini = ['id', 'hostname', 'ip', 'platform', 'protocols']
         fields_small = fields_mini + [
             'protocol', 'port', 'protocols', 'is_active', 'public_ip',
             'number', 'vendor', 'model', 'sn', 'cpu_model', 'cpu_count',

apps/notifications/ws.py

@@ -1,6 +1,6 @@
 import threading
 import json
-
+from redis.exceptions import ConnectionError
 from channels.generic.websocket import JsonWebsocketConsumer
 
 from common.utils import get_logger
@@ -12,13 +12,14 @@ logger = get_logger(__name__)
 
 
 class SiteMsgWebsocket(JsonWebsocketConsumer):
-    disconnected = False
     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.start()
@@ -48,9 +49,8 @@ class SiteMsgWebsocket(JsonWebsocketConsumer):
         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():
+        try:
+            for message in self.chan.listen():
                 if message['type'] != 'message':
                     continue
                 try:
@@ -64,7 +64,10 @@ class SiteMsgWebsocket(JsonWebsocketConsumer):
                         self.send_unread_msg_count()
                 except json.JSONDecoder as e:
                     logger.debug('Decode json error: ', e)
+        except ConnectionError:
+            logger.debug('Redis chan closed')
 
     def disconnect(self, close_code):
-        self.disconnected = True
+        if self.chan is not None:
+            self.chan.close()
         self.close()

apps/users/serializers/user.py

@@ -127,7 +127,7 @@ class UserSerializer(CommonBulkSerializerMixin, serializers.ModelSerializer):
     def validate_password(self, password):
         from ..utils import check_password_rules
         password_strategy = self.initial_data.get('password_strategy')
-        if self.instance is None and password_strategy != self.PasswordStrategy.custom:
+        if self.instance is None and password_strategy != PasswordStrategy.custom:
             # 创建用户，使用邮件设置密码
             return
         if self.instance and not password:

utils/check_celery.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-if [[ "$(ps axu | grep 'celery' | grep -v 'grep' | grep -cv 'defunct')" == "5" ]];then
+if [[ "$(ps axu | grep 'celery' | grep -v 'grep' | grep -cv 'defunct')" -gt "2" ]];then
   exit 0
 else
   exit 1