feat: add live listener count

server.py

@@ -3,10 +3,11 @@ import os
 
 import websocket
 from contextlib import asynccontextmanager
-from fastapi import FastAPI
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 from logger import log_info, log_warn
+from websocket_connection_manager import WebSocketConnectionManager
 
 # the index of the current audio track from 0 to 9
 current_index = -1
@@ -15,6 +16,8 @@ t = None
 # websocket connection to the inference server
 ws = None
 ws_url = ""
+ws_connection_manager = WebSocketConnectionManager()
+active_listeners = set()
 
 
 @asynccontextmanager
@@ -83,7 +86,6 @@ def advance():
         current_index = 0
     else:
         current_index = current_index + 1
-
     threading.Thread(target=generate_new_audio).start()
 
     t = threading.Timer(60, advance)
@@ -98,4 +100,35 @@ def get_current_audio():
     return FileResponse(f"{current_index}.mp3")
 
 
+@app.websocket("/ws")
+async def ws_endpoint(ws: WebSocket):
+    await ws_connection_manager.connect(ws)
+
+    addr = ""
+    if ws.client:
+        addr, _ = ws.client
+    else:
+        await ws.close()
+        ws_connection_manager.disconnect(ws)
+
+    try:
+        while True:
+            msg = await ws.receive_text()
+
+            if msg == "playing":
+                active_listeners.add(addr)
+                await ws_connection_manager.broadcast(f"{len(active_listeners)}")
+            elif msg == "paused":
+                active_listeners.remove(addr)
+                await ws_connection_manager.broadcast(f"{len(active_listeners)}")
+
+    except WebSocketDisconnect:
+        if ws.client:
+            addr, _ = ws.client
+            active_listeners.discard(addr)
+        ws_connection_manager.disconnect(ws)
+
+        await ws_connection_manager.broadcast(f"{len(active_listeners)}")
+
+
 app.mount("/", StaticFiles(directory="web", html=True), name="web")