{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns #-}
{- |
Module : OpenTelemetry.Instrumentation.Kafka
Description : OpenTelemetry instrumentation for hw-kafka-client
This module provides OpenTelemetry instrumentation for the hw-kafka-client library.
It adds distributed tracing capabilities to Kafka producer and consumer operations,
automatically propagating context between services via Kafka message headers.
-}
module OpenTelemetry.Instrumentation.Kafka (
-- * Producer
produceMessage,
-- * Consumer
pollMessage,
) where
import Control.Monad (void)
import Control.Monad.IO.Class (MonadIO)
import Control.Monad.IO.Unlift (MonadUnliftIO)
import Data.Bifunctor (first)
import Data.ByteString (ByteString)
import qualified Data.CaseInsensitive as CI
import qualified Data.Map as M
import Data.String (IsString)
import Data.Text.Encoding (decodeUtf8')
import GHC.Stack.Types (HasCallStack)
import Kafka.Consumer (
ConsumerProperties (cpProps),
ConsumerRecord (crHeaders, crKey, crOffset, crPartition, crTopic),
KafkaConsumer,
Offset (unOffset),
)
import qualified Kafka.Consumer as KC
import Kafka.Producer (
KafkaError,
KafkaProducer,
ProducePartition (SpecifiedPartition, UnassignedPartition),
ProducerRecord (prHeaders, prKey, prPartition, prTopic),
)
import qualified Kafka.Producer as KP
import Kafka.Types (
Headers,
PartitionId (unPartitionId),
Timeout,
TopicName (unTopicName),
headersFromList,
headersToList,
)
import Network.HTTP.Types (RequestHeaders)
import OpenTelemetry.Attributes.Map (AttributeMap, insertAttributeByKey)
import qualified OpenTelemetry.Context as Context
import OpenTelemetry.Context.ThreadLocal (attachContext, getContext)
import OpenTelemetry.Propagator (extract, inject)
import OpenTelemetry.SemanticConventions (
messaging_destination_name,
messaging_kafka_consumer_group,
messaging_kafka_destination_partition,
messaging_kafka_message_key,
messaging_kafka_message_offset,
messaging_operation,
)
import OpenTelemetry.Trace.Core (
SpanArguments (kind),
SpanKind (Consumer, Producer),
Tracer,
addAttributesToSpanArguments,
callerAttributes,
defaultSpanArguments,
detectInstrumentationLibrary,
getGlobalTracerProvider,
getTracerProviderPropagators,
inSpan'',
makeTracer,
toAttribute,
tracerOptions,
)
producerOperationName :: IsString a => a
producerOperationName = "send"
consumerOperationName :: IsString a => a
consumerOperationName = "process"
-- | Span arguments for producer operations
producerSpanArgs :: SpanArguments
producerSpanArgs =
defaultSpanArguments {kind = Producer}
rightToMaybe :: Either a b -> Maybe b
rightToMaybe (Right b) = Just b
rightToMaybe _ = Nothing
-- | Span attributes with caller information and kafka specifics
producerAttributes :: HasCallStack => ProducerRecord -> AttributeMap
producerAttributes record =
let
addOperationName =
insertAttributeByKey messaging_operation producerOperationName
addDestination =
insertAttributeByKey messaging_destination_name $ toAttribute . unTopicName . prTopic $ record
addPartition =
case prPartition record of
SpecifiedPartition p64 ->
insertAttributeByKey messaging_kafka_destination_partition $ toAttribute $ p64
UnassignedPartition ->
id
addKey =
case prKey record >>= rightToMaybe . decodeUtf8' of
Just key ->
insertAttributeByKey messaging_kafka_message_key $ toAttribute key
Nothing ->
id
in
(addOperationName . addDestination . addPartition . addKey)
callerAttributes
-- | Span arguments for consumer operations
consumerSpanArgs :: SpanArguments
consumerSpanArgs = defaultSpanArguments {kind = Consumer}
-- | Span attributes for consumer with caller information and kafka specifics
consumerAttributes
:: HasCallStack
=> ConsumerProperties
-> ConsumerRecord (Maybe ByteString) (Maybe ByteString)
-> AttributeMap
consumerAttributes consumerProperties record =
let
addOperationName =
insertAttributeByKey messaging_operation consumerOperationName
addDestination =
insertAttributeByKey messaging_destination_name $ toAttribute . unTopicName . crTopic $ record
addConsumerGroup =
-- NOTE: unfortunately, hw-kafka-client does not expose an API to get the consumer, this a flaky workaround
case M.lookup "group.id" $ cpProps consumerProperties of
Just groupId -> insertAttributeByKey messaging_kafka_consumer_group $ toAttribute groupId
Nothing -> id
addPartition =
insertAttributeByKey messaging_kafka_destination_partition $ toAttribute . unPartitionId . crPartition $ record
addOffset =
insertAttributeByKey messaging_kafka_message_offset $ toAttribute . unOffset . crOffset $ record
addKey =
case crKey record >>= rightToMaybe . decodeUtf8' of
Just key ->
insertAttributeByKey messaging_kafka_message_key $ toAttribute key
Nothing ->
id
in
(addOperationName . addDestination . addConsumerGroup . addPartition . addOffset . addKey)
callerAttributes
-- | Get the tracer for rdkafka instrumentation
rdkafkaTracer :: (MonadIO m) => m Tracer
rdkafkaTracer = do
provider <- getGlobalTracerProvider
return $ makeTracer provider $detectInstrumentationLibrary tracerOptions
-- | Convert Kafka headers to HTTP headers format
kafkaHeadersToHttpHeaders :: Headers -> RequestHeaders
kafkaHeadersToHttpHeaders = map (first CI.mk) . headersToList
-- | Convert HTTP headers to Kafka headers format
httpHeadersToKafkaHeaders :: RequestHeaders -> Headers
httpHeadersToKafkaHeaders = headersFromList . map (first CI.foldedCase)
{- | Produce a message to Kafka with OpenTelemetry instrumentation.
This function wraps the standard Kafka producer with OpenTelemetry tracing.
It creates a new span for the produce operation and injects the current context
into the message headers.
-}
produceMessage
:: (MonadUnliftIO m, HasCallStack)
=> KafkaProducer
-> ProducerRecord
-> m (Maybe KafkaError)
produceMessage producer record =
let
headers = prHeaders record
topicName = prTopic record
spanName = producerOperationName <> " " <> unTopicName topicName
attributes = producerAttributes record
spanArguments = addAttributesToSpanArguments attributes producerSpanArgs
in
do
tracer <- rdkafkaTracer
ctxt <- getContext
inSpan'' tracer spanName spanArguments $ \newSpan -> do
propagator <- getTracerProviderPropagators <$> getGlobalTracerProvider
extraHeaders <- inject propagator (Context.insertSpan newSpan ctxt) []
let newKafkaHeaders = headers <> httpHeadersToKafkaHeaders extraHeaders
let newKafkaRecord = record {prHeaders = newKafkaHeaders}
KP.produceMessage producer newKafkaRecord
{- | Poll for a single message from Kafka with OpenTelemetry instrumentation.
This function wraps the standard Kafka consumer with OpenTelemetry tracing.
It creates a new span for the poll operation and extracts any tracing context
from the message headers.
-}
pollMessage
:: (MonadUnliftIO m, HasCallStack)
=> ConsumerProperties
-> KafkaConsumer
-> Timeout
-> m (Either KafkaError (ConsumerRecord (Maybe ByteString) (Maybe ByteString)))
-- ^ Returns either an error or the consumed record
pollMessage consumerProperties consumer timeout =
do
KC.pollMessage consumer timeout >>= \case
Left err -> pure $ Left err
Right cr ->
let
attributes = consumerAttributes consumerProperties cr
topicName = crTopic cr
spanName = consumerOperationName <> " " <> unTopicName topicName
in
do
tracer <- rdkafkaTracer
ctxt <- getContext
propagator <- getTracerProviderPropagators <$> getGlobalTracerProvider
ctx <- extract propagator (kafkaHeadersToHttpHeaders $ crHeaders cr) ctxt
void $ attachContext ctx
inSpan'' tracer spanName (addAttributesToSpanArguments attributes consumerSpanArgs) $ \_span -> do
return $ Right cr