module Network.Statsd (
StatsdClient,
client,
fromURI,
Hostname,
Port,
Stat,
Key,
increment,
decrement,
count,
gauge,
timing,
histogram,
) where
import Control.Monad
import Control.Exception
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BLazy
import qualified Data.ByteString.Char8 as BC
import Data.ByteString.Lazy.Builder (int64LE, toLazyByteString)
import Data.Word
import Data.Byteable
import Data.Maybe
import System.Time
import System.IO.Error
import Crypto.Hash
import Crypto.Random.DRBG
import Text.Printf
import Data.Time.Units
import qualified Network.Socket as Net hiding (send, sendTo, recv, recvFrom)
import qualified Network.Socket.ByteString as Net
import Network.URI
type Stat = String
data StatsdClient = StatsdClient { getSocket :: Net.Socket
, getNamespace :: Stat
, getSigningKey :: Maybe Key
}
type Hostname = String
type Port = Int
fromURI :: URI -> IO StatsdClient
fromURI uri = case uriAuthority uri of
Nothing -> fail "invalid URI"
Just auth -> let hostname = uriRegName' auth
port = fromMaybe 8126 $ (actualPort . uriPort) auth
prefix = replace '/' '.' (uriPath uri)
key = let userInfo = uriUserInfo auth
init' = if null userInfo
then ""
else init userInfo
(user, pass) = case break (==':') init' of
(u, ':':p) -> (u, p)
_ -> ("", "")
in if null pass
then Nothing
else Just pass
in client hostname port prefix key
where
replace :: Eq a => a -> a -> [a] -> [a]
replace from to list = (\a -> if a == from then to else a) <$> list
actualPort :: String -> Maybe Int
actualPort (':':xs) = (Just . read) xs
actualPort _ = Nothing
uriRegName' :: URIAuth -> String
uriRegName' auth = let hostname = uriRegName auth
in (takeWhile (/=']') . dropWhile (=='[')) hostname
client :: Hostname -> Port -> Stat -> Maybe Key -> IO StatsdClient
client host port namespace key = do
(addr:_) <- Net.getAddrInfo Nothing (Just host) (Just $ show port)
sock <- Net.socket (Net.addrFamily addr) Net.Datagram Net.defaultProtocol
Net.connect sock (Net.addrAddress addr)
return $ StatsdClient sock namespace key
increment :: StatsdClient -> Stat -> IO ()
increment client stat = count client stat 1
decrement :: StatsdClient -> Stat -> IO ()
decrement client stat = count client stat (1)
count :: StatsdClient -> Stat -> Int -> IO ()
count client stat value = void . send client $ encode (getNamespace client) stat value Count
gauge :: StatsdClient -> Stat -> Int -> IO ()
gauge client stat value = void . send client $ encode (getNamespace client) stat value Gauge
timing :: StatsdClient -> Stat -> Millisecond -> IO ()
timing client stat value = void . send client $ encode (getNamespace client) stat (fromIntegral value) Timing
histogram :: StatsdClient -> Stat -> Int -> IO ()
histogram client stat value = void . send client $ encode (getNamespace client) stat value Histogram
encode :: Stat -> Stat -> Int -> Type -> Payload
encode namespace stat value stat_type =
let prefix = if null namespace
then ""
else namespace ++ "."
message = printf "%s%s:%s|%s" prefix stat (show value) (show stat_type)
in BC.pack message
type Payload = B.ByteString
send :: StatsdClient -> Payload -> IO (Either IOError ())
send client payload = do
signedPayload <- signed (getSigningKey client) payload
tryIOError . void $ Net.send (getSocket client) signedPayload
type Nonce = B.ByteString
type Key = String
signed :: Maybe Key -> Payload -> IO Payload
signed Nothing payload = return payload
signed (Just key) payload = do
(TOD sec _) <- getClockTime
let timestamp = B.concat . BLazy.toChunks . toLazyByteString . int64LE $ fromIntegral sec
gen <- newGenIO :: IO CtrDRBG
let (nonce, _) = throwLeft $ genBytes 4 gen
let newPayload = B.concat [timestamp, nonce, payload]
return $ sign key newPayload
sign :: Key -> Payload -> Payload
sign key payload = let keyBytes = BC.pack key
signature = toBytes (hmac keyBytes payload :: HMAC SHA256)
in B.append signature payload
data Type = Count | Gauge | Timing | Histogram
instance Show Type where
show Count = "c"
show Gauge = "g"
show Timing = "ms"
show Histogram = "h"