{-|
Module: OpenTracing.Propagation

Types and functions for serializing and deserializing `SpanContext`s across
process boundaries.

One of the big motiviating use cases for propagation is for tracing distributed
executions through RPC calls.
-}
{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE LambdaCase             #-}
{-# LANGUAGE OverloadedStrings      #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TupleSections          #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeOperators          #-}

module OpenTracing.Propagation
    ( TextMap
    , Headers
--  , Binary

    , Propagation
    , HasPropagation(..)

    , Carrier(..)
    , HasCarrier
    , HasCarriers
    , carrier

    , inject
    , extract

    , otPropagation
    , b3Propagation

    , _OTTextMap
    , _OTHeaders
    , _B3TextMap
    , _B3Headers

    , _HeadersTextMap

    -- * Re-exports from 'Data.Vinyl'
    , Rec ((:&), RNil)
    , rappend, (<+>)
    , rcast
    )
where

import           Control.Applicative     ((<|>))
import           Control.Lens
import           Data.Bool               (bool)
import           Data.ByteString.Builder (toLazyByteString)
import qualified Data.CaseInsensitive    as CI
import           Data.HashMap.Strict     (HashMap)
import qualified Data.HashMap.Strict     as HashMap
import           Data.Maybe              (catMaybes)
import           Data.Proxy
import           Data.Text               (Text, isPrefixOf, toLower)
import           Data.Text.Encoding      (decodeUtf8, encodeUtf8)
import qualified Data.Text.Read          as Text
import           Data.Vinyl
import           Data.Word
import           Network.HTTP.Types      (Header)
import           OpenTracing.Span
import           OpenTracing.Types
import           URI.ByteString          (urlDecodeQuery, urlEncodeQuery)


type TextMap = HashMap Text Text
type Headers = [Header]
--type Binary  = Lazy.ByteString

-- | A `Propagation` contains the different ways that a `SpanContext` can be
-- serialized and deserialized. For example @Propagation '[TextMap, Headers]@ indicates
-- support for serializing to `Header` or to `TextMap`.
--
-- @since 0.1.0.0
type Propagation carriers = Rec Carrier carriers

-- | A typeclass for application environments that contain a `Propagation`.
--
-- @since 0.1.0.0
class HasPropagation a p | a -> p where
    propagation :: Getting r a (Propagation p)

instance HasPropagation (Propagation p) p where
    propagation :: forall r. Getting r (Propagation p) (Propagation p)
propagation = forall a. a -> a
id

-- | `Carrier a` is a way to convert a `SpanContext` into or from an `a`.
--
-- @since 0.1.0.0
newtype Carrier a = Carrier { forall a. Carrier a -> Prism' a SpanContext
fromCarrier :: Prism' a SpanContext }

type HasCarrier  c  cs = c  ∈ cs
type HasCarriers cs ds = cs ⊆ ds

-- | Retrieve a (de)serialization lens from the application context for
-- format @c@.
--
-- @since 0.1.0.0
carrier
    :: ( HasCarrier     c cs
       , HasPropagation r cs
       )
    => proxy c -- ^ Proxy for the carrier type @c@.
    -> r -- ^ The application context
    -> Prism' c SpanContext
carrier :: forall c (cs :: [*]) r (proxy :: * -> *).
(HasCarrier c cs, HasPropagation r cs) =>
proxy c -> r -> Prism' c SpanContext
carrier proxy c
_c r
r = forall a. Carrier a -> Prism' a SpanContext
fromCarrier forall a b. (a -> b) -> a -> b
$ forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view (forall a (p :: [*]) r.
HasPropagation a p =>
Getting r a (Propagation p)
propagation forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {k} (r :: k) (record :: (k -> *) -> [k] -> *) (rs :: [k])
       (f :: k -> *) (g :: * -> *).
(RecElem record r r rs rs (RIndex r rs), RecElemFCtx record f,
 Functor g) =>
(f r -> g (f r)) -> record f rs -> g (record f rs)
rlens) r
r

-- | Serialize a `SpanContext` into the format `c` using a serializer from
-- the application context.
--
-- @since 0.1.0.0
inject
    :: forall c r p.
       ( HasCarrier     c p
       , HasPropagation r p
       )
    => r
    -> SpanContext
    -> c
inject :: forall c r (p :: [*]).
(HasCarrier c p, HasPropagation r p) =>
r -> SpanContext -> c
inject r
r = forall b (m :: * -> *) t. MonadReader b m => AReview t b -> m t
review (forall c (cs :: [*]) r (proxy :: * -> *).
(HasCarrier c cs, HasPropagation r cs) =>
proxy c -> r -> Prism' c SpanContext
carrier (forall {k} (t :: k). Proxy t
Proxy @c) r
r)

-- | Attempt to deserialize a `SpanContext` from the format @c@ using a deserializer
-- from the application context
--
-- @since 0.1.0.0
extract
    :: forall c r p.
       ( HasCarrier     c p
       , HasPropagation r p
       )
    => r
    -> c
    -> Maybe SpanContext
extract :: forall c r (p :: [*]).
(HasCarrier c p, HasPropagation r p) =>
r -> c -> Maybe SpanContext
extract r
r = forall s (m :: * -> *) a.
MonadReader s m =>
Getting (First a) s a -> m (Maybe a)
preview (forall c (cs :: [*]) r (proxy :: * -> *).
(HasCarrier c cs, HasPropagation r cs) =>
proxy c -> r -> Prism' c SpanContext
carrier (forall {k} (t :: k). Proxy t
Proxy @c) r
r)


-- | A propagation using an "ot" prefix.
-- No parent span id is propagated in OT.
otPropagation :: Propagation '[TextMap, Headers]
otPropagation :: Propagation '[TextMap, Headers]
otPropagation = forall a. Prism' a SpanContext -> Carrier a
Carrier Prism' TextMap SpanContext
_OTTextMap 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
_OTHeaders forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& forall {u} (a :: u -> *). Rec a '[]
RNil

-- | A propagation using an "x-b3" prefix for use with Zipkin.
b3Propagation :: Propagation '[TextMap, Headers]
b3Propagation :: Propagation '[TextMap, Headers]
b3Propagation = forall a. Prism' a SpanContext -> Carrier a
Carrier Prism' TextMap SpanContext
_B3TextMap 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
_B3Headers forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& forall {u} (a :: u -> *). Rec a '[]
RNil


_OTTextMap :: Prism' TextMap SpanContext
_OTTextMap :: Prism' TextMap SpanContext
_OTTextMap = 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 c :: SpanContext
c@SpanContext{Maybe Word64
Word64
TextMap
TraceID
Sampled
_ctxBaggage :: SpanContext -> TextMap
_ctxSampled :: SpanContext -> Sampled
ctxParentSpanID :: SpanContext -> Maybe Word64
ctxSpanID :: SpanContext -> Word64
ctxTraceID :: SpanContext -> TraceID
_ctxBaggage :: TextMap
_ctxSampled :: Sampled
ctxParentSpanID :: Maybe Word64
ctxSpanID :: Word64
ctxTraceID :: TraceID
..} = forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList forall a b. (a -> b) -> a -> b
$
          (Text
"ot-tracer-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)
        forall a. a -> [a] -> [a]
: (Text
"ot-tracer-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)
        forall a. a -> [a] -> [a]
: (Text
"ot-tracer-sampled", 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 Prism' Text Sampled
_OTSampled) 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
"ot-baggage-" forall a. Semigroup a => a -> a -> a
<>)) (forall k v. HashMap k v -> [(k, v)]
HashMap.toList TextMap
_ctxBaggage)

    toCtx :: TextMap -> Maybe SpanContext
toCtx TextMap
m = TraceID
-> Word64 -> Maybe Word64 -> Sampled -> TextMap -> SpanContext
SpanContext
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Text
"ot-tracer-traceid" 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 forall a. AsHex a => Prism' Hex a
_Hex forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Hex
knownHex)
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Text
"ot-tracer-spanid"  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 forall a. AsHex a => Prism' Hex a
_Hex forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Hex
knownHex)
        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 -- nb. parent is not propagated in OT
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Text
"ot-tracer-sampled" 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 Sampled
_OTSampled)
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k v. (k -> v -> Bool) -> HashMap k v -> HashMap k v
HashMap.filterWithKey (\Text
k Text
_ -> Text
"ot-baggage-" Text -> Text -> Bool
`isPrefixOf` Text
k) TextMap
m)


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

