{-# LANGUAGE DataKinds         #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}

module OpenTracing.Jaeger.Propagation
    ( jaegerPropagation

    , _JaegerTextMap
    , _JaegerHeaders

    , _UberTraceId
    )
where

import           Control.Lens
import           Data.Bits
import qualified Data.HashMap.Strict     as HashMap
import           Data.Text               (Text, isPrefixOf)
import qualified Data.Text               as Text
import qualified Data.Text.Read          as Text
import           OpenTracing.Propagation
import           OpenTracing.Span
import           OpenTracing.Types


jaegerPropagation :: Propagation '[TextMap, Headers]
jaegerPropagation :: Propagation '[TextMap, Headers]
jaegerPropagation = forall a. Prism' a SpanContext -> Carrier a
Carrier Prism' TextMap SpanContext
_JaegerTextMap forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& forall a. Prism' a SpanContext -> Carrier a
Carrier Prism' Headers SpanContext
_JaegerHeaders forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& forall {u} (a :: u -> *). Rec a '[]
RNil


_JaegerTextMap :: Prism' TextMap SpanContext
_JaegerTextMap :: Prism' TextMap SpanContext
_JaegerTextMap = forall b s a. (b -> s) -> (s -> Maybe a) -> Prism s s a b
prism' SpanContext -> TextMap
fromCtx TextMap -> Maybe SpanContext
toCtx
  where
    fromCtx :: SpanContext -> TextMap
fromCtx SpanContext
c = forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList forall a b. (a -> b) -> a -> b
$
          (Text
"uber-trace-id", forall b (m :: * -> *) t. MonadReader b m => AReview t b -> m t
review Prism' Text SpanContext
_UberTraceId SpanContext
c)
        forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map (forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over forall s t a b. Field1 s t a b => Lens s t a b
_1 (Text
"uberctx-" forall a. Semigroup a => a -> a -> a
<>)) (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view (Lens' SpanContext TextMap
ctxBaggage forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (p :: * -> * -> *) (f :: * -> *) s a.
(Profunctor p, Contravariant f) =>
(s -> a) -> Optic' p f s a
to forall k v. HashMap k v -> [(k, v)]
HashMap.toList) SpanContext
c)

    toCtx :: TextMap -> Maybe SpanContext
toCtx TextMap
m =
          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall s t a b. ASetter s t a b -> b -> s -> t
set Lens' SpanContext TextMap
ctxBaggage
                    (forall k v. (k -> v -> Bool) -> HashMap k v -> HashMap k v
HashMap.filterWithKey (\Text
k Text
_ -> Text
"uberctx-" Text -> Text -> Bool
`isPrefixOf` Text
k) TextMap
m))
        forall a b. (a -> b) -> a -> b
$ forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Text
"uber-trace-id" TextMap
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview Prism' Text SpanContext
_UberTraceId

_JaegerHeaders :: Prism' Headers SpanContext
_JaegerHeaders :: Prism' Headers SpanContext
_JaegerHeaders = Iso' Headers TextMap
_HeadersTextMap forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prism' TextMap SpanContext
_JaegerTextMap

_UberTraceId :: Prism' Text SpanContext
_UberTraceId :: Prism' Text SpanContext
_UberTraceId = forall b s a. (b -> s) -> (s -> Maybe a) -> Prism s s a b
prism' SpanContext -> Text
fromCtx Text -> Maybe SpanContext
toCtx
  where
    fromCtx :: SpanContext -> Text
fromCtx c :: SpanContext
c@SpanContext{Maybe Word64
Word64
TextMap
Sampled
TraceID
ctxTraceID :: SpanContext -> TraceID
ctxSpanID :: SpanContext -> Word64
ctxParentSpanID :: SpanContext -> Maybe Word64
_ctxSampled :: SpanContext -> Sampled
_ctxBaggage :: SpanContext -> TextMap
_ctxBaggage :: TextMap
_ctxSampled :: Sampled
ctxParentSpanID :: Maybe Word64
ctxSpanID :: Word64
ctxTraceID :: TraceID
..} =
        let traceid :: Text
traceid = forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall a. AsHex a => Getter a Text
hexText TraceID
ctxTraceID
            spanid :: Text
spanid  = forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall a. AsHex a => Getter a Text
hexText Word64
ctxSpanID
            parent :: Text
parent  = forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. Monoid a => a
mempty (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall a. AsHex a => Getter a Text
hexText) Maybe Word64
ctxParentSpanID
            flags :: Text
flags   = if forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view (Lens' SpanContext Sampled
ctxSampled forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall t b. AReview t b -> Getter b t
re Iso' Bool Sampled
_IsSampled) SpanContext
c then Text
"1" else Text
"0"
         in Text -> [Text] -> Text
Text.intercalate Text
":" [Text
traceid, Text
spanid, Text
parent, Text
flags]

    toCtx :: Text -> Maybe SpanContext
toCtx Text
t =
        let sampledFlag :: Word
sampledFlag = Word
1 :: Word
            debugFlag :: Word
debugFlag   = Word
2 :: Word
            shouldSample :: Word -> Bool
shouldSample Word
fs = Word
fs forall a. Bits a => a -> a -> a
.&. Word
sampledFlag forall a. Ord a => a -> a -> Bool
> Word
0 Bool -> Bool -> Bool
|| Word
fs forall a. Bits a => a -> a -> a
.&. Word
debugFlag forall a. Ord a => a -> a -> Bool
> Word
0
         in case (Char -> Bool) -> Text -> [Text]
Text.split (forall a. Eq a => a -> a -> Bool
==Char
':') Text
t of
                [Text
traceid, Text
spanid, Text
_, Text
flags] -> TraceID
-> Word64 -> Maybe Word64 -> Sampled -> TextMap -> SpanContext
SpanContext
                    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview forall a. AsHex a => Prism' Hex a
_Hex (Text -> Hex
knownHex Text
traceid)
                    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview forall a. AsHex a => Prism' Hex a
_Hex (Text -> Hex
knownHex Text
spanid)
                    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
                    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just Sampled
NotSampled)
                               (forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Iso' Bool Sampled
_IsSampled forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Bool
shouldSample forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst)
                               (forall a. Integral a => Reader a
Text.decimal Text
flags)
                    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty

                [Text]
_ -> forall a. Maybe a
Nothing