kafka.net: support socket_connection_setup_timeout_ms w/ exp backoff (KIP-601) (#3027)

Author: dpkp

kafka/net/connection.py

@@ -269,7 +269,6 @@ def connection_made(self, transport):
             client_id=self.config['client_id'],
             receive_message_max_bytes=self.config['receive_message_max_bytes'],
             ident=log_prefix)
-        self.net.call_soon(self._check_version)
 
     def pause(self, v):
         self.paused.add(v)
@@ -345,31 +344,36 @@ def _maybe_unthrottle(self):
             self._throttle_time = 0
             self.unpause('throttle')
 
-    async def _check_version(self, timeout_ms=None):
+    async def initialize(self, timeout_at=None):
+        if timeout_at is None:
+            timeout_at = self._timeout_at()
+        try:
+            await self._get_api_versions(timeout_at)
+            if self.sasl_enabled:
+                await self._sasl_authenticate(timeout_at)
+        except Exception as error:
+            self.close(error)
+        else:
+            self._init_complete()
+
+    async def _get_api_versions(self, timeout_at=None):
+        if timeout_at is None:
+            timeout_at = self._timeout_at()
         if self.broker_version_data is not None:
             try:
                 self._api_versions_idx = self.broker_version_data.api_version(ApiVersionsRequest)
             except Errors.IncompatibleBrokerVersion:
                 log.debug('%s: Using pre-configured api_version %s for ApiVersions', self, self.broker_version)
-                self._init_complete()
                 return
 
-        if timeout_ms is not None:
-            timeout_ms = self.config['api_version_auto_timeout_ms']
-        timeout_at = self._timeout_at(timeout_ms=timeout_ms)
         while timeout_at > time.monotonic():
             version = self._api_versions_idx
             request = ApiVersionsRequest(
                 version=version,
                 client_software_name=self.config['client_software_name'],
                 client_software_version=self.config['client_software_version'],
             )
-            try:
-                response = await self._send_request(request, timeout_at=timeout_at)
-            except Exception as exc:
-                self.close(exc)
-                return
-
+            response = await self._send_request(request, timeout_at=timeout_at)
             error_type = Errors.for_code(response.error_code)
             if error_type is Errors.NoError:
                 break
@@ -382,92 +386,74 @@ async def _check_version(self, timeout_ms=None):
                     self._api_versions_idx = 0
                 continue
             else:
-                self.close(error_type())
-                return
+                raise error_type()
+        else:
+            raise Errors.KafkaTimeoutError('Timeout during ApiVersions check')
 
         api_versions = {api_version.api_key: (api_version.min_version, api_version.max_version)
                         for api_version in response.api_keys}
         self.broker_version_data = BrokerVersionData(api_versions=api_versions)
         log.info('%s: Broker version identified as %s', self, '.'.join(map(str, self.broker_version)))
-        if self.sasl_enabled:
-            await self._sasl_authenticate()
-        if self.initializing:
-            self._init_complete()
 
     @property
     def sasl_enabled(self):
         return self.config['security_protocol'] in ('SASL_PLAINTEXT', 'SASL_SSL')
 
-    async def _sasl_authenticate(self):
+    async def _sasl_authenticate(self, timeout_at=None):
+        if timeout_at is None:
+            timeout_at = self._timeout_at()
         # Step 1: SaslHandshake to negotiate mechanism
         request = SaslHandshakeRequest(
             mechanism=self.config['sasl_mechanism'],
             max_version=1)
-        try:
-            response = await self._send_request(request)
-        except Exception as exc:
-            self.close(Errors.KafkaConnectionError('SaslHandshake failed: %s' % exc))
-            return
-
+        response = await self._send_request(request, timeout_at=timeout_at)
         error_type = Errors.for_code(response.error_code)
         if error_type is not Errors.NoError:
             log.error('%s: SaslHandshake failed: %s', self, error_type.__name__)
-            self.close(error_type())
-            return
+            raise error_type()
 
         if self.config['sasl_mechanism'] not in response.mechanisms:
