feat: Allow customization of socket options

This allows clients to configure (e.g.) socket keep alive probes

Author: dustinbyrne

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,

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"

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