From 3eb16109009431cbe90d3e8aa63e863c7c63b779 Mon Sep 17 00:00:00 2001 From: Dustin Byrne Date: Mon, 8 Dec 2025 23:23:21 -0500 Subject: [PATCH 1/5] feat: Allow customization of socket options This allows clients to configure (e.g.) socket keep alive probes --- posthog/__init__.py | 5 +++ posthog/request.py | 85 +++++++++++++++++++++++++++++++----- posthog/test/test_request.py | 19 ++++++++ 3 files changed, 98 insertions(+), 11 deletions(-) diff --git a/posthog/__init__.py b/posthog/__init__.py index 98286dd7..445c17c6 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -22,6 +22,11 @@ InconclusiveMatchError as InconclusiveMatchError, RequiresServerEvaluation as RequiresServerEvaluation, ) +from posthog.request import ( + enable_keep_alive as enable_keep_alive, + set_socket_options as set_socket_options, + SocketOptions as SocketOptions, +) from posthog.types import ( FeatureFlag, FlagsAndPayloads, diff --git a/posthog/request.py b/posthog/request.py index 2540f0e7..dbbeff87 100644 --- a/posthog/request.py +++ b/posthog/request.py @@ -1,19 +1,47 @@ import json import logging import re +import socket from dataclasses import dataclass from datetime import date, datetime from gzip import GzipFile from io import BytesIO -from typing import Any, Optional, Union +from typing import Any, List, Optional, Tuple, Union + import requests from dateutil.tz import tzutc +from requests.adapters import HTTPAdapter +from urllib3.connection import HTTPConnection from urllib3.util.retry import Retry from posthog.utils import remove_trailing_slash from posthog.version import VERSION +SocketOptions = List[Tuple[int, int, Union[int, bytes]]] + +KEEPALIVE_IDLE_SECONDS = 60 +KEEPALIVE_INTERVAL_SECONDS = 60 +KEEPALIVE_PROBE_COUNT = 3 + +# TCP keepalive probes idle connections to prevent them from being dropped. +# SO_KEEPALIVE is cross-platform, but timing options vary: +# - Linux: TCP_KEEPIDLE, TCP_KEEPINTVL, TCP_KEEPCNT +# - macOS: only SO_KEEPALIVE (uses system defaults) +# - Windows: TCP_KEEPIDLE, TCP_KEEPINTVL (since Windows 10 1709) +KEEP_ALIVE_SOCKET_OPTIONS: SocketOptions = list( + HTTPConnection.default_socket_options +) + [ + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), +] +for attr, value in [ + ("TCP_KEEPIDLE", KEEPALIVE_IDLE_SECONDS), + ("TCP_KEEPINTVL", KEEPALIVE_INTERVAL_SECONDS), + ("TCP_KEEPCNT", KEEPALIVE_PROBE_COUNT), +]: + if hasattr(socket, attr): + KEEP_ALIVE_SOCKET_OPTIONS.append((socket.SOL_TCP, getattr(socket, attr), value)) + def _mask_tokens_in_url(url: str) -> str: """Mask token values in URLs for safe logging, keeping first 10 chars visible.""" @@ -29,17 +57,52 @@ class GetResponse: not_modified: bool = False -# Retry on both connect and read errors -# by default read errors will only retry idempotent HTTP methods (so not POST) -adapter = requests.adapters.HTTPAdapter( - max_retries=Retry( - total=2, - connect=2, - read=2, +class HTTPAdapterWithSocketOptions(HTTPAdapter): + """HTTPAdapter with configurable socket options.""" + + def __init__(self, *args, socket_options: Optional[SocketOptions] = None, **kwargs): + self.socket_options = socket_options + super().__init__(*args, **kwargs) + + def init_poolmanager(self, *args, **kwargs): + if self.socket_options is not None: + kwargs["socket_options"] = self.socket_options + super().init_poolmanager(*args, **kwargs) + + +def _build_session(socket_options: Optional[SocketOptions] = None) -> requests.Session: + adapter = HTTPAdapterWithSocketOptions( + max_retries=Retry( + total=2, + connect=2, + read=2, + ), + socket_options=socket_options, ) -) -_session = requests.sessions.Session() -_session.mount("https://", adapter) + session = requests.sessions.Session() + session.mount("https://", adapter) + return session + + +_session = _build_session() + + +def set_socket_options(socket_options: Optional[SocketOptions]) -> None: + """ + Configure socket options for all HTTP connections. + + Example: + from posthog import set_socket_options + set_socket_options([(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]) + """ + global _session + _session = _build_session(socket_options) + + +def enable_keep_alive() -> None: + """Enable TCP keepalive to prevent idle connections from being dropped.""" + set_socket_options(KEEP_ALIVE_SOCKET_OPTIONS) + US_INGESTION_ENDPOINT = "https://us.i.posthog.com" EU_INGESTION_ENDPOINT = "https://eu.i.posthog.com" diff --git a/posthog/test/test_request.py b/posthog/test/test_request.py index 7eee835f..7464a612 100644 --- a/posthog/test/test_request.py +++ b/posthog/test/test_request.py @@ -10,12 +10,15 @@ APIError, DatetimeSerializer, GetResponse, + KEEP_ALIVE_SOCKET_OPTIONS, QuotaLimitError, _mask_tokens_in_url, batch_post, decide, determine_server_host, + enable_keep_alive, get, + set_socket_options, ) from posthog.test.test_utils import TEST_API_KEY @@ -344,3 +347,19 @@ def test_get_removes_trailing_slash_from_host(self, mock_get): ) def test_routing_to_custom_host(host, expected): assert determine_server_host(host) == expected + + +def test_enable_keep_alive_sets_socket_options(): + enable_keep_alive() + from posthog.request import _session + + adapter = _session.get_adapter("https://example.com") + assert adapter.socket_options == KEEP_ALIVE_SOCKET_OPTIONS + + +def test_set_socket_options_clears_with_none(): + set_socket_options(None) + from posthog.request import _session + + adapter = _session.get_adapter("https://example.com") + assert adapter.socket_options is None From 99bacfbe802ffbcd74675899df1aa8cac29a5098 Mon Sep 17 00:00:00 2001 From: Dustin Byrne Date: Tue, 9 Dec 2025 00:06:23 -0500 Subject: [PATCH 2/5] test: Prevent leaking modified _session --- posthog/test/test_request.py | 25 ++++++++++++++++--------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/posthog/test/test_request.py b/posthog/test/test_request.py index 7464a612..52dbc157 100644 --- a/posthog/test/test_request.py +++ b/posthog/test/test_request.py @@ -350,16 +350,23 @@ def test_routing_to_custom_host(host, expected): def test_enable_keep_alive_sets_socket_options(): - enable_keep_alive() - from posthog.request import _session + try: + enable_keep_alive() + from posthog.request import _session - adapter = _session.get_adapter("https://example.com") - assert adapter.socket_options == KEEP_ALIVE_SOCKET_OPTIONS + adapter = _session.get_adapter("https://example.com") + assert adapter.socket_options == KEEP_ALIVE_SOCKET_OPTIONS + finally: + set_socket_options(None) def test_set_socket_options_clears_with_none(): - set_socket_options(None) - from posthog.request import _session - - adapter = _session.get_adapter("https://example.com") - assert adapter.socket_options is None + try: + enable_keep_alive() + set_socket_options(None) + from posthog.request import _session + + adapter = _session.get_adapter("https://example.com") + assert adapter.socket_options is None + finally: + set_socket_options(None) From 157258099b8f9577a2d39af3162fa672cf6958a6 Mon Sep 17 00:00:00 2001 From: Dustin Byrne Date: Tue, 9 Dec 2025 00:06:56 -0500 Subject: [PATCH 3/5] chore: Add a type ignore --- posthog/request.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/posthog/request.py b/posthog/request.py index dbbeff87..6ea36052 100644 --- a/posthog/request.py +++ b/posthog/request.py @@ -11,7 +11,7 @@ import requests from dateutil.tz import tzutc -from requests.adapters import HTTPAdapter +from requests.adapters import HTTPAdapter # type: ignore[import-untyped] from urllib3.connection import HTTPConnection from urllib3.util.retry import Retry From 5d9b8e826348adcc9d4eebca073390acf6528ff6 Mon Sep 17 00:00:00 2001 From: Dustin Byrne Date: Tue, 9 Dec 2025 01:19:38 -0500 Subject: [PATCH 4/5] feat: Add `disable_connection_reuse` config method Disables connection pooling --- posthog/__init__.py | 1 + posthog/request.py | 21 ++++++++++++++++++--- posthog/test/test_request.py | 12 ++++++++++++ 3 files changed, 31 insertions(+), 3 deletions(-) diff --git a/posthog/__init__.py b/posthog/__init__.py index 445c17c6..b513d8de 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -23,6 +23,7 @@ RequiresServerEvaluation as RequiresServerEvaluation, ) from posthog.request import ( + disable_connection_reuse as disable_connection_reuse, enable_keep_alive as enable_keep_alive, set_socket_options as set_socket_options, SocketOptions as SocketOptions, diff --git a/posthog/request.py b/posthog/request.py index 6ea36052..9cc1c51a 100644 --- a/posthog/request.py +++ b/posthog/request.py @@ -85,6 +85,14 @@ def _build_session(socket_options: Optional[SocketOptions] = None) -> requests.S _session = _build_session() +_socket_options: Optional[SocketOptions] = None +_pooling_enabled = True + + +def _get_session() -> requests.Session: + if _pooling_enabled: + return _session + return _build_session(_socket_options) def set_socket_options(socket_options: Optional[SocketOptions]) -> None: @@ -95,7 +103,8 @@ def set_socket_options(socket_options: Optional[SocketOptions]) -> None: from posthog import set_socket_options set_socket_options([(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]) """ - global _session + global _session, _socket_options + _socket_options = socket_options _session = _build_session(socket_options) @@ -104,6 +113,12 @@ def enable_keep_alive() -> None: set_socket_options(KEEP_ALIVE_SOCKET_OPTIONS) +def disable_connection_reuse() -> None: + """Disable connection reuse, creating a fresh connection for each request.""" + global _pooling_enabled + _pooling_enabled = False + + US_INGESTION_ENDPOINT = "https://us.i.posthog.com" EU_INGESTION_ENDPOINT = "https://eu.i.posthog.com" DEFAULT_HOST = US_INGESTION_ENDPOINT @@ -148,7 +163,7 @@ def post( gz.write(data.encode("utf-8")) data = buf.getvalue() - res = _session.post(url, data=data, headers=headers, timeout=timeout) + res = _get_session().post(url, data=data, headers=headers, timeout=timeout) if res.status_code == 200: log.debug("data uploaded successfully") @@ -263,7 +278,7 @@ def get( if etag: headers["If-None-Match"] = etag - res = _session.get(full_url, headers=headers, timeout=timeout) + res = _get_session().get(full_url, headers=headers, timeout=timeout) masked_url = _mask_tokens_in_url(full_url) diff --git a/posthog/test/test_request.py b/posthog/test/test_request.py index 52dbc157..56415277 100644 --- a/posthog/test/test_request.py +++ b/posthog/test/test_request.py @@ -6,6 +6,7 @@ import pytest import requests +import posthog.request as request_module from posthog.request import ( APIError, DatetimeSerializer, @@ -16,6 +17,7 @@ batch_post, decide, determine_server_host, + disable_connection_reuse, enable_keep_alive, get, set_socket_options, @@ -370,3 +372,13 @@ def test_set_socket_options_clears_with_none(): assert adapter.socket_options is None finally: set_socket_options(None) + + +def test_disable_connection_reuse_creates_fresh_sessions(): + try: + disable_connection_reuse() + session1 = request_module._get_session() + session2 = request_module._get_session() + assert session1 is not session2 + finally: + request_module._pooling_enabled = True From c8e3bd5421875379594db528a39d851c466ae544 Mon Sep 17 00:00:00 2001 From: Dustin Byrne Date: Tue, 9 Dec 2025 01:25:04 -0500 Subject: [PATCH 5/5] set_socket_options is idempotent --- posthog/request.py | 2 ++ posthog/test/test_request.py | 11 +++++++++++ 2 files changed, 13 insertions(+) diff --git a/posthog/request.py b/posthog/request.py index 9cc1c51a..7199f3a8 100644 --- a/posthog/request.py +++ b/posthog/request.py @@ -104,6 +104,8 @@ def set_socket_options(socket_options: Optional[SocketOptions]) -> None: set_socket_options([(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]) """ global _session, _socket_options + if socket_options == _socket_options: + return _socket_options = socket_options _session = _build_session(socket_options) diff --git a/posthog/test/test_request.py b/posthog/test/test_request.py index 56415277..128123fe 100644 --- a/posthog/test/test_request.py +++ b/posthog/test/test_request.py @@ -382,3 +382,14 @@ def test_disable_connection_reuse_creates_fresh_sessions(): assert session1 is not session2 finally: request_module._pooling_enabled = True + + +def test_set_socket_options_is_idempotent(): + try: + enable_keep_alive() + session1 = request_module._session + enable_keep_alive() + session2 = request_module._session + assert session1 is session2 + finally: + set_socket_options(None)