Repository navigation
feat(confluent-kafka): patch AIOProducer and AIOConsumer - #5181
Heisenbrick wants to merge 32 commits into
Conversation
|
|
Pull request dashboard statusWaiting on the author · refreshed 2026-11-11 14:18 UTC Respond to 6 review items (e.g. link a commit, explain why not, ask a follow-up): Status above doesn't look right?
|
|
@Heisenbrick please fix the failing jobs |
- Remove trailing whitespace in __init__.py (ruff W291) - Fix blank line counts: E302 (2 before top-level class) and E303 (too many blank lines inside class body) - Add blank line before _uninstrument after wrapt block (E301) - Remove extra blank line before MockedAIOProducer in tests/utils.py - Fix changelog fragment line-ending (LF, not CRLF) - Add AutoInstrumentedAIOProducer/AIOConsumer to docs/nitpick-exceptions.ini so Sphinx --nitpicky does not fail on unresolved class references
|
@xrmx Thank you for taking a look! I have pushed a fix commit addressing the lint, formatting, and documentation checks:
Please let me know if anything else needs adjustments! |
- Replace public AIOProducer and AIOConsumer on module and restore on uninstrument - Do not force kwargs['headers'] = [] in wrap_aio_produce if headers is None, avoiding NotImplementedError - Add test_aio_instrument_and_uninstrument and update test_aio_produce
|
@JacksonWeber Thank you for the insightful review and catching these points! I have pushed a new commit addressing both items:
All unit tests and linters are passing locally. Please let me know if any further adjustments are needed! |
…1113 and CRLF line endings
|
Hi @JacksonWeber, thank you for the feedback! I have addressed your review comment — wrap_aio_poll is now a pass-through since AIOConsumer already delegates to an inner Consumer that is instrumented via AutoInstrumentedConsumer, so adding a second span layer at the AIO level was producing duplicates. I have also removed _current_consume_span from AutoInstrumentedAIOConsumer since span lifecycle is now fully managed by the inner Consumer layer. Additionally, I have updated test_aio_poll to mirror the real AIOConsumer → Consumer structure — MockedAIOConsumer now wraps an inner MockConsumer whose poll is instrumented via wrap_poll, and the test asserts that exactly one process span is created with no duplicates from the AIO wrapper. The CI failures from the previous commit (license header position, pylint W1113/C0328, and CRLF line endings) have also been fixed in this push. |
|
I have addressed your review comment — wrap_aio_poll is now a pass-through since AIOConsumer already delegates to an inner Consumer that is instrumented via AutoInstrumentedConsumer, so adding a second span layer at the AIO level was producing duplicates. I have also removed _current_consume_span from AutoInstrumentedAIOConsumer since span lifecycle is now fully managed by the inner Consumer layer. Additionally, I have updated test_aio_poll to mirror the real AIOConsumer → Consumer structure — MockedAIOConsumer now wraps an inner MockConsumer whose poll is instrumented via wrap_poll, and the test asserts that exactly one process span is created with no duplicates from the AIO wrapper. The CI failures from the previous commit (license header position, pylint W1113/C0328, and CRLF line endings) have also been fixed in this push. |
This comment has been minimized.
This comment has been minimized.
|
@emdneto @JacksonWeber — I've pushed an update in commit 43c2693 addressing the async context feedback. wrap_aio_poll now re-attaches the process span to the caller's async task context after the inner Consumer.poll() returns from the ThreadPoolExecutor, rather than leaving it confined to the worker thread's context copy. The thread-level token is cleared beforehand to prevent cross-thread detach errors, and wrap_aio_close cleans up the async context token on close. MockedAIOConsumer now uses a real ThreadPoolExecutor via loop.run_in_executor() to match the actual AIOConsumer behaviour, and test_aio_poll asserts trace.get_current_span().name == "test-topic process" from the caller's async task — confirming context is correctly propagated out of the worker thread. Happy to make any further changes if needed. Thanks for the reviews! |
|
@emdneto @JacksonWeber — pushed one more fix in commit 8967bfc to resolve the CI / lint / instrumentation-confluent-kafka failure. Pylint was flagging four issues in the AIO-related code: W0406 (import-self) on import confluent_kafka and from confluent_kafka import ... — false positive caused by pylint confusing the instrumentation module's directory name with the installed confluent_kafka package. Suppressed with inline # pylint: disable=import-self. |
|
@JacksonWeber @emdneto — I have updated the branch with the fixes for review feedback and CI checks:
All 19 tests, linters, and documentation builds are verified green locally. Whenever you have a chance, could you please approve the workflow run? Thank you! |
There was a problem hiding this comment.
🟡 Changes recommended
The async consumer leaves ended process spans attached to reused executor workers, corrupting context for subsequent polls, while batch-consume context handling and producer trace propagation also remain incomplete.
5 open issues
What changed in this PR
Adds asynchronous Confluent Kafka instrumentation for AIOProducer and AIOConsumer, extending the package’s existing synchronous client support toward issue #5115.
Changes:
- Adds optional async-client imports, instrumented subclasses, and wrapper registration and removal.
- Adds producer spans and consumer caller-task context handling, with async mocks and tests.
- Updates documentation reference exceptions and adds a changelog entry.
| File | Description |
|---|---|
instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/utils.py |
Adds async producer and executor-backed consumer mocks. |
instrumentation/opentelemetry-instrumentation-confluent-kafka/tests/test_instrumentation.py |
Tests async produce, poll, and instrumentation registration. |
instrumentation/opentelemetry-instrumentation-confluent-kafka/src/opentelemetry/instrumentation/confluent_kafka/__init__.py |
Adds optional async-client patching and telemetry wrappers. |
docs/nitpick-exceptions.ini |
Adds reference exceptions for async client classes. |
.changelog/5181.added |
Announces async Confluent Kafka support. |
Review effort: Balanced
| class AutoInstrumentedAIOProducer(AIOProducer): | ||
| def __init__(self, *args, **kwargs): | ||
| super().__init__(*args, **kwargs) | ||
| self.config = _capture_config(args, kwargs) |
| wrapt.wrap_function_wrapper( | ||
| AutoInstrumentedAIOConsumer, | ||
| "poll", | ||
| _inner_wrap_aio_poll, | ||
| ) |
| headers = KafkaPropertiesExtractor.extract_produce_headers(args, kwargs) | ||
| if headers is not None: |
| def instrumented_inner_poll(timeout=None): | ||
| return ConfluentKafkaInstrumentor.wrap_poll(raw_poll, inner, tracer, (timeout,), {}) | ||
|
|
||
| inner.poll = instrumented_inner_poll |
|
@aabmass @JacksonWeber — I have pushed commit 29a8ecb addressing all items raised in the review:
Could you please approve and trigger the workflow run when you have a moment? Thank you! |
| from collections.abc import Collection | ||
|
|
||
| import confluent_kafka | ||
| import confluent_kafka # pylint: disable=import-self |
| 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 |
There was a problem hiding this comment.
Do we really need to try the experimental imports if some version has already the stable one? I mean it sounds legit to add support for something that is considered stable and not experimental
|
|
||
|
|
||
| # Instrument asynchronous AIOProducer and AIOConsumer if available in confluent-kafka | ||
| if AIOProducer is not None: |
There was a problem hiding this comment.
Given the import is shared can we use a shared if _aio_module instead?
Description
This PR adds asynchronous instrumentation support for
AIOProducerandAIOConsumerinopentelemetry-instrumentation-confluent-kafka.Previously, the instrumentation only patched the synchronous
ProducerandConsumerclasses. When using asynchronous Kafka clients (confluent_kafka.aio), no spans were generated and distributed trace contexts were not propagated across message boundaries.This change:
AIOProducerandAIOConsumer(with fallback/graceful handling if.aiois not available in olderconfluent-kafkaversions).AutoInstrumentedAIOProducerandAutoInstrumentedAIOConsumerto wrap async calls.wrap_aio_produce,wrap_aio_poll, andwrap_aio_closeto start spans and inject/extract trace headers.Fixes #5115
Type of change
How Has This Been Tested?
test_aio_produce: verifies async produce generates a PRODUCER span and injects trace headers into the Kafka message.test_aio_poll: verifies async poll receives message and creates the CONSUMER process span.Does This PR Require a Core Repo Change?
Checklist: