diff --git a/.changelog/5181.added b/.changelog/5181.added new file mode 100644 index 0000000000..a9b9de1c3a --- /dev/null +++ b/.changelog/5181.added @@ -0,0 +1 @@ +`opentelemetry-instrumentation-confluent-kafka`: add support for `AIOProducer` and `AIOConsumer` diff --git a/docs/nitpick-exceptions.ini b/docs/nitpick-exceptions.ini index ac6e29690a..d1e9009986 100644 --- a/docs/nitpick-exceptions.ini +++ b/docs/nitpick-exceptions.ini @@ -13,6 +13,18 @@ py-class= opentelemetry.sdk.trace.id_generator.IdGenerator opentelemetry.instrumentation.confluent_kafka.ProxiedProducer opentelemetry.instrumentation.confluent_kafka.ProxiedConsumer + opentelemetry.instrumentation.confluent_kafka.AutoInstrumentedAIOProducer + opentelemetry.instrumentation.confluent_kafka.AutoInstrumentedAIOConsumer + confluent_kafka.aio.producer._AIOProducer.AIOProducer + confluent_kafka.aio._AIOConsumer.AIOConsumer + confluent_kafka.aio.AIOProducer + confluent_kafka.aio.AIOConsumer + confluent_kafka.experimental.aio.producer._AIOProducer.AIOProducer + confluent_kafka.experimental.aio._AIOConsumer.AIOConsumer + confluent_kafka.experimental.aio.AIOProducer + confluent_kafka.experimental.aio.AIOConsumer + AIOProducer + AIOConsumer opentelemetry.instrumentation.instrumentor.BaseInstrumentor ; - AwsXRayIdGenerator TextMapPropagator diff --git a/instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/__init__.py b/instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/__init__.py index c08aa2f8cc..fccc881e39 100644 --- a/instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/__init__.py +++ b/instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/__init__.py @@ -99,9 +99,24 @@ def instrument_consumer(consumer: Consumer, tracer_provider=None): ... from collections.abc import Collection -import confluent_kafka +import confluent_kafka # pylint: disable=import-self import wrapt -from confluent_kafka import Consumer, Producer +from confluent_kafka import Consumer, Producer # pylint: disable=import-self + +try: + import confluent_kafka.aio as _aio_module # pylint: disable=no-name-in-module + from confluent_kafka.aio import AIOConsumer, AIOProducer # pylint: disable=no-name-in-module +except ImportError: + try: + import confluent_kafka.experimental.aio as _aio_module # pylint: disable=no-name-in-module + from confluent_kafka.experimental.aio import ( # pylint: disable=no-name-in-module + AIOConsumer, + AIOProducer, + ) + except ImportError: + AIOProducer = None + AIOConsumer = None + _aio_module = None from opentelemetry import context, propagate, trace from opentelemetry.instrumentation.instrumentor import BaseInstrumentor @@ -130,9 +145,10 @@ def _capture_config(args, kwargs): ``conf=`` kwarg, or (for Consumer) expanded as **kwargs.""" if args and isinstance(args[0], dict): return args[0] - conf = kwargs.get("conf") - if isinstance(conf, dict): - return conf + for key in ("conf", "producer_conf", "consumer_conf"): + conf = kwargs.get(key) + if isinstance(conf, dict): + return conf # confluent_kafka.Consumer also supports Consumer(**conf) — in that case # the kwargs themselves are the config. if kwargs: @@ -169,6 +185,52 @@ def close(self): # pylint: disable=useless-super-delegation return super().close() +# Instrument asynchronous AIOProducer and AIOConsumer if available in confluent-kafka +if AIOProducer is not None: + + class AutoInstrumentedAIOProducer(AIOProducer): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.config = _capture_config(args, kwargs) + + # Deliberately implemented so wrapt can intercept the async produce call + async def produce(self, topic, value=None, *args, **kwargs): # pylint: disable=keyword-arg-before-vararg + return await super().produce(topic, value, *args, **kwargs) +else: + AutoInstrumentedAIOProducer = None + + +if AIOConsumer is not None: + + class AutoInstrumentedAIOConsumer(AIOConsumer): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.config = _capture_config(args, kwargs) + # Token for the span re-attached to the caller's async task context. + # Managed by wrap_aio_poll / wrap_aio_consume / wrap_aio_close. + self._aio_context_token = None + inner = getattr(self, "_consumer", None) + if inner is not None: + try: + inner._is_aio = True + except AttributeError: + pass + + # Deliberately implemented so wrapt can intercept the async poll call + async def poll(self, timeout=-1): + return await super().poll(timeout) + + # Deliberately implemented so wrapt can intercept the async consume call + async def consume(self, *args, **kwargs): # pylint: disable=useless-super-delegation + return await super().consume(*args, **kwargs) + + # Deliberately implemented so wrapt can intercept the async close call + async def close(self): + return await super().close() +else: + AutoInstrumentedAIOConsumer = None + + class ProxiedProducer(Producer): def __init__(self, producer: Producer, tracer: Tracer): self._producer = producer @@ -297,6 +359,21 @@ def _instrument(self, **kwargs): confluent_kafka.Producer = AutoInstrumentedProducer confluent_kafka.Consumer = AutoInstrumentedConsumer + if _aio_module is not None: + if hasattr(_aio_module, "AIOProducer") and AutoInstrumentedAIOProducer is not None: + self._original_kafka_aio_producer = _aio_module.AIOProducer + _aio_module.AIOProducer = AutoInstrumentedAIOProducer + if hasattr(_aio_module, "AIOConsumer") and AutoInstrumentedAIOConsumer is not None: + self._original_kafka_aio_consumer = _aio_module.AIOConsumer + _aio_module.AIOConsumer = AutoInstrumentedAIOConsumer + + if hasattr(confluent_kafka, "AIOProducer") and AutoInstrumentedAIOProducer is not None: + self._original_kafka_top_aio_producer = confluent_kafka.AIOProducer + confluent_kafka.AIOProducer = AutoInstrumentedAIOProducer + if hasattr(confluent_kafka, "AIOConsumer") and AutoInstrumentedAIOConsumer is not None: + self._original_kafka_top_aio_consumer = confluent_kafka.AIOConsumer + confluent_kafka.AIOConsumer = AutoInstrumentedAIOConsumer + tracer_provider = kwargs.get("tracer_provider") tracer = trace.get_tracer( __name__, @@ -342,13 +419,67 @@ def _inner_wrap_close(func, instance, args, kwargs): "close", _inner_wrap_close, ) + if AutoInstrumentedAIOProducer is not None: + + def _inner_wrap_aio_produce(func, instance, args, kwargs): + return ConfluentKafkaInstrumentor.wrap_aio_produce(func, instance, self._tracer, args, kwargs) + + wrapt.wrap_function_wrapper( + AutoInstrumentedAIOProducer, + "produce", + _inner_wrap_aio_produce, + ) + + if AutoInstrumentedAIOConsumer is not None: + + def _inner_wrap_aio_poll(func, instance, args, kwargs): + return ConfluentKafkaInstrumentor.wrap_aio_poll(func, instance, self._tracer, args, kwargs) + + def _inner_wrap_aio_consume(func, instance, args, kwargs): + return ConfluentKafkaInstrumentor.wrap_aio_consume(func, instance, self._tracer, args, kwargs) + + def _inner_wrap_aio_close(func, instance, args, kwargs): + return ConfluentKafkaInstrumentor.wrap_aio_close(func, instance, args, kwargs) + + wrapt.wrap_function_wrapper( + AutoInstrumentedAIOConsumer, + "poll", + _inner_wrap_aio_poll, + ) + wrapt.wrap_function_wrapper( + AutoInstrumentedAIOConsumer, + "consume", + _inner_wrap_aio_consume, + ) + wrapt.wrap_function_wrapper( + AutoInstrumentedAIOConsumer, + "close", + _inner_wrap_aio_close, + ) def _uninstrument(self, **kwargs): confluent_kafka.Producer = self._original_kafka_producer confluent_kafka.Consumer = self._original_kafka_consumer + if _aio_module is not None: + if hasattr(self, "_original_kafka_aio_producer") and self._original_kafka_aio_producer is not None: + _aio_module.AIOProducer = self._original_kafka_aio_producer + if hasattr(self, "_original_kafka_aio_consumer") and self._original_kafka_aio_consumer is not None: + _aio_module.AIOConsumer = self._original_kafka_aio_consumer + + if hasattr(self, "_original_kafka_top_aio_producer") and self._original_kafka_top_aio_producer is not None: + confluent_kafka.AIOProducer = self._original_kafka_top_aio_producer + if hasattr(self, "_original_kafka_top_aio_consumer") and self._original_kafka_top_aio_consumer is not None: + confluent_kafka.AIOConsumer = self._original_kafka_top_aio_consumer + unwrap(AutoInstrumentedProducer, "produce") unwrap(AutoInstrumentedConsumer, "poll") + if AutoInstrumentedAIOProducer is not None: + unwrap(AutoInstrumentedAIOProducer, "produce") + if AutoInstrumentedAIOConsumer is not None: + unwrap(AutoInstrumentedAIOConsumer, "poll") + unwrap(AutoInstrumentedAIOConsumer, "consume") + unwrap(AutoInstrumentedAIOConsumer, "close") @staticmethod def wrap_produce(func, instance, tracer, args, kwargs): @@ -399,7 +530,10 @@ def wrap_poll(func, instance, tracer, args, kwargs): "recv", end_on_exit=True, kind=trace.SpanKind.CONSUMER, attributes=receive_attributes ): _create_new_consume_span(instance, tracer, [record], attributes) - instance._current_context_token = context.attach(trace.set_span_in_context(instance._current_consume_span)) + if not getattr(instance, "_is_aio", False): + instance._current_context_token = context.attach( + trace.set_span_in_context(instance._current_consume_span) + ) return record @@ -424,12 +558,98 @@ def wrap_consume(func, instance, tracer, args, kwargs): "recv", end_on_exit=True, kind=trace.SpanKind.CONSUMER, attributes=receive_attributes ): _create_new_consume_span(instance, tracer, records, attributes) - instance._current_context_token = context.attach(trace.set_span_in_context(instance._current_consume_span)) + if not getattr(instance, "_is_aio", False): + instance._current_context_token = context.attach( + trace.set_span_in_context(instance._current_consume_span) + ) return records @staticmethod def wrap_close(func, instance, args, kwargs): + if instance._current_consume_span: _end_current_consume_span(instance) func(*args, **kwargs) + + @staticmethod + async def wrap_aio_produce(func, instance, tracer, args, kwargs): + topic = kwargs.get("topic") + if not topic: + topic = args[0] + + span_name = _get_span_name("send", topic) + topic = KafkaPropertiesExtractor.extract_produce_topic(args, kwargs) + bootstrap_servers = KafkaPropertiesExtractor.extract_bootstrap_servers(instance) + attributes = _get_attributes( + topic, + operation=MessagingOperationTypeValues.PUBLISH, + bootstrap_servers=bootstrap_servers, + ) + with tracer.start_as_current_span(name=span_name, kind=trace.SpanKind.PRODUCER, attributes=attributes): + # Note: confluent-kafka's AIOProducer does not currently support the + # headers keyword argument and raises NotImplementedError if headers are + # passed when none were provided. Trace context propagation via headers + # is therefore applied only when the caller explicitly passes headers. + headers = KafkaPropertiesExtractor.extract_produce_headers(args, kwargs) + if headers is not None: + propagate.inject( + headers, + setter=_kafka_setter, + ) + + return await func(*args, **kwargs) + + @staticmethod + async def wrap_aio_poll(func, instance, tracer, args, kwargs): + # Detach any span attached to the caller's async context from a + # previous poll call. Keeping the token lifetime to a single poll + # prevents stale context if the caller never calls close(). + if getattr(instance, "_aio_context_token", None): + context.detach(instance._aio_context_token) + instance._aio_context_token = None + + # The real AIOConsumer runs Consumer.poll() inside a ThreadPoolExecutor + # via loop.run_in_executor(). Any context.attach() done in that worker + # thread is confined to the thread's context copy and is NOT visible in + # the caller's async task. We therefore await first and then + # re-attach the span to the caller's async task context ourselves. + record = await func(*args, **kwargs) + + if record: + # Read the span created by the instrumented inner Consumer on the + # worker thread and attach it to the caller's async task context. + inner = getattr(instance, "_consumer", None) + if inner is not None: + span = getattr(inner, "_current_consume_span", None) + if span is not None: + instance._aio_context_token = context.attach(trace.set_span_in_context(span)) + + return record + + @staticmethod + async def wrap_aio_consume(func, instance, tracer, args, kwargs): + # Detach any span attached to the caller's async context from a + # previous consume or poll call. + if getattr(instance, "_aio_context_token", None): + context.detach(instance._aio_context_token) + instance._aio_context_token = None + + records = await func(*args, **kwargs) + + if records and len(records) > 0: + inner = getattr(instance, "_consumer", None) + if inner is not None: + span = getattr(inner, "_current_consume_span", None) + if span is not None: + instance._aio_context_token = context.attach(trace.set_span_in_context(span)) + + return records + + @staticmethod + async def wrap_aio_close(func, instance, args, kwargs): + # Detach any span still attached to the async task before closing. + if getattr(instance, "_aio_context_token", None): + context.detach(instance._aio_context_token) + instance._aio_context_token = None + return await func(*args, **kwargs) diff --git a/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py b/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py index 7b7b7c3dee..8a05a9d618 100644 --- a/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py +++ b/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py @@ -1,10 +1,10 @@ # Copyright The OpenTelemetry Authors # SPDX-License-Identifier: Apache-2.0 - -# pylint: disable=no-name-in-module,import-outside-toplevel - +# pylint: disable=no-name-in-module,import-outside-toplevel,too-many-public-methods +import asyncio from unittest import mock +from opentelemetry import trace from opentelemetry.instrumentation.confluent_kafka import ( AutoInstrumentedConsumer, AutoInstrumentedProducer, @@ -606,3 +606,224 @@ def test_consumer_sets_bootstrap_servers_attributes(self) -> None: process_span = next(s for s in self.memory_exporter.get_finished_spans() if s.name == "topic-1 process") self.assertEqual(process_span.attributes[SERVER_ADDRESS], "broker-1") self.assertEqual(process_span.attributes[SERVER_PORT], 9092) + + def test_aio_produce(self) -> None: + from .utils import MockedAIOProducer # noqa: PLC0415 + + async def _run(): + instrumentation = ConfluentKafkaInstrumentor() + tracer = self.tracer_provider.get_tracer(__name__) + producer = MockedAIOProducer({"bootstrap.servers": "localhost:9092"}) + + # Without headers: verify produce succeeds and headers kwarg is NOT forced + # (confluent-kafka AIOProducer raises NotImplementedError if headers is passed) + call_kwargs = {"topic": "test-topic", "value": b"test-value"} + await instrumentation.wrap_aio_produce( + producer.produce, + producer, + tracer, + args=(), + kwargs=call_kwargs, + ) + self.assertNotIn("headers", call_kwargs) + spans = self.memory_exporter.get_finished_spans() + self.assertEqual(len(spans), 1) + self.assertEqual(spans[0].name, "test-topic send") + + # With headers: verify traceparent is injected into supplied headers + self.memory_exporter.clear() + supplied_headers = [] + await instrumentation.wrap_aio_produce( + producer.produce, + producer, + tracer, + args=(), + kwargs={"topic": "test-topic", "value": b"test-value", "headers": supplied_headers}, + ) + self.assertTrue(any("traceparent" in str(h) for h in supplied_headers)) + + asyncio.run(_run()) + + def test_aio_poll(self) -> None: + from .utils import MockedAIOConsumer # noqa: PLC0415 + + async def _run(): + instrumentation = ConfluentKafkaInstrumentor() + tracer = self.tracer_provider.get_tracer(__name__) + + # MockedAIOConsumer dispatches the inner poll to a ThreadPoolExecutor, + # matching the real AIOConsumer's behaviour. This exercises the path + # where the span is created on a worker thread and must be re-attached + # to the caller's async-task context by wrap_aio_poll. + queue = [MockedMessage("test-topic", 0, 0, [])] + consumer = MockedAIOConsumer( + queue, + {"bootstrap.servers": "localhost:9092", "group.id": "g", "auto.offset.reset": "earliest"}, + ) + + # Instrument the inner consumer's poll to simulate AutoInstrumentedConsumer. + inner = consumer._consumer + raw_poll = inner.poll + + def instrumented_inner_poll(timeout=None): + return ConfluentKafkaInstrumentor.wrap_poll(raw_poll, inner, tracer, (timeout,), {}) + + inner.poll = instrumented_inner_poll + + # First poll: inner Consumer creates a span on a worker thread; + # wrap_aio_poll re-attaches it to the caller's async task context. + msg = await instrumentation.wrap_aio_poll( + consumer.poll, + consumer, + tracer, + args=(), + kwargs={}, + ) + self.assertIsNotNone(msg) + + # The span must now be current in the caller's async task. + active_span = trace.get_current_span() + self.assertTrue(active_span.is_recording()) + self.assertEqual(active_span.name, "test-topic process") + + # Verify worker thread context is not leaked to executor workers. + loop = asyncio.get_running_loop() + worker_span = await loop.run_in_executor(consumer._executor, trace.get_current_span) + self.assertFalse(worker_span.is_recording()) + + # Second (empty) poll detaches the previous context and ends the span. + await instrumentation.wrap_aio_poll( + consumer.poll, + consumer, + tracer, + args=(), + kwargs={}, + ) + + spans = self.memory_exporter.get_finished_spans() + process_spans = [s for s in spans if s.name == "test-topic process"] + # Exactly one process span — no duplicates from the AIO wrapper layer. + self.assertEqual(len(process_spans), 1) + + # Context detached in async task: + self.assertFalse(trace.get_current_span().is_recording()) + + await consumer.close() + + asyncio.run(_run()) + + def test_aio_consume(self) -> None: + from .utils import MockedAIOConsumer # noqa: PLC0415 + + async def _run(): + instrumentation = ConfluentKafkaInstrumentor() + tracer = self.tracer_provider.get_tracer(__name__) + + queue = [ + MockedMessage("test-topic", 0, 0, []), + MockedMessage("test-topic", 0, 1, []), + ] + consumer = MockedAIOConsumer( + queue, + {"bootstrap.servers": "localhost:9092", "group.id": "g", "auto.offset.reset": "earliest"}, + ) + + inner = consumer._consumer + raw_consume = inner.consume + + def instrumented_inner_consume(*args, **kwargs): + return ConfluentKafkaInstrumentor.wrap_consume(raw_consume, inner, tracer, args, kwargs) + + inner.consume = instrumented_inner_consume + + # First batch consume + msgs = await instrumentation.wrap_aio_consume( + consumer.consume, + consumer, + tracer, + args=(2,), + kwargs={}, + ) + self.assertEqual(len(msgs), 2) + + # The span must now be current in the caller's async task. + active_span = trace.get_current_span() + self.assertTrue(active_span.is_recording()) + self.assertEqual(active_span.name, "test-topic process") + + # Subsequent empty consume ends the previous span and detaches context. + await instrumentation.wrap_aio_consume( + consumer.consume, + consumer, + tracer, + args=(1,), + kwargs={}, + ) + + spans = self.memory_exporter.get_finished_spans() + process_spans = [s for s in spans if s.name == "test-topic process"] + self.assertEqual(len(process_spans), 1) + + # Context detached in async task: + self.assertFalse(trace.get_current_span().is_recording()) + + await consumer.close() + + asyncio.run(_run()) + + def test_aio_config_kwargs(self) -> None: + from opentelemetry.instrumentation.confluent_kafka import ( # noqa: PLC0415 + AutoInstrumentedAIOConsumer, + AutoInstrumentedAIOProducer, + ) + + async def _run(): + instrumentor = ConfluentKafkaInstrumentor() + instrumentor.instrument() + try: + if AutoInstrumentedAIOProducer is not None: + prod = AutoInstrumentedAIOProducer(producer_conf={"bootstrap.servers": "broker-p:9092"}) + self.assertEqual(prod.config.get("bootstrap.servers"), "broker-p:9092") + await prod.close() + + if AutoInstrumentedAIOConsumer is not None: + cons = AutoInstrumentedAIOConsumer( + consumer_conf={"bootstrap.servers": "broker-c:9092", "group.id": "g"} + ) + self.assertEqual(cons.config.get("bootstrap.servers"), "broker-c:9092") + inner = getattr(cons, "_consumer", None) + if inner is not None: + self.assertTrue(getattr(inner, "_is_aio", False)) + await cons.close() + finally: + instrumentor.uninstrument() + + asyncio.run(_run()) + + def test_aio_instrument_and_uninstrument(self) -> None: + from opentelemetry.instrumentation.confluent_kafka import ( # noqa: PLC0415 + AutoInstrumentedAIOConsumer, + AutoInstrumentedAIOProducer, + _aio_module, + ) + + if _aio_module is None or AutoInstrumentedAIOProducer is None: + return + + instrumentor = ConfluentKafkaInstrumentor() + orig_producer = getattr(_aio_module, "AIOProducer", None) + orig_consumer = getattr(_aio_module, "AIOConsumer", None) + + instrumentor.instrument() + try: + if orig_producer is not None: + self.assertIs(_aio_module.AIOProducer, AutoInstrumentedAIOProducer) + if orig_consumer is not None: + self.assertIs(_aio_module.AIOConsumer, AutoInstrumentedAIOConsumer) + finally: + instrumentor.uninstrument() + + if orig_producer is not None: + self.assertIs(_aio_module.AIOProducer, orig_producer) + if orig_consumer is not None: + self.assertIs(_aio_module.AIOConsumer, orig_consumer) diff --git a/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/utils.py b/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/utils.py index 86c360eeb9..937f0c7fe2 100644 --- a/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/utils.py +++ b/instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/utils.py @@ -1,6 +1,8 @@ # Copyright The OpenTelemetry Authors # SPDX-License-Identifier: Apache-2.0 +import asyncio +import concurrent.futures from confluent_kafka import Consumer, Producer @@ -9,6 +11,8 @@ class MockConsumer(Consumer): def __init__(self, queue, config): self._queue = queue self.config = config + self._current_consume_span = None + self._current_context_token = None super().__init__(config) def consume(self, num_messages=1, *args, **kwargs): # pylint: disable=keyword-arg-before-vararg @@ -81,3 +85,48 @@ def poll(self, *args, **kwargs): def flush(self, *args, **kwargs): return len(self._queue) + + +class MockedAIOProducer: + def __init__(self, config=None): + self._queue = [] + self.config = config or {} + + async def produce(self, topic, value=None, *args, **kwargs): # pylint: disable=keyword-arg-before-vararg + msg = MockedMessage( + topic=topic, + partition=0, + offset=0, + headers=kwargs.get("headers", []), + key=kwargs.get("key"), + value=value, + ) + self._queue.append(msg) + return msg + + +class MockedAIOConsumer: + """Mock of confluent_kafka.AIOConsumer. + + Dispatches the inner Consumer poll to a ThreadPoolExecutor to replicate + the thread-pool behaviour of the real AIOConsumer, so that tests cover + the async-task context-management in wrap_aio_poll. + """ + + def __init__(self, queue, config=None): + self._consumer = MockConsumer(queue, config or {}) + self._consumer._is_aio = True + self.config = config or {} + self._aio_context_token = None + self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + + async def poll(self, timeout=None): + loop = asyncio.get_running_loop() + return await loop.run_in_executor(self._executor, self._consumer.poll, timeout) + + async def consume(self, num_messages=1, *args, **kwargs): # pylint: disable=keyword-arg-before-vararg + loop = asyncio.get_running_loop() + return await loop.run_in_executor(self._executor, self._consumer.consume, num_messages, *args, **kwargs) + + async def close(self): + self._executor.shutdown(wait=False)