{-# 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.Core (
  SpanContext (SpanContext, isRemote, spanId, traceFlags, traceId, traceState),
  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 (ShortByteString -> TraceId)
-> (ByteString -> ShortByteString) -> ByteString -> TraceId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ShortByteString
newTraceIdFromHeader (ByteString -> TraceId) -> Maybe ByteString -> Maybe TraceId
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HeaderName -> RequestHeaders -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup HeaderName
forall s. IsString s => s
traceIdKey RequestHeaders
hs
              SpanId
parentId <- ShortByteString -> SpanId
SpanId (ShortByteString -> SpanId)
-> (ByteString -> ShortByteString) -> ByteString -> SpanId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ShortByteString
newSpanIdFromHeader (ByteString -> SpanId) -> Maybe ByteString -> Maybe SpanId
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HeaderName -> RequestHeaders -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup HeaderName
forall s. IsString s => s
parentIdKey RequestHeaders
hs
              Text
samplingPriority <- String -> Text
T.pack (String -> Text) -> (ByteString -> String) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String
BC.unpack (ByteString -> Text) -> Maybe ByteString -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HeaderName -> RequestHeaders -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup HeaderName
forall s. IsString s => s
samplingPriorityKey RequestHeaders
hs
              SpanContext -> Maybe SpanContext
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SpanContext -> Maybe SpanContext)
-> SpanContext -> Maybe SpanContext
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 Text
forall s. IsString s => s
samplingPriorityKey, Text -> Value
TS.Value Text
samplingPriority)]
                  }
        case Maybe SpanContext
spanContext' of
          Maybe SpanContext
Nothing -> Context -> IO Context
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Context
c
          Just SpanContext
spanContext -> Context -> IO Context
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Context -> IO Context) -> Context -> IO Context
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 -> RequestHeaders -> IO RequestHeaders
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure RequestHeaders
hs
          Just Span
span' -> do
            SpanContext {TraceId
traceId :: SpanContext -> TraceId
traceId :: TraceId
traceId, SpanId
spanId :: SpanContext -> SpanId
spanId :: SpanId
spanId, traceState :: SpanContext -> TraceState
traceState = TraceState [(Key, Value)]
traceState} <- Span -> IO SpanContext
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 Key -> [(Key, Value)] -> Maybe Value
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (Text -> Key
TS.Key Text
forall s. IsString s => s
samplingPriorityKey) [(Key, Value)]
traceState of
                Maybe Value
Nothing -> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
"1" -- when an origin of the trace
                Just (TS.Value Text
p) -> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ String -> ByteString
BC.pack (String -> ByteString) -> String -> ByteString
forall a b. (a -> b) -> a -> b
$ Text -> String
T.unpack Text
p
            RequestHeaders -> IO RequestHeaders
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RequestHeaders -> IO RequestHeaders)
-> RequestHeaders -> IO RequestHeaders
forall a b. (a -> b) -> a -> b
$
              (HeaderName
forall s. IsString s => s
traceIdKey, ByteString
traceIdValue)
                (HeaderName, ByteString) -> RequestHeaders -> RequestHeaders
forall a. a -> [a] -> [a]
: (HeaderName
forall s. IsString s => s
parentIdKey, ByteString
parentIdValue)
                (HeaderName, ByteString) -> RequestHeaders -> RequestHeaders
forall a. a -> [a] -> [a]
: (HeaderName
forall s. IsString s => s
samplingPriorityKey, ByteString
samplingPriority)
                (HeaderName, ByteString) -> RequestHeaders -> RequestHeaders
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