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

module OpenTracing.Jaeger.CollectorReporter
    ( JaegerCollectorOptions
    , jaegerCollectorOptions
    , jcoManager
    , jcoServiceName
    , jcoServiceTags
    , jcoAddr
    , jcoErrorLog

    , defaultJaegerCollectorAddr

    , JaegerCollector
    , newJaegerCollector
    , closeJaegerCollector
    , withJaegerCollector

    , jaegerCollectorReporter

    , jaegerPropagation

    , newManager
    , defaultManagerSettings
    )
where

import           Control.Lens                   (makeLenses, set, view)
import           Control.Monad                  (unless)
import           Control.Monad.Catch
import           Control.Monad.IO.Class
import           Data.ByteString.Builder
import           Data.Monoid
import           Data.Text                      (Text)
import           Data.Vector                    (fromList)
import qualified Jaeger_Types                   as Thrift
import           Network.HTTP.Client
import           Network.HTTP.Types.Status
import           OpenTracing.Jaeger.Propagation (jaegerPropagation)
import           OpenTracing.Jaeger.Thrift
import           OpenTracing.Reporting
import           OpenTracing.Span
import           OpenTracing.Tags
import           OpenTracing.Types
import           Thrift.Protocol.Binary
import           Thrift.Transport.Empty


newtype JaegerCollector = JaegerCollector { JaegerCollector -> BatchEnv
fromJaegerCollector :: BatchEnv }

data JaegerCollectorOptions = JaegerCollectorOptions
    { JaegerCollectorOptions -> Manager
_jcoManager     :: Manager
    , JaegerCollectorOptions -> Text
_jcoServiceName :: Text
    , JaegerCollectorOptions -> Tags
_jcoServiceTags :: Tags
    , JaegerCollectorOptions -> Addr 'HTTP
_jcoAddr        :: Addr 'HTTP
    , JaegerCollectorOptions -> Builder -> IO ()
_jcoErrorLog    :: Builder -> IO ()
    }

makeLenses ''JaegerCollectorOptions

jaegerCollectorOptions :: Manager -> Text -> JaegerCollectorOptions
jaegerCollectorOptions :: Manager -> Text -> JaegerCollectorOptions
jaegerCollectorOptions Manager
mgr Text
srv = JaegerCollectorOptions :: Manager
-> Text
-> Tags
-> Addr 'HTTP
-> (Builder -> IO ())
-> JaegerCollectorOptions
JaegerCollectorOptions
    { _jcoManager :: Manager
_jcoManager     = Manager
mgr
    , _jcoServiceName :: Text
_jcoServiceName = Text
srv
    , _jcoServiceTags :: Tags
_jcoServiceTags = Tags
forall a. Monoid a => a
mempty
    , _jcoAddr :: Addr 'HTTP
_jcoAddr        = Addr 'HTTP
defaultJaegerCollectorAddr
    , _jcoErrorLog :: Builder -> IO ()
_jcoErrorLog    = Builder -> IO ()
defaultErrorLog
    }

defaultJaegerCollectorAddr :: Addr 'HTTP
defaultJaegerCollectorAddr :: Addr 'HTTP
defaultJaegerCollectorAddr = HostName -> Port -> Bool -> Addr 'HTTP
HTTPAddr HostName
"127.0.0.1" Port
14268 Bool
False

newJaegerCollector :: JaegerCollectorOptions -> IO JaegerCollector
newJaegerCollector :: JaegerCollectorOptions -> IO JaegerCollector
newJaegerCollector opt :: JaegerCollectorOptions
opt@JaegerCollectorOptions{Text
Manager
Tags
Addr 'HTTP
Builder -> IO ()
_jcoErrorLog :: Builder -> IO ()
_jcoAddr :: Addr 'HTTP
_jcoServiceTags :: Tags
_jcoServiceName :: Text
_jcoManager :: Manager
_jcoErrorLog :: JaegerCollectorOptions -> Builder -> IO ()
_jcoAddr :: JaegerCollectorOptions -> Addr 'HTTP
_jcoServiceTags :: JaegerCollectorOptions -> Tags
_jcoServiceName :: JaegerCollectorOptions -> Text
_jcoManager :: JaegerCollectorOptions -> Manager
..} = do
    Request
rq <- IO Request
mkReq
    (BatchEnv -> JaegerCollector) -> IO BatchEnv -> IO JaegerCollector
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap BatchEnv -> JaegerCollector
JaegerCollector
        (IO BatchEnv -> IO JaegerCollector)
-> (([FinishedSpan] -> IO ()) -> IO BatchEnv)
-> ([FinishedSpan] -> IO ())
-> IO JaegerCollector
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BatchOptions -> IO BatchEnv
newBatchEnv
        (BatchOptions -> IO BatchEnv)
-> (([FinishedSpan] -> IO ()) -> BatchOptions)
-> ([FinishedSpan] -> IO ())
-> IO BatchEnv
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASetter
  BatchOptions BatchOptions (Builder -> IO ()) (Builder -> IO ())
-> (Builder -> IO ()) -> BatchOptions -> BatchOptions
forall s t a b. ASetter s t a b -> b -> s -> t
set ASetter
  BatchOptions BatchOptions (Builder -> IO ()) (Builder -> IO ())
Lens' BatchOptions (Builder -> IO ())
boptErrorLog Builder -> IO ()
_jcoErrorLog (BatchOptions -> BatchOptions)
-> (([FinishedSpan] -> IO ()) -> BatchOptions)
-> ([FinishedSpan] -> IO ())
-> BatchOptions
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([FinishedSpan] -> IO ()) -> BatchOptions
batchOptions
        (([FinishedSpan] -> IO ()) -> IO JaegerCollector)
-> ([FinishedSpan] -> IO ()) -> IO JaegerCollector
forall a b. (a -> b) -> a -> b
$ Manager
-> (Builder -> IO ())
-> Request
-> Process
-> [FinishedSpan]
-> IO ()
reporter Manager
_jcoManager Builder -> IO ()
_jcoErrorLog Request
rq Process
tproc
  where
    mkReq :: IO Request
mkReq = do
        Request
rq <- HostName -> IO Request
forall (m :: * -> *). MonadThrow m => HostName -> m Request
parseRequest
                    (HostName -> IO Request) -> HostName -> IO Request
forall a b. (a -> b) -> a -> b
$ HostName
"http://" HostName -> HostName -> HostName
forall a. Semigroup a => a -> a -> a
<> Getting HostName JaegerCollectorOptions HostName
-> JaegerCollectorOptions -> HostName
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view ((Addr 'HTTP -> Const HostName (Addr 'HTTP))
-> JaegerCollectorOptions -> Const HostName JaegerCollectorOptions
Lens' JaegerCollectorOptions (Addr 'HTTP)
jcoAddr ((Addr 'HTTP -> Const HostName (Addr 'HTTP))
 -> JaegerCollectorOptions -> Const HostName JaegerCollectorOptions)
-> ((HostName -> Const HostName HostName)
    -> Addr 'HTTP -> Const HostName (Addr 'HTTP))
-> Getting HostName JaegerCollectorOptions HostName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (HostName -> Const HostName HostName)
-> Addr 'HTTP -> Const HostName (Addr 'HTTP)
forall (a :: Protocol). Lens' (Addr a) HostName
addrHostName) JaegerCollectorOptions
opt
                   HostName -> HostName -> HostName
forall a. Semigroup a => a -> a -> a
<> HostName
":"
                   HostName -> HostName -> HostName
forall a. Semigroup a => a -> a -> a
<> Port -> HostName
forall a. Show a => a -> HostName
show (Getting Port JaegerCollectorOptions Port
-> JaegerCollectorOptions -> Port
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view ((Addr 'HTTP -> Const Port (Addr 'HTTP))
-> JaegerCollectorOptions -> Const Port JaegerCollectorOptions
Lens' JaegerCollectorOptions (Addr 'HTTP)
jcoAddr ((Addr 'HTTP -> Const Port (Addr 'HTTP))
 -> JaegerCollectorOptions -> Const Port JaegerCollectorOptions)
-> ((Port -> Const Port Port)
    -> Addr 'HTTP -> Const Port (Addr 'HTTP))
-> Getting Port JaegerCollectorOptions Port
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Port -> Const Port Port) -> Addr 'HTTP -> Const Port (Addr 'HTTP)
forall (a :: Protocol). Lens' (Addr a) Port
addrPort) JaegerCollectorOptions
opt)
                   HostName -> HostName -> HostName
forall a. Semigroup a => a -> a -> a
<> HostName
"/api/traces?format=jaeger.thrift"
        Request -> IO Request
forall (f :: * -> *) a. Applicative f => a -> f a
pure Request
rq { method :: Method
method = Method
"POST", secure :: Bool
secure = Getting Bool JaegerCollectorOptions Bool
-> JaegerCollectorOptions -> Bool
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view ((Addr 'HTTP -> Const Bool (Addr 'HTTP))
-> JaegerCollectorOptions -> Const Bool JaegerCollectorOptions
Lens' JaegerCollectorOptions (Addr 'HTTP)
jcoAddr ((Addr 'HTTP -> Const Bool (Addr 'HTTP))
 -> JaegerCollectorOptions -> Const Bool JaegerCollectorOptions)
