feat: support batch (getmany) in aiokafka instrumentation

Author: dimastbk

instrumentation/opentelemetry-instrumentation-aiokafka/src/opentelemetry/instrumentation/aiokafka/__init__.py

@@ -78,6 +78,7 @@ async def async_consume_hook(span, record, args, kwargs):
 from opentelemetry import trace
 from opentelemetry.instrumentation.aiokafka.package import _instruments
 from opentelemetry.instrumentation.aiokafka.utils import (
+    _wrap_getmany,
     _wrap_getone,
     _wrap_send,
 )
@@ -131,7 +132,13 @@ def _instrument(self, **kwargs):
             "getone",
             _wrap_getone(tracer, async_consume_hook),
         )
+        wrap_function_wrapper(
+            aiokafka.AIOKafkaConsumer,
+            "getmany",
+            _wrap_getmany(tracer, async_consume_hook),
+        )
 
     def _uninstrument(self, **kwargs):
         unwrap(aiokafka.AIOKafkaProducer, "send")
         unwrap(aiokafka.AIOKafkaConsumer, "getone")
+        unwrap(aiokafka.AIOKafkaConsumer, "getmany")

instrumentation/opentelemetry-instrumentation-aiokafka/src/opentelemetry/instrumentation/aiokafka/utils.py

@@ -1,9 +1,22 @@
+from __future__ import annotations
+
+import asyncio
 import json
 from logging import getLogger
-from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Protocol,
+    Tuple,
+    Union,
+)
 
 import aiokafka
-from aiokafka import ConsumerRecord
 
 from opentelemetry import context, propagate, trace
 from opentelemetry.context import Context
@@ -13,6 +26,54 @@
 from opentelemetry.trace import Tracer
 from opentelemetry.trace.span import Span
 
+if TYPE_CHECKING:
+    from aiokafka.structs import RecordMetadata
+
+    class AIOKafkaGetOneProto(Protocol):
+        async def __call__(
+            self, *partitions: aiokafka.TopicPartition
+        ) -> aiokafka.ConsumerRecord[object, object]: ...
+
+    class AIOKafkaGetManyProto(Protocol):
+        async def __call__(
+            self,
+            *partitions: aiokafka.TopicPartition,
+            timeout_ms: int = 0,
+            max_records: int | None = None,
+        ) -> dict[
+            aiokafka.TopicPartition,
+            list[aiokafka.ConsumerRecord[object, object]],
+        ]: ...
+
+    class AIOKafkaSendProto(Protocol):
+        async def __call__(
+            self,
+            topic: str,
+            value: Any | None = None,
+            key: Any | None = None,
+            partition: int | None = None,
+            timestamp_ms: int | None = None,
+            headers: HeadersT | None = None,
+        ) -> asyncio.Future[RecordMetadata]: ...
+
+
+ProduceHookT = Optional[
+    Callable[[Span, Tuple[Any, ...], Dict[str, Any]], Awaitable[None]]
+]
+ConsumeHookT = Optional[
+    Callable[
+        [
+            Span,
+            aiokafka.ConsumerRecord[object, object],
+            Tuple[aiokafka.TopicPartition, ...],
+            Dict[str, Any],
+        ],
+        Awaitable[None],
+    ]
+]
+
+HeadersT = List[Tuple[str, Optional[bytes]]]
+
 _LOG = getLogger(__name__)
 
 