-            self.close(Errors.UnsupportedSaslMechanismError(
+            raise Errors.UnsupportedSaslMechanismError(
                 'Kafka broker does not support %s sasl mechanism. Enabled mechanisms: %s'
-                % (self.config['sasl_mechanism'], response.mechanisms)))
-            return
+                % (self.config['sasl_mechanism'], response.mechanisms))
 
         # Step 2: SASL authentication exchange
         version = response.API_VERSION
         # Prefer the configured hostname (stored on the transport) so that
         # mechanisms like GSSAPI construct service principals against the
         # user-supplied name, not whichever IP getaddrinfo handed us.
         sasl_host = self.transport.host if self.transport.host else self.transport.getPeer()[0]
-        try:
-            mechanism = get_sasl_mechanism(self.config['sasl_mechanism'])(
-                host=sasl_host, **self.config)
-        except Exception as exc:
-            self.close(exc)
-            return
+        mechanism = get_sasl_mechanism(self.config['sasl_mechanism'])(
+            host=sasl_host, **self.config)
 
-        while not mechanism.is_done():
+        while not mechanism.is_done() and timeout_at > time.monotonic():
             token = mechanism.auth_bytes()
             if version == 1:
                 auth_request = SaslAuthenticateRequest(token, version=0)
             else:
                 auth_request = SaslBytesRequest(token)
-
-            try:
-                auth_response = await self._send_request(auth_request)
-            except Exception as exc:
-                self.close(Errors.KafkaConnectionError('SaslAuthenticate failed: %s' % exc))
-                return
-
+            auth_response = await self._send_request(auth_request, timeout_at=timeout_at)
             error_type = Errors.for_code(auth_response.error_code)
             if error_type is not Errors.NoError:
-                self.close(Errors.SaslAuthenticationFailedError(
-                    '%s: %s' % (error_type.__name__, auth_response.error_message)))
-                return
+                raise Errors.SaslAuthenticationFailedError(
+                    '%s: %s' % (error_type.__name__, auth_response.error_message))
 
             # GSSAPI does not get a final recv in v0 unframed mode
             if version == 0 and mechanism.is_done():
                 break
-
             mechanism.receive(auth_response.auth_bytes)
 
-        if not mechanism.is_authenticated():
-            self.close(Errors.SaslAuthenticationFailedError(
-                'Failed to authenticate via SASL %s' % self.config['sasl_mechanism']))
-            return
+        if time.monotonic() > timeout_at:
+            raise Errors.KafkaTimeoutError('SASL Authentication timed out')
+        elif not mechanism.is_authenticated():
+            raise Errors.SaslAuthenticationFailedError(
+                'Failed to authenticate via SASL %s' % self.config['sasl_mechanism'])
 
         log.info('%s: %s', self, mechanism.auth_details())
 
     def _init_complete(self):
-        self.initializing = False
-        self.connected = True
-        self.send_buffered()
-        self._init_future.success(True)
+        if self.initializing:
+            self.initializing = False
+            self.connected = True
+            self.send_buffered()
+            self._init_future.success(True)

kafka/net/inet.py

@@ -1,6 +1,7 @@
 import errno
 import logging
 import socket
+import time
 from urllib.parse import urlparse
 
 import kafka.errors as Errors
@@ -9,18 +10,20 @@
 log = logging.getLogger(__name__)
 
 
-async def create_connection(net, host, port, socket_options=(), proxy_url=None):
+async def create_connection(net, host, port, socket_options=(), proxy_url=None, timeout_at=None):
     """Connect to host:port; raises KafkaConnectionError on failure"""
     socket_factory = KafkaNetSocket(proxy_url)
     addrs = socket_factory.dns_lookup(host, port)
     exceptions = [Errors.KafkaConnectionError('DNS Resolution failure')]
     for res in addrs:
         try:
             log.debug('%s: Attempting to connect to %s (options: %s)', socket_factory, res, socket_options)
-            sock = await socket_factory.connect(net, res, socket_options)
+            sock = await socket_factory.connect(net, res, socket_options, timeout_at=timeout_at)
         except (socket.error, OSError) as e:
             exceptions.append(Errors.KafkaConnectionError('unable to connect: %s' % (e,)))
             continue
+        except Errors.KafkaTimeoutError:
+            raise Errors.KafkaConnectionError('Connection timed out')
         except Errors.KafkaConnectionError as e:
             exceptions.append(e)
             continue
