diff --git a/kafka/net/backend/abstract.py b/kafka/net/backend/abstract.py index 183de2632..f4fdb8ce6 100644 --- a/kafka/net/backend/abstract.py +++ b/kafka/net/backend/abstract.py @@ -27,7 +27,7 @@ * ``wait_read`` / ``wait_write`` / ``unregister_event`` -- the low-level fd-readiness primitives. They are the *selector's* private mechanism (used - only inside ``kafka/net/backend/transport.py`` + ``inet.py``, zero core callers) and + only inside ``kafka/net/backend/transport.py``, zero core callers) and do not port to asyncio/Twisted. The connection seam replaces them. * ``poll(timeout_ms, future=...)`` -- the legacy single-tick driver. Its only remaining caller is the ``KafkaNetClient`` compat shim @@ -142,6 +142,21 @@ def pause_writing(self) -> None: ... def resume_writing(self) -> None: ... +import socket +from typing import List, Tuple, Union, Any + +# Complete signature representation +AddrInfoResult = List[ + Tuple[ + socket.AddressFamily, # 0. family (e.g., AF_INET, AF_INET6) pylint: disable=no-member + socket.SocketKind, # 1. type (e.g., SOCK_STREAM, SOCK_DGRAM) pylint: disable=no-member + int, # 2. proto (protocol number) + str, # 3. canonname (canonical name string) + Union[Tuple[str, int], Tuple[str, int, int, int]] # 4. sockaddr + ] +] + + @runtime_checkable class NetBackend(Protocol): """Structural contract for a pluggable async event-loop backend. @@ -188,6 +203,9 @@ def sleep(self, delay: float) -> Any: """Awaitable that resolves after ``delay`` seconds.""" # --- connection seam -------------------------------------------------- + async def getaddrinfo(self, host: str, port: int) -> AddrInfoResult: + """Resolve host/port via DNS""" + async def create_connection( self, protocol: NetProtocol, diff --git a/kafka/net/backend/asyncio_backend.py b/kafka/net/backend/asyncio_backend.py index 15ee91de8..e268ef935 100644 --- a/kafka/net/backend/asyncio_backend.py +++ b/kafka/net/backend/asyncio_backend.py @@ -338,6 +338,9 @@ async def waiter(): return state['value'] # --- connection seam -------------------------------------------------- + async def getaddrinfo(self, host, port): + return await self._loop.getaddrinfo(host, port) + async def create_connection(self, protocol, host, port, *, ssl=None, proxy_url=None, socket_options=(), timeout_at=None): if proxy_url is not None: diff --git a/kafka/net/backend/inet.py b/kafka/net/backend/inet.py index f6d752cb1..02d6adcbb 100644 --- a/kafka/net/backend/inet.py +++ b/kafka/net/backend/inet.py @@ -10,28 +10,6 @@ log = logging.getLogger(__name__) -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, 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 - else: - return sock - raise exceptions[-1] - - class KafkaNetSocket: # scheme => handling class _registry = {} diff --git a/kafka/net/backend/selector.py b/kafka/net/backend/selector.py index 1e23f3d04..e56626add 100644 --- a/kafka/net/backend/selector.py +++ b/kafka/net/backend/selector.py @@ -1,6 +1,7 @@ import collections import copy import enum +import errno import inspect import logging import heapq @@ -11,7 +12,6 @@ import kafka.errors as Errors from kafka.future import Future -from kafka.net.backend.inet import create_connection as _inet_create_connection from kafka.net.backend.transport import KafkaTCPTransport from kafka.net.ssl import KafkaSSLTransport from kafka.version import __version__ @@ -560,22 +560,85 @@ def create_future(self): """ return SelectorFuture() + async def getaddrinfo(self, host, port): + # XXX: all DNS functions in Python are blocking. If we really + # want to be non-blocking here, we need to use a 3rd-party + # library like python-adns, or move resolution onto its + # own thread. This will be subject to the default libc + # name resolution timeout (5s on most Linux boxes) + try: + return socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM) + except socket.gaierror as ex: + err_str = "DNS lookup failed for %s:%d, %r" % (host, port, ex) + raise Errors.KafkaConnectionError(err_str) + + async def connect_host(self, host, port, socket_options=(), timeout_at=None): + """Connect to host:port; raises KafkaConnectionError on failure""" + addrs = await self.getaddrinfo(host, port) + exceptions = [Errors.KafkaConnectionError('DNS Resolution failure')] + for addrinfo in addrs: + try: + return await self.connect_addrinfo(addrinfo, 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 + raise exceptions[-1] + + async def connect_addrinfo(self, addrinfo, socket_options=(), timeout_at=None): + """Create non-blocking socket (with options) and connect to addrinfo tuple""" + log.debug('%s: Attempting to connect to %s (options: %s)', self, addrinfo, socket_options) + family, sock_type, proto, _canonname, sockaddr = addrinfo + sock = socket.socket(family, sock_type, proto) + sock.setblocking(False) + for option in socket_options: + sock.setsockopt(*option) + while timeout_at is None or time.monotonic() < timeout_at: + ret = None + try: + ret = sock.connect_ex(sockaddr) + except BlockingIOError: + ret = errno.EWOULDBLOCK + except socket.error as err: + ret = err.errno + + # Connection succeeded + if not ret or ret == errno.EISCONN: + log.debug('Connected: %s', sock) + return sock + + # 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 self.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') + async def create_connection(self, protocol, host, port, *, ssl=None, proxy_url=None, socket_options=(), timeout_at=None): """Establish a connected transport to host:port and wire ``protocol``. - The selector owns the raw socket: DNS + non-blocking connect (with - optional SOCKS5/HTTP-CONNECT proxy via KafkaNetSocket), then wraps it - in a TCP or SSL transport, runs the TLS handshake, and calls + The selector owns the raw socket: DNS + non-blocking connect, then + wraps it in a TCP or SSL transport, runs the TLS handshake, and calls ``protocol.connection_made(transport)`` -- mirroring asyncio/Twisted, which own the socket and wire the protocol at connect time. On any failure (handshake error, or a ``protocol`` that refuses the transport because it closed mid-connect) the transport is closed before raising, so the caller never handles a transport instance directly. """ - sock = await _inet_create_connection(self, host, port, socket_options, - proxy_url=proxy_url, timeout_at=timeout_at) + sock = await self.connect_host(host, port, + socket_options=socket_options, + timeout_at=timeout_at) transport = KafkaTCPTransport(self, sock, host=host) if ssl is not None: ssl_wrapper = KafkaSSLTransport(self, ssl, host=host) diff --git a/test/net/backend/test_inet.py b/test/net/backend/test_inet.py deleted file mode 100644 index 12dd3f569..000000000 --- a/test/net/backend/test_inet.py +++ /dev/null @@ -1,344 +0,0 @@ -import errno -import socket -from unittest.mock import MagicMock, call, patch - -import pytest - -from kafka.net.backend.inet import create_connection, KafkaNetSocket -from kafka.net.socks5 import Socks5Proxy -from kafka.net.http_connect import HttpConnectProxy -import kafka.errors as Errors - - -class TestDnsLookup: - def test_valid_host(self): - results = KafkaNetSocket().dns_lookup('localhost', 9092) - assert len(results) > 0 - for res in results: - assert len(res) == 5 - - def test_invalid_host(self): - with patch('kafka.net.backend.inet.socket.getaddrinfo', side_effect=socket.gaierror): - results = KafkaNetSocket().dns_lookup('invalid.host', 9092) - assert results == [] - - def test_numeric_host(self): - results = KafkaNetSocket().dns_lookup('127.0.0.1', 9092) - assert len(results) > 0 - assert results[0][4][0] == '127.0.0.1' - - -class TestSockConnect: - def test_immediate_connect(self, net): - factory = KafkaNetSocket() - sock = MagicMock() - sock.connect_ex.return_value = 0 - result = net.run(factory.sock_connect(net, sock, ('127.0.0.1', 9092))) - assert result is sock - sock.connect_ex.assert_called_once_with(('127.0.0.1', 9092)) - - def test_eisconn(self, net): - factory = KafkaNetSocket() - sock = MagicMock() - sock.connect_ex.return_value = errno.EISCONN - result = net.run(factory.sock_connect(net, sock, ('127.0.0.1', 9092))) - assert result is sock - - def test_connection_refused(self, net): - factory = KafkaNetSocket() - sock = MagicMock() - sock.connect_ex.return_value = errno.ECONNREFUSED - with pytest.raises(Errors.KafkaConnectionError): - net.run(factory.sock_connect(net, sock, ('127.0.0.1', 9092))) - - def test_socket_error_uses_errno(self, net): - factory = KafkaNetSocket() - sock = MagicMock() - sock.connect_ex.side_effect = socket.error(errno.ECONNREFUSED, 'refused') - with pytest.raises(Errors.KafkaConnectionError): - net.run(factory.sock_connect(net, sock, ('127.0.0.1', 9092))) - - def test_error_after_wait_write(self, net): - """connect_ex returns EINPROGRESS, then after wait_write fires the - second connect_ex returns the real error.""" - factory = KafkaNetSocket() - # socketpair endpoints are always immediately writable, so wait_write - # fires on the first poll and we re-enter the loop. - rsock, wsock = socket.socketpair() - rsock.setblocking(False) - wsock.setblocking(False) - try: - sock = MagicMock() - sock.connect_ex.side_effect = [errno.EINPROGRESS, errno.ECONNREFUSED] - sock.fileno.return_value = wsock.fileno() - with pytest.raises(Errors.KafkaConnectionError): - net.run(factory.sock_connect(net, sock, ('127.0.0.1', 9092))) - assert sock.connect_ex.call_count == 2 - finally: - rsock.close() - wsock.close() - - -class TestCreateConnection: - def test_dns_failure(self, net): - with patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=[]): - with pytest.raises(Errors.KafkaConnectionError, match='DNS'): - net.run(create_connection(net, 'badhost', 9092)) - - def test_socket_init_failure(self, net): - fake_addr = [(socket.AF_INET, socket.SOCK_STREAM, 6, '', ('127.0.0.1', 9092))] - with patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=fake_addr), \ - patch('kafka.net.backend.inet.socket.socket', side_effect=OSError('no socket')): - with pytest.raises(Errors.KafkaConnectionError): - net.run(create_connection(net, 'host', 9092)) - - def test_successful_connection(self, net): - fake_addr = [(socket.AF_INET, socket.SOCK_STREAM, 6, '', ('127.0.0.1', 9092))] - mock_sock = MagicMock() - mock_sock.connect_ex.return_value = 0 - with patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=fake_addr), \ - patch('kafka.net.backend.inet.socket.socket', return_value=mock_sock): - result = net.run( - create_connection(net, 'host', 9092)) - assert result is mock_sock - mock_sock.setblocking.assert_called_with(False) - - def test_tries_multiple_addresses(self, net): - addr1 = (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('10.0.0.1', 9092)) - addr2 = (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('10.0.0.2', 9092)) - mock_sock1 = MagicMock() - mock_sock1.connect_ex.return_value = errno.ECONNREFUSED - mock_sock2 = MagicMock() - mock_sock2.connect_ex.return_value = 0 - sockets = iter([mock_sock1, mock_sock2]) - with patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=[addr1, addr2]), \ - patch('kafka.net.backend.inet.socket.socket', side_effect=lambda *a: next(sockets)): - result = net.run( - create_connection(net, 'host', 9092)) - assert result is mock_sock2 - - def test_socket_options_applied(self, net): - fake_addr = [(socket.AF_INET, socket.SOCK_STREAM, 6, '', ('127.0.0.1', 9092))] - mock_sock = MagicMock() - mock_sock.connect_ex.return_value = 0 - opts = [ - (socket.IPPROTO_TCP, socket.TCP_NODELAY, 1), - (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), - ] - with patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=fake_addr), \ - patch('kafka.net.backend.inet.socket.socket', return_value=mock_sock): - net.run(create_connection(net, 'host', 9092, socket_options=opts)) - mock_sock.setsockopt.assert_has_calls([ - call(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1), - call(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), - ]) - assert mock_sock.setsockopt.call_count == 2 - - -class TestCreateConnectionWithProxy: - def test_proxy_creates_socket(self, net): - mock_sock = MagicMock() - mock_sock.connect_ex.return_value = 0 - fake_addr = (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('127.0.0.1', 9092)) - with patch('kafka.net.socks5.Socks5Proxy._get_proxy_addr'), \ - patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=[fake_addr]), \ - patch('kafka.net.socks5.Socks5Proxy.connect', return_value=mock_sock) as mock_connect: - result = net.run( - create_connection(net, 'broker', 9092, proxy_url='socks5://proxy:1080')) - mock_connect.assert_called_once_with(net, fake_addr, (), timeout_at=None) - assert result is mock_sock - - def test_proxy_remote_dns_skips_local_lookup(self, net): - mock_sock = MagicMock() - mock_sock.connect_ex.return_value = 0 - with patch('kafka.net.socks5.Socks5Proxy._get_proxy_addr'), \ - patch('kafka.net.socks5.Socks5Proxy.socket', return_value=mock_sock), \ - patch('kafka.net.socks5.Socks5Proxy.connect_ex', return_value=0), \ - patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup') as mock_dns: - result = net.run( - create_connection(net, 'broker', 9092, proxy_url='socks5h://proxy:1080')) - mock_dns.assert_not_called() - - def test_no_proxy_uses_direct_socket(self, net): - fake_addr = (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('127.0.0.1', 9092)) - mock_sock = MagicMock() - mock_sock.connect_ex.return_value = 0 - with patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=[fake_addr]), \ - patch('kafka.net.backend.inet.socket.socket', return_value=mock_sock), \ - patch('kafka.net.socks5.Socks5Proxy.connect') as mock_connect: - result = net.run( - create_connection(net, 'host', 9092)) - mock_connect.assert_not_called() - assert result is mock_sock - - def test_socks5h_does_dns_for_proxy_not_target(self, net): - """Companion to test_proxy_remote_dns_skips_local_lookup: with - _get_proxy_addr running normally, exactly one dns_lookup is made and - it is for the proxy hostname, not the target.""" - proxy_addr = (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('1.2.3.4', 1080)) - mock_sock = MagicMock() - with patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=[proxy_addr]) as mock_dns, \ - patch('kafka.net.socks5.Socks5Proxy.socket', return_value=mock_sock), \ - patch('kafka.net.socks5.Socks5Proxy.connect_ex', return_value=0): - net.run( - create_connection(net, 'broker', 9092, proxy_url='socks5h://proxy:1080')) - assert mock_dns.call_count == 1 - assert mock_dns.call_args.args[:2] == ('proxy', 1080) - - def test_socks5_proxy_dns_gaierror_raises(self): - with patch('kafka.net.backend.inet.socket.getaddrinfo', side_effect=socket.gaierror): - with pytest.raises(Errors.KafkaConnectionError): - KafkaNetSocket('socks5://bogus.proxy:1080') - - def test_socks5_proxy_dns_empty_raises(self): - with patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=[]): - with pytest.raises(Errors.KafkaConnectionError): - KafkaNetSocket('socks5://proxy:1080') - - def test_proxy_connect_dispatches_through_inherited_connect(self, net): - """create_connection -> Socks5Proxy.connect (inherited from - KafkaNetSocket) -> Socks5Proxy.socket + Socks5Proxy.connect_ex.""" - fake_addr = (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('127.0.0.1', 9092)) - mock_sock = MagicMock() - with patch('kafka.net.socks5.Socks5Proxy._get_proxy_addr'), \ - patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=[fake_addr]), \ - patch('kafka.net.socks5.Socks5Proxy.socket', return_value=mock_sock) as mock_socket, \ - patch('kafka.net.socks5.Socks5Proxy.connect_ex', return_value=0) as mock_connect_ex: - result = net.run( - create_connection(net, 'broker', 9092, proxy_url='socks5://proxy:1080')) - mock_socket.assert_called_once_with(socket.AF_INET, socket.SOCK_STREAM, 6) - mock_connect_ex.assert_called_once_with(mock_sock, ('127.0.0.1', 9092)) - assert result is mock_sock - - -class TestKafkaNetSocketRegistry: - def test_socks5(self): - assert 'socks5' in KafkaNetSocket._registry - with patch('kafka.net.socks5.Socks5Proxy._get_proxy_addr'): - factory = KafkaNetSocket('socks5://foo.bar') - assert isinstance(factory, Socks5Proxy) - - def test_socks5h(self): - assert 'socks5h' in KafkaNetSocket._registry - with patch('kafka.net.socks5.Socks5Proxy._get_proxy_addr'): - factory = KafkaNetSocket('socks5h://foo.bar') - assert isinstance(factory, Socks5Proxy) - - def test_http(self): - assert 'http' in KafkaNetSocket._registry - with patch('kafka.net.http_connect.HttpConnectProxy._get_proxy_addr'): - factory = KafkaNetSocket('http://proxy:8080') - assert isinstance(factory, HttpConnectProxy) - - def test_default(self): - factory = KafkaNetSocket() - assert type(factory) is KafkaNetSocket - - def test_unknown_scheme_raises(self): - with pytest.raises(ValueError, match='Unsupported proxy url scheme'): - KafkaNetSocket('ftp://proxy:8080') - - def test_no_scheme_raises(self): - with pytest.raises(ValueError, match='scheme'): - KafkaNetSocket('no-scheme') - - def test_empty_string_raises(self): - with pytest.raises(ValueError, match='scheme'): - KafkaNetSocket('') - - def test_kafka_net_import_registers_socks5(self): - # Importing kafka.net must register Socks5Proxy. Regression guard - # against accidentally dropping the import from kafka/net/__init__.py. - import kafka.net # noqa: F401 - assert KafkaNetSocket._registry['socks5'] is Socks5Proxy - assert KafkaNetSocket._registry['socks5h'] is Socks5Proxy - - def test_subclass_auto_registers(self): - class _TestProxy(KafkaNetSocket): - SCHEMES = ('test-autoregister',) - def __init__(self, proxy_url): - self.proxy_url = proxy_url - try: - assert KafkaNetSocket._registry['test-autoregister'] is _TestProxy - sock = KafkaNetSocket('test-autoregister://x') - assert isinstance(sock, _TestProxy) - assert sock.proxy_url == 'test-autoregister://x' # pylint: disable=no-member - finally: - KafkaNetSocket._registry.pop('test-autoregister', None) - - def test_duplicate_scheme_last_wins(self): - prior = KafkaNetSocket._registry.get('test-dup') - class _First(KafkaNetSocket): - SCHEMES = ('test-dup',) - class _Second(KafkaNetSocket): - SCHEMES = ('test-dup',) - try: - assert KafkaNetSocket._registry['test-dup'] is _Second - finally: - if prior is None: - KafkaNetSocket._registry.pop('test-dup', None) - else: - KafkaNetSocket._registry['test-dup'] = prior - - -class TestKafkaNetSocketExtensionPattern: - """Validates the contract a third-party scheme handler must implement. - - Two flavours: a SOCKS-style class that overrides only connect_ex (and - inherits the default socket()/sock_connect() machinery), and an - asyncio-style class that overrides connect() entirely (e.g. to delegate - to an existing asyncio socket implementation). - """ - - def test_connect_ex_only_subclass(self, net): - """An HTTP CONNECT-style handler that only overrides connect_ex.""" - ex_calls = [] - - class _HttpConnect(KafkaNetSocket): - SCHEMES = ('test-httpconnect',) - def __init__(self, proxy_url): - self.proxy_url = proxy_url - def connect_ex(self, sock, sockaddr): - ex_calls.append((sock, sockaddr)) - return 0 - - try: - fake_addr = (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('10.0.0.1', 9092)) - mock_sock = MagicMock() - with patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=[fake_addr]), \ - patch('kafka.net.backend.inet.socket.socket', return_value=mock_sock): - result = net.run( - create_connection(net, 'broker', 9092, - proxy_url='test-httpconnect://proxy:8080')) - assert result is mock_sock - assert ex_calls == [(mock_sock, ('10.0.0.1', 9092))] - mock_sock.setblocking.assert_called_with(False) - finally: - KafkaNetSocket._registry.pop('test-httpconnect', None) - - def test_connect_override_subclass(self, net): - """An asyncio-style handler that overrides connect() entirely; the - default socket()/sock_connect()/connect_ex() flow is bypassed.""" - connect_calls = [] - - class _AsyncIoSock(KafkaNetSocket): - SCHEMES = ('test-asyncio',) - def __init__(self, proxy_url): - self.proxy_url = proxy_url - async def connect(self, net, addrinfo, socket_options=(), timeout_at=None): - connect_calls.append((addrinfo, tuple(socket_options), timeout_at)) - return 'asyncio-stream-handle' - - try: - fake_addr = (socket.AF_INET, socket.SOCK_STREAM, 6, '', ('10.0.0.1', 9092)) - with patch('kafka.net.backend.inet.KafkaNetSocket.dns_lookup', return_value=[fake_addr]), \ - patch('kafka.net.backend.inet.socket.socket') as mock_sock_cls: - result = net.run( - create_connection(net, 'broker', 9092, - proxy_url='test-asyncio://x', - timeout_at=123)) - assert result == 'asyncio-stream-handle' - assert connect_calls == [(fake_addr, (), 123)] - mock_sock_cls.assert_not_called() - finally: - KafkaNetSocket._registry.pop('test-asyncio', None) diff --git a/test/net/backend/test_selector.py b/test/net/backend/test_selector.py index bc1c61a4b..88b4989c3 100644 --- a/test/net/backend/test_selector.py +++ b/test/net/backend/test_selector.py @@ -1,11 +1,13 @@ +import errno import socket import selectors import threading import time +from unittest.mock import MagicMock, call, patch import pytest -from kafka.errors import KafkaTimeoutError +import kafka.errors as Errors from kafka.future import Future from kafka.net.backend.selector import ( KernelEvent, @@ -331,7 +333,7 @@ async def reader(): try: await net.wait_read(rsock, timeout_at=time.monotonic() + 0.02) outcome.append(('resumed', rsock.recv(1024))) - except KafkaTimeoutError: + except Errors.KafkaTimeoutError: outcome.append(('timeout',)) net.call_soon(reader) @@ -988,7 +990,7 @@ async def work(): self._wedge(net, release) th, outcome = self._run_in_thread(net, work) # no timeout_ms assert not th.is_alive(), 'run() hung past the default deadline (#3121)' - assert isinstance(outcome.get('exc'), KafkaTimeoutError), ( + assert isinstance(outcome.get('exc'), Errors.KafkaTimeoutError), ( 'expected KafkaTimeoutError, got: %r' % outcome) assert 0.15 <= outcome['elapsed'] < 2.0 assert 'may be stalled' in str(outcome['exc']) @@ -1009,7 +1011,7 @@ async def work(): self._wedge(net, release) th, outcome = self._run_in_thread(net, work, timeout_ms=100) assert not th.is_alive() - assert isinstance(outcome.get('exc'), KafkaTimeoutError) + assert isinstance(outcome.get('exc'), Errors.KafkaTimeoutError) assert 0.15 <= outcome['elapsed'] < 2.0 finally: release.set() @@ -1054,7 +1056,7 @@ async def waits_forever(): try: start = time.monotonic() - with pytest.raises(KafkaTimeoutError): + with pytest.raises(Errors.KafkaTimeoutError): net.run(waits_forever) # no start() -> fallback path elapsed = time.monotonic() - start assert 0.1 <= elapsed < 2.0 @@ -1075,7 +1077,7 @@ async def work(): self._wedge(net, release) with caplog.at_level('WARNING', logger='kafka.net.backend.selector'): th, outcome = self._run_in_thread(net, work) - assert isinstance(outcome.get('exc'), KafkaTimeoutError) + assert isinstance(outcome.get('exc'), Errors.KafkaTimeoutError) assert any('did not complete within' in r.message and 'work' in r.message for r in caplog.records), ( 'expected a backstop WARNING, got: %r' @@ -1099,7 +1101,7 @@ async def work(): self._wedge(net, release) with caplog.at_level('WARNING', logger='kafka.net.backend.selector'): th, outcome = self._run_in_thread(net, work) # caller times out - assert isinstance(outcome.get('exc'), KafkaTimeoutError) + assert isinstance(outcome.get('exc'), Errors.KafkaTimeoutError) # Release the wedge so the abandoned coroutine now completes. release.set() deadline = time.monotonic() + 2.0 @@ -1115,3 +1117,145 @@ async def work(): finally: release.set() net.close() + + +class TestGetaddrinfo: + def test_valid_host(self, net): + results = net.run(net.getaddrinfo('localhost', 9092)) + assert len(results) > 0 + for res in results: + assert len(res) == 5 + + def test_invalid_host(self, net): + # getaddrinfo now raises (rather than returning []) so connect_host + # surfaces DNS failures directly. + with patch('kafka.net.backend.selector.socket.getaddrinfo', side_effect=socket.gaierror): + with pytest.raises(Errors.KafkaConnectionError, match='DNS'): + net.run(net.getaddrinfo('invalid.host', 9092)) + + def test_numeric_host(self, net): + results = net.run(net.getaddrinfo('127.0.0.1', 9092)) + assert len(results) > 0 + assert results[0][4][0] == '127.0.0.1' + + +class TestConnectAddrinfo: + ADDRINFO = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, '', ('127.0.0.1', 9092)) + + def test_immediate_connect(self, net): + sock = MagicMock() + sock.connect_ex.return_value = 0 + with patch('kafka.net.backend.selector.socket.socket', return_value=sock): + result = net.run(net.connect_addrinfo(self.ADDRINFO)) + assert result is sock + sock.connect_ex.assert_called_once_with(('127.0.0.1', 9092)) + sock.setblocking.assert_called_with(False) + + def test_eisconn(self, net): + sock = MagicMock() + sock.connect_ex.return_value = errno.EISCONN + with patch('kafka.net.backend.selector.socket.socket', return_value=sock): + result = net.run(net.connect_addrinfo(self.ADDRINFO)) + assert result is sock + + def test_connection_refused(self, net): + sock = MagicMock() + sock.connect_ex.return_value = errno.ECONNREFUSED + with patch('kafka.net.backend.selector.socket.socket', return_value=sock): + with pytest.raises(Errors.KafkaConnectionError): + net.run(net.connect_addrinfo(self.ADDRINFO)) + + def test_socket_error_uses_errno(self, net): + sock = MagicMock() + sock.connect_ex.side_effect = socket.error(errno.ECONNREFUSED, 'refused') + with patch('kafka.net.backend.selector.socket.socket', return_value=sock): + with pytest.raises(Errors.KafkaConnectionError): + net.run(net.connect_addrinfo(self.ADDRINFO)) + + def test_error_after_wait_write(self, net): + """connect_ex returns EINPROGRESS, then after wait_write fires the + second connect_ex returns the real error.""" + # socketpair endpoints are always immediately writable, so wait_write + # fires on the first poll and we re-enter the loop. + rsock, wsock = socket.socketpair() + rsock.setblocking(False) + wsock.setblocking(False) + try: + sock = MagicMock() + sock.connect_ex.side_effect = [errno.EINPROGRESS, errno.ECONNREFUSED] + sock.fileno.return_value = wsock.fileno() + with patch('kafka.net.backend.selector.socket.socket', return_value=sock): + with pytest.raises(Errors.KafkaConnectionError): + net.run(net.connect_addrinfo(self.ADDRINFO)) + assert sock.connect_ex.call_count == 2 + finally: + rsock.close() + wsock.close() + + def test_socket_options_applied(self, net): + sock = MagicMock() + sock.connect_ex.return_value = 0 + opts = [ + (socket.IPPROTO_TCP, socket.TCP_NODELAY, 1), + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), + ] + with patch('kafka.net.backend.selector.socket.socket', return_value=sock): + net.run(net.connect_addrinfo(self.ADDRINFO, socket_options=opts)) + sock.setsockopt.assert_has_calls([ + call(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1), + call(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), + ]) + assert sock.setsockopt.call_count == 2 + + +class TestConnectHost: + FAKE_ADDR = [(socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, '', ('127.0.0.1', 9092))] + + def test_dns_failure(self, net): + with patch('kafka.net.backend.selector.socket.getaddrinfo', side_effect=socket.gaierror): + with pytest.raises(Errors.KafkaConnectionError, match='DNS'): + net.run(net.connect_host('badhost', 9092)) + + def test_socket_init_failure(self, net): + with patch('kafka.net.backend.selector.socket.getaddrinfo', return_value=self.FAKE_ADDR), \ + patch('kafka.net.backend.selector.socket.socket', side_effect=OSError('no socket')): + with pytest.raises(Errors.KafkaConnectionError): + net.run(net.connect_host('host', 9092)) + + def test_successful_connection(self, net): + sock = MagicMock() + sock.connect_ex.return_value = 0 + with patch('kafka.net.backend.selector.socket.getaddrinfo', return_value=self.FAKE_ADDR), \ + patch('kafka.net.backend.selector.socket.socket', return_value=sock): + result = net.run(net.connect_host('host', 9092)) + assert result is sock + sock.setblocking.assert_called_with(False) + + def test_tries_multiple_addresses(self, net): + addr1 = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, '', ('10.0.0.1', 9092)) + addr2 = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, '', ('10.0.0.2', 9092)) + sock1 = MagicMock() + sock1.connect_ex.return_value = errno.ECONNREFUSED + sock2 = MagicMock() + sock2.connect_ex.return_value = 0 + sockets = iter([sock1, sock2]) + with patch('kafka.net.backend.selector.socket.getaddrinfo', return_value=[addr1, addr2]), \ + patch('kafka.net.backend.selector.socket.socket', side_effect=lambda *a: next(sockets)): + result = net.run(net.connect_host('host', 9092)) + assert result is sock2 + + def test_socket_options_applied(self, net): + sock = MagicMock() + sock.connect_ex.return_value = 0 + opts = [ + (socket.IPPROTO_TCP, socket.TCP_NODELAY, 1), + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), + ] + with patch('kafka.net.backend.selector.socket.getaddrinfo', return_value=self.FAKE_ADDR), \ + patch('kafka.net.backend.selector.socket.socket', return_value=sock): + net.run(net.connect_host('host', 9092, socket_options=opts)) + sock.setsockopt.assert_has_calls([ + call(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1), + call(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), + ]) + assert sock.setsockopt.call_count == 2