fix: Close streaming connection pool on FDv1 data source shutdown

ldclient/impl/datasource/streaming.py

@@ -5,6 +5,7 @@
 from typing import Optional
 from urllib import parse
 
+import urllib3
 from ld_eventsource import SSEClient
 from ld_eventsource.actions import Event, Fault
 from ld_eventsource.config import (
@@ -41,6 +42,32 @@
 ParsedPath = namedtuple('ParsedPath', ['kind', 'key'])
 
 
+def _close_pool_manager(pool: Optional[urllib3.PoolManager]) -> None:
+    """Close every pooled connection in ``pool`` so the underlying TCP sockets
+    are torn down. ``HTTPConnectionPool.close()`` drains its queue and calls
+    ``conn.close()`` on each connection, which sends the FIN the server is
+    waiting on. ``SSEClient.close()`` only releases the stream connection back
+    to the pool; under urllib3 >= 1.26.16 / 2.x, ``PoolManager.clear()`` no
+    longer closes sockets synchronously (teardown is deferred to GC), so the
+    streaming socket would otherwise linger ESTABLISHED after shutdown."""
+    if pool is None:
+        return
+    try:
+        # ``RecentlyUsedContainer`` deliberately disallows iteration; ``keys()``
+        # returns a thread-safe snapshot. We look each one up to close its
+        # underlying ``HTTPConnectionPool``.
+        for key in list(pool.pools.keys()):
+            try:
+                connection_pool = pool.pools.get(key)
+                if connection_pool is not None:
+                    connection_pool.close()
+            except Exception:  # pylint: disable=broad-except
+                log.warning("Error closing streaming connection pool", exc_info=True)
+        pool.clear()
+    except Exception:  # pylint: disable=broad-except
+        log.warning("Error closing streaming pool manager", exc_info=True)
+
+
 class StreamingUpdateProcessor(Thread, UpdateProcessor):
     def __init__(self, config, store, ready, diagnostic_accumulator):
         Thread.__init__(self, name="ldclient.datasource.streaming")
@@ -52,6 +79,8 @@ def __init__(self, config, store, ready, diagnostic_accumulator):
         self._data_source_update_sink = config.data_source_update_sink
         self._store = store
         self._running = False
+        self._sse: Optional[SSEClient] = None
+        self._sse_pool: Optional[urllib3.PoolManager] = None
         self._ready = ready
         self._diagnostic_accumulator = diagnostic_accumulator
         self._connection_attempt_start_time = None
@@ -99,6 +128,10 @@ def run(self):
                 if not self._handle_error(action.error):
                     break
         self._sse.close()
+        # See _close_pool_manager: SSEClient.close() only releases the connection
+        # back to the pool, so we force-close the pool to actually sever the TCP
+        # socket rather than leaving it for GC.
+        _close_pool_manager(self._sse_pool)
 
     def _record_stream_init(self, failed: bool):
         if self._diagnostic_accumulator and self._connection_attempt_start_time:
@@ -110,9 +143,12 @@ def _create_sse_client(self) -> SSEClient:
         # We don't want the stream to use the same read timeout as the rest of the SDK.
         http_factory = _http_factory(self._config)
         stream_http_factory = HTTPFactory(http_factory.base_headers, http_factory.http_config, override_read_timeout=stream_read_timeout)
+        # Retain the pool so we can force-close it on shutdown; SSEClient.close()
+        # won't close a caller-supplied pool. See _close_pool_manager.
+        self._sse_pool = stream_http_factory.create_pool_manager(1, self._uri)
         return SSEClient(
             connect=ConnectStrategy.http(
-                url=self._uri, headers=http_factory.base_headers, pool=stream_http_factory.create_pool_manager(1, self._uri), urllib3_request_options={"timeout": stream_http_factory.timeout}
+                url=self._uri, headers=http_factory.base_headers, pool=self._sse_pool, urllib3_request_options={"timeout": stream_http_factory.timeout}
             ),
             error_strategy=ErrorStrategy.always_continue(),  # we'll make error-handling decisions when we see a Fault
             initial_retry_delay=self._config.initial_reconnect_delay,
@@ -129,6 +165,10 @@ def __stop_with_error_info(self, error: Optional[DataSourceErrorInfo]):
         self._running = False
         if self._sse:
             self._sse.close()
+        # See _close_pool_manager: this is what actually severs the TCP
+        # connection. stop() may run on a different thread than run(); the pool
+        # close is idempotent, so calling it from both is safe.
+        _close_pool_manager(self._sse_pool)
 
         if self._data_source_update_sink is None:
             return
@@ -217,7 +257,7 @@ def _handle_error(self, error: Exception) -> bool:
             if self._data_source_update_sink is not None:
                 self._data_source_update_sink.update_status(DataSourceState.INTERRUPTED, DataSourceErrorInfo(DataSourceErrorKind.UNKNOWN, 0, time.time(), str(error)))
             # no stacktrace here because, for a typical connection error, it'll just be a lengthy tour of urllib3 internals
-        self._connection_attempt_start_time = time.time() + self._sse.next_retry_delay
+        self._connection_attempt_start_time = time.time() + self._sse.next_retry_delay  # type: ignore
         return True
 
     @staticmethod

ldclient/impl/datasourcev2/streaming.py

@@ -143,10 +143,10 @@ def _close_pool_manager(pool: Optional[urllib3.PoolManager]) -> None:
                 if connection_pool is not None:
                     connection_pool.close()
             except Exception:  # pylint: disable=broad-except
-                log.debug("Error closing streaming connection pool", exc_info=True)
+                log.warning("Error closing streaming connection pool", exc_info=True)
         pool.clear()
     except Exception:  # pylint: disable=broad-except
-        log.debug("Error closing streaming pool manager", exc_info=True)
+        log.warning("Error closing streaming pool manager", exc_info=True)
 
 
 class StreamingDataSource(Synchronizer, DiagnosticSource):

ldclient/testing/impl/datasource/test_streaming.py

@@ -438,6 +438,65 @@ def test_failure_transitions_from_valid():
             assert spy.statuses[1].error.status_code == 401
 
 
+def test_stop_closes_underlying_pool():
+    """On shutdown the underlying urllib3 connection pool must be torn down so
+    the streaming TCP socket is actually closed. SSEClient.close() only releases
+    the connection back to the pool; under urllib3 >= 1.26.16 / 2.x that leaves
+    the socket open until GC, so the processor force-closes the pool itself."""
+
+    class TrackingConnectionPool:
+        def __init__(self):
+            self.closed = False
+
+        def close(self):
+            self.closed = True
+
+    class TrackingPoolDict:
+        def __init__(self, items):
+            self._items = items
+
+        def keys(self):
+            return list(self._items.keys())
+
+        def get(self, key):
+            return self._items.get(key)
+
+    class TrackingPool:
+        """Stand-in PoolManager that records clear() and exposes a keys()-iterable
+        pools attribute matching urllib3's RecentlyUsedContainer."""
+
+        def __init__(self):
+            self.cleared = False
+            self.connection_pool = TrackingConnectionPool()
+            self.pools = TrackingPoolDict({"key": self.connection_pool})
+
+        def clear(self):
+            self.cleared = True
+
+    class FakeSSEClient:
+        def __init__(self):
+            self.closed = False
+
+        def close(self):
+            self.closed = True
+
+    store = InMemoryFeatureStore()
+    ready = Event()
+    config = Config(sdk_key='sdk-key', stream_uri='http://localhost')
+
+    sp = StreamingUpdateProcessor(config, store, ready, None)
+    tracking_pool = TrackingPool()
+    fake_sse = FakeSSEClient()
+    sp._sse = fake_sse
+    sp._sse_pool = tracking_pool
+
+    sp.stop()
+
+    assert fake_sse.closed is True
+    assert tracking_pool.cleared is True
+    assert tracking_pool.connection_pool.closed is True
+
+
 def expect_item(store, kind, item):
     assert store.get(kind, item['key'], lambda x: x) == item