Merge pull request #438 from boxerab/thread-safety-client-manager

Add thread safety to client manager with threading lock

tests/test_server_extended.py

@@ -1,5 +1,6 @@
 import json
 import time
+import threading
 import unittest
 from unittest import mock
 from unittest.mock import MagicMock, patch
@@ -35,6 +36,70 @@ class TestClientManagerAddRemove(unittest.TestCase):
         self.cm.remove_client(ws)  # should not raise
 
 
+class TestClientManagerThreadSafety(unittest.TestCase):
+    def test_concurrent_add_remove(self):
+        cm = ClientManager(max_clients=100, max_connection_time=600)
+        errors = []
+
+        def add_clients(start_idx):
+            try:
+                for i in range(50):
+                    ws = MagicMock(name=f"ws-{start_idx}-{i}")
+                    client = MagicMock(name=f"client-{start_idx}-{i}")
+                    cm.add_client(ws, client)
+            except Exception as e:
+                errors.append(e)
+
+        def remove_clients():
+            try:
+                for _ in range(25):
+                    with cm.lock:
+                        if cm.clients:
+                            ws = next(iter(cm.clients))
+                        else:
+                            continue
+                    cm.remove_client(ws)
+            except Exception as e:
+                errors.append(e)
+
+        threads = [
+            threading.Thread(target=add_clients, args=(0,)),
+            threading.Thread(target=add_clients, args=(1,)),
+            threading.Thread(target=remove_clients),
+            threading.Thread(target=remove_clients),
+        ]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        self.assertEqual(errors, [])
+
+    def test_concurrent_get_client(self):
+        cm = ClientManager(max_clients=100, max_connection_time=600)
+        ws = MagicMock()
+        client = MagicMock()
+        cm.add_client(ws, client)
+        errors = []
+        results = []
+
+        def get_many():
+            try:
+                for _ in range(100):
+                    results.append(cm.get_client(ws))
+            except Exception as e:
+                errors.append(e)
+
+        threads = [threading.Thread(target=get_many) for _ in range(4)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        self.assertEqual(errors, [])
+        self.assertTrue(all(r is client for r in results))
+
+
 class TestClientManagerServerFull(unittest.TestCase):
     def setUp(self):
         self.cm = ClientManager(max_clients=1, max_connection_time=60)

whisper_live/server.py

@@ -39,6 +39,7 @@ class ClientManager:
         self.start_times = {}
         self.max_clients = max_clients
         self.max_connection_time = max_connection_time
+        self.lock = threading.Lock()
 
     def add_client(self, websocket, client):
         """
@@ -48,6 +49,7 @@ class ClientManager:
             websocket: The websocket associated with the client to add.
             client: The client object to be added and tracked.
         """
+        with self.lock:
             self.clients[websocket] = client
             self.start_times[websocket] = time.time()
 
@@ -61,6 +63,7 @@ class ClientManager:
         Returns:
             The client object if found, False otherwise.
         """
+        with self.lock:
             if websocket in self.clients:
                 return self.clients[websocket]
             return False
@@ -73,10 +76,11 @@ class ClientManager:
         Args:
             websocket: The websocket associated with the client to be removed.
         """
+        with self.lock:
             client = self.clients.pop(websocket, None)
+            self.start_times.pop(websocket, None)
         if client:
             client.cleanup()
-        self.start_times.pop(websocket, None)
 
     def get_wait_time(self):
         """
@@ -85,6 +89,7 @@ class ClientManager:
         Returns:
             The estimated wait time in minutes for new clients to connect. Returns 0 if there are available slots.
         """
+        with self.lock:
             wait_time = None
             for start_time in self.start_times.values():
                 current_client_time_remaining = self.max_connection_time - (time.time() - start_time)
@@ -103,9 +108,15 @@ class ClientManager:
         Returns:
             True if the server is full, False otherwise.
         """
+        with self.lock:
             if len(self.clients) >= self.max_clients:
-            wait_time = self.get_wait_time()
+                wait_time = None
-            response = {"uid": options["uid"], "status": "WAIT", "message": wait_time}
+                for start_time in self.start_times.values():
+                    remaining = self.max_connection_time - (time.time() - start_time)
+                    if wait_time is None or remaining < wait_time:
+                        wait_time = remaining
+                wait_time_minutes = wait_time / 60 if wait_time is not None else 0
+                response = {"uid": options["uid"], "status": "WAIT", "message": wait_time_minutes}
                 websocket.send(json.dumps(response))
                 return True
             return False
@@ -120,10 +131,12 @@ class ClientManager:
         Returns:
             True if the client's connection time has exceeded the maximum limit, False otherwise.
         """
+        with self.lock:
             elapsed_time = time.time() - self.start_times[websocket]
-        if elapsed_time >= self.max_connection_time:
+            client = self.clients.get(websocket)
-            self.clients[websocket].disconnect()
+        if elapsed_time >= self.max_connection_time and client:
-            logging.warning(f"Client with uid '{self.clients[websocket].client_uid}' disconnected due to overtime.")
+            client.disconnect()
+            logging.warning(f"Client with uid '{client.client_uid}' disconnected due to overtime.")
             return True
         return False