fix(traces): classify the export response without reading its body

A requests timeout bounds read inactivity, not the whole response, so a
proxy answering 503 with a body that keeps dripping held the exporter's
single flight open indefinitely. Stream the response, classify it from
status and headers, and close it unread.

Author: turnipdabeets

posthog/test/tracing/test_transport.py

@@ -1,6 +1,9 @@
 import gzip
 import json
+import threading
+import time
 from datetime import datetime, timezone
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from types import SimpleNamespace
 from unittest import mock
 
@@ -199,3 +202,64 @@ def test_survives_a_throwing_headers_object(self):
         session.post.return_value = response
         outcome, _ = send(session=session)
         assert outcome == SendOutcome("retry-later", None)
+
+
+class _ChunkedHandler(BaseHTTPRequestHandler):
+    def do_POST(self):
+        self.rfile.read(int(self.headers.get("Content-Length", 0)))
+        self.send_response(self.server.status)
+        self.send_header("Transfer-Encoding", "chunked")
+        self.end_headers()
+        if self.server.status < 300:
+            self.wfile.write(b"0\r\n\r\n")
+            return
+        # An error body that drips a chunk every 10 ms and never finishes.
+        while not self.server.stop.is_set():
+            try:
+                self.wfile.write(b"1\r\nx\r\n")
+                self.wfile.flush()
+            except OSError:
+                return
+            time.sleep(0.01)
+
+    def log_message(self, *args):
+        pass
+
+
+@pytest.fixture
+def local_server():
+    servers = []
+
+    def start(status):
+        server = ThreadingHTTPServer(("127.0.0.1", 0), _ChunkedHandler)
+        server.status = status
+        server.stop = threading.Event()
+        thread = threading.Thread(target=server.serve_forever, daemon=True)
+        thread.start()
+        servers.append((server, thread))
+        return "http://127.0.0.1:{}".format(server.server_port)
+
+    yield start
+    for server, thread in servers:
+        server.stop.set()
+        server.shutdown()
+        server.server_close()
+        thread.join(2)
+
+
+class TestResponseBody:
+    def test_closes_the_response_without_reading_the_body(self):
+        _, session = send()
+        assert session.post.call_args[1]["stream"] is True
+        assert session.post.return_value.close.called
+
+    def test_does_not_wait_for_a_dripping_error_body(self, local_server):
+        client = fake_client(host=local_server(503), timeout=0.5)
+        started = time.monotonic()
+        outcome = send_traces_batch(client, PAYLOAD)
+        assert outcome == SendOutcome("retry-later", None)
+        assert time.monotonic() - started < 2
+
+    def test_a_completed_response_is_still_ok(self, local_server):
+        client = fake_client(host=local_server(200), timeout=0.5)
+        assert send_traces_batch(client, PAYLOAD) == SendOutcome("ok")

posthog/tracing/_transport.py

@@ -108,11 +108,21 @@ def send_traces_batch(client: Any, payload: dict) -> SendOutcome:
                 "User-Agent": USER_AGENT,
             },
             timeout=timeout,
+            stream=True,
         )
     except requests.exceptions.RequestException as e:
         log.debug("Span batch request failed: %s", e)
         return SendOutcome("retry-later")
+    # Status and headers alone classify the response, so the body is never
+    # read: the timeout bounds read inactivity, and a body that keeps dripping
+    # would otherwise hold the exporter's single flight open indefinitely.
+    try:
+        return _classify(response)
+    finally:
+        response.close()
+
 
+def _classify(response: requests.Response) -> SendOutcome:
     status = response.status_code
     if status < 300:
         return OK

typings/requests/__init__.pyi

@@ -8,6 +8,7 @@ class Response:
     text: str
     headers: dict[str, str]
     def json(self) -> Any: ...
+    def close(self) -> None: ...
 
 class Session:
     def mount(self, prefix: str, adapter: adapters.HTTPAdapter) -> None: ...
@@ -19,6 +20,7 @@ class Session:
         data: str | bytes,
         headers: dict[str, str],
         timeout: int,
+        stream: bool = ...,
     ) -> Response: ...
     def get(
         self,