From 7746491e19b4641d4fc7d3b0d726be59651842fb Mon Sep 17 00:00:00 2001
From: feng <1304903146@qq.com>
Date: Wed, 8 Nov 2023 10:26:08 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E5=9C=A8=E7=BA=BF=E7=94=A8=E6=88=B7?=
 =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=98=AF=E5=90=A6=E6=B4=BB=E8=B7=83=E7=9A=84?=
 =?UTF-8?q?=E5=B1=9E=E6=80=A7?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/audits/models.py    |  5 +++++
 apps/notifications/ws.py | 22 ++++++++++++++--------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/apps/audits/models.py b/apps/audits/models.py
index e1ae029f1..11f3d93f9 100644
--- a/apps/audits/models.py
+++ b/apps/audits/models.py
@@ -12,6 +12,7 @@ from django.utils.translation import gettext, gettext_lazy as _
 
 from common.db.encoder import ModelJSONFieldEncoder
 from common.utils import lazyproperty, i18n_trans
+from notifications.ws import WS_SESSION_KEY
 from ops.models import JobExecution
 from orgs.mixins.models import OrgModelMixin, Organization
 from orgs.utils import current_org
@@ -275,6 +276,10 @@ class UserSession(models.Model):
     def backend_display(self):
         return gettext(self.backend)
 
+    @property
+    def is_active(self):
+        return caches.sismember(WS_SESSION_KEY, self.key)
+
     @property
     def date_expired(self):
         session_store_cls = import_module(settings.SESSION_ENGINE).SessionStore
diff --git a/apps/notifications/ws.py b/apps/notifications/ws.py
index 391bb659d..bf9e04bdd 100644
--- a/apps/notifications/ws.py
+++ b/apps/notifications/ws.py
@@ -1,23 +1,27 @@
-import threading
 import json
-from channels.generic.websocket import JsonWebsocketConsumer
 
-from common.utils import get_logger
+from channels.generic.websocket import JsonWebsocketConsumer
+from django.core.cache import caches
+
 from common.db.utils import safe_db_connection
-from .site_msg import SiteMessageUtil
+from common.utils import get_logger
 from .signal_handlers import new_site_msg_chan
+from .site_msg import SiteMessageUtil
 
 logger = get_logger(__name__)
+WS_SESSION_KEY = 'ws_session_key'
 
 
 class SiteMsgWebsocket(JsonWebsocketConsumer):
-    refresh_every_seconds = 10
     sub = None
+    refresh_every_seconds = 10
 
     def connect(self):
         user = self.scope["user"]
         if user.is_authenticated:
             self.accept()
+            session = self.scope['session']
+            caches.sadd(WS_SESSION_KEY, session.session_key)
             self.sub = self.watch_recv_new_site_msg()
         else:
             self.close()
@@ -58,6 +62,8 @@ class SiteMsgWebsocket(JsonWebsocketConsumer):
         return new_site_msg_chan.subscribe(handle_new_site_msg_recv)
 
     def disconnect(self, code):
-        if self.sub:
-            self.sub.unsubscribe()
-
+        if not self.sub:
+            return
+        self.sub.unsubscribe()
+        session = self.scope['session']
+        caches.srem(WS_SESSION_KEY, session.session_key)