Socket.hs

module Network.StatsD.Socket
    ( connectStatsD, sendStatsDIO, StatsD
    , withStatsD, HasStatsD(..)
    , statsd
    ) where

import           Network.StatsD.Datagram

import           Control.Monad.Base
import           Control.Monad.Trans.Control
import qualified Network.Socket as S
import qualified Network.Socket.ByteString as SB

-- * Basic types and actions

-- | Socket container.
data StatsD = StatsD
    { sdSocket   :: S.Socket
    , sdSockAddr :: S.SockAddr
    }

-- | Initialize a StatsD container.
connectStatsD :: String -> String -> IO StatsD
connectStatsD hostname port = do
    ais <- S.getAddrInfo Nothing (Just hostname) (Just port)
    addr <- case ais of
                ai:_ -> return ai
                _    -> fail "StatsD connection failed: bad address."

    sock <- S.socket (S.addrFamily addr) S.Datagram S.defaultProtocol

    return $ StatsD sock (S.addrAddress addr)

-- | Send a metric or an event to a connected statsd.
sendStatsDIO :: ToDatagram a => StatsD -> a -> IO ()
sendStatsDIO sd x = do
    let payload = renderDatagram $ toDatagram x
    SB.sendAllTo (sdSocket sd) payload (sdSockAddr sd)

-- * StatsD-providing monads

class (MonadBaseControl IO m) => HasStatsD m where
    getStatsD :: m StatsD

-- | Extract a StatsD state from application monad stack.
withStatsD :: (HasStatsD m, MonadBaseControl IO m)
           => (StatsD -> m a)
           -> m a
withStatsD action = getStatsD >>= action

-- | Send a metric or an event from application monad.
statsd :: (HasStatsD m, ToDatagram a)
        => a
        -> m ()
statsd payload = withStatsD $ \sd ->
    liftBase (sendStatsDIO sd payload)