_OTSampled :: Prism' Text Sampled
_OTSampled :: Prism' Text Sampled
_OTSampled = forall b s a. (b -> s) -> (s -> Maybe a) -> Prism s s a b
prism' Sampled -> Text
enc Text -> Maybe Sampled
dec
    where
      enc :: Sampled -> Text
enc = \case Sampled
Sampled -> Text
"1"
                  Sampled
_       -> Text
"0"

      dec :: Text -> Maybe Sampled
dec = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> b -> a
const forall a. Maybe a
Nothing) forall a. a -> a
id
          forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Word8
x,Text
_) -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ if Word8
x forall a. Eq a => a -> a -> Bool
== (Word8
1 :: Word8) then Sampled
Sampled else Sampled
NotSampled)
          forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Integral a => Reader a
Text.decimal

_B3TextMap :: Prism' TextMap SpanContext
_B3TextMap :: Prism' TextMap SpanContext
_B3TextMap = 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 ctx :: SpanContext
ctx@SpanContext{Maybe Word64
Word64
TextMap
TraceID
Sampled
_ctxBaggage :: TextMap
_ctxSampled :: Sampled
ctxParentSpanID :: Maybe Word64
ctxSpanID :: Word64
ctxTraceID :: TraceID
_ctxBaggage :: SpanContext -> TextMap
_ctxSampled :: SpanContext -> Sampled
ctxParentSpanID :: SpanContext -> Maybe Word64
ctxSpanID :: SpanContext -> Word64
ctxTraceID :: SpanContext -> TraceID
..} = forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [Maybe a] -> [a]
catMaybes forall a b. (a -> b) -> a -> b
$
          forall a. a -> Maybe a
