CollectorReporter.hs

{-# 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 { fromJaegerCollector :: BatchEnv }

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

makeLenses ''JaegerCollectorOptions

jaegerCollectorOptions :: Manager -> Text -> JaegerCollectorOptions
jaegerCollectorOptions mgr srv = JaegerCollectorOptions
    { _jcoManager     = mgr
    , _jcoServiceName = srv
    , _jcoServiceTags = mempty
    , _jcoAddr        = defaultJaegerCollectorAddr
    , _jcoErrorLog    = defaultErrorLog
    }

defaultJaegerCollectorAddr :: Addr 'HTTP
defaultJaegerCollectorAddr = HTTPAddr "127.0.0.1" 14268 False

newJaegerCollector :: JaegerCollectorOptions -> IO JaegerCollector
newJaegerCollector opt@JaegerCollectorOptions{..} = do
    rq <- mkReq
    fmap JaegerCollector
        . newBatchEnv
        . set boptErrorLog _jcoErrorLog . batchOptions
        $ reporter _jcoManager _jcoErrorLog rq tproc
  where
    mkReq = do
        rq <- parseRequest
                    $ "http://" <> view (jcoAddr . addrHostName) opt
                   <> ":"
                   <> show (view (jcoAddr . addrPort) opt)
                   <> "/api/traces?format=jaeger.thrift"
        pure rq { method = "POST", secure = view (jcoAddr . addrSecure) opt }

    tproc = toThriftProcess _jcoServiceName _jcoServiceTags


closeJaegerCollector :: JaegerCollector -> IO ()
closeJaegerCollector = closeBatchEnv . fromJaegerCollector

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


jaegerCollectorReporter :: MonadIO m => JaegerCollector -> FinishedSpan -> m ()
jaegerCollectorReporter = batchReporter . fromJaegerCollector


reporter
    :: Manager
    -> (Builder -> IO ())
    -> Request
    -> Thrift.Process
    -> [FinishedSpan]
    -> IO ()
reporter mgr errlog rq tproc (fromList -> spans) = do
    rs <- responseStatus <$> httpLbs rq { requestBody = body } mgr
    unless (statusIsSuccessful rs) $
        errlog $ shortByteString "Error from Jaeger Collector: "
              <> intDec (statusCode rs)
              <> char8 '\n'
  where
    body = RequestBodyLBS . serializeBatch $ toThriftBatch tproc spans

    -- nb. collector accepts 'BinaryProtocol', but agent 'CompactProtocol'
    serializeBatch = Thrift.encode_Batch (BinaryProtocol EmptyTransport)