@@ -79,17 +82,17 @@ def dns_lookup(self, host, port, raise_error=False):
     def socket(self, family=socket.AF_UNSPEC, sock_type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP):
         return socket.socket(family, sock_type, proto)
 
-    async def connect(self, net, addrinfo, socket_options=()):
+    async def connect(self, net, addrinfo, socket_options=(), timeout_at=None):
         """Create non-blocking socket (with options) and connect to addrinfo tuple"""
         family, sock_type, proto, _canonname, sockaddr = addrinfo
         sock = self.socket(family, sock_type, proto)
         sock.setblocking(False)
         for option in socket_options:
             sock.setsockopt(*option)
-        return await self.sock_connect(net, sock, sockaddr)
+        return await self.sock_connect(net, sock, sockaddr, timeout_at=timeout_at)
 
-    async def sock_connect(self, net, sock, sockaddr):
-        while True:
+    async def sock_connect(self, net, sock, sockaddr, timeout_at=None):
+        while timeout_at is None or time.monotonic() < timeout_at:
             ret = None
             try:
                 ret = self.connect_ex(sock, sockaddr)
@@ -106,12 +109,14 @@ async def sock_connect(self, net, sock, sockaddr):
             # Needs retry
             # WSAEINVAL == 10022, but errno.WSAEINVAL is not available on non-win systems
             elif ret in (errno.EINPROGRESS, errno.EALREADY, errno.EWOULDBLOCK, 10022):
-                await net.wait_write(sock)
+                await net.wait_write(sock, timeout_at=timeout_at)
 
             # Connection failed
             else:
                 errstr = errno.errorcode.get(ret, 'UNKNOWN')
                 raise Errors.KafkaConnectionError('{} {}'.format(ret, errstr))
+        else:
+            raise Errors.KafkaTimeoutError('Connection timed out')
 
     def connect_ex(self, sock, sockaddr):
         return sock.connect_ex(sockaddr)

kafka/net/manager.py

@@ -30,7 +30,8 @@ class KafkaConnectionManager:
         'reconnect_backoff_ms': 50,
         'reconnect_backoff_max_ms': 30000,
         'request_timeout_ms': 30000,
-        'socket_connection_timeout_ms': 5000,
+        'socket_connection_setup_timeout_ms': 10000,
+        'socket_connection_setup_timeout_max_ms': 30000,
         'socket_options': [
             (socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),
             (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
@@ -86,7 +87,7 @@ def __init__(self, net, **configs):
         )
         self.cluster.attach(self)
         self._conns = {}
-        self._backoff = dict() # node_id => (failures, backoff_until)
+        self._backoff = dict() # node_id => (failures, backoff_until, socket_connect_setup_timeout_ms)
         # Cache the most recent SASL / SSL / auth failure per node so we can
         # surface it to the user instead of silently retrying forever.
         # Cleared on successful connect.
@@ -119,7 +120,9 @@ async def _do_bootstrap(self, deadline):
             bootstrap_broker = random.choice(self.cluster.bootstrap_brokers())
             log.debug('Attempting bootstrap with %s', bootstrap_broker)
             try:
+                timeout_ms = (deadline - time.monotonic()) * 1000 if deadline is not None else None
                 conn = self.get_connection(bootstrap_broker.node_id,
+                                           timeout_ms=timeout_ms,
                                            pop_on_close=False,
                                            refresh_metadata_on_err=False,
                                            reset_backoff_on_connect=False)
@@ -218,10 +221,11 @@ def _build_ssl_context(self):
             ctx.verify_flags |= ssl.VERIFY_CRL_CHECK_LEAF
         return ctx
 
-    async def _build_transport(self, node):
+    async def _build_transport(self, node, timeout_at=None):
         sock = await create_connection(self._net, node.host, node.port,
                                        self.config['socket_options'],
-                                       proxy_url=self.config['proxy_url'])
+                                       proxy_url=self.config['proxy_url'],
+                                       timeout_at=timeout_at)
         if self.ssl_enabled:
             transport = KafkaSSLTransport(self._net, sock, self._build_ssl_context(),
                                           host=node.host, ssl_check_hostname=self.config['ssl_check_hostname'])
