{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}

module OpenTelemetry.Propagator.Datadog (
  datadogTraceContextPropagator,
  convertOpenTelemetrySpanIdToDatadogSpanId,
  convertOpenTelemetryTraceIdToDatadogTraceId,
) where

import qualified Data.ByteString.Char8 as BC
import qualified Data.ByteString.Short.Internal as SBI
import Data.Primitive (ByteArray (ByteArray))
import Data.String (IsString)
import qualified Data.Text as T
import Data.Word (Word64)
import Network.HTTP.Types (
  RequestHeaders,
  ResponseHeaders,
 )
import OpenTelemetry.Common (TraceFlags (TraceFlags))
import OpenTelemetry.Context (
  Context,
  insertSpan,
  lookupSpan,
 )
import OpenTelemetry.Internal.Trace.Id (
  SpanId (SpanId),
  TraceId (TraceId),
 )
import OpenTelemetry.Propagator (Propagator (Propagator, extractor, injector, propagatorNames))
import OpenTelemetry.Propagator.Datadog.Internal (
  indexByteArrayNbo,
  newHeaderFromSpanId,
  newHeaderFromTraceId,
  newSpanIdFromHeader,
  newTraceIdFromHeader,
 )
import OpenTelemetry.Trace (SpanContext (SpanContext, isRemote, spanId, traceFlags, traceId, traceState))
import OpenTelemetry.Trace.Core (
  getSpanContext,
  wrapSpanContext,
 )
import OpenTelemetry.Trace.TraceState (TraceState (TraceState))
import qualified OpenTelemetry.Trace.TraceState as TS


-- Reference: bi-directional conversion of IDs of Open Telemetry and ones of Datadog
-- - English: https://docs.datadoghq.com/tracing/other_telemetry/connect_logs_and_traces/opentelemetry/
-- - Japanese: https://docs.datadoghq.com/ja/tracing/connect_logs_and_traces/opentelemetry/
datadogTraceContextPropagator :: Propagator Context RequestHeaders ResponseHeaders
datadogTraceContextPropagator :: Propagator Context RequestHeaders RequestHeaders
datadogTraceContextPropagator =
  Propagator
    { propagatorNames :: [Text]
propagatorNames = [Text
"datadog trace context"]
    , extractor :: RequestHeaders -> Context -> IO Context
extractor = \RequestHeaders
hs Context
c -> do
        let spanContext' :: Maybe SpanContext
spanContext' = do
              TraceId
traceId <- ShortByteString -> TraceId
TraceId forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ShortByteString
newTraceIdFromHeader forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup forall s. IsString s => s
traceIdKey RequestHeaders
hs
              SpanId
parentId <- ShortByteString -> SpanId
SpanId forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ShortByteString
newSpanIdFromHeader forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup forall s. IsString s => s
parentIdKey RequestHeaders
hs
              Text
samplingPriority <- String -> Text
T.pack forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String
BC.unpack forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup forall s. IsString s => s
samplingPriorityKey RequestHeaders
hs
              forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
                SpanContext
                  { TraceId
traceId :: TraceId
traceId :: TraceId
traceId
                  , spanId :: SpanId
spanId = SpanId
parentId
                  , isRemote :: Bool
isRemote = Bool
True
                  , -- when 0, not sampled
                    -- refer: OpenTelemetry.Internal.Trace.Types.isSampled
                    traceFlags :: TraceFlags
traceFlags = Word8 -> TraceFlags
TraceFlags Word8
1
                  , traceState :: TraceState
traceState = [(Key, Value)] -> TraceState
TraceState [(Text -> Key
TS.Key forall s. IsString s => s
samplingPriorityKey, Text -> Value
TS.Value Text
samplingPriority)]
                  }
        case Maybe SpanContext
spanContext' of
          Maybe SpanContext
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Context
c
          Just SpanContext
spanContext -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Span -> Context -> Context
insertSpan (SpanContext -> Span
wrapSpanContext SpanContext
spanContext) Context
c
    , injector :: Context -> RequestHeaders -> IO RequestHeaders
injector = \Context
c RequestHeaders
hs ->
        case Context -> Maybe Span
lookupSpan Context
c of
          Maybe Span
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure RequestHeaders
hs
          Just Span
span' -> do
            SpanContext {TraceId
traceId :: TraceId
traceId :: SpanContext -> TraceId
traceId, SpanId
spanId :: SpanId
spanId :: SpanContext -> SpanId
spanId, traceState :: SpanContext -> TraceState
traceState = TraceState [(Key, Value)]
traceState} <- forall (m :: * -> *). MonadIO m => Span -> m SpanContext
getSpanContext Span
span'
            let traceIdValue :: ByteString
traceIdValue = (\(TraceId ShortByteString
b) -> ShortByteString -> ByteString
newHeaderFromTraceId ShortByteString
b) TraceId
traceId
                parentIdValue :: ByteString
parentIdValue = (\(SpanId ShortByteString
b) -> ShortByteString -> ByteString
newHeaderFromSpanId ShortByteString
b) SpanId
spanId
            ByteString
samplingPriority <-
              case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (Text -> Key
TS.Key forall s. IsString s => s
samplingPriorityKey) [(Key, Value)]
traceState of
                Maybe Value
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
"1" -- when an origin of the trace
                Just (TS.Value Text
p) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ String -> ByteString
BC.pack forall a b. (a -> b) -> a -> b
$ Text -> String
T.unpack Text
p
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
              (forall s. IsString s => s
traceIdKey, ByteString
traceIdValue)
                forall a. a -> [a] -> [a]
: (forall s. IsString s => s
parentIdKey, ByteString
parentIdValue)
                forall a. a -> [a] -> [a]
: (forall s. IsString s => s
samplingPriorityKey, ByteString
samplingPriority)
                forall a. a -> [a] -> [a]
: RequestHeaders
hs
    }
  where
    traceIdKey, parentIdKey, samplingPriorityKey :: (IsString s) => s
    traceIdKey :: forall s. IsString s => s
traceIdKey = s
"x-datadog-trace-id"
    parentIdKey :: forall s. IsString s => s
parentIdKey = s
"x-datadog-parent-id"
    samplingPriorityKey :: forall s. IsString s => s
samplingPriorityKey = s
"x-datadog-sampling-priority"


convertOpenTelemetrySpanIdToDatadogSpanId :: SpanId -> Word64
convertOpenTelemetrySpanIdToDatadogSpanId :: SpanId -> Word64
convertOpenTelemetrySpanIdToDatadogSpanId (SpanId (SBI.SBS ByteArray#
a)) = ByteArray -> Int -> Word64
indexByteArrayNbo (ByteArray# -> ByteArray
ByteArray ByteArray#
a) Int
0


convertOpenTelemetryTraceIdToDatadogTraceId :: TraceId -> Word64
convertOpenTelemetryTraceIdToDatadogTraceId :: TraceId -> Word64
convertOpenTelemetryTraceIdToDatadogTraceId (TraceId (SBI.SBS ByteArray#
a)) = ByteArray -> Int -> Word64
indexByteArrayNbo (ByteArray# -> ByteArray
ByteArray ByteArray#
a) Int
1