{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
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]
}
b3Extractor :: [(HeaderName, ByteString)] -> Maybe Core.SpanContext
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
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
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
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)]
}