@@ -235,11 +239,11 @@ async def _build_transport(self, node):
         else:
             return transport
 
-    async def _connect(self, node, conn, reset_backoff_on_connect=True):
+    async def _connect(self, node, conn, reset_backoff_on_connect=True, timeout_at=None):
         try:
-            transport = await self._build_transport(node)
+            transport = await self._build_transport(node, timeout_at=timeout_at)
             conn.connection_made(transport)
-            await conn.init_future
+            await conn.initialize(timeout_at=timeout_at)
         except Exception as exc:
             log.error('Connection failed: %s', exc)
             conn.connection_lost(exc)
@@ -280,12 +284,10 @@ def get_connection(self, node_id, timeout_ms=None,
         if refresh_metadata_on_err:
             conn.close_future.add_errback(lambda _: self.cluster.request_update())
         self._conns[node_id] = conn
-        self._net.call_soon(lambda: self._connect(node, conn, reset_backoff_on_connect=reset_backoff_on_connect))
         if timeout_ms is None:
-            timeout_ms = self.config['socket_connection_timeout_ms']
-        self._net.call_later(timeout_ms / 1000,
-                             lambda: conn.close(Errors.KafkaConnectionError('Connection timed out'))
-                                     if not conn.init_future.is_done else None)
+            timeout_ms = self.socket_connection_setup_timeout_ms(node_id)
+        timeout_at = time.monotonic() + timeout_ms / 1000
+        self._net.call_soon(lambda: self._connect(node, conn, reset_backoff_on_connect=reset_backoff_on_connect, timeout_at=timeout_at))
         return conn
 
     def send(self, request, node_id=None, request_timeout_ms=None):
@@ -335,18 +337,29 @@ def reset_backoff(self, node_id):
         except KeyError:
             pass
 
-    def reconnect_jitter_pct(self):
+    def jitter_pct(self):
         return random.uniform(0.8, 1.2)
 
+    def _calculate_exp_timeout(self, key, failures):
+        max_keys = {
+            'reconnect_backoff_ms': 'reconnect_backoff_max_ms',
+            'socket_connection_setup_timeout_ms': 'socket_connection_setup_timeout_max_ms',
+        }
+        timeout_ms = self.config[key] * 2 ** (failures - 1)
+        if key in max_keys:
+            max_ms = self.config[max_keys[key]]
+            timeout_ms = min(max_ms, timeout_ms)
+        return timeout_ms * self.jitter_pct()
+
     def update_backoff(self, node_id):
-        failures, _ = self._backoff.get(node_id, (0, 0))
+        failures, _, _ = self._backoff.get(node_id, (0, 0, 0))
         failures += 1
-        backoff_ms = self.config['reconnect_backoff_ms'] * 2 ** (failures - 1)
-        backoff_ms = min(backoff_ms, self.config['reconnect_backoff_max_ms'])
-        backoff_ms *= self.reconnect_jitter_pct()
-        log.debug('%s reconnect backoff %d ms after %s failures', node_id, backoff_ms, failures)
+        backoff_ms = self._calculate_exp_timeout('reconnect_backoff_ms', failures)
+        connect_ms = self._calculate_exp_timeout('socket_connection_setup_timeout_ms', failures)
+        log.debug('%s reconnect backoff %d ms / connect timeout %d ms after %s failures',
+                  node_id, backoff_ms, connect_ms, failures)
         backoff_until_time = time.monotonic() + (backoff_ms / 1000)
-        self._backoff[node_id] = (failures, backoff_until_time)
+        self._backoff[node_id] = (failures, backoff_until_time, connect_ms)
 
     def connection_delay(self, node_id):
         """Connection delay in seconds.
@@ -357,6 +370,11 @@ def connection_delay(self, node_id):
             return 0
         return max(0, self._backoff[node_id][1] - time.monotonic())
 
+    def socket_connection_setup_timeout_ms(self, node_id):
+        if node_id not in self._backoff:
+            return self.config['socket_connection_setup_timeout_ms']
+        return self._backoff[node_id][2]
+
     def auth_failure(self, node_id):
         """Return the most recent auth-class failure for ``node_id``,
         or None if there is no sticky failure on record."""