{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE NamedFieldPuns             #-}
{-# LANGUAGE OverloadedStrings          #-}
{-# LANGUAGE RecordWildCards            #-}
{-# LANGUAGE StrictData                 #-}
{-# LANGUAGE TemplateHaskell            #-}

module OpenTracing.Jaeger.AgentReporter
    ( JaegerAgentOptions
    , jaegerAgentOptions
    , jaoServiceName
    , jaoServiceTags
    , jaoAddr
    , jaoErrorLog

    , defaultJaegerAgentAddr

    , JaegerAgent
    , newJaegerAgent
    , closeJaegerAgent
    , withJaegerAgent

    , jaegerAgentReporter

    , jaegerPropagation
    )
where

import qualified Agent.Client                   as Thrift
import           Control.Exception.Safe
import           Control.Lens                   (makeLenses, view)
import           Control.Monad.IO.Class
import           Data.ByteString.Builder
import           Data.Text                      (Text)
import qualified Data.Vector                    as Vector
import qualified Jaeger.Types                   as Thrift
import           Network.Socket
import           OpenTracing.Jaeger.Propagation (jaegerPropagation)
import           OpenTracing.Jaeger.Thrift
import           OpenTracing.Reporting          (defaultErrorLog)
import           OpenTracing.Span
import           OpenTracing.Tags
import           OpenTracing.Types
import qualified Pinch
import qualified Pinch.Client as Pinch
import qualified Pinch.Transport as Pinch

data JaegerAgent = JaegerAgent
    { JaegerAgent -> Process
envLocalProcess :: Thrift.Process
    , JaegerAgent -> Builder -> IO ()
envErrorLog     :: Builder -> IO ()
    , JaegerAgent -> JaegerClient
envClient       :: JaegerClient
    }

data JaegerClient = JaegerClient
  {
    JaegerClient -> Client
jclClient :: Pinch.Client
  , JaegerClient -> Socket
jclSocket :: Socket
  }

instance Pinch.ThriftClient JaegerClient where
  call :: forall a. JaegerClient -> ThriftCall a -> IO a
call JaegerClient{Client
jclClient :: Client
jclClient :: JaegerClient -> Client
jclClient} = forall c a. ThriftClient c => c -> ThriftCall a -> IO a
Pinch.call Client
jclClient

data JaegerAgentOptions = JaegerAgentOptions
    { JaegerAgentOptions -> Text
_jaoServiceName :: Text
    , JaegerAgentOptions -> Tags
_jaoServiceTags :: Tags
    , JaegerAgentOptions -> Addr 'UDP
_jaoAddr        :: Addr 'UDP
    , JaegerAgentOptions -> Builder -> IO ()
_jaoErrorLog    :: Builder -> IO ()
    }

jaegerAgentOptions :: Text -> JaegerAgentOptions
jaegerAgentOptions :: Text -> JaegerAgentOptions
jaegerAgentOptions Text
srv = JaegerAgentOptions
    { _jaoServiceName :: Text
_jaoServiceName = Text
srv
    , _jaoServiceTags :: Tags
_jaoServiceTags = forall a. Monoid a => a
mempty
    , _jaoAddr :: Addr 'UDP
_jaoAddr        = Addr 'UDP
defaultJaegerAgentAddr
    , _jaoErrorLog :: Builder -> IO ()
_jaoErrorLog    = Builder -> IO ()
defaultErrorLog
    }

defaultJaegerAgentAddr :: Addr 'UDP
defaultJaegerAgentAddr :: Addr 'UDP
defaultJaegerAgentAddr = HostName -> Port -> Addr 'UDP
UDPAddr HostName
"127.0.0.1" Port
6831


newJaegerAgent :: JaegerAgentOptions -> IO JaegerAgent
newJaegerAgent :: JaegerAgentOptions -> IO JaegerAgent
newJaegerAgent JaegerAgentOptions{Text
Tags
Addr 'UDP
Builder -> IO ()
_jaoErrorLog :: Builder -> IO ()
_jaoAddr :: Addr 'UDP
_jaoServiceTags :: Tags
_jaoServiceName :: Text
_jaoErrorLog :: JaegerAgentOptions -> Builder -> IO ()
_jaoAddr :: JaegerAgentOptions -> Addr 'UDP
_jaoServiceTags :: JaegerAgentOptions -> Tags
_jaoServiceName :: JaegerAgentOptions -> Text
..} =
    let tproc :: Process
tproc = Text -> Tags -> Process
toThriftProcess Text
_jaoServiceName Tags
_jaoServiceTags
     in Process -> (Builder -> IO ()) -> JaegerClient -> JaegerAgent
JaegerAgent Process
tproc Builder -> IO ()
_jaoErrorLog forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Addr 'UDP -> IO JaegerClient
openAgentTransport Addr 'UDP
_jaoAddr

closeJaegerAgent :: JaegerAgent -> IO ()
closeJaegerAgent :: JaegerAgent -> IO ()
closeJaegerAgent JaegerAgent{envClient :: JaegerAgent -> JaegerClient
envClient=JaegerClient{Socket
jclSocket :: Socket
jclSocket :: JaegerClient -> Socket
jclSocket}} =
  forall (m :: * -> *) a.
MonadCatch m =>
(SomeException -> m a) -> m a -> m a
handleAny (forall a b. a -> b -> a
const (forall (m :: * -> *) a. Monad m => a -> m a
return ())) forall a b. (a -> b) -> a -> b
$
    Socket -> IO ()
close Socket
jclSocket

withJaegerAgent
    :: ( MonadIO   m
       , MonadMask m
       )
    => JaegerAgentOptions
    -> (JaegerAgent -> m a)
    -> m a
withJaegerAgent :: forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
JaegerAgentOptions -> (JaegerAgent -> m a) -> m a
withJaegerAgent JaegerAgentOptions
opts =
    forall (m :: * -> *) a b c.
MonadMask m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ JaegerAgentOptions -> IO JaegerAgent
newJaegerAgent JaegerAgentOptions
opts) (forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall b c a. (b -> c) -> (a -> b) -> a -> c
. JaegerAgent -> IO ()
closeJaegerAgent)

openAgentTransport :: Addr 'UDP -> IO JaegerClient
openAgentTransport :: Addr 'UDP -> IO JaegerClient
openAgentTransport Addr 'UDP
addr = do
    AddrInfo{[AddrInfoFlag]
Maybe HostName
ProtocolNumber
SockAddr
SocketType
Family
addrFlags :: AddrInfo -> [AddrInfoFlag]
addrFamily :: AddrInfo -> Family
addrSocketType :: AddrInfo -> SocketType
addrProtocol :: AddrInfo -> ProtocolNumber
addrAddress :: AddrInfo -> SockAddr
addrCanonName :: AddrInfo -> Maybe HostName
addrCanonName :: Maybe HostName
addrAddress :: SockAddr
addrProtocol :: ProtocolNumber
addrSocketType :: SocketType
addrFamily :: Family
addrFlags :: [AddrInfoFlag]
..} : [AddrInfo]
_ <- Maybe AddrInfo -> Maybe HostName -> Maybe HostName -> IO [AddrInfo]
getAddrInfo (forall a. a -> Maybe a
Just AddrInfo
defaultHints { addrSocketType :: SocketType
addrSocketType = SocketType
Datagram })
                                    (forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall (a :: Protocol). Lens' (Addr a) HostName
addrHostName forall a b. (a -> b) -> a -> b
$ Addr 'UDP
addr)
                                    (forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> HostName
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall (a :: Protocol). Lens' (Addr a) Port
addrPort forall a b. (a -> b) -> a -> b
$ Addr 'UDP
addr)
    Socket
sock <- Family -> SocketType -> ProtocolNumber -> IO Socket
socket Family
addrFamily SocketType
addrSocketType ProtocolNumber
addrProtocol
    Socket -> SockAddr -> IO ()
connect Socket
sock SockAddr
addrAddress
    Channel
channel <- forall c.
Connection c =>
c -> (c -> IO Transport) -> Protocol -> IO Channel
Pinch.createChannel Socket
sock forall c. Connection c => c -> IO Transport
Pinch.unframedTransport Protocol
Pinch.compactProtocol
    forall (m :: * -> *) a. Monad m => a -> m a
return JaegerClient
      {
        jclClient :: Client
jclClient = Channel -> Client
Pinch.client Channel
channel
      , jclSocket :: Socket
jclSocket = Socket
sock
      }

jaegerAgentReporter :: MonadIO m => JaegerAgent -> FinishedSpan -> m ()
jaegerAgentReporter :: forall (m :: * -> *).
MonadIO m =>
JaegerAgent -> FinishedSpan -> m ()
jaegerAgentReporter JaegerAgent{Process
JaegerClient
Builder -> IO ()
envClient :: JaegerClient
envErrorLog :: Builder -> IO ()
envLocalProcess :: Process
envClient :: JaegerAgent -> JaegerClient
envErrorLog :: JaegerAgent -> Builder -> IO ()
envLocalProcess :: JaegerAgent -> Process
..} FinishedSpan
s = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ IO ()
emit forall (m :: * -> *) a.
MonadCatch m =>
m a -> (SomeException -> m a) -> m a
`catchAny` forall {a}. Show a => a -> IO ()
err
  where
    emit :: IO ()
emit = forall c a. ThriftClient c => c -> ThriftCall a -> IO a
Pinch.call JaegerClient
envClient (Batch -> ThriftCall ()
Thrift.emitBatch Batch
batch)
    batch :: Batch
batch = Process -> Vector FinishedSpan -> Batch
toThriftBatch Process
envLocalProcess (forall a. a -> Vector a
Vector.singleton FinishedSpan
s)
    err :: a -> IO ()
err a
e = Builder -> IO ()
envErrorLog forall a b. (a -> b) -> a -> b
$ ShortByteString -> Builder
shortByteString ShortByteString
"Jaeger Agent Thrift error: "
                       forall a. Semigroup a => a -> a -> a
<> HostName -> Builder
string8 (forall a. Show a => a -> HostName
show a
e)
                       forall a. Semigroup a => a -> a -> a
<> Char -> Builder
char8 Char
'\n'

makeLenses ''JaegerAgentOptions