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