{-# LANGUAGE RankNTypes #-}
module Servant.Tracing (
    WithTracing,
    TracingInstructions(..),
    instructionsToHeader,
    getInstructions
    ) where

import Tracing.Core (Tracer, TraceId(..), SpanId(..), MonadTracer, TracingInstructions(..))

import Control.Arrow (first)
import Control.Monad.Trans (liftIO, MonadIO)
import qualified Data.Text as T
import qualified Data.ByteString.Char8 as BS
import qualified Data.ByteString.Lex.Integral as BS
import Data.Text.Read(hexadecimal)
import Data.Bits (testBit, (.|.))
import Data.Maybe (catMaybes, fromMaybe)
import Data.Monoid ((<>))
import Servant.API.Header (Header)
import System.Random (randomRIO)
import Web.HttpApiData (FromHttpApiData(..))

-- | Constrain the 'ServerT''s base monad such that it provides an instance of 'MonadTracer'
type WithTracing = Header "uber-trace-id" TracingInstructions


-- | Jaeger format: http://jaeger.readthedocs.io/en/latest/client_libraries/#propagation-format
-- This allows the trace backend to reassemble downstream traces.
instructionsToHeader :: TracingInstructions -> T.Text
instructionsToHeader :: TracingInstructions -> Text
instructionsToHeader TracingInstructions {$sel:traceId:TracingInstructions :: TracingInstructions -> TraceId
traceId=(TraceId Int64
tid), SpanId
$sel:spanId:TracingInstructions :: TracingInstructions -> SpanId
spanId :: SpanId
spanId, Maybe SpanId
$sel:parentSpanId:TracingInstructions :: TracingInstructions -> Maybe SpanId
parentSpanId :: Maybe SpanId
parentSpanId, Bool
$sel:sample:TracingInstructions :: TracingInstructions -> Bool
sample :: Bool
sample, Bool
$sel:debug:TracingInstructions :: TracingInstructions -> Bool
debug :: Bool
debug} =
    Int64 -> Text
toField Int64
tidText -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>Text
":"Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>
    (Int64 -> Text
toField (Int64 -> Text) -> Int64 -> Text
forall a b. (a -> b) -> a -> b
$ SpanId -> Int64
unSpanId SpanId
spanId) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":"Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>
    (Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
"" (Maybe Text -> Text) -> Maybe Text -> Text
forall a b. (a -> b) -> a -> b
$ (Int64 -> Text
toField (Int64 -> Text) -> (SpanId -> Int64) -> SpanId -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SpanId -> Int64
unSpanId) (SpanId -> Text) -> Maybe SpanId -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe SpanId
parentSpanId) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>
    (String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
setFlags)
    where
        unSpanId :: SpanId -> Int64
unSpanId (SpanId Int64
sid) = Int64
sid
        toField :: Int64 -> Text
toField = String -> Text
T.pack (String -> Text) -> (Int64 -> String) -> Int64 -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String
BS.unpack (ByteString -> String) -> (Int64 -> ByteString) -> Int64 -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Maybe ByteString -> ByteString
forall a. a -> Maybe a -> a
fromMaybe ByteString
"" (Maybe ByteString -> ByteString)
-> (Int64 -> Maybe ByteString) -> Int64 -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Maybe ByteString
forall a. Integral a => a -> Maybe ByteString
BS.packHexadecimal
        setFlags :: Int
        setFlags :: Int
setFlags = (if Bool
debug then Int
2 else Int
0) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. (if Bool
sample then Int
1 else Int
0) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.|. Int
0


instance FromHttpApiData TracingInstructions where
    parseUrlPiece ::
        T.Text
        -> Either T.Text TracingInstructions
    parseUrlPiece :: Text -> Either Text TracingInstructions
parseUrlPiece Text
raw =
        case Text -> Text -> [Text]
T.splitOn Text
":" Text
raw of
            [Text
rawTraceId, Text
rawSpanId, Text
rawParentId, Text
flags] -> let
                res :: Either String TracingInstructions
res = do
                    TraceId
traceId <- Int64 -> TraceId
TraceId (Int64 -> TraceId)
-> ((Integer, Text) -> Int64) -> (Integer, Text) -> TraceId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int64)
-> ((Integer, Text) -> Integer) -> (Integer, Text) -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, Text) -> Integer
forall a b. (a, b) -> a
fst ((Integer, Text) -> TraceId)
-> Either String (Integer, Text) -> Either String TraceId
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Reader Integer
forall a. Integral a => Reader a
hexadecimal Text
rawTraceId
                    SpanId
spanId <- Int64 -> SpanId
SpanId (Int64 -> SpanId)
-> ((Integer, Text) -> Int64) -> (Integer, Text) -> SpanId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int64)
-> ((Integer, Text) -> Integer) -> (Integer, Text) -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, Text) -> Integer
forall a b. (a, b) -> a
fst ((Integer, Text) -> SpanId)
-> Either String (Integer, Text) -> Either String SpanId
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Reader Integer
forall a. Integral a => Reader a
hexadecimal Text
rawSpanId
                    let resolvedPid :: Either String (Maybe Integer, Text)
resolvedPid = if Text -> Bool
T.null Text
rawParentId
                                      then (Maybe Integer, Text) -> Either String (Maybe Integer, Text)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Integer
forall a. Maybe a
Nothing, Text
"")
                                      else (Integer -> Maybe Integer)