@@ -97,14 +158,6 @@ async def _extract_send_partition(
         return None
 
 
-ProduceHookT = Optional[Callable[[Span, Tuple, Dict], Awaitable[None]]]
-ConsumeHookT = Optional[
-    Callable[[Span, ConsumerRecord, Tuple, Dict], Awaitable[None]]
-]
-
-HeadersT = List[Tuple[str, Optional[bytes]]]
-
-
 class AIOKafkaContextGetter(textmap.Getter[HeadersT]):
     def get(self, carrier: HeadersT, key: str) -> Optional[List[str]]:
         if carrier is None:
@@ -198,7 +251,7 @@ def _enrich_send_span(
     )
 
 
-def _enrich_anext_span(
+def _enrich_getone_span(
     span: Span,
     *,
     bootstrap_servers: Union[str, List[str]],
@@ -247,19 +300,93 @@ def _enrich_anext_span(
         )
 
 
+def _enrich_getmany_poll_span(
+    span: Span,
+    *,
+    bootstrap_servers: Union[str, List[str]],
+    client_id: str,
+    consumer_group: Optional[str],
+    message_count: int,
+) -> None:
+    if not span.is_recording():
+        return
+
+    span.set_attribute(
+        messaging_attributes.MESSAGING_SYSTEM,
+        messaging_attributes.MessagingSystemValues.KAFKA.value,
+    )
+    span.set_attribute(
+        server_attributes.SERVER_ADDRESS, json.dumps(bootstrap_servers)
+    )
+    span.set_attribute(messaging_attributes.MESSAGING_CLIENT_ID, client_id)
+
+    if consumer_group is not None:
+        span.set_attribute(
+            messaging_attributes.MESSAGING_CONSUMER_GROUP_NAME, consumer_group
+        )
+
+    span.set_attribute(
+        messaging_attributes.MESSAGING_BATCH_MESSAGE_COUNT, message_count
+    )
+
+    span.set_attribute(messaging_attributes.MESSAGING_OPERATION_NAME, "poll")
+    span.set_attribute(
+        messaging_attributes.MESSAGING_OPERATION_TYPE,
+        messaging_attributes.MessagingOperationTypeValues.RECEIVE.value,
+    )
+
+
+def _enrich_getmany_topic_span(
+    span: Span,
+    *,
+    bootstrap_servers: Union[str, List[str]],
+    client_id: str,
+    consumer_group: Optional[str],
+    topic: str,
+    partition: int,
+    message_count: int,
+) -> None:
+    if not span.is_recording():
+        return
+
+    _enrich_base_span(
+        span,
+        bootstrap_servers=bootstrap_servers,
+        client_id=client_id,
+        topic=topic,
+        partition=partition,
+        key=None,
+    )
+
+    if consumer_group is not None:
+        span.set_attribute(
+            messaging_attributes.MESSAGING_CONSUMER_GROUP_NAME, consumer_group
+        )
+
+    span.set_attribute(
+        messaging_attributes.MESSAGING_BATCH_MESSAGE_COUNT, message_count
+    )
+
+    span.set_attribute(messaging_attributes.MESSAGING_OPERATION_NAME, "poll")
+    span.set_attribute(
+        messaging_attributes.MESSAGING_OPERATION_TYPE,
+        messaging_attributes.MessagingOperationTypeValues.RECEIVE.value,
+    )
+
+
 def _get_span_name(operation: str, topic: str):
     return f"{topic} {operation}"
 
 
 def _wrap_send(
     tracer: Tracer, async_produce_hook: ProduceHookT
-) -> Callable[..., Awaitable[None]]:
+) -> Callable[..., Awaitable[asyncio.Future[RecordMetadata]]]:
     async def _traced_send(
-        func: Callable[..., Awaitable[None]],
+        func: AIOKafkaSendProto,
         instance: aiokafka.AIOKafkaProducer,
         args: Tuple[Any],
         kwargs: Dict[str, Any],
-    ) -> None:
+    ) -> asyncio.Future[RecordMetadata]:
         headers = _extract_send_headers(args, kwargs)
         if headers is None:
             headers = []
@@ -301,14 +428,14 @@ async def _traced_send(
 async def _create_consumer_span(
     tracer: Tracer,
     async_consume_hook: ConsumeHookT,
-    record: ConsumerRecord,
+    record: aiokafka.ConsumerRecord[object, object],
     extracted_context: Context,
     bootstrap_servers: Union[str, List[str]],
     client_id: str,
     consumer_group: Optional[str],
     args: Tuple[Any],
     kwargs: Dict[str, Any],
-):
+) -> trace.Span:
     span_name = _get_span_name("receive", record.topic)
     with tracer.start_as_current_span(
         span_name,
@@ -317,7 +444,7 @@ async def _create_consumer_span(
     ) as span:
         new_context = trace.set_span_in_context(span, extracted_context)
         token = context.attach(new_context)
-        _enrich_anext_span(
+        _enrich_getone_span(
             span,
             bootstrap_servers=bootstrap_servers,
             client_id=client_id,
@@ -334,16 +461,18 @@ async def _create_consumer_span(
             _LOG.exception(hook_exception)
         context.detach(token)
 
+    return span
+
 
 def _wrap_getone(
     tracer: Tracer, async_consume_hook: ConsumeHookT
-) -> Callable[..., Awaitable[aiokafka.ConsumerRecord]]:
-    async def _traced_next(
-        func: Callable[..., Awaitable[aiokafka.ConsumerRecord]],
+) -> Callable[..., Awaitable[aiokafka.ConsumerRecord[object, object]]]:
+    async def _traced_getone(
+        func: AIOKafkaGetOneProto,
         instance: aiokafka.AIOKafkaConsumer,
         args: Tuple[Any],
         kwargs: Dict[str, Any],
-    ) -> aiokafka.ConsumerRecord:
+    ) -> aiokafka.ConsumerRecord[object, object]:
         record = await func(*args, **kwargs)
 
         if record:
@@ -367,4 +496,80 @@ async def _traced_next(
             )
         return record
 
-    return _traced_next
+    return _traced_getone
+
+
+def _wrap_getmany(
+    tracer: Tracer, async_consume_hook: ConsumeHookT
+) -> Callable[
+    ...,
+    Awaitable[
+        dict[
+            aiokafka.TopicPartition,
+            list[aiokafka.ConsumerRecord[object, object]],
+        ]
+    ],
+]:
+    async def _traced_getmany(
+        func: AIOKafkaGetManyProto,
+        instance: aiokafka.AIOKafkaConsumer,
+        args: Tuple[Any],
+        kwargs: Dict[str, Any],
+    ) -> dict[
+        aiokafka.TopicPartition, list[aiokafka.ConsumerRecord[object, object]]
+    ]:
+        records = await func(*args, **kwargs)
+
+        if records:
+            bootstrap_servers = _extract_bootstrap_servers(instance._client)
+            client_id = _extract_client_id(instance._client)
+            consumer_group = _extract_consumer_group(instance)
+
+            span_name = _get_span_name(
+                "poll", ", ".join([topic.topic for topic in records.keys()])
+            )
+            with tracer.start_as_current_span(
+                span_name, kind=trace.SpanKind.CLIENT
+            ) as poll_span:
+                _enrich_getmany_poll_span(
+                    poll_span,
+                    bootstrap_servers=bootstrap_servers,
+                    client_id=client_id,
+                    consumer_group=consumer_group,
+                    message_count=sum(len(r) for r in records.values()),
+                )
+
+                for topic, topic_records in records.items():
+                    span_name = _get_span_name("poll", topic.topic)
+                    with tracer.start_as_current_span(
+                        span_name, kind=trace.SpanKind.CLIENT
+                    ) as topic_span:
+                        _enrich_getmany_topic_span(
+                            topic_span,
+                            bootstrap_servers=bootstrap_servers,
+                            client_id=client_id,
+                            consumer_group=consumer_group,
+                            topic=topic.topic,
+                            partition=topic.partition,
+                            message_count=len(topic_records),
+                        )
+
+                        for record in topic_records:
+                            extracted_context = propagate.extract(
+                                record.headers, getter=_aiokafka_getter
+                            )
+                            record_span = await _create_consumer_span(
+                                tracer,
+                                async_consume_hook,
+                                record,
+                                extracted_context,
+                                bootstrap_servers,
+                                client_id,
+                                consumer_group,
+                                args,
+                                kwargs,
+                            )
+                            topic_span.add_link(record_span.get_span_context())
+        return records
+
+    return _traced_getmany