diff --git a/mixpanel/flags/local_feature_flags.py b/mixpanel/flags/local_feature_flags.py index 50518de..7409e1b 100644 --- a/mixpanel/flags/local_feature_flags.py +++ b/mixpanel/flags/local_feature_flags.py @@ -23,6 +23,7 @@ ) from .utils import ( REQUEST_HEADERS, + close_async_client_from_sync, dispatch_exposure, generate_traceparent, normalized_hash, @@ -549,8 +550,11 @@ async def __aenter__(self): return self def shutdown(self): + # SDK-85: close both clients. close_async_client_from_sync raises + # if a loop is already running — async callers should use __aexit__. self.stop_polling_for_definitions() self._sync_client.close() + close_async_client_from_sync(self._async_client) def __enter__(self): return self @@ -559,8 +563,8 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): logger.info("Exiting the LocalFeatureFlagsProvider and cleaning up resources") await self.astop_polling_for_definitions() await self._async_client.aclose() + self._sync_client.close() def __exit__(self, exc_type, exc_val, exc_tb): logger.info("Exiting the LocalFeatureFlagsProvider and cleaning up resources") - self.stop_polling_for_definitions() - self._sync_client.close() + self.shutdown() diff --git a/mixpanel/flags/remote_feature_flags.py b/mixpanel/flags/remote_feature_flags.py index 0fe73b5..e31945b 100644 --- a/mixpanel/flags/remote_feature_flags.py +++ b/mixpanel/flags/remote_feature_flags.py @@ -21,6 +21,7 @@ from .utils import ( EXPOSURE_EVENT, REQUEST_HEADERS, + close_async_client_from_sync, dispatch_exposure, generate_traceparent, prepare_common_query_params, @@ -423,7 +424,10 @@ def _lookup_flag_in_response( return fallback_value.as_fallback(FallbackReason.flag_not_found()), True def shutdown(self): + # SDK-85: close both clients. close_async_client_from_sync raises + # if a loop is already running — async callers should use __aexit__. self._sync_client.close() + close_async_client_from_sync(self._async_client) def __enter__(self): return self @@ -433,8 +437,9 @@ async def __aenter__(self): def __exit__(self, exc_type, exc_val, exc_tb): logger.info("Exiting the RemoteFeatureFlagsProvider and cleaning up resources") - self._sync_client.close() + self.shutdown() async def __aexit__(self, exc_type, exc_val, exc_tb): logger.info("Exiting the RemoteFeatureFlagsProvider and cleaning up resources") await self._async_client.aclose() + self._sync_client.close() diff --git a/mixpanel/flags/test_local_feature_flags.py b/mixpanel/flags/test_local_feature_flags.py index 2de7546..6f5554b 100644 --- a/mixpanel/flags/test_local_feature_flags.py +++ b/mixpanel/flags/test_local_feature_flags.py @@ -996,8 +996,7 @@ async def test_local_flags_async_with_service_account_credentials(): assert provider._async_client.auth is not None assert isinstance(provider._async_client.auth, httpx.BasicAuth) - await provider._async_client.aclose() - provider.shutdown() + await provider.__aexit__(None, None, None) def test_local_flags_fallback_to_token_without_credentials(): @@ -1029,3 +1028,31 @@ def test_local_flags_fallback_to_token_without_credentials(): assert provider._request_params["lib_version"] == "1.0.0" provider.shutdown() + + +# SDK-85: sync __exit__ and shutdown() historically only closed +# _sync_client, leaking _async_client's connection pool + background +# transport. The two tests below fail on the pre-fix code. + + +def test_sync_shutdown_closes_both_clients(): + config = LocalFlagsConfig(enable_polling=False) + provider = LocalFeatureFlagsProvider("test-token", config, "1.0.0", Mock()) + + assert not provider._sync_client.is_closed + assert not provider._async_client.is_closed + + provider.shutdown() + + assert provider._sync_client.is_closed + assert provider._async_client.is_closed + + +def test_sync_context_manager_exit_closes_both_clients(): + config = LocalFlagsConfig(enable_polling=False) + with LocalFeatureFlagsProvider("test-token", config, "1.0.0", Mock()) as provider: + assert not provider._sync_client.is_closed + assert not provider._async_client.is_closed + + assert provider._sync_client.is_closed + assert provider._async_client.is_closed diff --git a/mixpanel/flags/test_remote_feature_flags.py b/mixpanel/flags/test_remote_feature_flags.py index 3f8bf78..c9ed4af 100644 --- a/mixpanel/flags/test_remote_feature_flags.py +++ b/mixpanel/flags/test_remote_feature_flags.py @@ -623,3 +623,31 @@ def test_remote_flags_fallback_to_token_without_credentials(): assert provider._request_params_base["lib_version"] == "1.0.0" provider.shutdown() + + +# SDK-85: sync __exit__ and shutdown() historically only closed +# _sync_client, leaking _async_client's connection pool + background +# transport. The two tests below fail on the pre-fix code. + + +def test_sync_shutdown_closes_both_clients(): + config = RemoteFlagsConfig() + provider = RemoteFeatureFlagsProvider("test-token", config, "1.0.0", Mock()) + + assert not provider._sync_client.is_closed + assert not provider._async_client.is_closed + + provider.shutdown() + + assert provider._sync_client.is_closed + assert provider._async_client.is_closed + + +def test_sync_context_manager_exit_closes_both_clients(): + config = RemoteFlagsConfig() + with RemoteFeatureFlagsProvider("test-token", config, "1.0.0", Mock()) as provider: + assert not provider._sync_client.is_closed + assert not provider._async_client.is_closed + + assert provider._sync_client.is_closed + assert provider._async_client.is_closed diff --git a/mixpanel/flags/test_utils.py b/mixpanel/flags/test_utils.py index e8fd0ce..767127a 100644 --- a/mixpanel/flags/test_utils.py +++ b/mixpanel/flags/test_utils.py @@ -5,10 +5,12 @@ from concurrent.futures import Future, ThreadPoolExecutor from unittest.mock import MagicMock +import httpx import pytest from .utils import ( _log_tracker_future_exception, + close_async_client_from_sync, dispatch_exposure, generate_traceparent, normalized_hash, @@ -81,3 +83,30 @@ def test_log_tracker_future_exception_ignores_cancelled_future(self, caplog): assert not any( "Exposure event failed" in rec.message for rec in caplog.records ), "cancelled futures must not produce an error log" + + +class TestCloseAsyncClientFromSync: + # SDK-85: shutdown() and __exit__ need to close the AsyncClient + # from sync context. Callers already inside a running event loop + # must use __aexit__ — see the raises test below. + + def test_closes_client_when_no_loop_running(self): + client = httpx.AsyncClient() + assert not client.is_closed + + close_async_client_from_sync(client) + + assert client.is_closed + + @pytest.mark.asyncio + async def test_raises_when_called_from_running_loop(self): + # Fire-and-forget scheduling on the running loop can be cancelled + # by loop teardown before aclose finishes, so the helper refuses + # the call and points async callers at __aexit__. + client = httpx.AsyncClient() + try: + with pytest.raises(RuntimeError, match="running event loop"): + close_async_client_from_sync(client) + assert not client.is_closed + finally: + await client.aclose() diff --git a/mixpanel/flags/utils.py b/mixpanel/flags/utils.py index 30f93eb..9acb814 100644 --- a/mixpanel/flags/utils.py +++ b/mixpanel/flags/utils.py @@ -1,16 +1,46 @@ from __future__ import annotations +import asyncio import logging import uuid from typing import TYPE_CHECKING, Any, Callable +from asgiref.sync import async_to_sync + if TYPE_CHECKING: from concurrent.futures import Executor, Future + import httpx + logger = logging.getLogger(__name__) EXPOSURE_EVENT = "$experiment_started" + +def close_async_client_from_sync(client: httpx.AsyncClient) -> None: + """SDK-85: close an ``httpx.AsyncClient`` from sync code. + + Bridges to sync via ``asgiref.async_to_sync`` and blocks until the + close completes. If a loop is already running on the current thread + (e.g. ``shutdown()`` was called from inside an ``async def``), raises + ``RuntimeError`` — scheduling a background ``aclose`` on the running + loop is unsafe because loop teardown can cancel the task before it + finishes, defeating the fix. Async callers should use ``__aexit__`` + or await ``aclose()`` directly. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + async_to_sync(client.aclose)() + return + + raise RuntimeError( + "close_async_client_from_sync() cannot be called from a running " + "event loop. Use 'async with provider:' or await the provider's " + "__aexit__ from async code." + ) + + REQUEST_HEADERS: dict[str, str] = { "X-Scheme": "https", "X-Forwarded-Proto": "https",