{-# 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.Lazy           (fromStrict)
import           Data.ByteString.Builder
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 qualified Pinch

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
. Method -> ByteString
fromStrict (Method -> ByteString) -> (Batch -> Method) -> Batch -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Batch -> Method
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 -> Method
serializeBatch = Protocol -> Batch -> Method
forall a. Pinchable a => Protocol -> a -> Method
Pinch.encode Protocol
Pinch.binaryProtocol