merging consumers changes

okbrandon · Oct 31, 2024 · 3cd09c1 · 3cd09c1
1 parent bd170a2
commit 3cd09c1
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 8 deletions.
diff --git a/backend/api/consumers/chat.py b/backend/api/consumers/chat.py
@@ -14,6 +14,7 @@
 
 from django.db.models import Q, Count
 from django.utils import timezone
+from django.core.cache import cache
 
 from ..models import Conversation, User, Relationship, Match, UserSettings, Tournament
 from ..util import generate_id, get_safe_profile, get_user_id_from_token
@@ -46,24 +47,38 @@ async def connect(self):
         self.user = await sync_to_async(User.objects.get)(userID=userID)
         self.user_group_name = f"chat_{self.user.userID}"
 
+        connection_count_key = f"chat_user_connections_{self.user.userID}"
+        connection_count = cache.get(connection_count_key, 0)
+        cache.set(connection_count_key, connection_count + 1, timeout=None)
+
         await self.channel_layer.group_add(
             self.user_group_name,
             self.channel_name
         )
 
-        logger.info(f"[{self.__class__.__name__}] User {self.user.username} connected")
-        await self.ensure_conversations_exist(self.user)
-        logger.info(f"[{self.__class__.__name__}] User {self.user.username} conversations ensured")
+        if connection_count <= 0:
+            await self.ensure_conversations_exist(self.user)
+            logger.info(f"[{self.__class__.__name__}] User {self.user.username} connected, conversations ensured")
+        else:
+            logger.info(f"[{self.__class__.__name__}] User {self.user.username} reconnected")
         await self.accept()
 
     async def disconnect(self, close_code):
         if self.user:
+            connection_count_key = f"chat_user_connections_{self.user.userID}"
+            connection_count = cache.get(connection_count_key, 0) - 1
+
             await self.channel_layer.group_discard(
                 self.user_group_name,
                 self.channel_name
             )
 
-            logger.info(f"[{self.__class__.__name__}] User {self.user.username} disconnected")
+            if connection_count > 0:
+                cache.set(connection_count_key, connection_count, timeout=None)
+                logger.info(f"[{self.__class__.__name__}] User {self.user.username} disconnected, {connection_count} connections remaining")
+            else:
+                cache.delete(connection_count_key)
+                logger.info(f"[{self.__class__.__name__}] User {self.user.username} disconnected")
 
     async def receive(self, text_data):
         try:
@@ -190,4 +205,4 @@ def ensure_conversations_exist(self, user):
                 new_conversation = Conversation.objects.create(conversationID=generate_id("conv"), conversationType='private_message')
                 new_conversation.receipientID = user.userID
                 new_conversation.participants.add(user, friend)
-                new_conversation.save()
+                new_conversation.save()
diff --git a/backend/api/consumers/status.py b/backend/api/consumers/status.py
@@ -10,6 +10,7 @@
 
 from django.db.models import Q
 from django.utils import timezone
+from django.core.cache import cache
 
 from ..models import User, Relationship
 from ..util import get_safe_profile, get_user_id_from_token
@@ -45,6 +46,10 @@ async def connect(self):
         self.user = await sync_to_async(User.objects.get)(userID=userID)
         self.user_group_name = f"status_{self.user.userID}"
 
+        connection_count_key = f"status_user_connections_{self.user.userID}"
+        connection_count = cache.get(connection_count_key, 0)
+        cache.set(connection_count_key, connection_count + 1, timeout=None)
+
         await self.channel_layer.group_add(
             self.user_group_name,
             self.channel_name
@@ -55,16 +60,26 @@ async def connect(self):
         logger.info(f"[{self.__class__.__name__}] User {self.user.username} connected")
 
     async def disconnect(self, close_code):
+
         if self.heartbeat_task and not self.heartbeat_task.done():
             self.heartbeat_task.cancel()
         if self.user is not None:
+            connection_count_key = f"status_user_connections_{self.user.userID}"
+            connection_count = cache.get(connection_count_key, 0) - 1
+
             await self.channel_layer.group_discard(
                 self.user_group_name,
                 self.channel_name
             )
-            await self.update_user_status(False, None)
-            await self.notify_friends_connection(self.user)
-            logger.info(f"[{self.__class__.__name__}] User {self.user.username} disconnected")
+
+            if connection_count > 0:
+                cache.set(connection_count_key, connection_count, timeout=None)
+                logger.info(f"[{self.__class__.__name__}] User {self.user.username} disconnected, {connection_count} connections remaining")
+            else:
+                cache.delete(connection_count_key)
+                await self.update_user_status(False, None)
+                await self.notify_friends_connection(self.user)
+                logger.info(f"[{self.__class__.__name__}] User {self.user.username} disconnected")
 
     async def receive(self, text_data):
         try: