{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ConstraintKinds #-}
module Data.Pickle ( Tags(..)
, StatsDConfig(..)
, MetricData
, defaultConfig
, setupPickle
, metric
, gauge
, counter
, timer
, showT
)
where
import qualified Data.Map.Strict as M
import Data.Maybe
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Text.IO as T
import Control.Exception
import Control.Monad
import Network.Socket hiding (send)
import Network.Socket.ByteString (send)
import Control.Monad.IO.Class
import System.IO.Unsafe
import Control.Concurrent.STM
type Tags = M.Map T.Text T.Text
data StatsDConfig = StatsDConfig { statsdHost :: T.Text
, statsdPort :: T.Text
, statsdPrefix :: T.Text
, statsdTags :: Tags
, statsdVerbose :: Bool
}
data Pickle = Pickle { pickleSock :: Socket
, pickleCfg :: StatsDConfig
}
type MetricData a = (Show a, Real a)
defaultConfig :: StatsDConfig
defaultConfig = StatsDConfig { statsdHost = "127.0.0.1"
, statsdPort = "8125"
, statsdPrefix = ""
, statsdTags = M.empty
, statsdVerbose = False
}
setupPickle :: StatsDConfig -> IO ()
setupPickle cfg = bracketOnError checkPickle (const finish) setPickle
where checkPickle = atomically $ do
gp <- readTVar pickle
when (gpSetupRunning gp) retry
writeTVar pickle (gp { gpSetupRunning = True})
pure (gpPickle gp)
finish = atomically $ do
modifyTVar' pickle (\gp -> gp { gpSetupRunning = False })
setPickle Nothing = do
pick <- initPickle cfg
atomically $ do
writeTVar pickle (GlobalPickle False (Just pick))
setPickle (Just oldPick) = do
newPick <- initPickle cfg
atomically $ do
writeTVar pickle (GlobalPickle False (Just newPick))
close (pickleSock oldPick)
gauge :: (MetricData a) => T.Text -> a -> Maybe Tags -> IO()
gauge name val mTags = metric "g" name val mTags Nothing
gage :: (MetricData a) => T.Text -> a -> Maybe Tags -> IO()
gage = gauge
guage :: (MetricData a) => T.Text -> a -> Maybe Tags -> IO()
guage = gauge
counter :: (MetricData a) => T.Text -> a -> Maybe Tags -> Maybe Float -> IO()
counter = metric "c"
timer :: (MetricData a) => T.Text -> a -> Maybe Tags -> Maybe Float -> IO()
timer = metric "ms"
metric :: (MetricData a)
=> T.Text
-> T.Text
-> a
-> Maybe Tags
-> Maybe Float
-> IO ()
metric kind n val mTags mSampling = do
mPick <- gpPickle <$> atomically (readTVar pickle)
case mPick of
Nothing -> pure ()
Just (Pickle sock cfg) -> do
let tags = parseTags $ (fromMaybe M.empty mTags) <> (statsdTags cfg)
sampling = maybe "" (\s -> "|@" <> showT s ) mSampling
name = (statsdPrefix cfg) <> n
msg = name <> ":" <> (showT val) <> "|" <> kind <> sampling
when (statsdVerbose cfg) (T.putStrLn $ "Sending metric: " <> msg)
void $ (try $ send sock $ T.encodeUtf8 msg :: IO(Either SomeException Int))
parseTags :: Tags -> T.Text
parseTags tags
| M.null tags = ""
| otherwise = parsed where
parsed = "#|" <> trimmed
trimmed = T.dropEnd 1 catted
catted = M.foldrWithKey (\k a b -> b <> k <> (if T.null a then "" else (":" <> a)) <> ",") "" tags
data GlobalPickle = GlobalPickle {
gpSetupRunning :: Bool
, gpPickle :: Maybe Pickle
}
pickle :: TVar GlobalPickle
pickle = unsafePerformIO $ newTVarIO $ GlobalPickle False Nothing
{-# NOINLINE pickle #-}
initPickle :: StatsDConfig -> IO Pickle
initPickle cfg = do
when (statsdVerbose cfg) $ putStrLn "Initializing Pickle StatsD Client.."
addrinfos <- getAddrInfo Nothing (Just $ T.unpack $ statsdHost cfg) (Just $ T.unpack $ statsdPort cfg)
let serveraddr = head addrinfos
sock <- socket (addrFamily serveraddr) Datagram defaultProtocol
connect sock (addrAddress serveraddr)
pure $ Pickle sock cfg
showT :: (Show a) => a -> T.Text
showT = T.pack . show