{-# LANGUAGE DataKinds         #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}
{-# LANGUAGE StrictData        #-}
{-# 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.Text                      (Text)
import           Data.Vector                    (fromList)
import qualified Jaeger.Types                   as Thrift
import           Network.HTTP.Client
import           Network.HTTP.Types             (hContentType)
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
    { _jcoManager :: Manager
_jcoManager     = Manager
mgr
    , _jcoServiceName :: Text
_jcoServiceName = Text
srv
    , _jcoServiceTags :: Tags
_jcoServiceTags = 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
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap BatchEnv -> JaegerCollector
JaegerCollector
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. BatchOptions -> IO BatchEnv
newBatchEnv
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s t a b. ASetter s t a b -> b -> s -> t
set Lens' BatchOptions (Builder -> IO ())
boptErrorLog Builder -> IO ()
_jcoErrorLog forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([FinishedSpan] -> IO ()) -> BatchOptions
batchOptions
        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 <- forall (m :: * -> *). MonadThrow m => HostName -> m Request
parseRequest
                    forall a b. (a -> b) -> a -> b
$ HostName
"http://" forall a. Semigroup a => a -> a -> a
<> forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view (Lens' JaegerCollectorOptions (Addr 'HTTP)
jcoAddr forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: Protocol). Lens' (Addr a) HostName
addrHostName) JaegerCollectorOptions
opt
                   forall a. Semigroup a => a -> a -> a
<> HostName
":"
                   forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> HostName
show (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view (Lens' JaegerCollectorOptions (Addr 'HTTP)
jcoAddr forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: Protocol). Lens' (Addr a) Port
addrPort) JaegerCollectorOptions
opt)
                   forall a. Semigroup a => a -> a -> a
<> HostName
"/api/traces?format=jaeger.thrift"
        forall (f :: * -> *) a. Applicative f => a -> f a
pure Request
rq
            { method :: Method
method = Method
"POST"
            , secure :: Bool
secure = forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view (Lens' JaegerCollectorOptions (Addr 'HTTP)
jcoAddr forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' (Addr 'HTTP) Bool
addrSecure) JaegerCollectorOptions
opt
            , requestHeaders :: RequestHeaders
requestHeaders = [(HeaderName
hContentType, Method
"application/x-thrift")]
            }

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


closeJaegerCollector :: JaegerCollector -> IO ()
closeJaegerCollector :: JaegerCollector -> IO ()
closeJaegerCollector = BatchEnv -> IO ()
closeBatchEnv 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 :: forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
JaegerCollectorOptions -> (JaegerCollector -> m a) -> m a
withJaegerCollector JaegerCollectorOptions
opts =
    forall (m :: * -> *) a c b.
MonadMask m =>
m a -> (a -> m c) -> (a -> m b) -> m b
bracket (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ JaegerCollectorOptions -> IO JaegerCollector
newJaegerCollector JaegerCollectorOptions
opts) (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. JaegerCollector -> IO ()
closeJaegerCollector)


jaegerCollectorReporter :: MonadIO m => JaegerCollector -> FinishedSpan -> m ()
jaegerCollectorReporter :: forall (m :: * -> *).
MonadIO m =>
JaegerCollector -> FinishedSpan -> m ()
jaegerCollectorReporter = forall (m :: * -> *). MonadIO m => BatchEnv -> FinishedSpan -> m ()
batchReporter 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 (forall a. [a] -> Vector a
fromList -> Vector FinishedSpan
spans) = do
    Status
rs <- forall body. Response body -> Status
responseStatus forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Request -> Manager -> IO (Response ())
httpNoBody Request
rq { requestBody :: RequestBody
requestBody = RequestBody
body } Manager
mgr
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Status -> Bool
statusIsSuccessful Status
rs) forall a b. (a -> b) -> a -> b
$
        Builder -> IO ()
errlog forall a b. (a -> b) -> a -> b
$ ShortByteString -> Builder
shortByteString ShortByteString
"Error from Jaeger Collector: "
              forall a. Semigroup a => a -> a -> a
<> Int -> Builder
intDec (Status -> Int
statusCode Status
rs)
              forall a. Semigroup a => a -> a -> a
<> Char -> Builder
char8 Char
'\n'
  where
    body :: RequestBody
body = Method -> RequestBody
RequestBodyBS forall b c a. (b -> c) -> (a -> b) -> a -> c
. Batch -> Method
serializeBatch 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 = forall a. Pinchable a => Protocol -> a -> Method
Pinch.encode Protocol
Pinch.binaryProtocol