fix: retire pooled HTTP connections after error or unread responses

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:

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}")

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

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):

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()