{-# 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 :: JaegerClient -> Client
jclClient :: Client
jclClient} = Client -> ThriftCall a -> IO a
forall c a. ThriftClient c => c -> ThriftCall a -> IO a
forall a. Client -> 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 = Tags
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
Addr 'UDP
Tags
Builder -> IO ()
_jaoServiceName :: JaegerAgentOptions -> Text
_jaoServiceTags :: JaegerAgentOptions -> Tags
_jaoAddr :: JaegerAgentOptions -> Addr 'UDP
_jaoErrorLog :: JaegerAgentOptions -> Builder -> IO ()
_jaoServiceName :: Text
_jaoServiceTags :: Tags
_jaoAddr :: Addr 'UDP
_jaoErrorLog :: Builder -> IO ()
..} =
    let tproc :: Process
tproc = Text -> Tags -> Process
toThriftProcess Text
_jaoServiceName Tags
_jaoServiceTags
     in Process -> (Builder -> IO ()) -> JaegerClient -> JaegerAgent
JaegerAgent Process
tproc Builder -> IO ()
_jaoErrorLog (JaegerClient -> JaegerAgent) -> IO JaegerClient -> IO JaegerAgent
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 :: JaegerClient -> Socket
jclSocket :: Socket
jclSocket}} =
  (SomeException -> IO ()) -> IO () -> IO ()
forall (m :: * -> *) a.
(HasCallStack, MonadCatch m) =>
(SomeException -> m a) -> m a -> m a
handleAny (IO () -> SomeException -> IO ()
forall a b. a -> b -> a
const (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())) (IO () -> IO ()) -> IO () -> IO ()
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 =
    m JaegerAgent
-> (JaegerAgent -> m ()) -> (JaegerAgent -> m a) -> m a
forall (m :: * -> *) a b c.
(HasCallStack, MonadMask m) =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket (IO JaegerAgent -> m JaegerAgent
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO JaegerAgent -> m JaegerAgent)
-> IO JaegerAgent -> m JaegerAgent
forall a b. (a -> b) -> a -> b
$ JaegerAgentOptions -> IO JaegerAgent
newJaegerAgent JaegerAgentOptions
opts) (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (JaegerAgent -> IO ()) -> JaegerAgent -> m ()
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
Family
SocketType
addrFlags :: [AddrInfoFlag]
addrFamily :: Family
addrSocketType :: SocketType
addrProtocol :: ProtocolNumber
addrAddress :: SockAddr
addrCanonName :: Maybe HostName
addrFlags :: AddrInfo -> [AddrInfoFlag]
addrFamily :: AddrInfo -> Family
addrSocketType :: AddrInfo -> SocketType
addrProtocol :: AddrInfo -> ProtocolNumber
addrAddress :: AddrInfo -> SockAddr
addrCanonName :: AddrInfo -> Maybe HostName
..} : [AddrInfo]
_ <- Maybe AddrInfo -> Maybe HostName -> Maybe HostName -> IO [AddrInfo]
forall (t :: * -> *).
GetAddrInfo t =>
Maybe AddrInfo
-> Maybe HostName -> Maybe HostName -> IO (t AddrInfo)
getAddrInfo (AddrInfo -> Maybe AddrInfo
forall a. a -> Maybe a
Just AddrInfo
defaultHints { addrSocketType = Datagram })
                                    (HostName -> Maybe HostName
forall a. a -> Maybe a
Just (HostName -> Maybe HostName)
-> (Addr 'UDP -> HostName) -> Addr 'UDP -> Maybe HostName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting HostName (Addr 'UDP) HostName -> Addr 'UDP -> HostName
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting HostName (Addr 'UDP) HostName
forall (a :: Protocol) (f :: * -> *).
Functor f =>
(HostName -> f HostName) -> Addr a -> f (Addr a)
addrHostName (Addr 'UDP -> Maybe HostName) -> Addr 'UDP -> Maybe HostName
forall a b. (a -> b) -> a -> b
$ Addr 'UDP
addr)
                                    (HostName -> Maybe HostName
forall a. a -> Maybe a
Just (HostName -> Maybe HostName)
-> (Addr 'UDP -> HostName) -> Addr 'UDP -> Maybe HostName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Port -> HostName
forall a. Show a => a -> HostName
show (Port -> HostName) -> (Addr 'UDP -> Port) -> Addr 'UDP -> HostName
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting Port (Addr 'UDP) Port -> Addr 'UDP -> Port
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Port (Addr 'UDP) Port
forall (a :: Protocol) (f :: * -> *).
Functor f =>
(Port -> f Port) -> Addr a -> f (Addr a)
addrPort (Addr 'UDP -> Maybe HostName) -> Addr 'UDP -> Maybe HostName
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 <- Socket -> (Socket -> IO Transport) -> Protocol -> IO Channel
forall c.
Connection c =>
c -> (c -> IO Transport) -> Protocol -> IO Channel
Pinch.createChannel Socket
sock Socket -> IO Transport
forall c. Connection c => c -> IO Transport
Pinch.unframedTransport Protocol
Pinch.compactProtocol
    JaegerClient -> IO JaegerClient
forall a. a -> IO a
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 ()
envLocalProcess :: JaegerAgent -> Process
envErrorLog :: JaegerAgent -> Builder -> IO ()
envClient :: JaegerAgent -> JaegerClient
envLocalProcess :: Process
envErrorLog :: Builder -> IO ()
envClient :: JaegerClient
..} FinishedSpan
s = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ IO ()
emit IO () -> (SomeException -> IO ()) -> IO ()
forall (m :: * -> *) a.
(HasCallStack, MonadCatch m) =>
m a -> (SomeException -> m a) -> m a
`catchAny` SomeException -> IO ()
forall {a}. Show a => a -> IO ()
err
  where
    emit :: IO ()
emit = JaegerClient -> ThriftCall () -> IO ()
forall c a. ThriftClient c => c -> ThriftCall a -> IO a
forall a. JaegerClient -> ThriftCall a -> IO a
Pinch.call JaegerClient
envClient (Batch -> ThriftCall ()
Thrift.emitBatch Batch
batch)
    batch :: Batch
batch = Process -> Vector FinishedSpan -> Batch
toThriftBatch Process
envLocalProcess (FinishedSpan -> Vector FinishedSpan
forall a. a -> Vector a
Vector.singleton FinishedSpan
s)
    err :: a -> IO ()
err a
e = Builder -> IO ()
envErrorLog (Builder -> IO ()) -> Builder -> IO ()
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Builder
shortByteString ShortByteString
"Jaeger Agent Thrift error: "
                       Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> HostName -> Builder
string8 (a -> HostName
forall a. Show a => a -> HostName
show a
e)
                       Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Char -> Builder
char8 Char
'\n'

makeLenses ''JaegerAgentOptions