Add flag to enable TCP keepalives (#163)

Author: BenjaminLowry

README.md

@@ -6,7 +6,7 @@
 
 ## Local development
 
-A light-weight HTTP client library.
+A light-weight HTTP client library for conjure.
 
 ### Overview
 

changelog/@unreleased/pr-161.v2.yml

@@ -0,0 +1,5 @@
+type: improvement
+improvement:
+  description: Add flag to enable TCP keepalives
+  links:
+  - https://github.com/palantir/conjure-python-client/pull/161

conjure_python_client/_http/requests_client.py

@@ -15,6 +15,7 @@
 from requests.adapters import HTTPAdapter, Response, CaseInsensitiveDict
 from typing import TypeVar, Type, List, Optional, Dict, Any, Union
 from requests.exceptions import HTTPError
+from requests.packages.urllib3.connection import HTTPConnection
 from requests.packages.urllib3.poolmanager import PoolManager
 from requests.packages.urllib3.util.ssl_ import create_urllib3_context
 from requests.packages.urllib3.util import Retry
@@ -24,6 +25,8 @@
 import os
 import random
 import requests
+import socket
+import sys
 
 
 T = TypeVar("T")
@@ -47,6 +50,24 @@
     "DHE-RSA-AES256-GCM-SHA384"
 )
 
+SOCKET_KEEP_ALIVE = (
+    socket.SOL_SOCKET,
+    socket.SO_KEEPALIVE,
+    1,
+)  # Enable keep alive.
+SOCKET_KEEP_INTVL = (
+    socket.SOL_TCP,
+    socket.TCP_KEEPINTVL,
+    120,
+)  # Interval of 120s between individual keepalive probes.
+KEEP_ALIVE_SOCKET_OPTIONS = [SOCKET_KEEP_ALIVE, SOCKET_KEEP_INTVL]
+if sys.platform != "darwin":
+    SOCKET_KEEP_IDLE = (
+        socket.SOL_TCP,
+        socket.TCP_KEEPIDLE,
+        120,
+    )  # After 120s of idle connection, start keepalive probes.
+    KEEP_ALIVE_SOCKET_OPTIONS.append(SOCKET_KEEP_IDLE)
 
 TRACE_ID_HEADER: str = "X-B3-TraceId"
 TRACE_ID_RANDOM_BYTES = 8
@@ -164,6 +185,7 @@ def create(
         user_agent: str,
         service_config: ServiceConfiguration,
         return_none_for_unknown_union_types: bool = False,
+        enable_keep_alive: bool = False,
     ) -> T:
         # setup retry to match java remoting
         # https://github.com/palantir/http-remoting/tree/3.12.0#quality-of-service-retry-failover-throttling
@@ -173,7 +195,9 @@ def create(
             status_forcelist=[308, 429, 503],
             backoff_factor=float(service_config.backoff_slot_size) / 1000,
         )
-        transport_adapter = TransportAdapter(max_retries=retry)
+        transport_adapter = TransportAdapter(
+            max_retries=retry, enable_keep_alive=enable_keep_alive
+        )
         # create a session, for shared connection polling, user agent, etc
         session = requests.Session()
         session.headers = CaseInsensitiveDict({"User-Agent": user_agent})
