{-# 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 = 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
Addr 'HTTP
Tags
Builder -> IO ()
_jcoManager :: JaegerCollectorOptions -> Manager
_jcoServiceName :: JaegerCollectorOptions -> Text
_jcoServiceTags :: JaegerCollectorOptions -> Tags
_jcoAddr :: JaegerCollectorOptions -> Addr 'HTTP
_jcoErrorLog :: JaegerCollectorOptions -> Builder -> IO ()
_jcoManager :: Manager
_jcoServiceName :: Text
_jcoServiceTags :: Tags
_jcoAddr :: Addr 'HTTP
_jcoErrorLog :: Builder -> IO ()
..} = do
    Request
rq <- IO Request
mkReq
    (BatchEnv -> JaegerCollector) -> IO BatchEnv -> IO JaegerCollector
forall a b. (a -> b) -> IO a -> IO b
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) (f :: * -> *).
Functor f =>
(HostName -> f HostName) -> Addr a -> f (Addr a)
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) (f :: * -> *).
Functor f =>
(Port -> f Port) -> Addr a -> f (Addr a)
addrPort) JaegerCollectorOptions
opt)
                   HostName -> HostName -> HostName
forall a. Semigroup a => a -> a -> a
<> HostName
"/api/traces?format=jaeger.thrift"
        Request -> IO Request
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Request
rq
            { method = "POST"
            , secure = view (jcoAddr . addrSecure) opt
            , requestHeaders = [(hContentType, "application/x-thrift")]
            }

    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 :: forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
JaegerCollectorOptions -> (JaegerCollector -> m a) -> m a
withJaegerCollector JaegerCollectorOptions
opts =
    m JaegerCollector
-> (JaegerCollector -> m ()) -> (JaegerCollector -> m a) -> m a
forall (m :: * -> *) a c b.
(HasCallStack, MonadMask m) =>
m a -> (a -> m c) -> (a -> m b) -> m b
bracket (IO JaegerCollector -> m JaegerCollector
forall a. IO a -> m a
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 a. IO a -> m a
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 :: forall (m :: * -> *).
MonadIO m =>
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 () -> Status
forall body. Response body -> Status
responseStatus (Response () -> Status) -> IO (Response ()) -> IO Status
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Request -> Manager -> IO (Response ())
httpNoBody Request
rq { 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 = Method -> RequestBody
RequestBodyBS (Method -> RequestBody)
-> (Batch -> Method) -> Batch -> RequestBody
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