Merge pull request #24 from PRUNplanner/improve-sse-pubsub-usertracking

improve(sse): remove pubsub leftover, add connection tracking

Author: jplacht

backend/analytics/admin.py

@@ -1,4 +1,5 @@
 import json
+import time
 
 import redis
 from django.apps import apps
@@ -71,6 +72,13 @@ def get_redis_stats():
             r = redis.from_url(settings.CACHES['default']['LOCATION'])
             info = r.info()
 
+            # sse metrics
+            stats_key = 'stream:active_connections'
+            # prune stale sesions
+            r.zremrangebyscore(stats_key, 0, time.time() - 30)
+            # get active users
+            sse_users = r.zcard(stats_key)
+
             hits = info.get('keyspace_hits', 0)
             misses = info.get('keyspace_misses', 0)
             total_reqs = hits + misses
@@ -79,6 +87,7 @@ def get_redis_stats():
             return {
                 'usage': info.get('used_memory_human', '0B'),
                 'hit_rate': hit_rate,
+                'active_stream_users': sse_users,
                 'active_connections': info.get('connected_clients', 0),
                 'blocked_clients': info.get('blocked_clients', 0),
                 'fragmentation': info.get('mem_fragmentation_ratio', 0),

backend/gamedata/api/sse.py

@@ -1,4 +1,6 @@
 import asyncio
+import time
+import uuid
 
 import structlog
 from django.conf import settings
@@ -18,19 +20,12 @@
 )
 @permission_classes([AllowAny])
 async def sse_stream_view(request):
-    raw_channels = request.GET.get('channels', 'beat')
-    channels = [c.strip() for c in raw_channels.split(',') if c.strip()]
+    channels = [c.strip() for c in request.GET.get('channels', 'beat').split(',') if c.strip()]
+    redis_keys = {f'stream:{c}': request.headers.get('Last-Event-ID', '0') for c in channels}
 
-    header_id = request.headers.get('Last-Event-ID')
-
-    if header_id:
-        # catch up to browsers last id
-        last_id = header_id
-    else:
-        # get existing history
-        last_id = '0'
-
-    redis_keys = {f'stream:{c}': last_id for c in channels}
+    # unique connection ID for stats
+    conn_id = str(uuid.uuid4())
+    stats_key = 'stream:active_connections'
 
     redis_url = settings.CACHES['default']['LOCATION']
 
@@ -40,13 +35,14 @@ async def event_generator():
 
         # initiate redis, pubsub and subscribe to channels
         redis = await aioredis.from_url(redis_url)
-        pubsub = redis.pubsub()
-        await pubsub.subscribe(*channels)
 
         total_sent_count = 0
 
-        # loop
         try:
+            # register connection with current timestamp
+            await redis.zadd(stats_key, {conn_id: time.time()})
+
+            # loop
             while True:
                 # XREAD expects a dict with { key: last_id }
                 events = await redis.xread(redis_keys, count=1, block=5000)
@@ -61,6 +57,8 @@ async def event_generator():
                             payload = data[b'payload'].decode('utf-8')
                             yield f'id: {last_id}\ndata: {payload}\n\n'
                 else:
+                    # keep-alive + heartbeat in redis
+                    await redis.zadd(stats_key, {conn_id: time.time()})
                     yield ': \n\n'
 
         # user left stream
@@ -69,7 +67,12 @@ async def event_generator():
 
         # cleanup
         finally:
-            await pubsub.unsubscribe(*channels)
+            # prune this connection
+            await redis.zrem(stats_key, conn_id)
+
+            # prune stale connections (> 30s)
+            await redis.zremrangebyscore(stats_key, 0, time.time() - 30)
+
             await redis.close()
 
             log.info(action='stream_close', messages_sent=total_sent_count)

backend/templates/admin/dashboard/redis_status.html

@@ -4,6 +4,10 @@
       <th>Active Connections</th>
       <td>{{ redis_stats.active_connections }}</td>
     </tr>
+    <tr>
+      <th>Stream Connections</th>
+      <td>{{ redis_stats.active_stream_users }}</td>
+    </tr>
     <tr>
       <th>Usage</th>
       <td>{{ redis_stats.usage }}</td>