@@ -196,18 +220,29 @@ def create(
 class TransportAdapter(HTTPAdapter):
     """Transport adapter that allows customising ssl things"""
 
+    __attrs__ = HTTPAdapter.__attrs__ + ["_enable_keep_alive"]
+
+    def __init__(self, *args, enable_keep_alive: bool = False, **kwargs):
+        self._enable_keep_alive = enable_keep_alive
+        super().__init__(*args, **kwargs)
+
     def init_poolmanager(
         self, connections, maxsize, block=False, **pool_kwargs
     ):
         self._pool_connections = connections
         self._pool_maxsize = maxsize
         self._pool_block = block
         ssl_context = create_urllib3_context(ciphers=CIPHERS)
+        keep_alive_pool_kwargs = {
+            "socket_options": HTTPConnection.default_socket_options
+            + KEEP_ALIVE_SOCKET_OPTIONS
+        }
+        if self._enable_keep_alive:
+            pool_kwargs = {**pool_kwargs, **keep_alive_pool_kwargs}
         self.poolmanager = PoolManager(
             num_pools=connections,
             maxsize=maxsize,
             block=block,
-            strict=True,
             ssl_context=ssl_context,
             **pool_kwargs
         )

conjure_python_client/_serde/decoder.py

@@ -141,8 +141,8 @@ def decode_conjure_union_type(
         # for backwards compatibility with conjure-python,
         # only pass in arg type_of_union if it is expected
         param_dict = inspect.signature(conjure_type.__init__).parameters
-        if 'type_of_union' in param_dict:
-            deserialized['type_of_union'] = type_of_union
+        if "type_of_union" in param_dict:
+            deserialized["type_of_union"] = type_of_union
         return conjure_type(**deserialized)
 
     @classmethod

setup.py

@@ -14,29 +14,63 @@
 #!/usr/bin/env python
 from setuptools import find_packages, setup, Command
 from os import path, makedirs, system
+import re
 import subprocess
 import sys
 
 VERSION_PY_PATH = "conjure_python_client/_version.py"
 
+GIT_REGEX_PATTERN = re.compile(
+    r"^"
+    + r"(?P<tag>[0-9]+\.[0-9]+\.[0-9]+)"
+    + r"(-rc(?P<rc>[0-9]+))?"
+    + r"(-(?P<distance>[0-9]+)-g(?P<hash>[a-f0-9]+))?"
+    + r"(\.(?P<dirty>dirty))?"
+    + r"$"
+)
+
+
+def convert_sls_version_to_python(sls_version: str) -> str:
+    match = GIT_REGEX_PATTERN.match(sls_version)
+    if not match:
+        raise RuntimeError(f"Invalid SLS version {sls_version}")
 
-if not path.exists(VERSION_PY_PATH):
-    try:
-        gitversion = (
-            subprocess.check_output(
-                "git describe --tags --always --first-parent".split()
+    python_version = match.group("tag")
+    rc_group = match.group("rc")
+    distance_group = match.group("distance")
+    hash_group = match.group("hash")
+    dirty_group = match.group("dirty")
+
+    if rc_group:
+        python_version += "rc" + rc_group
+    if distance_group:
+        if not hash_group:
+            raise RuntimeError(
+                f"Cannot specify commit distance without hash for version {sls_version}"
             )
-            .decode()
-            .strip()
-            .replace("-", "_")
+        python_version += "+" + distance_group + ".g" + hash_group
+    if dirty_group:
+        python_version += "." + dirty_group
+    return python_version
+
+
+try:
+    gitversion = (
+        subprocess.check_output(
+            "git describe --tags --always --first-parent".split()
         )
-        open(VERSION_PY_PATH, "w").write(
-            '__version__ = "{}"\n'.format(gitversion)
+        .decode()
+        .strip()
+    )
+    open(VERSION_PY_PATH, "w").write(
+        '__version__ = "{}"\n'.format(
+            convert_sls_version_to_python(gitversion)
         )
-        if not path.exists("build"):
-            makedirs("build")
-    except subprocess.CalledProcessError:
-        print("outside git repo, not generating new version string")
+    )
+    if not path.exists("build"):
+        makedirs("build")
+except subprocess.CalledProcessError:
+    print("outside git repo, not generating new version string")
 exec(open(VERSION_PY_PATH).read())
 
 
@@ -83,12 +117,12 @@ def blackCheck(self):
     # The project's main homepage.
     url="https://github.com/palantir/conjure-python-client",
     author="Palantir Technologies, Inc.",
-    classifiers=[ 
+    classifiers=[
         "License :: OSI Approved :: Apache Software License",
         "Programming Language :: Python :: 3",
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
-        "Programming Language :: Python :: 3.10"
+        "Programming Language :: Python :: 3.10",
     ],
     # You can just specify the packages manually here if your project is
     # simple. Or you can use find_packages().

test/http/__init__.py

@@ -0,0 +1,40 @@
+import sys
+
+from conjure_python_client._http.requests_client import (
+    SOCKET_KEEP_ALIVE,
+    SOCKET_KEEP_INTVL,
+    TransportAdapter,
+)
+
+if sys.platform != "darwin":
+    from conjure_python_client._http.requests_client import SOCKET_KEEP_IDLE
+
+
+def test_can_enable_keep_alives_in_transport_adapter():
+    assert (
+        TransportAdapter(
+            max_retries=12, enable_keep_alive=True
+        )._enable_keep_alive
+        is True
+    )
+    assert TransportAdapter(max_retries=12)._enable_keep_alive is False
+    assert TransportAdapter()._enable_keep_alive is False
+
+
+def test_keep_alive_passes_correct_options():
+    socket_options = TransportAdapter(
+        max_retries=12, enable_keep_alive=True
+    ).poolmanager.connection_pool_kw["socket_options"]
+    assert SOCKET_KEEP_ALIVE in socket_options
+    assert SOCKET_KEEP_INTVL in socket_options
+    if sys.platform != "darwin":
+        assert SOCKET_KEEP_IDLE in socket_options
+
+
+def test_keep_alive_passed_in_state_in_transport_adapter():
+    ta = TransportAdapter()
+    ta.__setstate__(
+        TransportAdapter(max_retries=12, enable_keep_alive=True).__getstate__()
+    )
+    assert ta._enable_keep_alive is True
+    assert ta.max_retries.total == 12