{-# 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.Semigroup
import           Data.Text                      (Text)
import qualified Data.Vector                    as Vector
import qualified Jaeger_Types                   as Thrift
import           Network.Socket
import qualified Network.Socket.ByteString.Lazy as Net
import           OpenTracing.Jaeger.Propagation (jaegerPropagation)
import           OpenTracing.Jaeger.Thrift
import           OpenTracing.Reporting          (defaultErrorLog)
import           OpenTracing.Span
import           OpenTracing.Tags
import           OpenTracing.Types
import qualified Thrift
import qualified Thrift.Protocol.Compact        as Thrift
import qualified Thrift.Transport.IOBuffer      as Thrift


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

data AgentTransport = AgentTransport
    { AgentTransport -> Socket
transSock :: Socket
    , AgentTransport -> WriteBuffer
transBuf  :: Thrift.WriteBuffer
    }

instance Thrift.Transport AgentTransport where
    tIsOpen :: AgentTransport -> IO Bool
tIsOpen = IO Bool -> AgentTransport -> IO Bool
forall a b. a -> b -> a
const (Bool -> IO Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True)
    tWrite :: AgentTransport -> ByteString -> IO ()
tWrite  = WriteBuffer -> ByteString -> IO ()
Thrift.writeBuf (WriteBuffer -> ByteString -> IO ())
-> (AgentTransport -> WriteBuffer)
-> AgentTransport
-> ByteString
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AgentTransport -> WriteBuffer
transBuf

    tFlush :: AgentTransport -> IO ()
tFlush AgentTransport{Socket
WriteBuffer
transBuf :: WriteBuffer
transSock :: Socket
transBuf :: AgentTransport -> WriteBuffer
transSock :: AgentTransport -> Socket
..} =
        WriteBuffer -> IO ByteString
Thrift.flushBuf WriteBuffer
transBuf IO ByteString -> (ByteString -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Socket -> ByteString -> IO ()
Net.sendAll Socket
transSock

    tClose :: AgentTransport -> IO ()
tClose   = [Char] -> AgentTransport -> IO ()
forall a. HasCallStack => [Char] -> a
error [Char]
"tClose undefined"
    tRead :: AgentTransport -> Int -> IO ByteString
tRead    = [Char] -> AgentTransport -> Int -> IO ByteString
forall a. HasCallStack => [Char] -> a
error [Char]
"tRead undefined"
    tPeek :: AgentTransport -> IO (Maybe Word8)
tPeek    = [Char] -> AgentTransport -> IO (Maybe Word8)
forall a. HasCallStack => [Char] -> a
error [Char]
"tPeek undefined"
    tReadAll :: AgentTransport -> Int -> IO ByteString
tReadAll = [Char] -> AgentTransport -> Int -> IO ByteString
forall a. HasCallStack => [Char] -> a
error [Char]
"tReadAll undefined"

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 :: Text
-> Tags -> Addr 'UDP -> (Builder -> IO ()) -> JaegerAgentOptions
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 = [Char] -> Port -> Addr 'UDP
UDPAddr [Char]
"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 ()) -> AgentTransport -> JaegerAgent
JaegerAgent Process
tproc Builder -> IO ()
_jaoErrorLog (AgentTransport -> JaegerAgent)
-> IO AgentTransport -> IO JaegerAgent
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Addr 'UDP -> IO AgentTransport
openAgentTransport Addr 'UDP
_jaoAddr

closeJaegerAgent :: JaegerAgent -> IO ()
closeJaegerAgent :: JaegerAgent -> IO ()
closeJaegerAgent JaegerAgent{AgentTransport
envTransport :: AgentTransport
envTransport :: JaegerAgent -> AgentTransport
envTransport} =
    (SomeException -> IO ()) -> IO () -> IO ()
forall (m :: * -> *) a.
MonadCatch m =>
(SomeException -> m a) -> m a -> m a
handleAny (IO () -> SomeException -> IO ()
forall a b. a -> b -> a
const (() -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        AgentTransport -> IO ()
forall a. Transport a => a -> IO ()
Thrift.tFlush AgentTransport
envTransport IO () -> IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> AgentTransport -> IO ()
forall a. Transport a => a -> IO ()
Thrift.tClose AgentTransport
envTransport

withJaegerAgent
    :: ( MonadIO   m
       , MonadMask m
       )
    => JaegerAgentOptions
    -> (JaegerAgent -> m a)
    -> m a
withJaegerAgent :: JaegerAgentOptions -> (JaegerAgent -> m a) -> m a
withJaegerAgent JaegerAgentOptions
opts =
    m JaegerAgent
-> (JaegerAgent -> m ()) -> (JaegerAgent -> m a) -> m a
forall (m :: * -> *) a b c.
MonadMask m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket (IO JaegerAgent -> m JaegerAgent
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 (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 AgentTransport
openAgentTransport :: Addr 'UDP -> IO AgentTransport
openAgentTransport Addr 'UDP
addr = do
    AddrInfo{[AddrInfoFlag]
Maybe [Char]
ProtocolNumber
SockAddr
SocketType
Family
addrFlags :: AddrInfo -> [AddrInfoFlag]
addrFamily :: AddrInfo -> Family
addrSocketType :: AddrInfo -> SocketType
addrProtocol :: AddrInfo -> ProtocolNumber
addrAddress :: AddrInfo -> SockAddr
addrCanonName :: AddrInfo -> Maybe [Char]
addrCanonName :: Maybe [Char]
addrAddress :: SockAddr
addrProtocol :: ProtocolNumber
addrSocketType :: SocketType
addrFamily :: Family
addrFlags :: [AddrInfoFlag]
..} : [AddrInfo]
_ <- Maybe AddrInfo -> Maybe [Char] -> Maybe [Char] -> IO [AddrInfo]
getAddrInfo (AddrInfo -> Maybe AddrInfo
forall a. a -> Maybe a
Just AddrInfo
defaultHints { addrSocketType :: SocketType
addrSocketType = SocketType
Datagram })
                                    ([Char] -> Maybe [Char]
forall a. a -> Maybe a
Just ([Char] -> Maybe [Char])
-> (Addr 'UDP -> [Char]) -> Addr 'UDP -> Maybe [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting [Char] (Addr 'UDP) [Char] -> Addr 'UDP -> [Char]
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting [Char] (Addr 'UDP) [Char]
forall (a :: Protocol). Lens' (Addr a) [Char]
addrHostName (Addr 'UDP -> Maybe [Char]) -> Addr 'UDP -> Maybe [Char]
forall a b. (a -> b) -> a -> b
$ Addr 'UDP
addr)
                                    ([Char] -> Maybe [Char]
forall a. a -> Maybe a
Just ([Char] -> Maybe [Char])
-> (Addr 'UDP -> [Char]) -> Addr 'UDP -> Maybe [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Port -> [Char]
forall a. Show a => a -> [Char]
show (Port -> [Char]) -> (Addr 'UDP -> Port) -> Addr 'UDP -> [Char]
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). Lens' (Addr a) Port
addrPort (Addr 'UDP -> Maybe [Char]) -> Addr 'UDP -> Maybe [Char]
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
    WriteBuffer
buf  <- IO WriteBuffer
Thrift.newWriteBuffer
    AgentTransport -> IO AgentTransport
forall (m :: * -> *) a. Monad m => a -> m a
return AgentTransport :: Socket -> WriteBuffer -> AgentTransport
AgentTransport
        { transSock :: Socket
transSock = Socket
sock
        , transBuf :: WriteBuffer
transBuf  = WriteBuffer
buf
        }

jaegerAgentReporter :: MonadIO m => JaegerAgent -> FinishedSpan -> m ()
jaegerAgentReporter :: JaegerAgent -> FinishedSpan -> m ()
jaegerAgentReporter JaegerAgent{Process
AgentTransport
Builder -> IO ()
envTransport :: AgentTransport
envErrorLog :: Builder -> IO ()
envLocalProcess :: Process
envTransport :: JaegerAgent -> AgentTransport
envErrorLog :: JaegerAgent -> Builder -> IO ()
envLocalProcess :: JaegerAgent -> Process
..} FinishedSpan
s = IO () -> m ()
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.
MonadCatch m =>
m a -> (SomeException -> m a) -> m a
`catchAny` SomeException -> IO ()
forall a. Show a => a -> IO ()
err
  where
    proto :: CompactProtocol AgentTransport
proto = AgentTransport -> CompactProtocol AgentTransport
forall a. a -> CompactProtocol a
Thrift.CompactProtocol AgentTransport
envTransport
    emit :: IO ()
emit  = (Any, CompactProtocol AgentTransport) -> Batch -> IO ()
forall p a. Protocol p => (a, p) -> Batch -> IO ()
Thrift.emitBatch (Any
forall a. HasCallStack => a
undefined, CompactProtocol AgentTransport
proto) 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
<> [Char] -> Builder
string8 (a -> [Char]
forall a. Show a => a -> [Char]
show a
e)
                       Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Char -> Builder
char8 Char
'\n'

makeLenses ''JaegerAgentOptions