Expose setsockopt in TCPConnector API (#10474)

Author: TimMenninger

CHANGES/10474.feature.rst

@@ -0,0 +1,2 @@
+Added ``tcp_sockopts`` to ``TCPConnector`` to allow specifying custom socket options
+-- by :user:`TimMenninger`.

CONTRIBUTORS.txt

@@ -341,6 +341,7 @@ Thanos Lefteris
 Thijs Vermeir
 Thomas Forbes
 Thomas Grainger
+Tim Menninger
 Tolga Tezel
 Tomasz Trebski
 Toshiaki Tanaka

aiohttp/connector.py

@@ -20,6 +20,7 @@
     DefaultDict,
     Deque,
     Dict,
+    Iterable,
     Iterator,
     List,
     Literal,
@@ -61,6 +62,11 @@
 )
 from .resolver import DefaultResolver
 
+if sys.version_info >= (3, 12):
+    from collections.abc import Buffer
+else:
+    Buffer = Union[bytes, bytearray, "memoryview[int]", "memoryview[bytes]"]
+
 if TYPE_CHECKING:
     import ssl
 
@@ -820,6 +826,8 @@ class TCPConnector(BaseConnector):
                            the happy eyeballs algorithm, set to None.
     interleave - “First Address Family Count” as defined in RFC 8305
     loop - Optional event loop.
+    tcp_sockopts - List of tuples of sockopts applied to underlying
+                   socket
     """
 
     allowed_protocol_schema_set = HIGH_LEVEL_SCHEMA_SET | frozenset({"tcp"})
@@ -841,6 +849,7 @@ def __init__(
         timeout_ceil_threshold: float = 5,
         happy_eyeballs_delay: Optional[float] = 0.25,
         interleave: Optional[int] = None,
+        tcp_sockopts: Iterable[Tuple[int, int, Union[int, Buffer]]] = [],
     ):
         super().__init__(
             keepalive_timeout=keepalive_timeout,
@@ -871,6 +880,7 @@ def __init__(
         self._happy_eyeballs_delay = happy_eyeballs_delay
         self._interleave = interleave
         self._resolve_host_tasks: Set["asyncio.Task[List[ResolveResult]]"] = set()
+        self._tcp_sockopts = tcp_sockopts
 
     def _close_immediately(self) -> List[Awaitable[object]]:
         for fut in chain.from_iterable(self._throttle_dns_futures.values()):
@@ -1113,6 +1123,8 @@ async def _wrap_create_connection(
                     interleave=self._interleave,
                     loop=self._loop,
                 )
+                for sockopt in self._tcp_sockopts:
+                    sock.setsockopt(*sockopt)
                 connection = await self._loop.create_connection(
                     *args, **kwargs, sock=sock
                 )

docs/client_advanced.rst

@@ -468,6 +468,21 @@ If your HTTP server uses UNIX domain sockets you can use
   session = aiohttp.ClientSession(connector=conn)
 
 
+Setting socket options
+^^^^^^^^^^^^^^^^^^^^^^
+
+Socket options passed to the :class:`~aiohttp.TCPConnector` will be passed
+to the underlying socket when creating a connection. For example, we may
+want to change the conditions under which we consider a connection dead.
+The following would change that to 9*7200 = 18 hours::
+
+  import socket
+
+  conn = aiohttp.TCPConnector(tcp_sockopts=[(socket.SOL_SOCKET,  socket.SO_KEEPALIVE,  True),
+                                            (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE,  7200),
+                                            (socket.IPPROTO_TCP, socket.TCP_KEEPCNT,      9) ])
+
+
 Named pipes in Windows
 ^^^^^^^^^^^^^^^^^^^^^^
 

docs/client_reference.rst

@@ -1128,7 +1128,8 @@ is controlled by *force_close* constructor's parameter).
                  resolver=None, keepalive_timeout=sentinel, \
                  force_close=False, limit=100, limit_per_host=0, \
                  enable_cleanup_closed=False, timeout_ceil_threshold=5, \
-                 happy_eyeballs_delay=0.25, interleave=None, loop=None)
+                 happy_eyeballs_delay=0.25, interleave=None, loop=None, \
+                 tcp_sockopts=[])
 
    Connector for working with *HTTP* and *HTTPS* via *TCP* sockets.
 
@@ -1249,6 +1250,12 @@ is controlled by *force_close* constructor's parameter).
 
         .. versionadded:: 3.10
 
+   :param list tcp_sockopts: options applied to the socket when a connection is
+      created. This should be a list of 3-tuples, each a ``(level, optname, value)``.
+      Each tuple is deconstructed and passed verbatim to ``<socket>.setsockopt``.
+
+        .. versionadded:: 3.12
+
    .. attribute:: family
 
       *TCP* socket family e.g. :data:`socket.AF_INET` or

tests/test_connector.py

@@ -3767,6 +3767,29 @@ def test_connect() -> Literal[True]:
     assert raw_response_list == [True, True]
 
 
+async def test_tcp_connector_setsockopts(
+    loop: asyncio.AbstractEventLoop, start_connection: mock.AsyncMock
+) -> None:
+    """Check that sockopts get passed to socket"""
+    conn = aiohttp.TCPConnector(
+        tcp_sockopts=[(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 2)]
+    )
+
+    with mock.patch.object(
+        conn._loop, "create_connection", autospec=True, spec_set=True
+    ) as create_connection:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            start_connection.return_value = s
+            create_connection.return_value = mock.Mock(), mock.Mock()
+
+            req = ClientRequest("GET", URL("https://127.0.0.1:443"), loop=loop)
+
+            with closing(await conn.connect(req, [], ClientTimeout())):
+                assert s.getsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT) == 2
+
+    await conn.close()
+
+
 def test_default_ssl_context_creation_without_ssl() -> None:
     """Verify _make_ssl_context does not raise when ssl is not available."""
     with mock.patch.object(connector_module, "ssl", None):