{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}

{- | B3 Propagation Requirements:
 https://github.com/openzipkin/b3-propagation
 https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/context/api-propagators.md#b3-requirements
-}
module OpenTelemetry.Propagator.B3 (
  b3TraceContextPropagator,
  b3MultiTraceContextPropagator,
) where

--------------------------------------------------------------------------------

import Control.Applicative ((<|>))
import Data.ByteString (ByteString)
import Data.List (intersperse)
import Data.Maybe (catMaybes, fromMaybe)
import qualified Data.Text.Encoding as Text
import Network.HTTP.Types (HeaderName, RequestHeaders, ResponseHeaders)
import OpenTelemetry.Common (TraceFlags (..))
import OpenTelemetry.Context (Context)
import qualified OpenTelemetry.Context as Context
import OpenTelemetry.Propagator (Propagator (..))
import OpenTelemetry.Propagator.B3.Internal
import qualified OpenTelemetry.Trace.Core as Core
import qualified OpenTelemetry.Trace.TraceState as TS
import Prelude


--------------------------------------------------------------------------------

b3TraceContextPropagator :: Propagator Context RequestHeaders ResponseHeaders
b3TraceContextPropagator :: Propagator Context RequestHeaders RequestHeaders
b3TraceContextPropagator =
  Propagator
    { propagatorNames :: [Text]
propagatorNames = [Text
"B3 Trace Context"]
    , extractor :: RequestHeaders -> Context -> IO Context
extractor = \RequestHeaders
hs Context
c ->
        case RequestHeaders -> Maybe SpanContext
b3Extractor RequestHeaders
hs 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
Context.insertSpan (SpanContext -> Span
Core.wrapSpanContext SpanContext
spanContext') Context
c
    , injector :: Context -> RequestHeaders -> IO RequestHeaders
injector = \Context
c RequestHeaders
hs ->
        case Context -> Maybe Span
Context.lookupSpan Context
c of
          Maybe Span
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure RequestHeaders
hs
          Just Span
span' -> do
            Core.SpanContext {TraceId
traceId :: SpanContext -> TraceId
traceId :: TraceId
traceId, SpanId
spanId :: SpanContext -> SpanId
spanId :: SpanId
spanId, traceState :: SpanContext -> TraceState
traceState = TS.TraceState [(Key, Value)]
traceState} <- forall (m :: * -> *). MonadIO m => Span -> m SpanContext
Core.getSpanContext Span
span'
            let traceIdValue :: ByteString
traceIdValue = TraceId -> ByteString
encodeTraceId TraceId
traceId
                spanIdValue :: ByteString
spanIdValue = SpanId -> ByteString
encodeSpanId SpanId
spanId
                samplingStateValue :: Maybe Text
samplingStateValue = forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (Text -> Key
TS.Key Text
"sampling-state") [(Key, Value)]
traceState forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Value -> Maybe SamplingState
samplingStateFromValue forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SamplingState -> Maybe Text
printSamplingStateSingle
                value :: ByteString
value = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a. a -> [a] -> [a]
intersperse ByteString
"-" forall a b. (a -> b) -> a -> b
$ [ByteString
traceIdValue, ByteString
spanIdValue] forall a. Semigroup a => a -> a -> a
<> forall a. [Maybe a] -> [a]
catMaybes [Text -> ByteString
Text.encodeUtf8 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Text
samplingStateValue]

            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (HeaderName
b3Header, ByteString
value) forall a. a -> [a] -> [a]
: RequestHeaders
hs
    }


b3MultiTraceContextPropagator :: Propagator Context RequestHeaders ResponseHeaders
b3MultiTraceContextPropagator :: Propagator Context RequestHeaders RequestHeaders
b3MultiTraceContextPropagator =
  Propagator
    { propagatorNames :: [Text]
propagatorNames = [Text
"B3 Multi Trace Context"]
    , extractor :: RequestHeaders -> Context -> IO Context
extractor = \RequestHeaders
hs Context
c -> do
        case RequestHeaders -> Maybe SpanContext
b3Extractor RequestHeaders
hs 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
Context.insertSpan (SpanContext -> Span
Core.wrapSpanContext SpanContext
spanContext') Context
c
    , injector :: Context -> RequestHeaders -> IO RequestHeaders
injector = \Context
c RequestHeaders
hs ->
        case Context -> Maybe Span
Context.lookupSpan Context
c of
          Maybe Span
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure RequestHeaders
hs
          Just Span
span' -> do
            Core.SpanContext {TraceId
traceId :: TraceId
traceId :: SpanContext -> TraceId
traceId, SpanId
spanId :: SpanId
spanId :: SpanContext -> SpanId
spanId, traceState :: SpanContext -> TraceState
traceState = TS.TraceState [(Key, Value)]
traceState} <- forall (m :: * -> *). MonadIO m => Span -> m SpanContext
Core.getSpanContext Span
span'
            let traceIdValue :: ByteString
traceIdValue = TraceId -> ByteString
encodeTraceId TraceId
traceId
                spanIdValue :: ByteString
spanIdValue = SpanId -> ByteString
encodeSpanId SpanId
spanId
                samplingStateValue :: Maybe (HeaderName, Text)
samplingStateValue = forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (Text -> Key
TS.Key Text
"sampling-state") [(Key, Value)]
traceState forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Value -> Maybe SamplingState
samplingStateFromValue forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SamplingState -> Maybe (HeaderName, Text)
printSamplingStateMulti

            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
              (HeaderName
xb3TraceIdHeader, ByteString
traceIdValue)
                forall a. a -> [a] -> [a]
: (HeaderName
xb3SpanIdHeader, ByteString
spanIdValue)
                forall a. a -> [a] -> [a]
: RequestHeaders
hs
                forall a. [a] -> [a] -> [a]
++ forall a. [Maybe a] -> [a]
catMaybes [forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Text -> ByteString
Text.encodeUtf8 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (HeaderName, Text)
samplingStateValue]
    }


--------------------------------------------------------------------------------

{- | For both @B3@ and @B3 Multi@ formats, we must attempt single and
 multi header extraction:
 https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/context/api-propagators.md#configuration
-}
b3Extractor :: [(HeaderName, ByteString)] -> Maybe Core.SpanContext
b3Extractor :: RequestHeaders -> Maybe SpanContext
b3Extractor RequestHeaders
hs = RequestHeaders -> Maybe SpanContext
b3SingleExtractor RequestHeaders
hs forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> RequestHeaders -> Maybe SpanContext
b3MultiExtractor RequestHeaders
hs


b3SingleExtractor :: [(HeaderName, ByteString)] -> Maybe Core.SpanContext
b3SingleExtractor :: RequestHeaders -> Maybe SpanContext
b3SingleExtractor RequestHeaders
hs = do
  B3SingleHeader {Maybe SpanId
TraceId
SpanId
SamplingState
parentSpanId :: B3SingleHeader -> Maybe SpanId
samplingState :: B3SingleHeader -> SamplingState
spanId :: B3SingleHeader -> SpanId
traceId :: B3SingleHeader -> TraceId
parentSpanId :: Maybe SpanId
samplingState :: SamplingState
spanId :: SpanId
traceId :: TraceId
..} <- ByteString -> Maybe B3SingleHeader
decodeB3SingleHeader forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup HeaderName
b3Header RequestHeaders
hs

  let traceFlags :: TraceFlags
traceFlags = if SamplingState
samplingState forall a. Eq a => a -> a -> Bool
== SamplingState
Accept Bool -> Bool -> Bool
|| SamplingState
samplingState forall a. Eq a => a -> a -> Bool
== SamplingState
Debug then Word8 -> TraceFlags
TraceFlags Word8
1 else Word8 -> TraceFlags
TraceFlags Word8
0

  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    Core.SpanContext
      { traceId :: TraceId
traceId = TraceId
traceId
      , spanId :: SpanId
spanId = SpanId
spanId
      , isRemote :: Bool
isRemote = Bool
True
      , traceFlags :: TraceFlags
traceFlags = TraceFlags
traceFlags
      , traceState :: TraceState
traceState = [(Key, Value)] -> TraceState
TS.TraceState [(Text -> Key
TS.Key Text
"sampling-state", SamplingState -> Value
samplingStateToValue SamplingState
samplingState)]
      }


b3MultiExtractor :: [(HeaderName, ByteString)] -> Maybe Core.SpanContext
b3MultiExtractor :: RequestHeaders -> Maybe SpanContext
b3MultiExtractor RequestHeaders
hs = do
  TraceId
traceId <- ByteString -> Maybe TraceId
decodeXb3TraceIdHeader forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup HeaderName
xb3TraceIdHeader RequestHeaders
hs
  SpanId
spanId <- ByteString -> Maybe SpanId
decodeXb3SpanIdHeader forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup HeaderName
xb3SpanIdHeader RequestHeaders
hs

  let sampled :: Maybe SamplingState
sampled = ByteString -> Maybe SamplingState
decodeXb3SampledHeader forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup HeaderName
xb3SampledHeader RequestHeaders
hs
      debug :: Maybe SamplingState
debug = ByteString -> Maybe SamplingState
decodeXb3FlagsHeader forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a b. Eq a => a -> [(a, b)] -> Maybe b
Prelude.lookup HeaderName
xb3FlagsHeader RequestHeaders
hs
      -- NOTE: Debug implies Accept (https://github.com/openzipkin/b3-propagation#debug-flag)
      samplingState :: SamplingState
samplingState = forall a. a -> Maybe a -> a
fromMaybe SamplingState
Defer forall a b. (a -> b) -> a -> b
$ Maybe SamplingState
sampled forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe SamplingState
debug
  let traceFlags :: TraceFlags
traceFlags = if SamplingState
samplingState forall a. Eq a => a -> a -> Bool
== SamplingState
Accept Bool -> Bool -> Bool
|| SamplingState
samplingState forall a. Eq a => a -> a -> Bool
== SamplingState
Debug then Word8 -> TraceFlags
TraceFlags Word8
1 else Word8 -> TraceFlags
TraceFlags Word8
0

  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    Core.SpanContext
      { traceId :: TraceId
traceId = TraceId
traceId
      , spanId :: SpanId
spanId = SpanId
spanId
      , isRemote :: Bool
isRemote = Bool
True
      , traceFlags :: TraceFlags
traceFlags = TraceFlags
traceFlags
      , traceState :: TraceState
traceState = [(Key, Value)] -> TraceState
TS.TraceState [(Text -> Key
TS.Key Text
"sampling-state", SamplingState -> Value
samplingStateToValue SamplingState
samplingState)]
      }