diff --git a/src/amplitude_experiment/cohort/cohort_download_api.py b/src/amplitude_experiment/cohort/cohort_download_api.py index e1a41a3..e441689 100644 --- a/src/amplitude_experiment/cohort/cohort_download_api.py +++ b/src/amplitude_experiment/cohort/cohort_download_api.py @@ -2,8 +2,9 @@ import logging import base64 import json +from contextlib import contextmanager from http.client import HTTPResponse -from typing import Optional +from typing import Iterator, Optional from ..version import __version__ from .cohort import Cohort @@ -36,27 +37,27 @@ def get_cohort(self, cohort_id: str, cohort: Optional[Cohort]) -> Cohort or None response = None try: last_modified = None if cohort is None else cohort.last_modified - response = self._get_cohort_members_request(cohort_id, last_modified) - self.logger.debug(f"getCohortMembers({cohort_id}): status={response.status}") - if response.status == 200: - cohort_info = json.loads(response.read().decode("utf8")) - self.logger.debug(f"getCohortMembers({cohort_id}): end - resultSize={cohort_info['size']}") - return Cohort( - id=cohort_info['cohortId'], - last_modified=cohort_info['lastModified'], - size=cohort_info['size'], - member_ids=set(cohort_info['memberIds']), - group_type=cohort_info['groupType'], - ) - elif response.status == 204: - self.logger.debug(f"getCohortMembers({cohort_id}): Cohort not modified") - return - elif response.status == 413: - raise CohortTooLargeException( - f"Cohort exceeds max cohort size of {self.max_cohort_size}: {response.status}") - elif response.status != 202: - raise HTTPErrorResponseException(response.status, - f"Unexpected response code: {response.status}") + with self._get_cohort_members_request(cohort_id, last_modified) as response: + self.logger.debug(f"getCohortMembers({cohort_id}): status={response.status}") + if response.status == 200: + cohort_info = json.loads(response.read().decode("utf8")) + self.logger.debug(f"getCohortMembers({cohort_id}): end - resultSize={cohort_info['size']}") + return Cohort( + id=cohort_info['cohortId'], + last_modified=cohort_info['lastModified'], + size=cohort_info['size'], + member_ids=set(cohort_info['memberIds']), + group_type=cohort_info['groupType'], + ) + elif response.status == 204: + self.logger.debug(f"getCohortMembers({cohort_id}): Cohort not modified") + return + elif response.status == 413: + raise CohortTooLargeException( + f"Cohort exceeds max cohort size of {self.max_cohort_size}: {response.status}") + elif response.status != 202: + raise HTTPErrorResponseException(response.status, + f"Unexpected response code: {response.status}") except Exception as e: if response and not (isinstance(e, HTTPErrorResponseException) and response.status == 429): errors += 1 @@ -65,19 +66,26 @@ def get_cohort(self, cohort_id: str, cohort: Optional[Cohort]) -> Cohort or None raise e time.sleep(COHORT_REQUEST_RETRY_DELAY_MILLIS / 1000) - def _get_cohort_members_request(self, cohort_id: str, last_modified: int) -> HTTPResponse: + @contextmanager + def _get_cohort_members_request(self, cohort_id: str, last_modified: int) -> Iterator[HTTPResponse]: headers = { 'Authorization': f'Basic {self._get_basic_auth()}', 'X-Amp-Exp-Library': f"experiment-python-server/{__version__}" } conn = self._connection_pool.acquire() + response = None try: url = f'/sdk/v1/cohort/{cohort_id}?maxCohortSize={self.max_cohort_size}' if last_modified is not None: url += f'&lastModified={last_modified}' response = conn.request('GET', url, headers=headers) - return response + yield response + except Exception: + conn.close() + raise finally: + if response is not None and (response.status != 200 or not response.isclosed()): + conn.close() self._connection_pool.release(conn) def _get_basic_auth(self) -> str: diff --git a/src/amplitude_experiment/flag/flag_config_api.py b/src/amplitude_experiment/flag/flag_config_api.py index 662f3cf..e33e51c 100644 --- a/src/amplitude_experiment/flag/flag_config_api.py +++ b/src/amplitude_experiment/flag/flag_config_api.py @@ -35,7 +35,11 @@ def _get_flag_configs(self) -> List[EvaluationFlag]: body = None try: response = conn.request('GET', '/sdk/v2/flags?v=0', body, headers) - response_body = response.read().decode("utf8") + try: + response_body = response.read().decode("utf8") + except Exception: + conn.close() + raise if response.status != 200: raise Exception( f"[Experiment] Get flagConfigs - received error response: ${response.status}: ${response_body}") diff --git a/src/amplitude_experiment/remote/client.py b/src/amplitude_experiment/remote/client.py index d4e8004..39bd8f6 100644 --- a/src/amplitude_experiment/remote/client.py +++ b/src/amplitude_experiment/remote/client.py @@ -149,6 +149,12 @@ def __do_fetch(self, user, fetch_options: FetchOptions = None): json.dumps(fetch_options.flagKeys, separators=(",", ":")).encode("utf-8") ).rstrip(b"=").decode("utf-8") + body = user_context.to_json().encode('utf8') + if len(body) > 8000: + self.logger.warning(f"[Experiment] encoded user object length ${len(body)} " + f"cannot be cached by CDN; must be < 8KB") + self.logger.debug(f"[Experiment] Fetch variants for user: {str(user_context)}") + acquire_timeout_millis = self.config.fetch_pool_acquire_timeout_millis try: conn = self._connection_pool.acquire( @@ -157,19 +163,21 @@ def __do_fetch(self, user, fetch_options: FetchOptions = 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)} " - f"cannot be cached by CDN; must be < 8KB") - self.logger.debug(f"[Experiment] Fetch variants for user: {str(user_context)}") try: response = conn.request('POST', '/sdk/v2/vardata?v=0', body, headers) elapsed = '%.3f' % ((time.time() - start) * 1000) self.logger.debug(f"[Experiment] Fetch complete in {elapsed} ms") if response.status != 200: + # The error body may be large or slow. Retire its connection instead of reading it. + conn.close() raise FetchException(response.status, f"Fetch error response: status={response.status} {response.reason}") - json_response = json.loads(response.read().decode("utf8")) + try: + response_as_bytes = response.read() + except Exception: + conn.close() + raise + json_response = json.loads(response_as_bytes.decode("utf8")) variants = evaluation_variants_json_to_variants(json_response) self.logger.debug(f"[Experiment] Fetched variants: {json.dumps(variants, default=str)}") return variants diff --git a/tests/cohort/cohort_download_api_test.py b/tests/cohort/cohort_download_api_test.py index 2a42394..e49f160 100644 --- a/tests/cohort/cohort_download_api_test.py +++ b/tests/cohort/cohort_download_api_test.py @@ -1,19 +1,20 @@ import json import logging import unittest +from contextlib import nullcontext from unittest import mock -from unittest.mock import MagicMock, patch +from unittest.mock import patch from src.amplitude_experiment.cohort.cohort import Cohort from src.amplitude_experiment.cohort.cohort_download_api import DirectCohortDownloadApi from src.amplitude_experiment.exception import CohortTooLargeException def response(code: int, body: dict = None): - mock_response = MagicMock() + mock_response = mock.MagicMock() mock_response.status = code if body is not None: mock_response.read.return_value = json.dumps(body).encode() - return mock_response + return nullcontext(mock_response) class CohortDownloadApiTest(unittest.TestCase): diff --git a/tests/pooled_response_test.py b/tests/pooled_response_test.py new file mode 100644 index 0000000..33457cd --- /dev/null +++ b/tests/pooled_response_test.py @@ -0,0 +1,164 @@ +import json +import logging +import socket +import threading +import time +import unittest +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from unittest.mock import patch + +from src.amplitude_experiment import RemoteEvaluationClient, RemoteEvaluationConfig, User +from src.amplitude_experiment.cohort.cohort_download_api import DirectCohortDownloadApi +from src.amplitude_experiment.flag.flag_config_api import FlagConfigApiV2 + + +class ResponseServer: + def __init__(self, responses): + self.responses = responses + self.requests = [] + self.lock = threading.Lock() + owner = self + + class Handler(BaseHTTPRequestHandler): + protocol_version = 'HTTP/1.1' + + def handle_one_request(self): + try: + super().handle_one_request() + except ConnectionResetError: + self.close_connection = True + + def do_POST(self): + self.reply() + + def do_GET(self): + self.reply() + + def reply(self): + length = int(self.headers.get('Content-Length', '0')) + if length: + self.rfile.read(length) + with owner.lock: + index = len(owner.requests) + owner.requests.append((self.path, self.client_address[1])) + status, body, stall = owner.responses[min(index, len(owner.responses) - 1)] + self.send_response(status) + self.send_header('Content-Length', str(len(body) + (99 if stall else 0))) + self.end_headers() + try: + self.wfile.write(body) + self.wfile.flush() + except BrokenPipeError: + return + if stall: + time.sleep(0.3) + + def log_message(self, *_): + pass + + self.server = ThreadingHTTPServer(('127.0.0.1', 0), Handler) + self.server.daemon_threads = True + self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) + + def __enter__(self): + self.thread.start() + return self + + def __exit__(self, *_): + self.server.shutdown() + self.server.server_close() + self.thread.join(timeout=2) + + @property + def url(self): + return f'http://127.0.0.1:{self.server.server_address[1]}' + + +class PooledResponseTest(unittest.TestCase): + def remote_client(self, url): + return RemoteEvaluationClient('test-key', RemoteEvaluationConfig( + server_url=url, fetch_retries=1, fetch_retry_backoff_min_millis=0, + fetch_timeout_millis=100, connection_pool_max_size=1)) + + def test_remote_error_body_then_retry(self): + with ResponseServer([(503, b'error', False), (200, b'{}', False)]) as server: + with self.remote_client(server.url) as client: + self.assertEqual({}, client.fetch_v2(User(user_id='test-user'))) + self.assertEqual(2, len(server.requests)) + + def test_remote_success_reuses_keep_alive_connection(self): + with ResponseServer([(200, b'{}', False), (200, b'{}', False)]) as server: + with self.remote_client(server.url) as client: + self.assertEqual({}, client.fetch_v2(User(user_id='first'))) + self.assertEqual({}, client.fetch_v2(User(user_id='second'))) + self.assertEqual(2, len(server.requests)) + self.assertEqual(server.requests[0][1], server.requests[1][1]) + + def test_remote_stalled_error_body_then_retry(self): + with ResponseServer([(503, b'e', True), (200, b'{}', False)]) as server: + with self.remote_client(server.url) as client: + self.assertEqual({}, client.fetch_v2(User(user_id='test-user'))) + self.assertEqual(2, len(server.requests)) + + def test_remote_stalled_success_body_then_retry(self): + with ResponseServer([(200, b'{', True), (200, b'{}', False)]) as server: + with self.remote_client(server.url) as client: + self.assertEqual({}, client.fetch_v2(User(user_id='test-user'))) + self.assertEqual(2, len(server.requests)) + + def test_remote_serialization_failure_releases_pool_slot(self): + with self.remote_client('http://127.0.0.1:1') as client: + with patch.object(User, 'to_json', side_effect=ValueError('bad user')): + with self.assertRaises(ValueError): + client._RemoteEvaluationClient__do_fetch(User(user_id='test-user')) + conn = client._connection_pool.acquire(blocking=False) + client._connection_pool.release(conn) + + def test_flag_config_stalled_body_does_not_poison_next_request(self): + with ResponseServer([(503, b'e', True), (200, b'[]', False)]) as server: + api = FlagConfigApiV2('test-key', server.url, 100) + try: + with self.assertRaises(socket.timeout): + api.get_flag_configs() + self.assertEqual([], api.get_flag_configs()) + self.assertEqual(2, len(server.requests)) + finally: + api._connection_pool.close() + + def test_cohort_202_body_then_success(self): + cohort_body = json.dumps({'cohortId': 'cohort', 'lastModified': 1, 'size': 1, + 'groupType': 'User', 'memberIds': ['user']}).encode() + with ResponseServer([(202, b'pending', False), (200, cohort_body, False)]) as server: + api = DirectCohortDownloadApi('key', 'secret', 100, server.url, logging.getLogger('test')) + api._connection_pool.max_size = 1 + try: + first = api._get_cohort_members_request('cohort', None) + with first as response: + self.assertEqual(202, response.status) + with api._get_cohort_members_request('cohort', None) as response: + self.assertEqual(200, response.status) + self.assertEqual(['user'], json.loads(response.read())['memberIds']) + self.assertEqual(2, len(server.requests)) + finally: + api._connection_pool.close() + + def test_cohort_stalled_success_body_does_not_poison_next_request(self): + cohort_body = json.dumps({'cohortId': 'cohort', 'lastModified': 1, 'size': 1, + 'groupType': 'User', 'memberIds': ['user']}).encode() + with ResponseServer([(200, b'{', True), (200, cohort_body, False)]) as server: + api = DirectCohortDownloadApi('key', 'secret', 100, server.url, logging.getLogger('test')) + api._connection_pool.max_size = 1 + api._connection_pool.read_timeout = 0.1 + try: + with self.assertRaises(socket.timeout): + with api._get_cohort_members_request('cohort', None) as response: + response.read() + with api._get_cohort_members_request('cohort', None) as response: + self.assertEqual(['user'], json.loads(response.read())['memberIds']) + self.assertEqual(2, len(server.requests)) + finally: + api._connection_pool.close() + + +if __name__ == '__main__': + unittest.main()