Fix thread leak in SessionShutdownObserver and WebRtcWorker

Author: kayush0712

streamlit_webrtc/shutdown.py

@@ -38,36 +38,36 @@ def _polling_thread_impl(
         app_session_ref: "weakref.ReferenceType[AppSession]",
         callback: Callback,
     ):
-        # Use polling because event-based methods are not available
-        # to observe the session lifecycle.
         while True:
             app_session = app_session_ref()
             if not app_session:
-                logger.debug("AppSession has removed.")
+                logger.debug("AppSession has been removed.")
                 break
             if app_session._state == AppSessionState.SHUTDOWN_REQUESTED:
-                logger.debug(
-                    "AppSession %s has been requested to shutdown.",
-                    app_session.id,
-                )
+                logger.debug("AppSession %s shutdown requested.", app_session.id)
                 break
             if self._polling_thread_stop_event.wait(1.0):
-                logger.debug(
-                    "The polling thread should be stopped. Exit the polling loop and return."
-                )
+                logger.debug("Polling thread stop requested. Exiting loop.")
                 return
 
-        # Ensure the flag is set
         self._polling_thread_stop_event.set()
-
-        logger.debug("AppSession shutdown has been detected.")
+        logger.debug("AppSession shutdown detected.")
         callback()
 
-    def stop(self):
-        if self._polling_thread_stop_event.is_set():
+    def stop(self, timeout: float = 1.0) -> None:
+        """
+        Stop the polling thread cleanly so we don't leak ShutdownPolling_* threads.
+        """
+        if not self._polling_thread:
             return
 
-        if self._polling_thread:
-            self._polling_thread_stop_event.set()
-            self._polling_thread.join()
-            self._polling_thread = None
+        self._polling_thread_stop_event.set()
+
+        if self._polling_thread.is_alive():
+            try:
+                self._polling_thread.join(timeout=timeout)
+            except Exception:
+                logger.exception("Failed to join ShutdownPolling thread")
+
+        self._polling_thread = None
+

streamlit_webrtc/webrtc.py

@@ -542,19 +542,10 @@ def on_track_created(track_type: TrackType, track: MediaStreamTrack):
         async def on_iceconnectionstatechange():
             logger.debug("ICE connection state is %s", self.pc.iceConnectionState)
             iceConnectionState = self.pc.iceConnectionState
-            
-            if iceConnectionState in ("failed","disconnect","closed"):
-                self._unset_processors()
-            if iceConnectionState in ("failed","disconnect","closed"):
-                try:
-                    if self.pc and self.pc.connectionState != "closed":
-                        await self.pc.close()
-                except Exception as e:
-                    logger.debug("Error occurred while closing the peer connection", e)
-                    
-                if self._session_shutdown_observer:
-                  self._session_shutdown_observer.stop()    
-                           
+
+            if iceConnectionState in ("closed", "failed", "disconnected"):
+                  logger.debug("Connection state is %s, stopping worker", iceConnectionState)
+                  self.stop()