fix(cncf-kubernetes): use Configuration.socket_options instead of monkey-patching urllib3 (#68396)

The `enable_tcp_keepalive` config option in the Kubernetes provider
relied on monkey-patching urllib3's default_socket_options.
In urllib3 v2.x, the socket_options parameter in HTTPConnection.__init__
is evaluated as a default argument at import time, so post-import
changes are never picked up by new connections.

Fix by passing socket options through the Kubernetes client's
Configuration.socket_options field, which is properly threaded
through ApiClient -> RESTClientObject -> urllib3.PoolManager.

Also includes TCP_NODELAY in the socket options to preserve the
default urllib3 behavior of disabling Nagle's algorithm.

Closes: #68396

Author: anmolxlight

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/kube_client.py

@@ -51,26 +51,33 @@ def _disable_verify_ssl() -> None:
     _import_err = e
 
 
-def _enable_tcp_keepalive() -> None:
+def _enable_tcp_keepalive(configuration: Configuration) -> None:
     """
-    Enable TCP keepalive mechanism.
+    Enable TCP keepalive mechanism on the provided Kubernetes client configuration.
 
-    This prevents urllib3 connection to hang indefinitely when idle connection
-    is time-outed on services like cloud load balancers or firewalls.
+    This prevents urllib3 connections from hanging indefinitely when an idle
+    connection is timed out by services like cloud load balancers or firewalls.
 
-    See https://github.com/apache/airflow/pull/11406 for detailed explanation.
+    Uses the ``socket_options`` field on the Kubernetes ``Configuration`` object,
+    which is threaded through to the underlying urllib3 ``PoolManager`` and
+    ``HTTPConnection``, rather than monkey-patching urllib3 connection defaults
+    (which no longer works with urllib3 v2.x).
+
+    See https://github.com/apache/airflow/pull/11406 for the original discussion
+    and https://github.com/apache/airflow/issues/68396 for the urllib3 v2 fix.
 
     Please ping @michalmisiewicz or @dimberman in the PR if you want to modify this function.
     """
     import socket
 
-    from urllib3.connection import HTTPConnection, HTTPSConnection
-
     tcp_keep_idle = conf.getint("kubernetes_executor", "tcp_keep_idle")
     tcp_keep_intvl = conf.getint("kubernetes_executor", "tcp_keep_intvl")
     tcp_keep_cnt = conf.getint("kubernetes_executor", "tcp_keep_cnt")
 
-    socket_options = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
+    socket_options: list[tuple[int, int, int | bytes]] = [
+        (socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),
+        (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
+    ]
 
     if hasattr(socket, "TCP_KEEPIDLE"):
         socket_options.append((socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, tcp_keep_idle))
@@ -87,17 +94,7 @@ def _enable_tcp_keepalive() -> None:
     else:
         log.debug("Unable to set TCP_KEEPCNT on this platform")
 
-    # Cast both the default options and our socket options
-    socket_options_cast: list[tuple[int, int, int | bytes]] = [
-        (level, opt, val) for level, opt, val in socket_options
-    ]
-    default_options_cast: list[tuple[int, int, int | bytes]] = [
-        (level, opt, val) for level, opt, val in HTTPSConnection.default_socket_options
-    ]
-
-    # Then use the cast versions for both HTTPS and HTTP
-    HTTPSConnection.default_socket_options = default_options_cast + socket_options_cast
-    HTTPConnection.default_socket_options = default_options_cast + socket_options_cast
+    configuration.socket_options = socket_options
 
 
 def get_kube_client(
@@ -118,10 +115,10 @@ def get_kube_client(
     if not has_kubernetes:
         raise _import_err
 
-    if conf.getboolean("kubernetes_executor", "enable_tcp_keepalive"):
-        _enable_tcp_keepalive()
-
     configuration = _get_default_configuration()
+
+    if conf.getboolean("kubernetes_executor", "enable_tcp_keepalive"):
+        _enable_tcp_keepalive(configuration)
     api_client_retry_configuration = conf.getjson(
         "kubernetes_executor", "api_client_retry_configuration", fallback={}
     )

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/test_client.py

@@ -16,13 +16,10 @@
 # under the License.
 from __future__ import annotations
 
-import socket
 from unittest import mock
 
 import pytest
 from kubernetes.client import Configuration
-from urllib3.connection import HTTPConnection, HTTPSConnection
-
 from airflow.providers.cncf.kubernetes.kube_client import (
     _disable_verify_ssl,
     _enable_tcp_keepalive,
@@ -65,19 +62,19 @@ def test_load_config_ssl_ca_cert(self, conf, config):
 
     @pytest.mark.platform("linux")
     def test_enable_tcp_keepalive(self):
-        socket_options = [
-            (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
-            (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 120),
-            (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 30),
-            (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 6),
-        ]
-        expected_http_connection_options = HTTPConnection.default_socket_options + socket_options
-        expected_https_connection_options = HTTPSConnection.default_socket_options + socket_options
-
-        _enable_tcp_keepalive()
-
-        assert HTTPConnection.default_socket_options == expected_http_connection_options
-        assert HTTPSConnection.default_socket_options == expected_https_connection_options
+        import socket
+
+        configuration = Configuration()
+        assert configuration.socket_options is None
+
+        _enable_tcp_keepalive(configuration)
+
+        assert configuration.socket_options is not None
+        assert configuration.socket_options[0] == (socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+        assert configuration.socket_options[1] == (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+        assert configuration.socket_options[2] == (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 120)
+        assert configuration.socket_options[3] == (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 30)
+        assert configuration.socket_options[4] == (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 6)
 
     def test_disable_verify_ssl(self):
         configuration = Configuration()

providers/google/src/airflow/providers/google/cloud/hooks/kubernetes_engine.py

@@ -76,7 +76,7 @@ def get_conn(self) -> client.ApiClient:
         configuration = self._get_config()
         configuration.refresh_api_key_hook = self._refresh_api_key_hook
         if self.enable_tcp_keepalive:
-            _enable_tcp_keepalive()
+            _enable_tcp_keepalive(configuration)
         return client.ApiClient(configuration)
 
     def _refresh_api_key_hook(self, configuration: client.configuration.Configuration):