-> ((Bool -> Const Bool Bool)
    -> Addr 'HTTP -> Const Bool (Addr 'HTTP))
-> Getting Bool JaegerCollectorOptions Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Bool -> Const Bool Bool) -> Addr 'HTTP -> Const Bool (Addr 'HTTP)
Lens' (Addr 'HTTP) Bool
addrSecure) JaegerCollectorOptions
opt }

    tproc :: Process
tproc = Text -> Tags -> Process
toThriftProcess Text
_jcoServiceName Tags
_jcoServiceTags


closeJaegerCollector :: JaegerCollector -> IO ()
closeJaegerCollector :: JaegerCollector -> IO ()
closeJaegerCollector = BatchEnv -> IO ()
closeBatchEnv (BatchEnv -> IO ())
-> (JaegerCollector -> BatchEnv) -> JaegerCollector -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JaegerCollector -> BatchEnv
fromJaegerCollector

withJaegerCollector
    :: ( MonadIO   m
       , MonadMask m
       )
    => JaegerCollectorOptions
    -> (JaegerCollector -> m a)
    -> m a
withJaegerCollector :: JaegerCollectorOptions -> (JaegerCollector -> m a) -> m a
withJaegerCollector JaegerCollectorOptions
opts =
    m JaegerCollector
-> (JaegerCollector -> m ()) -> (JaegerCollector -> m a) -> m a
forall (m :: * -> *) a c b.
MonadMask m =>
m a -> (a -> m c) -> (a -> m b) -> m b
bracket (IO JaegerCollector -> m JaegerCollector
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO JaegerCollector -> m JaegerCollector)
-> IO JaegerCollector -> m JaegerCollector
forall a b. (a -> b) -> a -> b
$ JaegerCollectorOptions -> IO JaegerCollector
newJaegerCollector JaegerCollectorOptions
opts) (IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ())
-> (JaegerCollector -> IO ()) -> JaegerCollector -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JaegerCollector -> IO ()
closeJaegerCollector)


jaegerCollectorReporter :: MonadIO m => JaegerCollector -> FinishedSpan -> m ()
jaegerCollectorReporter :: JaegerCollector -> FinishedSpan -> m ()
jaegerCollectorReporter = BatchEnv -> FinishedSpan -> m ()
forall (m :: * -> *). MonadIO m => BatchEnv -> FinishedSpan -> m ()
batchReporter (BatchEnv -> FinishedSpan -> m ())
-> (JaegerCollector -> BatchEnv)
-> JaegerCollector
-> FinishedSpan
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JaegerCollector -> BatchEnv
fromJaegerCollector


reporter
    :: Manager
    -> (Builder -> IO ())
    -> Request
    -> Thrift.Process
    -> [FinishedSpan]
    -> IO ()
reporter :: Manager
-> (Builder -> IO ())
-> Request
-> Process
-> [FinishedSpan]
-> IO ()
reporter Manager
mgr Builder -> IO ()
errlog Request
rq Process
tproc ([FinishedSpan] -> Vector FinishedSpan
forall a. [a] -> Vector a
fromList -> Vector FinishedSpan
spans) = do
    Status
rs <- Response ByteString -> Status
forall body. Response body -> Status
responseStatus (Response ByteString -> Status)
-> IO (Response ByteString) -> IO Status
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Request -> Manager -> IO (Response ByteString)
httpLbs Request
rq { requestBody :: RequestBody
requestBody = RequestBody
body } Manager
mgr
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Status -> Bool
statusIsSuccessful Status
rs) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        Builder -> IO ()
errlog (Builder -> IO ()) -> Builder -> IO ()
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Builder
shortByteString ShortByteString
"Error from Jaeger Collector: "
              Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Int -> Builder
intDec (Status -> Int
statusCode Status
rs)
              Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Char -> Builder
char8 Char
'\n'
  where
    body :: RequestBody
body = ByteString -> RequestBody
RequestBodyLBS (ByteString -> RequestBody)
-> (Batch -> ByteString) -> Batch -> RequestBody
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Batch -> ByteString
serializeBatch (Batch -> RequestBody) -> Batch -> RequestBody
forall a b. (a -> b) -> a -> b
$ Process -> Vector FinishedSpan -> Batch
toThriftBatch Process
tproc Vector FinishedSpan
spans

    -- nb. collector accepts 'BinaryProtocol', but agent 'CompactProtocol'
    serializeBatch :: Batch -> ByteString
serializeBatch = BinaryProtocol EmptyTransport -> Batch -> ByteString
forall p. StatelessProtocol p => p -> Batch -> ByteString
Thrift.encode_Batch (EmptyTransport -> BinaryProtocol EmptyTransport
forall a. Transport a => a -> BinaryProtocol a
BinaryProtocol EmptyTransport
EmptyTransport)