diff --git a/src/amplitude_experiment/remote/client.py b/src/amplitude_experiment/remote/client.py index 8db2653..cbba315 100644 --- a/src/amplitude_experiment/remote/client.py +++ b/src/amplitude_experiment/remote/client.py @@ -172,8 +172,8 @@ def __do_fetch(self, user, fetch_options: FetchOptions = None): def __setup_connection_pool(self): scheme, _, host = self.config.server_url.split('/', 3) timeout = self.config.fetch_timeout_millis / 1000 - self._connection_pool = HTTPConnectionPool(host, max_size=1, idle_timeout=30, - read_timeout=timeout, scheme=scheme) + self._connection_pool = HTTPConnectionPool(host, max_size=self.config.connection_pool_max_size, + idle_timeout=30, read_timeout=timeout, scheme=scheme) def close(self) -> None: """ diff --git a/src/amplitude_experiment/remote/config.py b/src/amplitude_experiment/remote/config.py index c4efdcc..5724e1e 100644 --- a/src/amplitude_experiment/remote/config.py +++ b/src/amplitude_experiment/remote/config.py @@ -18,6 +18,7 @@ def __init__(self, debug=False, fetch_retry_backoff_scalar=1.5, fetch_retry_timeout_millis=10000, server_zone: ServerZone = ServerZone.US, + connection_pool_max_size=1, logger=None): """ Initialize a config @@ -34,6 +35,10 @@ def __init__(self, debug=False, fetch_retry_backoff_scalar (float): Scales the minimum backoff exponentially. fetch_retry_timeout_millis (int): The request timeout for retrying fetch requests. server_zone (str): Select the Amplitude data center to get flags and variants from, US or EU. + connection_pool_max_size (int): The maximum number of HTTP connections kept in the fetch connection + pool, and therefore the maximum number of concurrent fetch requests. Additional concurrent fetches + beyond this limit block waiting for a connection to be released. Defaults to 1 (all fetches in the + process share a single keep-alive connection). 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. @@ -49,6 +54,7 @@ def __init__(self, debug=False, self.fetch_retry_backoff_scalar = fetch_retry_backoff_scalar self.fetch_retry_timeout_millis = fetch_retry_timeout_millis self.server_zone = server_zone + self.connection_pool_max_size = connection_pool_max_size 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..c59f506 100644 --- a/tests/remote/client_test.py +++ b/tests/remote/client_test.py @@ -124,3 +124,15 @@ def test_fetch_retry_with_response(self, response_code, error_message, should_ca if __name__ == '__main__': unittest.main() + + +class RemoteEvaluationClientConnectionPoolTestCase(unittest.TestCase): + + def test_connection_pool_defaults_to_single_connection(self): + with RemoteEvaluationClient(API_KEY) as client: + self.assertEqual(client._connection_pool.max_size, 1) + + def test_connection_pool_size_is_configurable(self): + config = RemoteEvaluationConfig(connection_pool_max_size=8) + with RemoteEvaluationClient(API_KEY, config) as client: + self.assertEqual(client._connection_pool.max_size, 8) diff --git a/tests/remote/config_test.py b/tests/remote/config_test.py index 2159b52..b9da77a 100644 --- a/tests/remote/config_test.py +++ b/tests/remote/config_test.py @@ -68,3 +68,15 @@ def test_default_logger_only_one_handler_added(self): if __name__ == '__main__': unittest.main() + + +class RemoteEvaluationConfigConnectionPoolTestCase(unittest.TestCase): + """Tests for RemoteEvaluationConfig connection pool configuration""" + + def test_connection_pool_max_size_defaults_to_one(self): + config = RemoteEvaluationConfig() + self.assertEqual(config.connection_pool_max_size, 1) + + def test_connection_pool_max_size_is_stored(self): + config = RemoteEvaluationConfig(connection_pool_max_size=8) + self.assertEqual(config.connection_pool_max_size, 8)