Fix thread leak for LOOPBACK workers in external worker pool.

Author: shunping

sdks/python/apache_beam/runners/portability/prism_runner_test.py

@@ -488,6 +488,63 @@ def test_singleton(self, enable_singleton):
       else:
         mock_prism_server.assert_called_once()
 
+  def test_loopback_worker_daemon_thread_accumulation(self):
+    """Verifies that in LOOPBACK mode, the external worker pool servicer properly
+    tracks active thread-based SdkHarness workers and cleanly shuts them down in
+    StopWorker via sentinel messages. This prevents background daemon threads from
+    accumulating across sequential pipeline executions and leaking resources.
+    """
+    import queue
+    import threading
+    import time
+    from apache_beam.portability.api import beam_fn_api_pb2
+    from apache_beam.runners.worker import worker_pool_main
+
+    servicer = worker_pool_main.BeamFnExternalWorkerPoolServicer(
+        use_process=False, state_cache_size=0, data_buffer_time_limit_ms=0)
+
+    active_workers = []
+    mock_responses = queue.Queue()
+
+    def mock_run(self_worker):
+      active_workers.append(self_worker)
+      mock_responses.get()
+      active_workers.remove(self_worker)
+
+    with mock.patch('apache_beam.runners.worker.sdk_worker.SdkHarness') as mock_harness:
+      mock_harness.return_value._responses = mock_responses
+      mock_harness.return_value.run = lambda: mock_run(mock_harness)
+
+      # Simulate starting Worker 1 for Pipeline 1
+      req1 = beam_fn_api_pb2.StartWorkerRequest(worker_id="worker_1")
+      req1.control_endpoint.url = "localhost:12345"
+      servicer.StartWorker(req1, None)
+
+      time.sleep(0.05)
+      self.assertEqual(len(active_workers), 1)
+
+      # Simulate stopping Worker 1 at the end of Pipeline 1
+      stop_req1 = beam_fn_api_pb2.StopWorkerRequest(worker_id="worker_1")
+      servicer.StopWorker(stop_req1, None)
+
+      time.sleep(0.05)
+      # Verify the fix: StopWorker successfully tells the thread harness to shut down,
+      # completely resolving the thread leak!
+      self.assertEqual(len(active_workers), 0)
+
+      # Simulate starting Worker 2 for Pipeline 2
+      req2 = beam_fn_api_pb2.StartWorkerRequest(worker_id="worker_2")
+      req2.control_endpoint.url = "localhost:12345"
+      servicer.StartWorker(req2, None)
+
+      time.sleep(0.05)
+      self.assertEqual(len(active_workers), 1)
+
+      # Clean up the second worker
+      servicer.StopWorker(beam_fn_api_pb2.StopWorkerRequest(worker_id="worker_2"), None)
+      time.sleep(0.05)
+      self.assertEqual(len(active_workers), 0)
+
 
 if __name__ == '__main__':
   # Run the tests.

sdks/python/apache_beam/runners/worker/worker_pool_main.py

@@ -82,6 +82,7 @@ def __init__(
     self._state_cache_size = state_cache_size
     self._data_buffer_time_limit_ms = data_buffer_time_limit_ms
     self._worker_processes: dict[str, subprocess.Popen] = {}
+    self._worker_threads: dict[str, sdk_worker.SdkHarness] = {}
 
   @classmethod
   def start(
@@ -166,6 +167,7 @@ def StartWorker(
             worker_id=start_worker_request.worker_id,
             state_cache_size=self._state_cache_size,
             data_buffer_time_limit_ms=self._data_buffer_time_limit_ms)
+        self._worker_threads[start_worker_request.worker_id] = worker
         worker_thread = threading.Thread(
             name='run_worker_%s' % start_worker_request.worker_id,
             target=worker.run)
@@ -188,6 +190,13 @@ def StopWorker(
       _LOGGER.info("Stopping worker %s" % stop_worker_request.worker_id)
       kill_process_gracefully(worker_process)
 
+    worker_thread_harness = self._worker_threads.pop(
+        stop_worker_request.worker_id, None)
+    if worker_thread_harness:
+      _LOGGER.info("Stopping thread worker %s" % stop_worker_request.worker_id)
+      from apache_beam.utils.sentinel import Sentinel
+      worker_thread_harness._responses.put(Sentinel.sentinel)
+
     return beam_fn_api_pb2.StopWorkerResponse()