GOVFOUN-408: Fix connection pool issue

Author: ankitpatnaik-atlan

pyatlan/client/atlan.py

@@ -214,7 +214,14 @@ def __init__(self, **data):
         # Pass self reference to transport for duplicate checking during retries
         self._session = httpx.Client(
             transport=PyatlanSyncTransport(
-                retry=self.retry, client=self, **transport_kwargs
+                retry=self.retry,
+                client=self,
+                limits=httpx.Limits(
+                    max_connections=50,
+                    max_keepalive_connections=10,
+                    keepalive_expiry=30.0,
+                ),
+                **transport_kwargs,
             ),
             headers={
                 "x-atlan-agent": "sdk",
@@ -554,7 +561,10 @@ def _call_api_internal(
         try:
             params["headers"]["X-Atlan-Request-Id"] = request_id_var.get()
             timeout = httpx.Timeout(
-                None, connect=self.connect_timeout, read=self.read_timeout
+                None,
+                connect=self.connect_timeout,
+                read=self.read_timeout,
+                pool=30.0,
             )
             if binary_data:
                 response = self._session.request(
@@ -2014,7 +2024,14 @@ def max_retries(  # type: ignore[misc]
             transport_kwargs["verify"] = self.verify
 
         new_transport = PyatlanSyncTransport(
-            retry=max_retries, client=self, **transport_kwargs
+            retry=max_retries,
+            client=self,
+            limits=httpx.Limits(
+                max_connections=50,
+                max_keepalive_connections=10,
+                keepalive_expiry=30.0,
+            ),
+            **transport_kwargs,
         )
         self._session._transport = new_transport
 

tests/integration/test_connection_pool_config.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 Atlan Pte. Ltd.
+"""
+Integration tests for GOVFOUN-408: httpcore connection pool deadlock fix.
+
+These tests run against a live Atlan tenant and verify:
+1. Pool limits and timeout config survive the full client init and real HTTP
+   connections — not just object construction.
+2. Concurrent requests complete without hanging under the configured pool size.
+3. httpx.PoolTimeout propagates through the full SDK stack instead of being
+   swallowed or causing an indefinite hang.
+4. max_retries context manager correctly installs and restores limits on a
+   client that is actively connected.
+"""
+
+import time
+from concurrent.futures import ThreadPoolExecutor, wait
+from unittest.mock import patch
+
+import httpx
+import pytest
+from httpx_retries import Retry
+
+from pyatlan.client.atlan import AtlanClient
+from pyatlan.client.transport import PyatlanSyncTransport
+
+
+def _get_httpcore_pool(client: AtlanClient):
+    return client._session._transport._transport._pool
+
+
+# ---------------------------------------------------------------------------
+# Pool limits — config survives real HTTP connections
+# ---------------------------------------------------------------------------
+
+
+def test_pool_limits_on_live_client(client: AtlanClient):
+    """
+    Pool limits are configured correctly on an AtlanClient built from the
+    integration environment (real base_url, real API key, real retry config).
+    The values are set at construction time and must survive unchanged.
+    """
+    pool = _get_httpcore_pool(client)
+    assert pool._max_connections == 50
+    assert pool._keepalive_expiry == 30.0
+    assert pool._max_keepalive_connections == 10
+
+
+# ---------------------------------------------------------------------------
+# Concurrent requests — happy path
+# ---------------------------------------------------------------------------
+
+
+def test_concurrent_requests_complete_without_deadlock(client: AtlanClient):
+    """
+    N concurrent requests (N << max_connections=50) must all complete.
+
+    Under the old config (pool=None, max_connections=100), if connections
+    accumulated CLOSE_WAIT sockets and filled all slots, threads waiting for
+    a connection would block on threading.Event.wait(timeout=None) forever.
+
+    With pool=30.0 and max_connections=50, threads either succeed or raise
+    PoolTimeout within 30s — they never hang indefinitely.
+    """
+    n_threads = 5
+    timeout_seconds = 60
+
+    def make_request():
+        return client.user.get_current()
+
+    with ThreadPoolExecutor(max_workers=n_threads) as executor:
+        futures = [executor.submit(make_request) for _ in range(n_threads)]
+        done, not_done = wait(futures, timeout=timeout_seconds)
+
+    # Only check for deadlock — individual futures may raise (auth errors, etc.)
+    assert len(not_done) == 0, (
+        f"{len(not_done)} of {n_threads} requests still pending after "
+        f"{timeout_seconds}s — possible connection pool deadlock"
+    )
+
+
+# ---------------------------------------------------------------------------
+# PoolTimeout propagation
+# ---------------------------------------------------------------------------
+
+
+def test_pool_timeout_propagates_through_sdk_stack(client: AtlanClient):
+    """
+    httpx.PoolTimeout injected at the transport layer must propagate up
+    through the SDK without being swallowed or stalling.
+
+    Before the fix, pool=None meant threads blocked indefinitely on
+    threading.Event.wait(timeout=None). With pool=30.0 the SDK raises an
+    exception quickly instead.
+    """
+    original_transport = client._session._transport
+
+    test_transport = PyatlanSyncTransport(
+        retry=Retry(total=0),
+        client=client,
+        limits=httpx.Limits(
+            max_connections=50,
+            max_keepalive_connections=10,
+            keepalive_expiry=30.0,
+        ),
+    )
+
+    def always_pool_timeout(request: httpx.Request) -> httpx.Response:
+        raise httpx.PoolTimeout("simulated pool exhaustion", request=request)
+
+    test_transport._transport.handle_request = always_pool_timeout  # type: ignore[method-assign]
+    client._session._transport = test_transport
+
+    try:
+        start = time.monotonic()
+        with pytest.raises(Exception):
+            client.user.get_current()
+        elapsed = time.monotonic() - start
+
+        assert elapsed < 5.0, (
+            f"PoolTimeout took {elapsed:.1f}s to propagate — "
+            "it may be getting swallowed or retried unexpectedly"
+        )
+    finally:
+        client._session._transport = original_transport
+
+
+# ---------------------------------------------------------------------------
+# max_retries context manager
+# ---------------------------------------------------------------------------
+
+
+def test_max_retries_transport_limits_on_live_client(client: AtlanClient):
+    """
+    max_retries context manager must install a transport with the same pool
+    limits on the real integration client (env-configured base_url, retry policy).
+    """
+    captured: dict = {}
+    original_init = PyatlanSyncTransport.__init__
+
+    def capturing_init(self_t, retry=None, client=None, **kwargs):  # noqa: ARG001
+        if "limits" in kwargs:
+            captured["limits"] = kwargs["limits"]
+        original_init(self_t, retry=retry, client=client, **kwargs)
+
+    with patch.object(PyatlanSyncTransport, "__init__", capturing_init):
+        with client.max_retries():
+            pass
+
+    limits = captured.get("limits")
+    assert limits is not None, "max_retries did not pass limits to PyatlanSyncTransport"
+    assert limits.max_connections == 50
+    assert limits.keepalive_expiry == 30.0
+    assert limits.max_keepalive_connections == 10
+
+
+def test_max_retries_restores_original_transport(client: AtlanClient):
+    """
+    max_retries must restore the original transport after the context exits,
+    even on the real integration client (env-configured base_url, retry policy).
+    """
+    original_transport = client._session._transport
+
+    with client.max_retries():
+        pass
+
+    assert client._session._transport is original_transport

tests/unit/test_connection_pool_config.py

@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2026 Atlan Pte. Ltd.
+"""
+Regression tests for GOVFOUN-408: httpcore connection pool deadlock fix.
+
+Root cause: csa-metadata-completeness hung indefinitely on masterc-prd because
+two missing configs caused all 100 httpcore connection slots to fill with
+CLOSE_WAIT zombies, then 109 worker threads blocked forever in
+wait_for_connection(timeout=None).
+
+These tests verify the two SDK-side fixes:
+1. httpx.Timeout has pool=30.0 — threads raise PoolTimeout after 30s instead
+   of waiting forever.
+2. httpx.Limits has keepalive_expiry=30.0 — client retires idle connections
+   before nginx's keepalive_timeout=75s FIN, preventing CLOSE_WAIT accumulation.
+"""
+
+from unittest.mock import Mock, patch
+
+import httpx
+import pytest
+
+from pyatlan.client.atlan import AtlanClient
+from pyatlan.client.transport import PyatlanSyncTransport
+from pyatlan.model.assets import AtlasGlossary
+
+
+@pytest.fixture(autouse=True)
+def set_env(monkeypatch):
+    monkeypatch.setenv("ATLAN_BASE_URL", "https://test.atlan.com")
+    monkeypatch.setenv("ATLAN_API_KEY", "test-api-key")
+
+
+@pytest.fixture()
+def client():
+    return AtlanClient()
+
+
+def _error_response():
+    r = Mock()
+    r.status_code = 500
+    r.text = "internal server error"
+    return r
+
+
+def _trigger_api_call(client: AtlanClient) -> None:
+    """Drive any outbound request through _call_api_internal."""
+    try:
+        client.asset.save(AtlasGlossary.creator(name="t"))
+    except Exception:
+        pass
+
+
+def _get_httpcore_pool(client: AtlanClient):
+    transport = client._session._transport
+    assert isinstance(transport, PyatlanSyncTransport)
+    return transport._transport._pool
+
+
+# ---------------------------------------------------------------------------
+# Pool timeout
+# ---------------------------------------------------------------------------
+
+
+@patch.object(AtlanClient, "_session")
+def test_pool_timeout_is_30_seconds(mock_session, client):
+    """pool=30.0 must be set — the live deadlock showed pool=None caused infinite blocking."""
+    mock_session.request.return_value = _error_response()
+    _trigger_api_call(client)
+    assert mock_session.request.called, "no HTTP request was made"
+    timeout = mock_session.request.call_args.kwargs["timeout"]
+    assert timeout.pool == 30.0
+
+
+@patch.object(AtlanClient, "_session")
+def test_pool_timeout_is_not_none(mock_session, client):
+    """pool timeout must never be None — threading.Event.wait(timeout=None) blocks forever."""
+    mock_session.request.return_value = _error_response()
+    _trigger_api_call(client)
+    assert mock_session.request.called
+    timeout = mock_session.request.call_args.kwargs["timeout"]
+    assert timeout.pool is not None
+
+
+@patch.object(AtlanClient, "_session")
+def test_connect_timeout_unchanged(mock_session, client):
+    """connect timeout must still equal client.connect_timeout (default 30s)."""
+    mock_session.request.return_value = _error_response()
+    _trigger_api_call(client)
+    assert mock_session.request.called
+    timeout = mock_session.request.call_args.kwargs["timeout"]
+    assert timeout.connect == client.connect_timeout
+
+
+@patch.object(AtlanClient, "_session")
+def test_read_timeout_unchanged(mock_session, client):
+    """read timeout must still equal client.read_timeout (default 900s)."""
+    mock_session.request.return_value = _error_response()
+    _trigger_api_call(client)
+    assert mock_session.request.called
+    timeout = mock_session.request.call_args.kwargs["timeout"]
+    assert timeout.read == client.read_timeout
+
+
+@patch.object(AtlanClient, "_session")
+def test_pool_timeout_propagates_not_hangs(mock_session, client):
+    """httpx.PoolTimeout must propagate — if it were swallowed the workflow would still hang."""
+    mock_session.request.side_effect = httpx.PoolTimeout(
+        "connection pool exhausted", request=None
+    )
+    with pytest.raises(Exception):
+        client.asset.save(AtlasGlossary.creator(name="t"))
+
+
+# ---------------------------------------------------------------------------
+# Transport limits
+# ---------------------------------------------------------------------------
+
+
+def test_transport_max_connections_is_50(client):
+    """max_connections=50 reduces blast radius when CLOSE_WAIT sockets accumulate."""
+    assert _get_httpcore_pool(client)._max_connections == 50
+
+
+def test_transport_keepalive_expiry_is_30_seconds(client):
+    """keepalive_expiry=30.0 — client closes idle connections before nginx's 75s FIN."""
+    assert _get_httpcore_pool(client)._keepalive_expiry == 30.0
+
+
+def test_transport_max_keepalive_connections_is_10(client):
+    """max_keepalive_connections=10 bounds idle connections held in the pool."""
+    assert _get_httpcore_pool(client)._max_keepalive_connections == 10
+
+
+# ---------------------------------------------------------------------------
+# max_retries context manager
+# ---------------------------------------------------------------------------
+
+
+def _capture_transport_limits(client: AtlanClient) -> httpx.Limits:
+    """Enter max_retries, capture the limits used to construct the new transport."""
+    captured: dict = {}
+    original_init = PyatlanSyncTransport.__init__
+
+    def capturing_init(self_t, retry=None, client=None, **kwargs):  # noqa: ARG001
+        if "limits" in kwargs:
+            captured["limits"] = kwargs["limits"]
+        original_init(self_t, retry=retry, client=client, **kwargs)
+
+    with patch.object(PyatlanSyncTransport, "__init__", capturing_init):
+        with client.max_retries():
+            pass
+
+    assert "limits" in captured, "max_retries did not pass limits to PyatlanSyncTransport"
+    return captured["limits"]
+
+
+def test_max_retries_transport_max_connections_is_50(client):
+    """max_retries context manager must use the same max_connections=50."""
+    assert _capture_transport_limits(client).max_connections == 50
+
+
+def test_max_retries_transport_keepalive_expiry_is_30_seconds(client):
+    """max_retries context manager must use keepalive_expiry=30.0."""
+    assert _capture_transport_limits(client).keepalive_expiry == 30.0
+
+
+def test_max_retries_transport_max_keepalive_connections_is_10(client):
+    """max_retries context manager must use max_keepalive_connections=10."""
+    assert _capture_transport_limits(client).max_keepalive_connections == 10
+
+
+def test_max_retries_replaces_transport_inside_context(client):
+    """max_retries installs a fresh transport for the duration of the context."""
+    original = client._session._transport
+    with client.max_retries():
+        assert client._session._transport is not original
+
+
+def test_max_retries_restores_original_transport_on_exit(client):
+    """max_retries restores the original transport when the context exits."""
+    original = client._session._transport
+    with client.max_retries():
+        pass
+    assert client._session._transport is original