diff --git a/kafka/admin/client.py b/kafka/admin/client.py index 3161fac70..824d0b284 100644 --- a/kafka/admin/client.py +++ b/kafka/admin/client.py @@ -142,9 +142,19 @@ class KafkaAdminClient( metadata before the configured timeout. Note that bootstrap is called eagerly from __init__(). Default: 30000 - selector (selectors.BaseSelector): Provide a specific selector - implementation to use for I/O multiplexing. - Default: selectors.DefaultSelector + net (str or kafka.net.backend.NetBackend): The async backend that runs + this client's network I/O event loop. One of: a NetBackend + instance; a registered name -- 'selector' (the built-in + NetworkSelector) or 'asyncio' (runs I/O on an asyncio loop); or None + to auto-detect. Auto-detect selects 'asyncio' when the client is + constructed inside a running asyncio event loop (via sniffio, or + asyncio's running-loop check), otherwise falls back to 'selector'. + Regardless of backend, the loop runs on its own dedicated daemon IO + thread and this client's public API stays synchronous (blocking); a + native awaitable API is a separate, later phase. Caveat: calling a + blocking client method from within your own running asyncio loop + will block that loop -- run such calls in a thread executor + (e.g. loop.run_in_executor). Default: None. metrics (kafka.metrics.Metrics): Optionally provide a metrics instance for capturing network IO stats. Default: None. metric_group_prefix (str): Prefix for metric names. Default: '' @@ -194,7 +204,7 @@ class KafkaAdminClient( 'ssl_crlfile': None, 'api_version': None, 'bootstrap_timeout_ms': 30000, - 'selector': selectors.DefaultSelector, + 'selector': None, # deprecated; use net instead 'sasl_mechanism': None, 'sasl_plain_username': None, 'sasl_plain_password': None, @@ -220,6 +230,7 @@ def __init__(self, **configs): raise KafkaConfigurationError("Unrecognized configs: {}".format(extra_configs)) self.config = copy.copy(self.DEFAULT_CONFIG) + self.config.pop('selector') self.config.update(configs) # Configure metrics diff --git a/kafka/consumer/group.py b/kafka/consumer/group.py index 84adfbbca..ed4530797 100644 --- a/kafka/consumer/group.py +++ b/kafka/consumer/group.py @@ -265,9 +265,19 @@ class KafkaConsumer: metrics. Default: 2 metrics_sample_window_ms (int): The maximum age in milliseconds of samples used to compute metrics. Default: 30000 - selector (selectors.BaseSelector): Provide a specific selector - implementation to use for I/O multiplexing. - Default: selectors.DefaultSelector + net (str or kafka.net.backend.NetBackend): The async backend that runs + this client's network I/O event loop. One of: a NetBackend + instance; a registered name -- 'selector' (the built-in + NetworkSelector) or 'asyncio' (runs I/O on an asyncio loop); or None + to auto-detect. Auto-detect selects 'asyncio' when the client is + constructed inside a running asyncio event loop (via sniffio, or + asyncio's running-loop check), otherwise falls back to 'selector'. + Regardless of backend, the loop runs on its own dedicated daemon IO + thread and this client's public API stays synchronous (blocking); a + native awaitable API is a separate, later phase. Caveat: calling a + blocking client method from within your own running asyncio loop + will block that loop -- run such calls in a thread executor + (e.g. loop.run_in_executor). Default: None. exclude_internal_topics (bool): Whether records from internal topics (such as offsets) should be exposed to the consumer. If set to True the only way to receive records from an internal topic is @@ -350,7 +360,7 @@ class KafkaConsumer: 'metrics_num_samples': 2, 'metrics_sample_window_ms': 30000, 'metric_group_prefix': 'consumer', - 'selector': selectors.DefaultSelector, + 'selector': None, # deprecated; use net instead 'exclude_internal_topics': True, 'sasl_mechanism': None, 'sasl_plain_username': None, @@ -373,6 +383,7 @@ def __init__(self, *topics, **configs): raise KafkaConfigurationError("Unrecognized configs: %s" % (extra_configs,)) self.config = copy.copy(self.DEFAULT_CONFIG) + self.config.pop('selector') # handled in manager init self.config.update(configs) deprecated = {'smallest': 'earliest', 'largest': 'latest'} diff --git a/kafka/net/asyncio_backend.py b/kafka/net/asyncio_backend.py new file mode 100644 index 000000000..24f7628ff --- /dev/null +++ b/kafka/net/asyncio_backend.py @@ -0,0 +1,489 @@ +"""An asyncio-backed NetBackend (Phase 1: own daemon thread). + +Runs a private ``asyncio`` event loop on a dedicated daemon thread, mirroring +``NetworkSelector``'s threading model, and implements the ``NetBackend`` +contract on top of asyncio primitives. Selected via ``net='asyncio'`` or +auto-detected when constructed inside a running asyncio loop (see +``kafka.net.backend.resolve_backend``). + +Phase 1 preserves the synchronous public API: ``run()`` blocks the calling +thread on the loop thread; it does not run on the caller's own loop. +""" +import asyncio +import inspect +import threading +import time + +import kafka.errors as Errors +from kafka.future import Future + + +class AsyncioFuture(Future): + """``create_future()`` result for the asyncio backend. + + Inherits ``kafka.future.Future``'s callback core and overrides only + ``__await__`` to bridge to an ``asyncio.Future`` so an asyncio Task can + await it. Per the BackendFuture contract this is created and resolved on + the loop thread only; a fresh asyncio.Future is minted per awaiter, so + fan-out (multiple awaiters / callbacks) is preserved. + """ + __slots__ = ('_loop',) + + def __init__(self, loop): + super().__init__() + self._loop = loop + + def __await__(self): + if not self.is_done: + aio = self._loop.create_future() + + def _resolve(_=None): + if aio.done(): + return + if self.exception is not None: + aio.set_exception(self.exception) + else: + aio.set_result(self.value) + + self.add_both(_resolve) + yield from aio.__await__() + if self.exception: + raise self.exception + return self.value + + +class _DeferredHandle: + """Cancelable handle for a timer/callback armed cross-thread. + + ``call_later``/``call_soon`` invoked off the loop thread schedule the real + handle via ``call_soon_threadsafe``; this box lets the caller ``cancel()`` + synchronously whether or not the real handle has been armed yet. + """ + __slots__ = ('_handle', '_cancelled') + + def __init__(self): + self._handle = None + self._cancelled = False + + def _arm(self, handle): + if self._cancelled: + handle.cancel() + else: + self._handle = handle + + def cancel(self): + self._cancelled = True + if self._handle is not None: + self._handle.cancel() + + +class AsyncioBackend: + def __init__(self, *, loop=None, loop_factory=None, **configs): + # Loop acquisition is a strategy, not a hardcode (forward seam for a + # native/Phase-2 mode). Default: own a fresh loop on our daemon thread. + # loop_factory lets callers pick the loop implementation (e.g. uvloop) + # while keeping the owned-thread model. loop= injects an existing loop + # (not owned -- close() won't close it); in Phase 1 it is still run on + # our own thread. True native reuse of a *running* caller loop is Phase 2. + if loop is not None: + self._loop = loop + self._owns_loop = False + else: + self._loop = (loop_factory or asyncio.new_event_loop)() + self._owns_loop = True + self._io_thread = None + self._closed = False + self._client_id = configs.get('client_id') or 'kafka-python' + # Strong refs to live tasks (asyncio only holds weak refs, so bare + # tasks can be GC'd mid-flight); mirrors NetworkSelector._pending_tasks. + self._pending = set() + # Cross-thread run() waiters, failed on stop() so callers don't hang. + self._pending_waiters = {} + self._pending_waiters_lock = threading.Lock() + + # --- lifecycle -------------------------------------------------------- + def start(self): + if self._io_thread is not None: + return + t = threading.Thread(target=self._run_forever, + name='kafka-io-%s' % self._client_id, daemon=True) + self._io_thread = t + t.start() + + def _run_forever(self): + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + def on_io_thread(self): + return self._io_thread is not None and threading.current_thread() is self._io_thread + + def stop(self, timeout_ms=None): + if self._io_thread is None: + return + self._loop.call_soon_threadsafe(self._loop.stop) + self._io_thread.join(timeout_ms / 1000 if timeout_ms is not None else None) + self._io_thread = None + self._fail_pending_waiters(Errors.KafkaConnectionError('Event loop stopped')) + + def close(self): + if self._closed: + return + self._closed = True + if self._io_thread is not None: + self.stop() + if not self._loop.is_closed(): + pending = [t for t in self._pending if not t.done()] + for t in pending: + t.cancel() + if pending: + try: + self._loop.run_until_complete( + asyncio.gather(*pending, return_exceptions=True)) + except RuntimeError: + pass + if self._owns_loop: + self._loop.close() + + def _fail_pending_waiters(self, exc): + with self._pending_waiters_lock: + waiters = list(self._pending_waiters.items()) + self._pending_waiters.clear() + for event, state in waiters: + if state['exception'] is None: + state['exception'] = exc + event.set() + + # --- scheduling ------------------------------------------------------- + def _spawn(self, coro): + task = self._loop.create_task(coro) + self._pending.add(task) + task.add_done_callback(self._pending.discard) + return task + + def _schedule(self, task, args=()): + """Run a coroutine / coroutine-function / callable on the loop thread. + + A plain callable that *returns* a coroutine (e.g. the manager's + ``call_soon(lambda: self._connect(...))``) has that coroutine run too, + mirroring NetworkSelector's Task, which steps into returned coroutines. + """ + if inspect.iscoroutine(task): + return self._spawn(task) + if inspect.iscoroutinefunction(task): + return self._spawn(task(*args)) + + def _call(): + result = task(*args) + if inspect.iscoroutine(result): + self._spawn(result) + return self._loop.call_soon(_call) + + def call_soon(self, task): + # On the loop thread: schedule directly. Off it (or before start()): + # route through call_soon_threadsafe so create_task/call_soon run on + # the loop thread as asyncio requires. + if self.on_io_thread(): + return self._schedule(task) + box = _DeferredHandle() + self._loop.call_soon_threadsafe(lambda: box._arm(self._schedule(task))) + return box + + def call_soon_threadsafe(self, callback): + if self._closed: + raise RuntimeError('AsyncioBackend closed!') + box = _DeferredHandle() + self._loop.call_soon_threadsafe(lambda: box._arm(self._schedule(callback))) + return box + + def _as_callback(self, task): + if inspect.iscoroutinefunction(task): + return lambda: self._spawn(task()) + if inspect.iscoroutine(task): + return lambda: self._spawn(task) + return task + + def call_at(self, when, task): + # Selector uses time.monotonic()-based absolute `when`; convert to a + # relative delay for asyncio's loop.time() base. + return self.call_later(max(0.0, when - time.monotonic()), task) + + def call_later(self, delay, task): + cb = self._as_callback(task) + if self.on_io_thread(): + return self._loop.call_later(delay, cb) + box = _DeferredHandle() + self._loop.call_soon_threadsafe( + lambda: box._arm(self._loop.call_later(delay, cb))) + return box + + def cancel(self, task): + if task is not None: + task.cancel() + + def sleep(self, delay): + return asyncio.sleep(delay) + + def wakeup(self): + # asyncio has no select() to interrupt from user code; a no-op + # threadsafe callback is enough to wake a blocked run_forever(). + try: + self._loop.call_soon_threadsafe(lambda: None) + except RuntimeError: + pass + + # --- futures ---------------------------------------------------------- + def create_future(self): + return AsyncioFuture(self._loop) + + async def _resolve_future(self, fut): + """Await any kafka.future.Future (plain or AsyncioFuture) to its value.""" + if fut.is_done: + if fut.exception is not None: + raise fut.exception + return fut.value + aio = self._loop.create_future() + + def _cb(_=None): + if aio.done(): + return + if fut.exception is not None: + aio.set_exception(fut.exception) + else: + aio.set_result(fut.value) + + fut.add_both(_cb) + return await aio + + async def _invoke(self, coro, *args): + """Invoke coro/awaitable/function and fully resolve the result. + + Mirrors NetworkSelector._invoke, but bridges any trailing kafka Future + through _resolve_future (a plain Future isn't awaitable under asyncio). + """ + if inspect.iscoroutinefunction(coro): + result = await coro(*args) + elif hasattr(coro, '__await__'): + result = await coro + else: + result = coro(*args) + if inspect.iscoroutine(result) or hasattr(result, '__await__'): + result = await result + while isinstance(result, Future): + result = await self._resolve_future(result) + return result + + def call_soon_with_future(self, coro, *args): + if hasattr(coro, '__await__') and args: + raise ValueError('initiated coroutine does not accept args') + future = AsyncioFuture(self._loop) + + async def wrapper(): + try: + future.success(await self._invoke(coro, *args)) + except BaseException as exc: + future.failure(exc) + + self.call_soon(wrapper) + return future + + # --- cross-thread bridge --------------------------------------------- + def run(self, coro, *args): + if self._closed: + raise RuntimeError('AsyncioBackend closed!') + if self._io_thread is None: + raise RuntimeError('AsyncioBackend not started; call start() first') + if self.on_io_thread(): + raise RuntimeError( + "Cannot block on net.run() from the IO thread itself. " + "This typically happens when a synchronous rebalance listener " + "(or another IO-thread callback) calls a blocking consumer/admin API. " + "Use AsyncConsumerRebalanceListener and await the async variant, " + "or move the blocking work to a worker thread.") + event = threading.Event() + state = {'value': None, 'exception': None} + + async def waiter(): + try: + state['value'] = await self._invoke(coro, *args) + except BaseException as exc: + if state['exception'] is None: + state['exception'] = exc + finally: + with self._pending_waiters_lock: + self._pending_waiters.pop(event, None) + event.set() + + with self._pending_waiters_lock: + self._pending_waiters[event] = state + self.call_soon(waiter) + event.wait() + if state['exception'] is not None: + raise state['exception'] # pylint: disable=raising-bad-type + return state['value'] + + # --- connection seam -------------------------------------------------- + async def create_connection(self, protocol, host, port, *, ssl=None, + proxy_url=None, socket_options=(), timeout_at=None): + if proxy_url is not None: + raise NotImplementedError( + 'The asyncio backend does not support proxy_url yet; use the ' + 'default selector backend for SOCKS5/HTTP-CONNECT proxying.') + server_hostname = host.rstrip('.') if ssl is not None else None + adapter = _AsyncioProtocolAdapter() + connect = self._loop.create_connection( + lambda: adapter, host, port, ssl=ssl, server_hostname=server_hostname) + try: + if timeout_at is not None: + connect = asyncio.wait_for(connect, max(0.0, timeout_at - time.monotonic())) + aio_transport, _ = await connect + except asyncio.TimeoutError: + raise Errors.KafkaConnectionError('Connection timed out') + except Errors.KafkaError: + raise + except Exception as exc: # noqa: BLE001 -- surface any connect error uniformly + raise Errors.KafkaConnectionError('unable to connect to %s:%s: %s' % (host, port, exc)) + sock = aio_transport.get_extra_info('socket') + if sock is not None: + for option in socket_options: + try: + sock.setsockopt(*option) + except OSError: + pass + return _AsyncioTransport(aio_transport, adapter, host, port) + + +class _AsyncioProtocolAdapter(asyncio.Protocol): + """Bridges asyncio's Protocol callbacks to a KafkaConnection. + + asyncio calls ``connection_made`` on this adapter during + ``create_connection`` -- before the manager wires the KafkaConnection via + ``transport.set_protocol(conn)`` (Option A: the caller runs connection_made + after its "closed during connect" check). Reading is paused until wired, so + no data is delivered early; a buffer/latched-loss is kept as a safety net. + """ + + def __init__(self): + self._kafka = None # the KafkaConnection, once wired + self._transport = None + self._buffer = [] + self._eof = False + self._lost = False + self._lost_exc = None + self._paused_writing = False + self._on_read = None # bumps the wrapper's last_read + + def connection_made(self, transport): + self._transport = transport + # Hold off delivery until the KafkaConnection is wired + resumes reading. + transport.pause_reading() + + def data_received(self, data): + if self._on_read is not None: + self._on_read() + if self._kafka is None: + self._buffer.append(data) + else: + self._kafka.data_received(data) + + def eof_received(self): + if self._kafka is not None: + return self._kafka.eof_received() + self._eof = True + return None + + def connection_lost(self, exc): + if self._kafka is not None: + self._kafka.connection_lost(exc) + else: + self._lost = True + self._lost_exc = exc + + def pause_writing(self): + if self._kafka is not None: + self._kafka.pause_writing() + else: + self._paused_writing = True + + def resume_writing(self): + if self._kafka is not None: + self._kafka.resume_writing() + else: + self._paused_writing = False + + def _wire(self, kafka_protocol, on_read): + self._kafka = kafka_protocol + self._on_read = on_read + if self._paused_writing: + kafka_protocol.pause_writing() + for data in self._buffer: + kafka_protocol.data_received(data) + self._buffer = [] + if self._eof: + kafka_protocol.eof_received() + if self._lost: + kafka_protocol.connection_lost(self._lost_exc) + + +class _AsyncioTransport: + """Transport returned by AsyncioBackend.create_connection. + + Wraps an asyncio transport + its protocol adapter and exposes the surface + KafkaConnection / the manager drive (a superset of the NetBackend Transport + protocol): write/close/abort/is_closing, pause/resume_reading, set/get + protocol, host/host_port/getPeer, and last_activity for idle sweeping. + """ + + def __init__(self, transport, adapter, host, port): + self._t = transport + self._adapter = adapter + self.host = host + self._port = port + self._protocol = None + self.last_write = time.monotonic() + self.last_read = time.monotonic() + + @property + def last_activity(self): + return max(self.last_write, self.last_read) + + def _bump_read(self): + self.last_read = time.monotonic() + + def get_protocol(self): + return self._protocol + + def set_protocol(self, protocol): + self._protocol = protocol + self._adapter._wire(protocol, self._bump_read) + + def write(self, data): + self.last_write = time.monotonic() + self._t.write(data) + + def close(self): + self._t.close() + + def abort(self, error=None): + self._t.abort() + + def is_closing(self): + return self._t.is_closing() + + def pause_reading(self): + try: + self._t.pause_reading() + except (RuntimeError, AttributeError): + pass + + def resume_reading(self): + try: + self._t.resume_reading() + except (RuntimeError, AttributeError): + pass + + def getPeer(self): + peer = self._t.get_extra_info('peername') + return peer if peer is not None else (self.host, self._port) + + def host_port(self): + return '%s:%s' % (self.host, self._port) diff --git a/kafka/net/backend.py b/kafka/net/backend.py index 9c77d26fe..0bbb94de3 100644 --- a/kafka/net/backend.py +++ b/kafka/net/backend.py @@ -284,3 +284,4 @@ def resolve_backend(net, config): register_backend_lazy('selector', 'kafka.net.selector', 'NetworkSelector') +register_backend_lazy('asyncio', 'kafka.net.asyncio_backend', 'AsyncioBackend') diff --git a/kafka/net/manager.py b/kafka/net/manager.py index aaa40107c..76feddf6c 100644 --- a/kafka/net/manager.py +++ b/kafka/net/manager.py @@ -57,6 +57,7 @@ class KafkaConnectionManager: 'metric_group_prefix': '', 'metadata_max_age_ms': 300000, 'client_dns_lookup': 'use_all_dns_ips', + 'selector': None, # deprecated; use net instead } _VALID_DNS_LOOKUP_MODES = ('use_all_dns_ips', 'resolve_canonical_bootstrap_servers_only') @@ -76,11 +77,16 @@ def __init__(self, net=None, **configs): log.warning('socks5_proxy is deprecated, use proxy_url instead') self.config['proxy_url'] = configs['socks5_proxy'] + if configs.get('selector') is None: + self.config.pop('selector') + else: + log.warning('selector is deprecated, use net instead') + # `net` is the raw backend selector: a NetBackend instance, a backend # name ('selector'/'asyncio'), or None to auto-detect / default to the # NetworkSelector. Resolved here (not in the legacy compat shim) so the # manager remains the durable entry point once compat.py is removed. - self._net = resolve_backend(net, configs) + self._net = resolve_backend(net, self.config) self.cluster = ClusterMetadata( bootstrap_servers=self.config['bootstrap_servers'], metadata_max_age_ms=self.config['metadata_max_age_ms'], diff --git a/kafka/producer/kafka.py b/kafka/producer/kafka.py index abf43479a..1eafbab61 100644 --- a/kafka/producer/kafka.py +++ b/kafka/producer/kafka.py @@ -375,9 +375,19 @@ class KafkaProducer: metrics. Default: 2 metrics_sample_window_ms (int): The maximum age in milliseconds of samples used to compute metrics. Default: 30000 - selector (selectors.BaseSelector): Provide a specific selector - implementation to use for I/O multiplexing. - Default: selectors.DefaultSelector + net (str or kafka.net.backend.NetBackend): The async backend that runs + this client's network I/O event loop. One of: a NetBackend + instance; a registered name -- 'selector' (the built-in + NetworkSelector) or 'asyncio' (runs I/O on an asyncio loop); or None + to auto-detect. Auto-detect selects 'asyncio' when the client is + constructed inside a running asyncio event loop (via sniffio, or + asyncio's running-loop check), otherwise falls back to 'selector'. + Regardless of backend, the loop runs on its own dedicated daemon IO + thread and this client's public API stays synchronous (blocking); a + native awaitable API is a separate, later phase. Caveat: calling a + blocking client method from within your own running asyncio loop + will block that loop -- run such calls in a thread executor + (e.g. loop.run_in_executor). Default: None. sasl_mechanism (str): Authentication mechanism when security_protocol is configured for SASL_PLAINTEXT or SASL_SSL. Valid values are: PLAIN, GSSAPI, OAUTHBEARER, SCRAM-SHA-256, SCRAM-SHA-512. @@ -447,7 +457,7 @@ class KafkaProducer: 'metrics_enabled': True, 'metrics_num_samples': 2, 'metrics_sample_window_ms': 30000, - 'selector': selectors.DefaultSelector, + 'selector': None, # deprecated; use net instead 'sasl_mechanism': None, 'sasl_plain_username': None, 'sasl_plain_password': None, @@ -461,8 +471,6 @@ class KafkaProducer: 'net': None, } - DEPRECATED_CONFIGS = () - _COMPRESSORS = { 'gzip': (has_gzip, LegacyRecordBatchBuilder.CODEC_GZIP), 'snappy': (has_snappy, LegacyRecordBatchBuilder.CODEC_SNAPPY), @@ -472,26 +480,20 @@ class KafkaProducer: } def __init__(self, **configs): + user_provided_configs = set(configs) + extra_configs = user_provided_configs.difference(self.DEFAULT_CONFIG) + if extra_configs: + raise Errors.KafkaConfigurationError("Unrecognized configs: {}".format(extra_configs)) + self.config = copy.copy(self.DEFAULT_CONFIG) - user_provided_configs = set(configs.keys()) - for key in self.config: - if key in configs: - self.config[key] = configs.pop(key) + self.config.pop('selector') + self.config.update(configs) for key in ('key_serializer', 'value_serializer'): if self.config[key] is not None and not isinstance(self.config[key], Serializer): warnings.warn('%s does not implement kafka.serializer.Serializer' % (key,), category=DeprecationWarning, stacklevel=3) self.config[key] = SerializeWrapper(self.config[key]) - for key in self.DEPRECATED_CONFIGS: - if key in configs: - configs.pop(key) - warnings.warn('Deprecated Producer config: %s' % (key,), category=DeprecationWarning) - - # Only check for extra config keys in top-level class - if configs: - raise ValueError('Unrecognized configs: %s' % (configs,)) - if self.config['client_id'] is None: self.config['client_id'] = 'kafka-python-producer-%s' % \ (PRODUCER_CLIENT_ID_SEQUENCE.increment(),) diff --git a/test/integration/fixtures.py b/test/integration/fixtures.py index dceab32f2..8ca5a5cf3 100644 --- a/test/integration/fixtures.py +++ b/test/integration/fixtures.py @@ -726,6 +726,12 @@ def _enrich_client_params(self, params, **defaults): if self.ssl_enabled: params.setdefault('ssl_cafile', os.path.join(self.ssl_dir, 'ca-cert')) params.setdefault('ssl_check_hostname', False) + # Run the whole integration suite against a chosen net backend, e.g. + # KAFKA_PYTHON_NET=asyncio make test. Unset -> default (selector). + # setdefault so a test that pins net= still wins. + net = os.environ.get('KAFKA_PYTHON_NET') + if net: + params.setdefault('net', net) return params diff --git a/test/integration/test_sasl_integration.py b/test/integration/test_sasl_integration.py index d36100bb1..bcfc7500a 100644 --- a/test/integration/test_sasl_integration.py +++ b/test/integration/test_sasl_integration.py @@ -1,13 +1,12 @@ import logging import os import uuid -import time import pytest from kafka import KafkaAdminClient, KafkaConsumer, KafkaProducer from kafka.admin import NewTopic -from kafka.net.compat import KafkaNetClient +from kafka.net.manager import KafkaConnectionManager from kafka.protocol.metadata import MetadataRequest from test.testutil import assert_message_count, env_kafka_version, random_string, special_to_underscore from test.integration.fixtures import client_params, create_topics @@ -76,16 +75,20 @@ def test_client(request, sasl_kafka): topic_name = special_to_underscore(request.node.name + random_string(4)) create_topics(sasl_kafka, [topic_name], num_partitions=1) - client = KafkaNetClient(**client_params(sasl_kafka, 'client')) - client._manager.run(client._manager.bootstrap_async) - request = MetadataRequest(topics=None, version=1) - timeout_at = time.time() + 1 - future = client.send(None, request) - client.poll(future=future, timeout_ms=10000) - if not future.is_done: - raise RuntimeError("Couldn't fetch topic response from Broker.") - elif future.failed(): - raise future.exception - result = future.value - assert topic_name in [t[1] for t in result.topics] - client.close() + # Low-level SASL round-trip via KafkaConnectionManager directly (no compat + # shim, no poll()): the started-loop + manager.run(coro) pattern the real + # clients use, so it runs on any net backend (selector or asyncio). + manager = KafkaConnectionManager(**client_params(sasl_kafka, 'client')) + manager._net.start() + try: + manager.bootstrap(timeout_ms=5000) + + async def fetch_metadata(): + future = manager.send(MetadataRequest(topics=None, version=1), node_id=None) + return await manager.wait_for(future, 10000) + + result = manager.run(fetch_metadata) + assert topic_name in [t[1] for t in result.topics] + finally: + manager.close() + manager._net.close() diff --git a/test/net/test_asyncio_backend.py b/test/net/test_asyncio_backend.py new file mode 100644 index 000000000..06c8c2971 --- /dev/null +++ b/test/net/test_asyncio_backend.py @@ -0,0 +1,352 @@ +"""Tests for the asyncio NetBackend (kafka/net/asyncio_backend.py). + +Covers backend-specific behavior (lifecycle, timers, cross-thread run), reuses +the shared BackendFuture conformance suite against the asyncio-backed future, +and drives a real protocol round-trip through a MockBroker on a started +AsyncioBackend -- the both-backends coverage for the async paths. +""" +import asyncio +import socket +import threading +import time + +import pytest + +from kafka.net.asyncio_backend import AsyncioBackend, AsyncioFuture +from kafka.net.backend import NetBackend +from kafka.net.manager import KafkaConnectionManager +from kafka.protocol.metadata import MetadataRequest +from test.mock_broker import MockBroker +from test.net.test_backend_future import BackendFutureContract + + +@pytest.fixture +def backend(): + b = AsyncioBackend(client_id='test') + try: + yield b + finally: + b.close() + + +@pytest.fixture +def started_backend(): + b = AsyncioBackend(client_id='test') + b.start() + try: + yield b + finally: + b.close() + + +class TestAsyncioBackendContract: + def test_satisfies_netbackend(self, backend): + assert isinstance(backend, NetBackend) + + def test_isinstance_after_start(self, started_backend): + assert isinstance(started_backend, NetBackend) + + +class TestLifecycle: + def test_start_idempotent(self, backend): + backend.start() + t = backend._io_thread + backend.start() + assert backend._io_thread is t + + def test_on_io_thread(self, started_backend): + async def where(): + return started_backend.on_io_thread() + assert started_backend.run(where) is True + assert started_backend.on_io_thread() is False + + def test_stop_is_idempotent(self, backend): + backend.start() + backend.stop() + backend.stop() # no raise + assert backend._io_thread is None + + def test_run_before_start_raises(self, backend): + async def noop(): + return 1 + with pytest.raises(RuntimeError, match='not started'): + backend.run(noop) + + def test_run_from_io_thread_raises(self, started_backend): + async def nested(): + started_backend.run(lambda: 1) + with pytest.raises(RuntimeError, match='IO thread'): + started_backend.run(nested) + + def test_stop_fails_pending_run_waiters(self, started_backend): + # A run() blocked on a never-resolving coroutine is released with an + # error when the loop stops, rather than hanging forever. + errors = [] + + def caller(): + async def forever(): + await asyncio.Event().wait() + try: + started_backend.run(forever) + except Exception as exc: # noqa: BLE001 + errors.append(exc) + + t = threading.Thread(target=caller) + t.start() + time.sleep(0.1) + started_backend.stop() + t.join(timeout=5) + assert not t.is_alive() + assert len(errors) == 1 + + +class TestRun: + def test_run_coroutine_function_with_args(self, started_backend): + async def add(x, y): + return x + y + assert started_backend.run(add, 2, 3) == 5 + + def test_run_propagates_exception(self, started_backend): + async def boom(): + raise ValueError('kaboom') + with pytest.raises(ValueError, match='kaboom'): + started_backend.run(boom) + + def test_run_resolves_returned_future(self, started_backend): + # _invoke must bridge a trailing kafka Future to its value. + def returns_future(): + fut = started_backend.create_future() + started_backend.call_soon(lambda: fut.success('deferred')) + return fut + assert started_backend.run(returns_future) == 'deferred' + + +class TestTimers: + def test_call_later_fires(self, started_backend): + fired = threading.Event() + started_backend.call_soon(lambda: started_backend.call_later(0.01, fired.set)) + assert fired.wait(timeout=2) + + def test_call_later_cancel_prevents_fire(self, started_backend): + fired = threading.Event() + + async def schedule_and_cancel(): + handle = started_backend.call_later(0.5, fired.set) + started_backend.cancel(handle) + started_backend.run(schedule_and_cancel) + assert not fired.wait(timeout=0.3) + + def test_call_at_converts_monotonic(self, started_backend): + fired = threading.Event() + when = time.monotonic() + 0.01 + started_backend.call_soon(lambda: started_backend.call_at(when, fired.set)) + assert fired.wait(timeout=2) + + def test_sleep_awaitable(self, started_backend): + async def nap(): + start = time.monotonic() + await started_backend.sleep(0.05) + return time.monotonic() - start + assert started_backend.run(nap) >= 0.04 + + +class TestCreateFuture: + def test_create_future_type(self, backend): + fut = backend.create_future() + assert isinstance(fut, AsyncioFuture) + assert not fut.is_done + + +class TestAsyncioBackendFuture(BackendFutureContract): + """Reuse the shared BackendFuture conformance suite for the asyncio future.""" + + @pytest.fixture(autouse=True) + def _net(self): + self.net = AsyncioBackend(client_id='future-contract') + self.net.start() + try: + yield + finally: + self.net.close() + + def make_future(self): + return self.net.create_future() + + def drive(self, coros): + async def _all(): + await asyncio.gather(*coros) + self.net.run(_all) + + +class _StubProtocol: + """Minimal KafkaConnection-shaped protocol for transport tests.""" + def __init__(self): + self.received = bytearray() + self.lost = False + def data_received(self, data): + self.received += data + def eof_received(self): + return None + def connection_lost(self, exc): + self.lost = True + def pause_writing(self): + pass + def resume_writing(self): + pass + + +def _echo_server(): + srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(('127.0.0.1', 0)) + srv.listen(1) + host, port = srv.getsockname() + + def serve(): + try: + conn, _ = srv.accept() + except OSError: + return + while True: + data = conn.recv(4096) + if not data: + break + conn.sendall(data) + conn.close() + + t = threading.Thread(target=serve, daemon=True) + t.start() + return srv, host, port + + +class TestCreateConnection: + def test_roundtrip_write_read_close(self, started_backend): + srv, host, port = _echo_server() + proto = _StubProtocol() + + async def do(): + transport = await started_backend.create_connection(proto, host, port) + # Wire like manager._connect -> conn.connection_made does. + transport.set_protocol(proto) + transport.resume_reading() + assert transport.host_port() == '%s:%s' % (host, port) + assert transport.getPeer()[0:2] == (host, port) + transport.write(b'ping') + for _ in range(100): + if proto.received: + break + await started_backend.sleep(0.01) + transport.close() + return bytes(proto.received) + + try: + assert started_backend.run(do) == b'ping' + finally: + srv.close() + + def test_early_data_is_buffered_until_wired(self, started_backend): + # Server sends immediately on connect; data must not be lost before + # the protocol is wired via set_protocol(). + srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(('127.0.0.1', 0)) + srv.listen(1) + host, port = srv.getsockname() + + def serve(): + conn, _ = srv.accept() + conn.sendall(b'hello-early') + conn.recv(4096) + conn.close() + + threading.Thread(target=serve, daemon=True).start() + proto = _StubProtocol() + + async def do(): + transport = await started_backend.create_connection(proto, host, port) + await started_backend.sleep(0.05) # let early bytes arrive (buffered) + transport.set_protocol(proto) # wiring flushes the buffer + transport.resume_reading() + for _ in range(100): + if proto.received: + break + await started_backend.sleep(0.01) + transport.close() + return bytes(proto.received) + + try: + assert started_backend.run(do) == b'hello-early' + finally: + srv.close() + + def test_connect_refused_raises(self, started_backend): + # Nothing listening -> KafkaConnectionError. + s = socket.socket(); s.bind(('127.0.0.1', 0)); host, port = s.getsockname(); s.close() + + async def do(): + await started_backend.create_connection(_StubProtocol(), host, port) + + with pytest.raises(Exception): # KafkaConnectionError + started_backend.run(do) + + def test_proxy_url_raises(self, started_backend): + async def do(): + await started_backend.create_connection( + _StubProtocol(), 'h', 1, proxy_url='socks5://proxy:1080') + with pytest.raises(NotImplementedError, match='proxy'): + started_backend.run(do) + + +class TestLoopConfig: + def test_loop_factory_used_and_owned(self): + calls = [] + + def factory(): + calls.append(1) + return asyncio.new_event_loop() + + b = AsyncioBackend(loop_factory=factory, client_id='lf') + try: + assert calls == [1] + assert b._owns_loop is True + finally: + b.close() + + def test_injected_loop_not_owned_or_closed(self): + loop = asyncio.new_event_loop() + b = AsyncioBackend(loop=loop, client_id='inj') + assert b._loop is loop + assert b._owns_loop is False + b.close() + assert not loop.is_closed() # injected loop must survive close() + loop.close() + + +class TestEndToEndMockBroker: + """Drive a real bootstrap + protocol round-trip through MockBroker on a + started AsyncioBackend -- proves the backend runs the full IO machinery.""" + + def test_bootstrap_and_send(self): + broker = MockBroker() + net = AsyncioBackend(client_id='e2e') + manager = KafkaConnectionManager( + net, + bootstrap_servers='%s:%d' % (broker.host, broker.port), + api_version=broker.broker_version, + request_timeout_ms=5000, + ) + broker.attach(manager) + net.start() + try: + manager.bootstrap(timeout_ms=5000) + assert manager.bootstrapped + assert manager.cluster.brokers() + + async def do_send(): + return await manager.send(MetadataRequest[0]([])) + resp = net.run(do_send) + assert resp is not None + assert broker.requests_received > 0 + finally: + manager.close() + net.close() diff --git a/test/net/test_backend.py b/test/net/test_backend.py index 70aa9b6eb..c86318d09 100644 --- a/test/net/test_backend.py +++ b/test/net/test_backend.py @@ -118,11 +118,12 @@ def test_unknown_name_raises(self): with pytest.raises(ValueError, match='Unknown net backend'): resolve_backend('bogus', {}) - def test_asyncio_name_unregistered_raises(self): - # In Phase-1/Step-3 the asyncio backend is not registered yet; an - # explicit request for it is a hard error (an auto-detect is not). - with pytest.raises(ValueError, match='Unknown net backend'): - resolve_backend('asyncio', {}) + def test_asyncio_name_resolves(self): + # net='asyncio' lazily imports + registers the asyncio backend. + from kafka.net.asyncio_backend import AsyncioBackend + b = resolve_backend('asyncio', {'client_id': 'x'}) + assert isinstance(b, AsyncioBackend) + b.close() def test_non_backend_instance_raises(self): with pytest.raises(TypeError): @@ -152,13 +153,24 @@ async def main(): assert asyncio.run(main()) is sentinel - def test_autodetect_falls_back_when_unregistered_in_loop(self, clean_registry): - _BACKENDS.pop('asyncio', None) # ensure not registered + def test_autodetect_asyncio_in_loop_returns_asyncio_backend(self): + # In a running asyncio loop with no explicit net, auto-detect lazily + # registers + selects the asyncio backend (Phase-1: still own thread). + from kafka.net.asyncio_backend import AsyncioBackend async def main(): - return resolve_backend(None, {}) + return resolve_backend(None, {'client_id': 'auto'}) - assert isinstance(asyncio.run(main()), NetworkSelector) + b = asyncio.run(main()) + assert isinstance(b, AsyncioBackend) + b.close() + + def test_autodetect_falls_back_for_unknown_framework(self, monkeypatch): + # A detected-but-unregistered framework (e.g. trio, no backend) falls + # back to the default selector rather than erroring. + import kafka.net.backend as backend_mod + monkeypatch.setattr(backend_mod, '_detect_async_library', lambda: 'trio') + assert isinstance(resolve_backend(None, {}), NetworkSelector) def test_no_running_loop_defaults_to_selector(self): assert isinstance(resolve_backend(None, {}), NetworkSelector)