From 81cdbbca95499537d402c39a3cdfb822dfad27fe Mon Sep 17 00:00:00 2001 From: Aaron Boxer Date: Fri, 17 Apr 2026 09:27:37 -0400 Subject: [PATCH] Add thread safety to ClientManager with threading.Lock - All ClientManager methods (add_client, get_client, remove_client, get_wait_time, is_server_full, is_client_timeout) now protected by a threading.Lock - cleanup() called outside the lock to avoid holding it during I/O - is_server_full() computes wait time inline under lock instead of calling get_wait_time() to avoid nested lock acquisition - Added concurrent thread safety tests for add/remove and get operations --- tests/test_server_extended.py | 65 +++++++++++++++++++++++++++++++++++ whisper_live/server.py | 57 ++++++++++++++++++------------ 2 files changed, 100 insertions(+), 22 deletions(-) diff --git a/tests/test_server_extended.py b/tests/test_server_extended.py index 699b3fb..3172f4e 100644 --- a/tests/test_server_extended.py +++ b/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) diff --git a/whisper_live/server.py b/whisper_live/server.py index 1bd2dfc..e3aaf6b 100644 --- a/whisper_live/server.py +++ b/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