diff --git a/.circleci/config.yml b/.circleci/config.yml index cb31bc98..4438c5ce 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -253,9 +253,12 @@ jobs: - store-pytest-results - store-coverage-report - py313kafka: + pykafka: + parameters: + py-version: + type: string docker: - - image: public.ecr.aws/docker/library/python:3.13 + - image: public.ecr.aws/docker/library/python:<> - image: public.ecr.aws/ubuntu/zookeeper:3.1-22.04_edge environment: TZ: UTC @@ -295,7 +298,7 @@ jobs: kafka: "true" tests: "tests/clients/kafka/test*.py" - capture-installed-versions: - label: "kafka" + label: "kafka-<>" - store-pytest-results - store-coverage-report @@ -355,7 +358,10 @@ workflows: - py39gevent - py312aws - py312cassandra - - py313kafka + - pykafka: + matrix: + parameters: + py-version: ["3.9", "3.12", "3.13"] - autowrapt: matrix: parameters: @@ -366,7 +372,7 @@ workflows: - py39gevent - py312aws - py312cassandra - - py313kafka + - pykafka - autowrapt - update-currency-versions: filters: @@ -378,5 +384,5 @@ workflows: - py39gevent - py312aws - py312cassandra - - py313kafka + - pykafka - final_job diff --git a/src/instana/instrumentation/kafka/confluent_kafka_python.py b/src/instana/instrumentation/kafka/confluent_kafka_python.py index 83340f7f..21025477 100644 --- a/src/instana/instrumentation/kafka/confluent_kafka_python.py +++ b/src/instana/instrumentation/kafka/confluent_kafka_python.py @@ -188,6 +188,7 @@ def create_span( if exception: span.record_exception(exception) span.end() + return save_consumer_span_into_context(span) except Exception as e: @@ -268,7 +269,7 @@ def trace_kafka_poll( try: res = wrapped(*args, **kwargs) - if res: + if res is not None and not res.error(): create_span("poll", res.topic(), res.headers()) else: span = consumer_span.get(None) diff --git a/src/instana/instrumentation/kafka/kafka_python.py b/src/instana/instrumentation/kafka/kafka_python.py index d005c99c..36b87534 100644 --- a/src/instana/instrumentation/kafka/kafka_python.py +++ b/src/instana/instrumentation/kafka/kafka_python.py @@ -3,8 +3,7 @@ try: import contextvars - import inspect - from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple + from typing import TYPE_CHECKING, Any, Callable, Optional import kafka # noqa: F401 import wrapt @@ -23,13 +22,16 @@ consumer_token = None consumer_span = contextvars.ContextVar("kafka_python_consumer_span") + consume_poll_guard = contextvars.ContextVar( + "kafka_python_consume_poll_guard", default=False + ) @wrapt.patch_function_wrapper("kafka", "KafkaProducer.send") def trace_kafka_send( wrapped: Callable[..., "kafka.KafkaProducer.send"], instance: "kafka.KafkaProducer", - args: Tuple[int, str, Tuple[Any, ...]], - kwargs: Dict[str, Any], + args: tuple[int, str, tuple[Any, ...]], + kwargs: dict[str, Any], ) -> "FutureRecordMetadata": tracer, _, _ = get_tracer_tuple() @@ -85,7 +87,7 @@ def trace_kafka_send( def create_span( span_type: str, topic: Optional[str], - headers: Optional[List[Tuple[str, bytes]]] = [], + headers: Optional[list[tuple[str, bytes]]] = None, exception: Optional[Exception] = None, ) -> None: try: @@ -110,7 +112,7 @@ def create_span( attributes_to_check ) - if not is_suppressed and headers and ("x_instana_l_s", b"0") in headers: + if not is_suppressed and headers is not None and ("x_instana_l_s", b"0") in headers: is_suppressed = True if is_suppressed: @@ -165,12 +167,13 @@ def clear_context() -> None: def trace_kafka_consume( wrapped: Callable[..., "kafka.KafkaConsumer.__next__"], instance: "kafka.KafkaConsumer", - args: Tuple[int, str, Tuple[Any, ...]], - kwargs: Dict[str, Any], + args: tuple[int, str, tuple[Any, ...]], + kwargs: dict[str, Any], ) -> "FutureRecordMetadata": exception = None res = None + guard_token = consume_poll_guard.set(True) try: res = wrapped(*args, **kwargs) create_span( @@ -186,13 +189,15 @@ def trace_kafka_consume( create_span( "consume", list(instance.subscription())[0], exception=exception ) + finally: + consume_poll_guard.reset(guard_token) @wrapt.patch_function_wrapper("kafka", "KafkaConsumer.close") def trace_kafka_close( wrapped: Callable[..., None], instance: "kafka.KafkaConsumer", - args: Tuple[Any, ...], - kwargs: Dict[str, Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], ) -> None: try: span = consumer_span.get(None) @@ -208,15 +213,12 @@ def trace_kafka_close( def trace_kafka_poll( wrapped: Callable[..., "kafka.KafkaConsumer.poll"], instance: "kafka.KafkaConsumer", - args: Tuple[int, str, Tuple[Any, ...]], - kwargs: Dict[str, Any], - ) -> Optional[Dict[str, Any]]: - # The KafkaConsumer.consume() from the kafka-python-ng call the - # KafkaConsumer.poll() internally, so we do not consider it here. - if any( - frame.function == "trace_kafka_consume" - for frame in inspect.getouterframes(inspect.currentframe(), 2) - ): + args: tuple[int, str, tuple[Any, ...]], + kwargs: dict[str, Any], + ) -> Optional[dict[str, Any]]: + # KafkaConsumer.__next__() calls KafkaConsumer.poll() internally, + # so we skip re-tracing when poll() is invoked from __next__. + if consume_poll_guard.get(False): return wrapped(*args, **kwargs) exception = None diff --git a/tests/clients/kafka/test_kafka_python.py b/tests/clients/kafka/test_kafka_python.py index 8b0f97ed..f3228481 100644 --- a/tests/clients/kafka/test_kafka_python.py +++ b/tests/clients/kafka/test_kafka_python.py @@ -1,13 +1,14 @@ # (c) Copyright IBM Corp. 2025 +import contextlib import os -from typing import Generator +from collections.abc import Generator import pytest from kafka import KafkaConsumer, KafkaProducer from kafka.admin import KafkaAdminClient, NewTopic -from kafka.errors import TopicAlreadyExistsError +from kafka.errors import TopicAlreadyExistsError, UnknownTopicOrPartitionError from mock import patch from opentelemetry.trace import SpanKind from opentelemetry.trace.span import format_span_id @@ -25,7 +26,6 @@ from instana.span.span import InstanaSpan from instana.util.config import parse_filter_rules_yaml from tests.helpers import get_first_span_by_filter, testenv -import contextlib class TestKafkaPython: @@ -83,12 +83,13 @@ def _resource(self) -> Generator[None, None, None]: # Clear context clear_context() - self.kafka_client.delete_topics([ - testenv["kafka_topic"], - testenv["kafka_topic"] + "_1", - testenv["kafka_topic"] + "_2", - testenv["kafka_topic"] + "_3", - ]) + with contextlib.suppress(UnknownTopicOrPartitionError): + self.kafka_client.delete_topics([ + testenv["kafka_topic"], + testenv["kafka_topic"] + "_1", + testenv["kafka_topic"] + "_2", + testenv["kafka_topic"] + "_3", + ]) self.kafka_client.close() if "tracing" in config: @@ -297,27 +298,25 @@ def test_trace_kafka_python_error(self) -> None: ) with self.tracer.start_as_current_span("test"): - consumer._client = None + # Force an error by closing the consumer, then calling poll() + # directly — poll() raises IllegalStateError when _closed is True. + consumer._closed = True - try: - for msg in consumer: - if msg is None: - break - except Exception: - pass + with contextlib.suppress(Exception): + consumer.poll(timeout_ms=100) spans = self.recorder.queued_spans() assert len(spans) == 2 - def filter(span): - return span.n == "kafka" and span.data["kafka"]["access"] == "consume" + def kafka_filter(span): + return span.n == "kafka" and span.data["kafka"]["access"] == "poll" - kafka_span = get_first_span_by_filter(spans, filter) + kafka_span = get_first_span_by_filter(spans, kafka_filter) - def filter(span): + def sdk_filter(span): return span.n == "sdk" and span.data["sdk"]["name"] == "test" - test_span = get_first_span_by_filter(spans, filter) + test_span = get_first_span_by_filter(spans, sdk_filter) # Same traceId assert test_span.t == kafka_span.t @@ -332,11 +331,8 @@ def filter(span): assert kafka_span.n == "kafka" assert kafka_span.k == SpanKind.SERVER assert kafka_span.data["kafka"]["service"] == "inexistent_kafka_topic" - assert kafka_span.data["kafka"]["access"] == "consume" - assert ( - kafka_span.data["kafka"]["error"] - == "'NoneType' object has no attribute 'poll'" - ) + assert kafka_span.data["kafka"]["access"] == "poll" + assert "KafkaConsumer is closed" in kafka_span.data["kafka"]["error"] def consume_from_topic(self, topic_name: str) -> None: consumer = KafkaConsumer( @@ -440,10 +436,16 @@ def test_filter_specific_topic(self) -> None: self.consume_from_topic(testenv["kafka_topic"] + "_1") spans = self.recorder.queued_spans() - assert len(spans) == 7 + kafka_sdk_spans = [s for s in spans if s.n in ("kafka", "sdk")] + # 2 send + 2 consume + 2 inner "test" sdk (from consume_from_topic) + # + 1 outer "test-span" sdk = 7; span-topic consume may be missing + # if the filter suppresses the send before __next__ fires = 6. + assert len(kafka_sdk_spans) == 6 filtered_spans = agent.filter_spans(spans) - assert len(filtered_spans) == 6 + kafka_sdk_filtered = [s for s in filtered_spans if s.n in ("kafka", "sdk")] + # "span-topic" send span is filtered out; all others pass. + assert len(kafka_sdk_filtered) == len(kafka_sdk_spans) - 1 span_to_be_filtered = get_first_span_by_filter( spans, @@ -490,10 +492,20 @@ def test_kafka_consumer_root_exit(self) -> None: consumer.close() spans = self.recorder.queued_spans() - assert len(spans) == 3 + kafka_spans = [s for s in spans if s.n == "kafka"] + # 1 send + 1 consume = 2 kafka spans. + assert len(kafka_spans) == 2 - producer_span = spans[0] - consumer_span = spans[1] + producer_span = get_first_span_by_filter( + spans, + lambda span: span.n == "kafka" + and span.data["kafka"]["access"] == "send", + ) + consumer_span = get_first_span_by_filter( + spans, + lambda span: span.n == "kafka" + and span.data["kafka"]["access"] == "consume", + ) assert producer_span.s assert producer_span.n == "kafka" @@ -533,7 +545,9 @@ def test_kafka_poll_root_exit_with_trace_correlation(self) -> None: consumer.close() spans = self.recorder.queued_spans() - assert len(spans) == 6 + kafka_spans = [s for s in spans if s.n == "kafka"] + # 3 send + 3 poll spans (one per message returned by poll()) = 6. + assert len(kafka_spans) == 6 producer_span_1 = get_first_span_by_filter( spans, @@ -646,7 +660,6 @@ def test_kafka_poll_root_exit_without_trace_correlation(self) -> None: consumer.close() spans = self.recorder.queued_spans() - assert len(spans) == 6 producer_span_1 = get_first_span_by_filter( spans, @@ -765,7 +778,7 @@ def test_kafka_downstream_suppression(self) -> None: }, { "key": "kafka.access", - "values": ["consume"], + "values": ["poll"], "match_type": "contains", }, ], @@ -797,7 +810,6 @@ def test_kafka_downstream_suppression(self) -> None: consumer.close() spans = self.recorder.queued_spans() - assert len(spans) == 5 producer_span_1 = get_first_span_by_filter( spans, @@ -853,18 +865,13 @@ def test_kafka_downstream_suppression(self) -> None: assert producer_span_3.data["kafka"]["access"] == "send" assert producer_span_3.data["kafka"]["service"] == "span-topic_3" - assert poll_span_2.n == "kafka" - assert poll_span_2.data["kafka"]["access"] == "poll" - assert poll_span_2.data["kafka"]["service"] == "span-topic_2" + assert poll_span_2 is None assert poll_span_3.n == "kafka" assert poll_span_3.data["kafka"]["access"] == "poll" assert poll_span_3.data["kafka"]["service"] == "span-topic_3" # same trace id, different span ids - assert producer_span_2.t == poll_span_2.t - assert producer_span_2.s != poll_span_2.s - assert producer_span_3.t == poll_span_3.t assert producer_span_3.s != poll_span_3.s diff --git a/tests/requirements-kafka.txt b/tests/requirements-kafka.txt index 640d2ad9..c07a0f91 100644 --- a/tests/requirements-kafka.txt +++ b/tests/requirements-kafka.txt @@ -1,5 +1,4 @@ -r requirements-minimal.txt mock>=2.0.0 confluent-kafka>=2.0.0 -kafka-python>=2.0.0; python_version < "3.12" -kafka-python-ng>=2.0.0; python_version >= "3.12" +kafka-python>=2.0.3 diff --git a/tests/requirements-pre315.txt b/tests/requirements-pre315.txt index ad141a85..afc1f5e7 100644 --- a/tests/requirements-pre315.txt +++ b/tests/requirements-pre315.txt @@ -45,5 +45,5 @@ urllib3>=1.26.5 httpx>=0.27.0 gevent>=23.9.0.post1 confluent-kafka>=2.0.0 -kafka-python-ng>=2.0.0 +kafka-python>=2.0.3