-> (Integer, Text) -> (Maybe Integer, Text)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first Integer -> Maybe Integer
forall a. a -> Maybe a
Just ((Integer, Text) -> (Maybe Integer, Text))
-> Either String (Integer, Text)
-> Either String (Maybe Integer, Text)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Reader Integer
forall a. Integral a => Reader a
hexadecimal Text
rawParentId
                    Maybe SpanId
parentId <- (Integer -> SpanId) -> Maybe Integer -> Maybe SpanId
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int64 -> SpanId
SpanId (Int64 -> SpanId) -> (Integer -> Int64) -> Integer -> SpanId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral) (Maybe Integer -> Maybe SpanId)
-> ((Maybe Integer, Text) -> Maybe Integer)
-> (Maybe Integer, Text)
-> Maybe SpanId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe Integer, Text) -> Maybe Integer
forall a b. (a, b) -> a
fst ((Maybe Integer, Text) -> Maybe SpanId)
-> Either String (Maybe Integer, Text)
-> Either String (Maybe SpanId)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either String (Maybe Integer, Text)
resolvedPid
                    Int
flagField <- Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int)
-> ((Integer, Text) -> Integer) -> (Integer, Text) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer, Text) -> Integer
forall a b. (a, b) -> a
fst ((Integer, Text) -> Int)
-> Either String (Integer, Text) -> Either String Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Reader Integer
forall a. Integral a => Reader a
hexadecimal Text
flags
                    let [Bool
sample, Bool
debug]= [Int -> Bool
sampleFlag, Int -> Bool
debugFlag] [Int -> Bool] -> [Int] -> [Bool]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Int
flagField]
                    TracingInstructions -> Either String TracingInstructions
forall (f :: * -> *) a. Applicative f => a -> f a
pure TracingInstructions :: TraceId
-> SpanId -> Maybe SpanId -> Bool -> Bool -> TracingInstructions
TracingInstructions {
                       $sel:traceId:TracingInstructions :: TraceId
traceId = TraceId
traceId,
                       $sel:spanId:TracingInstructions :: SpanId
spanId = SpanId
spanId,
                       $sel:parentSpanId:TracingInstructions :: Maybe SpanId
parentSpanId = Maybe SpanId
parentId,
                       $sel:sample:TracingInstructions :: Bool
sample = Bool
sample,
                       $sel:debug:TracingInstructions :: Bool
debug = Bool
debug
                    }
                in case Either String TracingInstructions
res of
                    Left String
err -> Text -> Either Text TracingInstructions
forall a b. a -> Either a b
Left (Text -> Either Text TracingInstructions)
-> Text -> Either Text TracingInstructions
forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack String
err
                    Right TracingInstructions
val -> TracingInstructions -> Either Text TracingInstructions
forall a b. b -> Either a b
Right  TracingInstructions
val
            [Text]
_ -> Text -> Either Text TracingInstructions
forall a b. a -> Either a b
Left (Text -> Either Text TracingInstructions)
-> Text -> Either Text TracingInstructions
forall a b. (a -> b) -> a -> b
$ Text
raw Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" is not a valid uber-trace-id header"
        where
            sampleFlag :: Int -> Bool
            sampleFlag :: Int -> Bool
sampleFlag = (Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
`testBit` Int
0)
            debugFlag :: Int -> Bool
            debugFlag :: Int -> Bool
debugFlag = (Int -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
`testBit` Int
1)


-- TODO write a monad that wraps servant & determines if it should sample or not. Takes a sampling determinant. Only evaluates if the header is not present

-- | In the event that there are no 'TracingInstructions' for this call, generate new instructions.
--
-- This has a
getInstructions :: MonadIO m =>
    Bool
    -> Maybe TracingInstructions
    -> m TracingInstructions
getInstructions :: Bool -> Maybe TracingInstructions -> m TracingInstructions
getInstructions Bool
debug Maybe TracingInstructions
Nothing = do
    Int64
newTraceId <- IO Int64 -> m Int64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int64 -> m Int64) -> IO Int64 -> m Int64
forall a b. (a -> b) -> a -> b
$ (Int64, Int64) -> IO Int64
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Int64
0, Int64
forall a. Bounded a => a
maxBound)
    Int64
newSpanId <- IO Int64 -> m Int64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int64 -> m Int64) -> IO Int64 -> m Int64
forall a b. (a -> b) -> a -> b
$ (Int64, Int64) -> IO Int64
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Int64
0, Int64
forall a. Bounded a => a
maxBound)
    Int
sample <- IO Int -> m Int
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> m Int) -> IO Int -> m Int
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> IO Int
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Int
0, Int
1000)
    TracingInstructions -> m TracingInstructions
forall (f :: * -> *) a. Applicative f => a -> f a
pure TracingInstructions :: TraceId
-> SpanId -> Maybe SpanId -> Bool -> Bool -> TracingInstructions
TracingInstructions {
        $sel:traceId:TracingInstructions :: TraceId
traceId = Int64 -> TraceId
TraceId Int64
newTraceId,
        $sel:spanId:TracingInstructions :: SpanId
spanId = Int64 -> SpanId
SpanId Int64
newSpanId,
        $sel:parentSpanId:TracingInstructions :: Maybe SpanId
parentSpanId = Maybe SpanId
forall a. Maybe a
Nothing,
        Bool
debug :: Bool
$sel:debug:TracingInstructions :: Bool
debug,
        $sel:sample:TracingInstructions :: Bool
sample = Int
sample Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (Int
1::Int)
        }
getInstructions Bool
_ (Just TracingInstructions
inst) = TracingInstructions -> m TracingInstructions
forall (f :: * -> *) a. Applicative f => a -> f a
pure TracingInstructions
inst