thread safety

app.py

@@ -34,7 +34,7 @@ max_retries = config.get("max_retries", 3)
 
 # Configure logging
 logging.basicConfig(
-    level=logging.INFO,
+    level=logging.DEBUG,
     format="%(asctime)s [%(levelname)s] %(message)s",
     handlers=[
         logging.FileHandler("app.log"),
@@ -49,6 +49,10 @@ os.makedirs("models", exist_ok=True)
 HEARTBEAT_INTERVAL = 2  # seconds
 WORKER_TIMEOUT_MS = 10000
 
+# Add a lock for thread-safe operations on shared resources
+streams_lock = threading.Lock()
+models_lock = threading.Lock()
+
 @app.websocket("/")
 async def detect(websocket: WebSocket):
     import asyncio
@@ -61,6 +65,7 @@ async def detect(websocket: WebSocket):
     # This function is user-modifiable
     # Save data you want to persist across frames in the persistent_data dictionary
     async def handle_detection(camera_id, stream, frame, websocket, model: YOLO, persistent_data):
+        try:
             boxes = []
             for r in model.track(frame, stream=False, persist=True):
                 for box in r.boxes:
@@ -87,10 +92,14 @@ async def detect(websocket: WebSocket):
             logging.debug(f"Sending detection data for camera {camera_id}: {detection_data}")
             await websocket.send_json(detection_data)
             return persistent_data
+        except Exception as e:
+            logging.error(f"Error in handle_detection for camera {camera_id}: {e}")
+            return persistent_data
 
     def frame_reader(camera_id, cap, buffer, stop_event):
         import time
         retries = 0
+        try:
             while not stop_event.is_set():
                 try:
                     ret, frame = cap.read()
@@ -132,6 +141,8 @@ async def detect(websocket: WebSocket):
                     logging.error(f"Unexpected error for camera {camera_id}: {e}")
                     cap.release()
                     break
+        except Exception as e:
+            logging.error(f"Error in frame_reader thread for camera {camera_id}: {e}")
 
     async def process_streams():
         global models
@@ -141,10 +152,13 @@ async def detect(websocket: WebSocket):
             while True:
                 start_time = time.time()
                 # Round-robin processing
-                for camera_id, stream in list(streams.items()):
+                with streams_lock:
+                    current_streams = list(streams.items())
+                for camera_id, stream in current_streams:
                     buffer = stream['buffer']
                     if not buffer.empty():
                         frame = buffer.get()
+                        with models_lock:
                             model = models.get(camera_id, {}).get(stream['modelId'])
                         key = (camera_id, stream['modelId'])
                         persistent_data = persistent_data_dict.get(key, {})
@@ -199,6 +213,7 @@ async def detect(websocket: WebSocket):
     async def on_message():
         global models
         while True:
+            try:
                 msg = await websocket.receive_text()
                 logging.debug(f"Received message: {msg}")
                 data = json.loads(msg)
@@ -213,6 +228,7 @@ async def detect(websocket: WebSocket):
                     modelName = payload.get("modelName")
     
                     if model_url:
+                        with models_lock:
                             if camera_id not in models:
                                 models[camera_id] = {}
                             if modelId not in models[camera_id]:
@@ -236,6 +252,7 @@ async def detect(websocket: WebSocket):
                                     logging.error(f"Failed to download model from {model_url}")
                                     continue
                     if camera_id and rtsp_url:
+                        with streams_lock:
                             if camera_id not in streams and len(streams) < max_streams:
                                 cap = cv2.VideoCapture(rtsp_url)
                                 if not cap.isOpened():
@@ -267,6 +284,8 @@ async def detect(websocket: WebSocket):
                 elif msg_type == "unsubscribe":
                     payload = data.get("payload", {})
                     camera_id = payload.get("cameraIdentifier")
+                    logging.debug(f"Unsubscribing from camera {camera_id}")
+                    with streams_lock:
                         if camera_id and camera_id in streams:
                             stream = streams.pop(camera_id)
                             stream['cap'].release()
@@ -307,76 +326,6 @@ async def detect(websocket: WebSocket):
                     await websocket.send_text(json.dumps(state_report))
                 else:
                     logging.error(f"Unknown message type: {msg_type}")
-
-    await websocket.accept()
-    task = asyncio.create_task(process_streams())
-    heartbeat_task = asyncio.create_task(send_heartbeat())
-    message_task = asyncio.create_task(on_message())
-
-    await asyncio.gather(heartbeat_task, message_task)
-
-    model = None
-    model_path = None
-
-    try:
-        while True:
-            try:
-                msg = await websocket.receive_text()
-                logging.debug(f"Received message: {msg}")
-                data = json.loads(msg)
-                camera_id = data.get("cameraIdentifier")
-                rtsp_url = data.get("rtspUrl")
-                model_url = data.get("modelUrl")
-                modelId = data.get("modelId")
-                modelName = data.get("modelName")
-    
-                if model_url:
-                    print(f"Downloading model from {model_url}")
-                    parsed_url = urlparse(model_url)
-                    filename = os.path.basename(parsed_url.path)    
-                    model_filename = os.path.join("models", filename)
-                    # Download the model
-                    response = requests.get(model_url, stream=True)
-                    if response.status_code == 200:
-                        with open(model_filename, 'wb') as f:
-                            for chunk in response.iter_content(chunk_size=8192):
-                                f.write(chunk)
-                        logging.info(f"Downloaded model from {model_url} to {model_filename}")
-                        model = YOLO(model_filename)
-                        if torch.cuda.is_available():
-                            model.to('cuda')
-                        class_names = model.names
-                    else:
-                        logging.error(f"Failed to download model from {model_url}")
-                        continue
-                if camera_id and rtsp_url:
-                    if camera_id not in streams and len(streams) < max_streams:
-                        cap = cv2.VideoCapture(rtsp_url)
-                        if not cap.isOpened():
-                            logging.error(f"Failed to open RTSP stream for camera {camera_id}")
-                            continue
-                        buffer = queue.Queue(maxsize=1)
-                        stop_event = threading.Event()
-                        thread = threading.Thread(target=frame_reader, args=(camera_id, cap, buffer, stop_event))
-                        thread.daemon = True
-                        thread.start()
-                        streams[camera_id] = {
-                            'cap': cap,
-                            'buffer': buffer,
-                            'thread': thread,
-                            'rtsp_url': rtsp_url,
-                            'stop_event': stop_event,
-                            'modelId': modelId,
-                            'modelName': modelName
-                        }
-                        logging.info(f"Subscribed to camera {camera_id} with modelId {modelId}, modelName {modelName} and URL {rtsp_url}")
-                    elif camera_id and camera_id in streams:
-                        stream = streams.pop(camera_id)
-                        stream['cap'].release()
-                        logging.info(f"Unsubscribed from camera {camera_id}")
-                elif data.get("command") == "stop":
-                    logging.info("Received stop command")
-                    break
             except json.JSONDecodeError:
                 logging.error("Received invalid JSON message")
             except (WebSocketDisconnect, ConnectionClosedError) as e:
@@ -385,11 +334,20 @@ async def detect(websocket: WebSocket):
             except Exception as e:
                 logging.error(f"Error handling message: {e}")
                 break
+
+    try:
+        await websocket.accept()
+        task = asyncio.create_task(process_streams())
+        heartbeat_task = asyncio.create_task(send_heartbeat())
+        message_task = asyncio.create_task(on_message())
+
+        await asyncio.gather(heartbeat_task, message_task)
     except Exception as e:
-        logging.error(f"Unexpected error in WebSocket connection: {e}")
+        logging.error(f"Error in detect websocket: {e}")
     finally:
         task.cancel()
         await task
+        with streams_lock:
             for camera_id, stream in streams.items():
                 stream['stop_event'].set()
                 stream['thread'].join()
@@ -397,5 +355,6 @@ async def detect(websocket: WebSocket):
                 stream['buffer'].queue.clear()
                 logging.info(f"Released camera {camera_id} and cleaned up resources")
             streams.clear()
+        with models_lock:
             models.clear()
         logging.info("WebSocket connection closed")