diff --git a/src/amplitude_experiment/remote/client.py b/src/amplitude_experiment/remote/client.py index 8db2653..a293edd 100644 --- a/src/amplitude_experiment/remote/client.py +++ b/src/amplitude_experiment/remote/client.py @@ -7,7 +7,7 @@ from .config import RemoteEvaluationConfig from .fetch_options import FetchOptions -from ..connection_pool import HTTPConnectionPool +from ..connection_pool import EmptyPoolError, HTTPConnectionPool from ..exception import FetchException from ..user import User from ..util.deprecated import deprecated @@ -149,7 +149,14 @@ def __do_fetch(self, user, fetch_options: FetchOptions = None): json.dumps(fetch_options.flagKeys, separators=(",", ":")).encode("utf-8") ).rstrip(b"=").decode("utf-8") - conn = self._connection_pool.acquire() + acquire_timeout_millis = self.config.fetch_pool_acquire_timeout_millis + try: + conn = self._connection_pool.acquire( + timeout=acquire_timeout_millis / 1000 if acquire_timeout_millis is not None else None + ) + except EmptyPoolError: + raise TimeoutError(f"Timed out waiting {acquire_timeout_millis}ms for a connection " + f"from the pool (max_size={self._connection_pool.max_size})") body = user_context.to_json().encode('utf8') if len(body) > 8000: self.logger.warning(f"[Experiment] encoded user object length ${len(body)} " diff --git a/src/amplitude_experiment/remote/config.py b/src/amplitude_experiment/remote/config.py index c4efdcc..900202a 100644 --- a/src/amplitude_experiment/remote/config.py +++ b/src/amplitude_experiment/remote/config.py @@ -17,6 +17,7 @@ def __init__(self, debug=False, fetch_retry_backoff_max_millis=10000, fetch_retry_backoff_scalar=1.5, fetch_retry_timeout_millis=10000, + fetch_pool_acquire_timeout_millis=None, server_zone: ServerZone = ServerZone.US, logger=None): """ @@ -33,6 +34,10 @@ def __init__(self, debug=False, greater than the max, the max is used for all subsequent retries. fetch_retry_backoff_scalar (float): Scales the minimum backoff exponentially. fetch_retry_timeout_millis (int): The request timeout for retrying fetch requests. + fetch_pool_acquire_timeout_millis (int | None): Maximum time, in milliseconds, a fetch may wait to + acquire a connection from the connection pool when all pooled connections are busy with concurrent + fetches. None (the default) preserves the existing behavior of waiting indefinitely. When set, an + exhausted wait raises TimeoutError, classified for retries exactly like a socket read timeout. server_zone (str): Select the Amplitude data center to get flags and variants from, US or EU. logger (logging.Logger): Optional logger instance. If provided, this logger will be used instead of creating a new one. The debug flag only applies when no logger is provided. @@ -48,6 +53,7 @@ def __init__(self, debug=False, self.fetch_retry_backoff_max_millis = fetch_retry_backoff_max_millis self.fetch_retry_backoff_scalar = fetch_retry_backoff_scalar self.fetch_retry_timeout_millis = fetch_retry_timeout_millis + self.fetch_pool_acquire_timeout_millis = fetch_pool_acquire_timeout_millis self.server_zone = server_zone if server_url == DEFAULT_SERVER_URL and server_zone == ServerZone.EU: self.server_url = EU_SERVER_URL diff --git a/tests/remote/client_test.py b/tests/remote/client_test.py index 45c45b8..f9871ec 100644 --- a/tests/remote/client_test.py +++ b/tests/remote/client_test.py @@ -1,4 +1,5 @@ import json +import time import unittest from unittest import mock @@ -41,6 +42,45 @@ def test_fetch_async(self): user = User(user_id='test_user') self.client.fetch_async(user, self.callback_for_async) + def test_pool_acquire_wait_unbounded_by_default(self): + assert RemoteEvaluationConfig().fetch_pool_acquire_timeout_millis is None + + def test_fetch_pool_exhausted_times_out(self): + client = RemoteEvaluationClient( + API_KEY, RemoteEvaluationConfig(fetch_pool_acquire_timeout_millis=100) + ) + # Occupy the pool's only connection so the next fetch must wait for it. + held = client._connection_pool.acquire() + try: + start = time.time() + with self.assertRaises(TimeoutError): + client._RemoteEvaluationClient__do_fetch(User(user_id='test_user')) + elapsed = time.time() - start + # Must fail promptly (bounded by fetch_timeout_millis), not hang indefinitely. + self.assertLess(elapsed, 5) + finally: + client._connection_pool.release(held) + client.close() + + def test_pool_acquire_timeout_classified_like_read_timeout(self): + # A pool-wait timeout must take the same retry path as a socket read timeout. + should_retry = RemoteEvaluationClient._RemoteEvaluationClient__should_retry_fetch + self.assertTrue(should_retry(TimeoutError("pool acquire timed out"))) + + def test_fetch_v2_pool_exhausted_matches_read_timeout_behavior(self): + # With fetch_retries=0, a pool-wait timeout must surface exactly like a read + # timeout: an empty variants dict, not None and not an exception. + client = RemoteEvaluationClient( + API_KEY, RemoteEvaluationConfig(fetch_pool_acquire_timeout_millis=100) + ) + held = client._connection_pool.acquire() + try: + variants = client.fetch_v2(User(user_id='test_user')) + self.assertEqual({}, variants) + finally: + client._connection_pool.release(held) + client.close() + def test_fetch_failed_with_retry(self): with RemoteEvaluationClient(API_KEY, RemoteEvaluationConfig(debug=False, fetch_retries=1, fetch_timeout_millis=1)) as client: @@ -53,7 +93,7 @@ def test_fetch_with_fetch_options(self): user = User(user_id='test_user') mock_conn = mock.MagicMock() - client._connection_pool.acquire = lambda: mock_conn + client._connection_pool.acquire = lambda **kwargs: mock_conn mock_conn.request.return_value = mock.MagicMock(status=200) mock_conn.request.return_value.read.return_value = json.dumps({ 'sdk-ci-test': {