{-# 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.Monoid
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 = Prism' TextMap SpanContext -> Carrier TextMap
forall a. Prism' a SpanContext -> Carrier a
Carrier Prism' TextMap SpanContext
_JaegerTextMap Carrier TextMap
-> Rec Carrier '[Headers] -> Propagation '[TextMap, Headers]
forall u (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& Prism' Headers SpanContext -> Carrier Headers
forall a. Prism' a SpanContext -> Carrier a
Carrier Prism' Headers SpanContext
_JaegerHeaders Carrier Headers -> Rec Carrier '[] -> Rec Carrier '[Headers]
forall u (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& Rec Carrier '[]
forall u (a :: u -> *). Rec a '[]
RNil


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

    toCtx :: TextMap -> Maybe SpanContext
toCtx TextMap
m =
          (SpanContext -> SpanContext)
-> Maybe SpanContext -> Maybe SpanContext
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ASetter SpanContext SpanContext TextMap TextMap
-> TextMap -> SpanContext -> SpanContext
forall s t a b. ASetter s t a b -> b -> s -> t
set ASetter SpanContext SpanContext TextMap TextMap
Lens' SpanContext TextMap
ctxBaggage
                    ((Text -> Text -> Bool) -> TextMap -> TextMap
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))
        (Maybe SpanContext -> Maybe SpanContext)
-> Maybe SpanContext -> Maybe SpanContext
forall a b. (a -> b) -> a -> b
$ Text -> TextMap -> Maybe Text
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Text
"uber-trace-id" TextMap
m Maybe Text -> (Text -> Maybe SpanContext) -> Maybe SpanContext
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Getting (First SpanContext) Text SpanContext
-> Text -> Maybe SpanContext
forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview Getting (First SpanContext) Text SpanContext
Prism' Text SpanContext
_UberTraceId

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

_UberTraceId :: Prism' Text SpanContext
_UberTraceId :: p SpanContext (f SpanContext) -> p Text (f Text)
_UberTraceId = (SpanContext -> Text)
-> (Text -> Maybe SpanContext) -> Prism' Text SpanContext
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 = Getting Text TraceID Text -> TraceID -> Text
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Text TraceID Text
forall a. AsHex a => Getter a Text
hexText TraceID
ctxTraceID
            spanid :: Text
spanid  = Getting Text Word64 Text -> Word64 -> Text
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Text Word64 Text
forall a. AsHex a => Getter a Text
hexText Word64
ctxSpanID
            parent :: Text
parent  = Text -> (Word64 -> Text) -> Maybe Word64 -> Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Text
forall a. Monoid a => a
mempty (Getting Text Word64 Text -> Word64 -> Text
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Text Word64 Text
forall a. AsHex a => Getter a Text
hexText) Maybe Word64
ctxParentSpanID
            flags :: Text
flags   = if Getting Bool SpanContext Bool -> SpanContext -> Bool
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view ((Sampled -> Const Bool Sampled)
-> SpanContext -> Const Bool SpanContext
Lens' SpanContext Sampled
ctxSampled ((Sampled -> Const Bool Sampled)
 -> SpanContext -> Const Bool SpanContext)
-> ((Bool -> Const Bool Bool) -> Sampled -> Const Bool Sampled)
-> Getting Bool SpanContext Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AReview Bool Sampled -> Getter Sampled Bool
forall t b. AReview t b -> Getter b t
re AReview Bool Sampled
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 Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Word
sampledFlag Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
> Word
0 Bool -> Bool -> Bool
|| Word
fs Word -> Word -> Word
forall a. Bits a => a -> a -> a
.&. Word
debugFlag Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
> Word
0
         in case (Char -> Bool) -> Text -> [Text]
Text.split (Char -> Char -> Bool
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
                    (TraceID
 -> Word64 -> Maybe Word64 -> Sampled -> TextMap -> SpanContext)
-> Maybe TraceID
-> Maybe
     (Word64 -> Maybe Word64 -> Sampled -> TextMap -> SpanContext)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Getting (First TraceID) Hex TraceID -> Hex -> Maybe TraceID
forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview Getting (First TraceID) Hex TraceID
forall a. AsHex a => Prism' Hex a
_Hex (Text -> Hex
knownHex Text
traceid)
                    Maybe (Word64 -> Maybe Word64 -> Sampled -> TextMap -> SpanContext)
-> Maybe Word64
-> Maybe (Maybe Word64 -> Sampled -> TextMap -> SpanContext)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Getting (First Word64) Hex Word64 -> Hex -> Maybe Word64
forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview Getting (First Word64) Hex Word64
forall a. AsHex a => Prism' Hex a
_Hex (Text -> Hex
knownHex Text
spanid)
                    Maybe (Maybe Word64 -> Sampled -> TextMap -> SpanContext)
-> Maybe (Maybe Word64)
-> Maybe (Sampled -> TextMap -> SpanContext)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Word64 -> Maybe (Maybe Word64)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Word64
forall a. Maybe a
Nothing
                    Maybe (Sampled -> TextMap -> SpanContext)
-> Maybe Sampled -> Maybe (TextMap -> SpanContext)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (String -> Maybe Sampled)
-> ((Word, Text) -> Maybe Sampled)
-> Either String (Word, Text)
-> Maybe Sampled
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Maybe Sampled -> String -> Maybe Sampled
forall a b. a -> b -> a
const (Maybe Sampled -> String -> Maybe Sampled)
-> Maybe Sampled -> String -> Maybe Sampled
forall a b. (a -> b) -> a -> b
$ Sampled -> Maybe Sampled
forall a. a -> Maybe a
Just Sampled
NotSampled)
                               (Sampled -> Maybe Sampled
forall a. a -> Maybe a
Just (Sampled -> Maybe Sampled)
-> ((Word, Text) -> Sampled) -> (Word, Text) -> Maybe Sampled
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting Sampled Bool Sampled -> Bool -> Sampled
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Sampled Bool Sampled
Iso' Bool Sampled
_IsSampled (Bool -> Sampled)
-> ((Word, Text) -> Bool) -> (Word, Text) -> Sampled
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word -> Bool
shouldSample (Word -> Bool) -> ((Word, Text) -> Word) -> (Word, Text) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word, Text) -> Word
forall a b. (a, b) -> a
fst)
                               (Reader Word
forall a. Integral a => Reader a
Text.decimal Text
flags)
                    Maybe (TextMap -> SpanContext)
-> Maybe TextMap -> Maybe SpanContext
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TextMap -> Maybe TextMap
forall (f :: * -> *) a. Applicative f => a -> f a
pure TextMap
forall a. Monoid a => a
mempty

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