{-# LANGUAGE ScopedTypeVariables #-}

module Tracing.Zipkin (
    publishZipkin,
    ZipkinSpan(..)
    ) where

import Tracing.Core (Span(..), SpanId(..), OpName(..), TraceId(..), SpanContext(..),
    SpanRelation(..))

import Control.Monad.Trans (liftIO, MonadIO)
import Control.Monad (void)
import Data.Monoid ((<>), mempty)
import Data.Aeson
import Data.Maybe (fromMaybe)
import Data.Int (Int64)
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.ByteString.Char8 as BS
import qualified Data.ByteString.Lazy as LBS
import qualified Data.ByteString.Lex.Integral as BS
import qualified Data.HashMap.Strict as HM
import Network.HTTP.Client


-- | Publish 'Span' in the <https://zipkin.io/pages/data_model.html Zipkin format> . No call is made
-- on an empty
publishZipkin :: MonadIO m =>
    String -- ^ The address of the backend server
    -> Manager
    -> [Span] -- ^ The traced spans to send to a Zipkin backend
    -> m (Maybe (Response T.Text))
publishZipkin :: String -> Manager -> [Span] -> m (Maybe (Response Text))
publishZipkin String
_ Manager
_ [] = Maybe (Response Text) -> m (Maybe (Response Text))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Response Text)
forall a. Maybe a
Nothing
publishZipkin String
destination Manager
manager [Span]
spans =
    IO (Maybe (Response Text)) -> m (Maybe (Response Text))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe (Response Text)) -> m (Maybe (Response Text)))
-> (IO (Response ByteString) -> IO (Maybe (Response Text)))
-> IO (Response ByteString)
-> m (Maybe (Response Text))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Response ByteString -> Maybe (Response Text))
-> IO (Response ByteString) -> IO (Maybe (Response Text))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Response Text -> Maybe (Response Text)
forall a. a -> Maybe a
Just (Response Text -> Maybe (Response Text))
-> (Response ByteString -> Response Text)
-> Response ByteString
-> Maybe (Response Text)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString -> Text) -> Response ByteString -> Response Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> Text
decode) (IO (Response ByteString) -> m (Maybe (Response Text)))
-> IO (Response ByteString) -> m (Maybe (Response Text))
forall a b. (a -> b) -> a -> b
$ Request -> Manager -> IO (Response ByteString)
httpLbs Request
zipkinReq Manager
manager
    where
        decode :: ByteString -> Text
decode = ByteString -> Text
T.decodeUtf8 (ByteString -> Text)
-> (ByteString -> ByteString) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
LBS.toStrict
        req :: Request
req = String -> Request
parseRequest_ String
destination
        body :: RequestBody
body = ByteString -> RequestBody
RequestBodyLBS (ByteString -> RequestBody)
-> ([ZipkinSpan] -> ByteString) -> [ZipkinSpan] -> RequestBody
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ZipkinSpan] -> ByteString
forall a. ToJSON a => a -> ByteString
encode ([ZipkinSpan] -> RequestBody) -> [ZipkinSpan] -> RequestBody
forall a b. (a -> b) -> a -> b
$ Span -> ZipkinSpan
ZipkinSpan (Span -> ZipkinSpan) -> [Span] -> [ZipkinSpan]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Span]
spans
        zipkinReq :: Request
zipkinReq = Request
req { method :: ByteString
method = ByteString
"POST", requestBody :: RequestBody
requestBody = RequestBody
body, requestHeaders :: RequestHeaders
requestHeaders = [(HeaderName
"content-type", ByteString
"application/json")]}

newtype ZipkinSpan = ZipkinSpan Span
instance ToJSON ZipkinSpan where
    toJSON :: ZipkinSpan -> Value