Just (Text
"x-b3-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)
        forall a. a -> [a] -> [a]
: forall a. a -> Maybe a
Just (Text
"x-b3-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)
        forall a. a -> [a] -> [a]
: forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Text
"x-b3-parentspanid",) forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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
        forall a. a -> [a] -> [a]
: forall a. a -> Maybe a
Just (Text
"x-b3-sampled", forall a. a -> a -> Bool -> a
bool Text
"false" Text
"true" forall a b. (a -> b) -> a -> b
$ 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
ctx)
        forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map (forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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
"ot-baggage-" forall a. Semigroup a => a -> a -> a
<>)) (forall k v. HashMap k v -> [(k, v)]
HashMap.toList TextMap
_ctxBaggage)

    toCtx :: TextMap -> Maybe SpanContext
toCtx TextMap
m = TraceID
-> Word64 -> Maybe Word64 -> Sampled -> TextMap -> SpanContext
SpanContext
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Text
"x-b3-traceid" 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 forall a. AsHex a => Prism' Hex a
_Hex forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Hex
knownHex)
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Text
"x-b3-spanid"  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 forall a. AsHex a => Prism' Hex a
_Hex forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Hex
knownHex)
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Text
"x-b3-parentspanid" 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 forall a. AsHex a => Prism' Hex a
_Hex forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Hex
knownHex)
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall {k} {a}.
(Hashable k, Eq a, IsString k, IsString a) =>
HashMap k a -> Maybe Sampled
b3Sampled TextMap
m forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> forall {k} {a}.
(Hashable k, Eq a, IsString k, IsString a) =>
HashMap k a -> Maybe Sampled
b3Debug TextMap
m forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> forall a. a -> Maybe a
Just Sampled
NotSampled)
        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall k v. (k -> v -> Bool) -> HashMap k v -> HashMap k v
HashMap.filterWithKey (\Text
k Text
_ -> Text
"ot-baggage-" Text -> Text -> Bool
`isPrefixOf` Text
k) TextMap
m)

    b3Sampled :: HashMap k a -> Maybe Sampled
b3Sampled HashMap k a
m = forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
"x-b3-sampled" HashMap k a
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        a
"true" -> forall a. a -> Maybe a
Just Sampled
Sampled
        a
_      -> forall a. Maybe a
Nothing

    b3Debug :: HashMap k a -> Maybe Sampled
b3Debug HashMap k a
m = forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup k
"x-b3-flags" HashMap k a
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        a
"1" -> forall a. a -> Maybe a
Just Sampled
Sampled
        a
_   -> forall a. Maybe a
Nothing

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

-- | Convert between a 'TextMap' and 'Headers'
--
-- Header field values are URL-encoded when converting from 'TextMap' to
-- 'Headers', and URL-decoded when converting the other way.
--
-- Note: validity of header fields is not checked (RFC 7230, 3.2.4)
_HeadersTextMap :: Iso' Headers TextMap
_HeadersTextMap :: Iso' Headers TextMap
_HeadersTextMap = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso Headers -> TextMap
toTextMap TextMap -> Headers
toHeaders
  where
    toHeaders :: TextMap -> Headers
toHeaders
        = forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (forall s. FoldCase s => s -> CI s
CI.mk forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
encodeUtf8)
                     (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall lazy strict. Strict lazy strict => Iso' lazy strict
strict forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> ByteString
toLazyByteString forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Builder
urlEncodeQuery forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
encodeUtf8))
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. HashMap k v -> [(k, v)]
HashMap.toList

    toTextMap :: Headers -> TextMap
toTextMap
        = forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap (Text -> Text
toLower forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
decodeUtf8 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s. CI s -> s
CI.original)
                     (ByteString -> Text
decodeUtf8 forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
urlDecodeQuery))