toJSON (ZipkinSpan Span
span) = [Pair] -> Value
object ([Pair] -> Value) -> [Pair] -> Value
forall a b. (a -> b) -> a -> b
$ [
        Key
"traceId" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= TraceId -> Text
unTrace (SpanContext -> TraceId
traceId (SpanContext -> TraceId) -> SpanContext -> TraceId
forall a b. (a -> b) -> a -> b
$ Span -> SpanContext
context Span
span),
        Key
"id" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.=  SpanId -> Text
unSpan (SpanContext -> SpanId
spanId (SpanContext -> SpanId) -> SpanContext -> SpanId
forall a b. (a -> b) -> a -> b
$ Span -> SpanContext
context Span
span),
        Key
"name" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.=  OpName -> Text
unOp (Span -> OpName
operationName Span
span),
        Key
"timestamp" Key -> Int64 -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= (POSIXTime -> Int64
forall a b. (RealFrac a, Integral b) => a -> b
floor (POSIXTime -> Int64)
-> (POSIXTime -> POSIXTime) -> POSIXTime -> Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. POSIXTime -> POSIXTime
toMicros (POSIXTime -> Int64) -> POSIXTime -> Int64
forall a b. (a -> b) -> a -> b
$ Span -> POSIXTime
timestamp Span
span :: Int64),
        Key
"kind" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= (Text
"CLIENT"::T.Text),
        Key
"duration" Key -> POSIXTime -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= (POSIXTime -> POSIXTime
toMicros (POSIXTime -> POSIXTime) -> POSIXTime -> POSIXTime
forall a b. (a -> b) -> a -> b
$ Span -> POSIXTime
duration Span
span),
        Key
"debug" Key -> Bool -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= (Span -> Bool
debug Span
span),
        Key
"localEndpoint" Key -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= ([Pair] -> Value
object [Key
"serviceName" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= (Span -> Text
serviceName Span
span)])
        ] [Pair] -> [Pair] -> [Pair]
forall a. Semigroup a => a -> a -> a
<>
        [SpanRelation] -> [Pair]
forall a. KeyValue a => [SpanRelation] -> [a]
parentId (Span -> [SpanRelation]
relations Span
span)
        where
            toMicros :: POSIXTime -> POSIXTime
toMicros = POSIXTime -> POSIXTime -> POSIXTime
forall a. Num a => a -> a -> a
(*) POSIXTime
1000000
            unOp :: OpName -> Text
unOp (OpName Text
n) = Text
n
            zipkinFormatId :: Maybe ByteString -> Text
zipkinFormatId = Int -> Text -> Text
padLeft Int
16 (Text -> Text)
-> (Maybe ByteString -> Text) -> Maybe ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
T.pack (String -> Text)
-> (Maybe ByteString -> String) -> Maybe ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String
BS.unpack (ByteString -> String)
-> (Maybe ByteString -> ByteString) -> Maybe ByteString -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Maybe ByteString -> ByteString
forall a. a -> Maybe a -> a
fromMaybe ByteString
"-1"
            unTrace :: TraceId -> Text
unTrace (TraceId Int64
t) = Maybe ByteString -> Text
zipkinFormatId (Maybe ByteString -> Text) -> Maybe ByteString -> Text
forall a b. (a -> b) -> a -> b
$ Int64 -> Maybe ByteString
forall a. Integral a => a -> Maybe ByteString
BS.packHexadecimal Int64
t
            unSpan :: SpanId -> Text
unSpan (SpanId Int64
s) = Maybe ByteString -> Text
zipkinFormatId (Maybe ByteString -> Text) -> Maybe ByteString -> Text
forall a b. (a -> b) -> a -> b
$ Int64 -> Maybe ByteString
forall a. Integral a => a -> Maybe ByteString
BS.packHexadecimal Int64
s
            parentId :: [SpanRelation] -> [a]
parentId (ChildOf SpanContext
ctx:[SpanRelation]
_) = [Key
"parentId" Key -> Text -> a
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= (SpanId -> Text
unSpan (SpanId -> Text) -> SpanId -> Text
forall a b. (a -> b) -> a -> b
$ SpanContext -> SpanId
spanId SpanContext
ctx)]
            parentId (FollowsFrom SpanContext
ctx:[SpanRelation]
_) = [Key
"parentId" Key -> Text -> a
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= (SpanId -> Text
unSpan (SpanId -> Text) -> SpanId -> Text
forall a b. (a -> b) -> a -> b
$ SpanContext -> SpanId
spanId SpanContext
ctx)]
            parentId [SpanRelation]
_ = []
            padLeft :: Int -> Text -> Text
padLeft Int
0 Text
txt = Text
txt
            padLeft Int
n Text
txt
                | Text -> Int
T.length Text
txt Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = Int -> Text -> Text
padLeft Int
n (Text
"0"Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>Text
txt)
                | Bool
otherwise = Text
txt