diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,7 @@ +Copyright 2017 Ilya Ostrovskiy + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,3 @@ +pls build your stuff with `-threaded` + +kthxbai diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/src/Control/Monad/Stats.hs b/src/Control/Monad/Stats.hs new file mode 100644 --- /dev/null +++ b/src/Control/Monad/Stats.hs @@ -0,0 +1,24 @@ +module Control.Monad.Stats + ( module Export + ) where + +import Control.Monad.Stats.MTL as Export +import Control.Monad.Stats.TH as Export +import Control.Monad.Stats.Types as Export + +-- sample usage +-- +-- defineCounter "vm.txs_processed" [("vm_type", "evm")] +-- defineTimer "vm.run_loop" [("vm_type", "evm")] +-- # quasiquotes out to +-- # vm_txs_processed = Counter { counterName = "vm.txs_processed", counterTags = [("vm_type", "evm")] } +-- +-- # then you just use it inside a transformed monad +-- flip runStatsT myConfig . forever $ do +-- start <- liftIO getTime +-- tickBy 20 vm_txs_processed -- whoo so fast +-- threadDelay 1000 -- crunchin real numbers here buddy! +-- tick vm_txs_processed +-- threadDelay 10000 -- oh noes a ddos +-- done <- liftIO getTime +-- time (done - start) vm_run_loop diff --git a/src/Control/Monad/Stats/Ethereal.hs b/src/Control/Monad/Stats/Ethereal.hs new file mode 100644 --- /dev/null +++ b/src/Control/Monad/Stats/Ethereal.hs @@ -0,0 +1,7 @@ +module Control.Monad.Stats.Ethereal + ( module Export + ) where + +import Control.Monad.Stats.Monad as Export +import Control.Monad.Stats.TH as Export +import Control.Monad.Stats.Types as Export diff --git a/src/Control/Monad/Stats/MTL.hs b/src/Control/Monad/Stats/MTL.hs new file mode 100644 --- /dev/null +++ b/src/Control/Monad/Stats/MTL.hs @@ -0,0 +1,64 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TemplateHaskell #-} +module Control.Monad.Stats.MTL + ( MonadStats + , StatsT(..) + , runStatsT + , runNoStatsT + , tick + , tickBy + , setCounter + , setGauge + , time + , histoSample + , addSetMember + , reportEvent + , reportServiceCheck + , MTLStatsT, mtlStatsT + ) where + +import Control.Monad.Ether +import Control.Monad.IO.Class +import qualified Control.Monad.Stats.Monad as Ethereal +import Control.Monad.Stats.Types +import Data.Time.Clock (NominalDiffTime) + +ethereal "MTLStatsT" "mtlStatsT" + +type MonadStats m = (Monad m, MonadIO m, Ethereal.MonadStats MTLStatsT m) +type StatsT = Ethereal.StatsT MTLStatsT + +runStatsT :: (MonadIO m) => StatsT m a -> StatsTConfig -> m a +runStatsT = Ethereal.runStatsT mtlStatsT + +runNoStatsT :: (MonadIO m) => StatsT m a -> m a +runNoStatsT = Ethereal.runNoStatsT mtlStatsT + +tick :: (MonadStats m) => Counter -> m () +tick = Ethereal.tick mtlStatsT + +tickBy :: (MonadStats m) => Int -> Counter -> m () +tickBy = Ethereal.tickBy mtlStatsT + +setCounter :: (MonadStats m) => Int -> Counter -> m () +setCounter = Ethereal.setCounter mtlStatsT + +setGauge :: (MonadStats m) => Int -> Gauge -> m () +setGauge = Ethereal.setGauge mtlStatsT + +time :: (MonadStats m) => NominalDiffTime -> Timer -> m () +time = Ethereal.time mtlStatsT + +histoSample :: (MonadStats m) => Int -> Histogram -> m () +histoSample = Ethereal.histoSample mtlStatsT + +addSetMember :: (MonadStats m) => Int -> Set -> m () +addSetMember = Ethereal.addSetMember mtlStatsT + +reportEvent :: (MonadStats m) => Event -> m () +reportEvent = Ethereal.reportEvent mtlStatsT + +reportServiceCheck :: (MonadStats m) => ServiceCheck -> ServiceCheckValue -> m () +reportServiceCheck = Ethereal.reportServiceCheck mtlStatsT diff --git a/src/Control/Monad/Stats/Monad.hs b/src/Control/Monad/Stats/Monad.hs new file mode 100644 --- /dev/null +++ b/src/Control/Monad/Stats/Monad.hs @@ -0,0 +1,235 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +module Control.Monad.Stats.Monad + ( MonadStats + , StatsT(..) + , runStatsT + , runNoStatsT + , borrowTMVar + , tick + , tickBy + , setCounter + , setGauge + , time + , histoSample + , addSetMember + , reportEvent + , reportServiceCheck + ) where + +import Control.Concurrent +import Control.Concurrent.STM +import Control.Exception (AsyncException (..), fromException) +import Control.Monad.Ether +import Control.Monad.IO.Class +import Control.Monad.Stats.Types +import Control.Monad.Stats.Util +import Data.ByteString (ByteString) +import Data.ByteString as ByteString +import qualified Data.ByteString.Char8 as Char8 +import Data.Dequeue +import Data.HashMap.Strict (HashMap) +import qualified Data.HashMap.Strict as HashMap +import Data.IORef +import Data.Time (NominalDiffTime) +import Data.Time.Clock.POSIX (POSIXTime, getPOSIXTime) +import qualified Network.Socket as Socket hiding (recv, recvFrom, + send, sendTo) +import qualified Network.Socket.ByteString as Socket +import System.Random (getStdRandom, random) + +type MonadStats t m = (Monad m, MonadIO m, MonadReader t StatsTEnvironment m) + +borrowTMVar :: (Monad m, MonadIO m) => TMVar a -> (a -> m b) -> m b +borrowTMVar tmvar m = do + var <- liftIO . atomically $ takeTMVar tmvar + v <- m var + liftIO . atomically $ putTMVar tmvar var + return v + +withSocket :: (MonadStats tag m) => proxy tag -> (Socket.Socket -> m ()) -> m () +withSocket tag m = withEnvironment tag $ \e -> borrowTMVar (envSocket e) m + +withEnvironment :: (MonadStats tag m) => proxy tag -> (StatsTEnvironment -> m ()) -> m () +withEnvironment tag f = ask tag >>= \case + NoStatsTEnvironment -> return () + env -> f env + +withSTS :: (MonadStats tag m) => proxy tag -> (StatsTState -> StatsTState) -> m () +withSTS tag f = withEnvironment tag $ \(StatsTEnvironment (_, _, state)) -> + liftIO $ atomicModifyIORef' state (\x -> (f x, ())) + +tick :: (MonadStats tag m) => proxy tag -> Counter -> m () +tick tag = tickBy tag 1 + +tickBy :: (MonadStats tag m) => proxy tag -> Int -> Counter -> m () +tickBy tag n c = withSTS tag f + where f (StatsTState m q) = flip StatsTState q $ case metricMapLookup c m of + Nothing -> metricMapInsert c MetricStore{metricValue = n} m + Just (MetricStore n') -> metricMapInsert c (MetricStore (n' + n)) m + +setRegularValue :: (MonadStats tag m, Metric m') => proxy tag -> Int -> m' -> m () +setRegularValue tag v c = withSTS tag f + where f (StatsTState m q) = StatsTState (metricMapInsert c (fromIntegral v) m) q + +enqueueNonMetric :: (MonadStats tag m) => proxy tag -> ByteString -> m () +enqueueNonMetric tag e = withSTS tag f + where f (StatsTState m q) = StatsTState m (pushBack q e) + +setCounter :: (MonadStats tag m) => proxy tag -> Int -> Counter -> m () +setCounter = setRegularValue + +setGauge :: (MonadStats tag m) => proxy tag -> Int -> Gauge -> m () +setGauge = setRegularValue + +time :: (Real n, Fractional n, MonadStats tag m) => proxy tag -> n -> Timer -> m () +time tag = setRegularValue tag . v + where v = round . (* 1000.0) . toDouble + +histoSample :: (MonadStats tag m) => proxy tag -> Int -> Histogram -> m () +histoSample tag v Histogram{..} = withEnvironment tag $ \env -> do + rando <- liftIO $ getStdRandom random -- per System.Random docs: for fractional types, + -- the range is normally the semi-closed interval [0,1). + when (rando < _histogramSampleRate) $ enqueueNonMetric tag (rendered env) + where rendered = renderSimpleMetric histogramName v "|h" histogramTags (Just _histogramSampleRate) + +renderTimestamp :: POSIXTime -> ByteString +renderTimestamp time = ByteString.concat ["|d:", Char8.pack . show $ posixToMillis time] + +renderKeyedField :: ByteString -> Maybe ByteString -> ByteString +renderKeyedField x = maybe "" (\y -> ByteString.concat [x,y]) + +renderSimpleMetric :: ByteString -> Int -> ByteString -> Tags -> Maybe SampleRate -> StatsTEnvironment -> ByteString +renderSimpleMetric name value kind tags sampleRate env = + ByteString.concat [ name, ":" + , Char8.pack (show value) + , kind + , sampleRateBit + , allTags + ] + where sampleRateBit = maybe "" (ByteString.append "|@" . Char8.pack . show) sampleRate + fullName = ByteString.concat [prefix cfg, name, suffix cfg] + cfg = envConfig env + allTags = renderAllTags [defaultTags cfg, tags] + +addSetMember :: (MonadStats tag m) => proxy tag -> Int -> Set -> m () +addSetMember tag member Set{..} = withEnvironment tag $ enqueueNonMetric tag . rendered + where rendered = renderSimpleMetric setName member "|s" setTags Nothing + +reportEvent :: (MonadStats tag m) => proxy tag -> Event -> m () +reportEvent tag Event{..} = withEnvironment tag $ \env -> do + timestamp <- case eventTimestamp of + Just x -> return $ renderTimestamp x + Nothing -> renderTimestamp <$> liftIO getPOSIXTime + enqueueNonMetric tag $ renderEvent timestamp (defaultTags $ envConfig env) + where renderEvent ts dt = ByteString.concat [ "_e{" + , Char8.pack . show $ ByteString.length eventName + , "," + , Char8.pack . show $ ByteString.length eventText + , "}:" + , eventName + , "|" + , eventText + , ts + , renderKeyedField "|h:" eventHostname + , renderKeyedField "|k:" eventAggKey + , renderKeyedField "|p:" (renderPriority <$> eventPriority) + , renderKeyedField "|s" eventSource + , renderKeyedField "|t:" (renderAlertType <$> eventAlertType) + , renderAllTags [dt, eventTags] + ] + +reportServiceCheck :: (MonadStats tag m) => proxy tag -> ServiceCheck -> ServiceCheckValue -> m () +reportServiceCheck t ServiceCheck{..} ServiceCheckValue{..} = ask t >>= \case + NoStatsTEnvironment -> return () + env -> do + timestamp <- case scvTimestamp of + Just x -> return $ renderTimestamp x + Nothing -> renderTimestamp <$> liftIO getPOSIXTime + enqueueNonMetric t $ renderSC timestamp + where renderSC ts = ByteString.concat [ "_sc|" + , serviceCheckName + , "|" + , renderServiceCheckStatus scvStatus + , ts + , renderKeyedField "|h:" scvHostname + , renderAllTags [defaultTags (envConfig env), serviceCheckTags] + , renderKeyedField "|m:" scvMessage + ] + +mkStatsDSocket :: (MonadIO m) => StatsTConfig -> m (TMVar Socket.Socket, Socket.SockAddr) +mkStatsDSocket cfg = do + addrInfos <- liftIO $ Socket.getAddrInfo opts' host' port' + case addrInfos of + [] -> error $ "Unsupported address: " ++ host cfg ++ ":" ++ show (port cfg) + (a:_) -> liftIO $ do + sock <- Socket.socket (Socket.addrFamily a) Socket.Datagram Socket.defaultProtocol + tmv <- atomically (newTMVar sock) + return (tmv, Socket.addrAddress a) + + where host' = Just $ host cfg + port' = Just . show $ port cfg + opts' = Nothing + +forkStatsThread :: (MonadIO m) => StatsTEnvironment -> m ThreadId +forkStatsThread env@(StatsTEnvironment (cfg, socket, state)) = do + me <- liftIO myThreadId + liftIO . forkFinally loop $ \e -> do + borrowTMVar socket $ \s -> do + isConnected <- Socket.isConnected s + when isConnected $ Socket.close s + case e of + Left e -> case (fromException e :: Maybe AsyncException) of + Just ThreadKilled -> return () + _ -> throwTo me e + Right () -> return () + where loop = forever $ do + let interval = flushInterval cfg + start <- getMicrotime + reportSamples env + end <- getMicrotime + threadDelay (interval * 1000 - fromIntegral (end - start)) + +reportSamples :: MonadIO m => StatsTEnvironment -> m () +reportSamples env@(StatsTEnvironment (cfg, socket, state)) = do + (samples, queuedEvents) <- getAndWipeStates + borrowTMVar socket $ \sock -> do + forM_ samples (reportSample sock) + liftIO $ forM_ queuedEvents (Socket.send sock) + + where reportSample :: (MonadIO m) => Socket.Socket -> (MetricStoreKey, MetricStore) -> m () + reportSample sock (key, MetricStore value) = void . liftIO $ Socket.send sock message + where message = if value < 0 + then ByteString.concat [msg' 0, "\n", msg' value] + else msg' value + msg' v = renderSimpleMetric (keyName key) v (keyKind key) (keyTags key) sampleRate env + value' = Char8.pack (show value) + sampleRate = if isHistogram key + then Just $ histogramSampleRate key + else Nothing + + getAndWipeStates :: (MonadIO m) => m ([(MetricStoreKey, MetricStore)], BankersDequeue ByteString) + getAndWipeStates = liftIO . atomicModifyIORef' state $ \(StatsTState m q) -> + (StatsTState HashMap.empty Data.Dequeue.empty, (HashMap.toList m, q)) + + +type StatsT t = ReaderT t StatsTEnvironment + +runStatsT :: (MonadIO m) => proxy t -> StatsT t m a -> StatsTConfig -> m a +runStatsT t m c = do + (socket, addr) <- mkStatsDSocket c + liftIO $ borrowTMVar socket (`Socket.connect` addr) + theEnv <- mkStatsTEnv c socket + flip (runReaderT t) theEnv $ do + tid <- forkStatsThread theEnv + ret <- m + reportSamples theEnv -- just in case our actions ran faster than 2x flush interval + liftIO $ killThread tid + return ret + +runNoStatsT :: (MonadIO m) => proxy t -> StatsT t m a -> m a +runNoStatsT t = flip (runReaderT t) NoStatsTEnvironment diff --git a/src/Control/Monad/Stats/TH.hs b/src/Control/Monad/Stats/TH.hs new file mode 100644 --- /dev/null +++ b/src/Control/Monad/Stats/TH.hs @@ -0,0 +1,127 @@ +module Control.Monad.Stats.TH + ( defineCounter + , defineGauge + , defineTimer + , defineHistogram + , defineSet + , defineServiceCheck + , __mkByteString + ) where + +import Control.Monad +import qualified Data.ByteString.Char8 as Char8 +import Data.Char +import Data.List +import Text.ParserCombinators.ReadP + +import qualified Language.Haskell.TH as TH + +-- +-- defineCounter "vm.txs_processed" [("vm_type", "evm")] +-- # quasiquotes out to +-- # vm_txs_processed = Counter { counterName = "vm.txs_processed", counterTags = [("vm_type", "evm")] } + +-- IMPORTANT: lets the quasiquoters run. pls keep +__mkByteString :: String -> Char8.ByteString +__mkByteString = Char8.pack + +defField :: String -> [TH.FieldExp] -> String -> [(String, String)] -> TH.DecsQ +defField typeName moreFields metricName metricTags = do + unless (isValidMetricName metricName) . fail $ metricName ++ " is not a valid name for a " ++ typeName + validatedTags <- sequence $ transformTag <$> metricTags + let type' = TH.mkName typeName + nameFieldName = TH.mkName $ snek typeName ++ "Name" + nameFieldExp = (nameFieldName, bsp metricName) + tagsFieldName = TH.mkName $ snek typeName ++ "Tags" + tagsFieldExp = (tagsFieldName, TH.ListE (tagE <$> validatedTags)) + allFieldExps = nameFieldExp : tagsFieldExp : moreFields + varName = mkMetricName metricName + varP = TH.VarP varName + sig = TH.SigD varName (TH.ConT type') + val = TH.ValD varP (TH.NormalB (TH.RecConE type' allFieldExps)) [] + char8Pack = TH.mkName "__mkByteString" + bsp = TH.AppE (TH.VarE char8Pack) . TH.LitE . TH.StringL + tagE (metName, metVal) = TH.TupE [bsp metName, bsp metVal] + return [sig, val] + + +defineCounter :: String -> [(String, String)] -> TH.DecsQ +defineCounter = defField "Counter" [] + +defineGauge :: String -> [(String, String)] -> TH.DecsQ +defineGauge = defField "Gauge" [] + +defineTimer :: String -> [(String, String)] -> TH.DecsQ +defineTimer = defField "Timer" [] + +defineHistogram :: String -> [(String, String)] -> Rational -> TH.DecsQ +defineHistogram metricName metricTags sampleRate = do + when (sampleRate < 0.0 || sampleRate > 1.0) . fail $ "Histogram sample rate must be between 0.0 and 1.0" + defField "Histogram" [sampleRateField] metricName metricTags + where sampleRateField = (TH.mkName "_histogramSampleRate", TH.LitE (TH.RationalL sampleRate)) + +defineSet :: String -> [(String, String)] -> TH.DecsQ +defineSet = defField "Set" [] + +defineServiceCheck :: String -> [(String, String)] -> TH.DecsQ +defineServiceCheck = defField "ServiceCheck" [] + +transformTag :: (String, String) -> TH.Q (String, String) +transformTag (name, value) = do + let strippedName = stripTrailingUnderscores name + strippedValue = stripTrailingUnderscores value + ret = (strippedName, strippedValue) + unless (isValidTagName name) . fail $ "Tag name `" ++ name ++ "` is not a valid name.`" + unless (isValidTagValueForm value) . fail $ "Tag value`" ++ name ++ "` is not a valid value.`" + unless ((length strippedValue + length strippedValue) <= 199) . fail $ "Tag `" ++ show ret ++ "` ends up longer than 200 chars`" + return ret + +satisfiesParser :: ReadP a -> String -> Bool +satisfiesParser p = predicate . readP_to_S p + where predicate hits = not (null hits) && length hits == 1 + +isValidMetricName :: String -> Bool +isValidMetricName = not . null . readP_to_S validateMetricName + +isValidTag :: (String, String) -> Bool +isValidTag (name, val) = isValidTagName name + +isValidTagName :: String -> Bool +isValidTagName name = isValidTagNameForm name && name /= "device" + +isValidTagNameForm :: String -> Bool +isValidTagNameForm = satisfiesParser validateTagName + +isValidTagValueForm :: String -> Bool +isValidTagValueForm = satisfiesParser validateTagValue + +isAsciiAlpha :: Char -> Bool +isAsciiAlpha c = isAscii c && isAlpha c +isAsciiAlphaNum :: Char -> Bool +isAsciiAlphaNum c = isAscii c && isAlphaNum c + +validateMetricName :: ReadP String +validateMetricName = do + first <- satisfy isAsciiAlpha + rest <- flip manyTill eof $ satisfy (\c -> c `elem` "_." || isAsciiAlphaNum c) + return $ first : rest + +validateTagName :: ReadP String +validateTagName = do + first <- satisfy isAsciiAlpha + rest <- flip manyTill eof $ satisfy (\c -> c `elem` "_-." || isAsciiAlphaNum c) + return $ first : rest + +validateTagValue :: ReadP String +validateTagValue = flip manyTill eof $ satisfy (\c -> c `elem` "_-." || isAsciiAlphaNum c) + +stripTrailingUnderscores :: String -> String +stripTrailingUnderscores = reverse . dropWhile (== '_') . reverse + +mkMetricName :: String -> TH.Name +mkMetricName = TH.mkName . fmap makeUnderscores + where makeUnderscores c = if isAsciiAlphaNum c then c else '_' + +snek :: String -> String +snek "" = "" +snek (c:cs) = toLower c : cs diff --git a/src/Control/Monad/Stats/Types.hs b/src/Control/Monad/Stats/Types.hs new file mode 100644 --- /dev/null +++ b/src/Control/Monad/Stats/Types.hs @@ -0,0 +1,219 @@ +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TupleSections #-} +module Control.Monad.Stats.Types where + +import Control.Concurrent.STM (TMVar) +import Control.Monad.Ether +import Control.Monad.IO.Class +import Data.ByteString (ByteString) +import qualified Data.ByteString as ByteString +import qualified Data.ByteString.Char8 as Char8 +import Data.Dequeue +import Data.Hashable +import Data.HashMap.Strict (HashMap) +import qualified Data.HashMap.Strict as HashMap +import Data.IORef +import Data.Time.Clock.POSIX (POSIXTime) +import Network.Socket (Socket) +import System.Random (Random) + +import Data.Typeable +import GHC.Generics + +type Tag = (ByteString, ByteString) +type Tags = [Tag] +newtype SampleRate = SampleRate Float deriving (Eq, Ord, Read, Show, Num, Fractional, Random, Generic, Typeable) + +data MetricStoreKey = CounterKey { ckMetric :: Counter } + | GaugeKey { gkMetric :: Gauge } + | TimerKey { tkMetric :: Timer } + | HistogramKey { hkMetric :: Histogram } + | SetKey { skMetric :: Set } + deriving (Eq, Ord, Read, Show, Generic, Typeable) + +isHistogram :: MetricStoreKey -> Bool +isHistogram (HistogramKey _) = True +isHistogram _ = False + +histogramSampleRate :: MetricStoreKey -> SampleRate +histogramSampleRate (HistogramKey h) = _histogramSampleRate h +histogramSampleRate _ = error "called histogramSampleRate on a non-HistogramKey. pls use `isHistogram` to avoid this" + +newtype MetricStore = MetricStore { metricValue :: Int } + deriving (Eq, Ord, Read, Show, Enum, Num, Real, Integral, Generic, Typeable) + +class (Eq m, Ord m, Read m, Show m) => Metric m where + metricStoreKey :: m -> MetricStoreKey + +instance Hashable MetricStoreKey where + hashWithSalt salt m = salt `hashWithSalt` keyName m `hashWithSalt` keyTags m `hashWithSalt` keyKind m + +keyName :: MetricStoreKey -> ByteString +keyName (CounterKey m) = counterName m -- this is literally a crime against humanity +keyName (GaugeKey m) = gaugeName m +keyName (TimerKey m) = timerName m +keyName (HistogramKey m) = histogramName m +keyName (SetKey m) = setName m + +keyTags :: MetricStoreKey -> Tags +keyTags (CounterKey m) = counterTags m +keyTags (GaugeKey m) = gaugeTags m +keyTags (TimerKey m) = timerTags m +keyTags (HistogramKey m) = histogramTags m +keyTags (SetKey m) = setTags m + +keyKind :: MetricStoreKey -> ByteString +keyKind (CounterKey m) = "|c" +keyKind (GaugeKey m) = "|g" +keyKind (TimerKey m) = "|ms" +keyKind (HistogramKey m) = "|h" +keyKind (SetKey m) = "|s" + +data Counter = Counter { counterName :: !ByteString, counterTags :: !Tags } + deriving (Eq, Ord, Read, Show, Generic, Typeable) + +data Gauge = Gauge { gaugeName :: !ByteString, gaugeTags :: !Tags } + deriving (Eq, Ord, Read, Show, Generic, Typeable) + +data Timer = Timer { timerName :: !ByteString, timerTags :: !Tags } + deriving (Eq, Ord, Read, Show, Generic, Typeable) + +data Histogram = Histogram { histogramName :: !ByteString , histogramTags :: !Tags, _histogramSampleRate :: !SampleRate } + deriving (Eq, Ord, Read, Show, Generic, Typeable) + +data Set = Set { setName :: !ByteString, setTags :: !Tags } + deriving (Eq, Ord, Read, Show, Generic, Typeable) + +data Event = + Event { eventName :: ByteString + , eventText :: ByteString + , eventTags :: Tags + , eventTimestamp :: Maybe POSIXTime + , eventHostname :: Maybe ByteString + , eventAggKey :: Maybe ByteString + , eventPriority :: Maybe Priority + , eventSource :: Maybe ByteString + , eventAlertType :: Maybe AlertType + } deriving (Eq, Ord, Show, Generic, Typeable) + +data ServiceCheck = + ServiceCheck { serviceCheckName :: ByteString + , serviceCheckTags :: Tags + } deriving (Eq, Ord, Read, Show, Generic, Typeable) + +data Priority = Normal | Low + deriving (Eq, Ord, Read, Show, Generic, Typeable) + +renderPriority :: Priority -> ByteString +renderPriority Normal = "normal" +renderPriority Low = "low" + +data AlertType = Error | Warning | Info | Success + deriving (Eq, Ord, Read, Show, Generic, Typeable) + +renderAlertType :: AlertType -> ByteString +renderAlertType Error = "error" +renderAlertType Warning = "warning" +renderAlertType Info = "info" +renderAlertType Success = "success" + +data ServiceCheckStatus = StatusOK | StatusWarning | StatusCritical | StatusUnknown + deriving (Eq, Ord, Read, Show, Generic, Typeable) + +data ServiceCheckValue = + ServiceCheckValue { scvStatus :: ServiceCheckStatus + , scvTimestamp :: Maybe POSIXTime + , scvHostname :: Maybe ByteString + , scvMessage :: Maybe ByteString + } deriving (Eq, Show, Generic, Typeable) + +renderServiceCheckStatus :: ServiceCheckStatus -> ByteString +renderServiceCheckStatus StatusOK = "0" +renderServiceCheckStatus StatusWarning = "1" +renderServiceCheckStatus StatusCritical = "2" +renderServiceCheckStatus StatusUnknown = "3" + +renderTags :: Tags -> ByteString +renderTags = ByteString.intercalate "," . map renderTag + where renderTag :: Tag -> ByteString + renderTag (k, v) = ByteString.concat [k, ":", v] + +renderAllTags :: [Tags] -> ByteString +renderAllTags tags = case concat tags of + [] -> "" + xs -> ByteString.concat ["|#", renderTags xs] + +instance Metric Counter where + metricStoreKey = CounterKey + +instance Metric Gauge where + metricStoreKey = GaugeKey + +instance Metric Timer where + metricStoreKey = TimerKey + +instance Metric Histogram where + metricStoreKey = HistogramKey + +instance Metric Set where + metricStoreKey = SetKey + +data StatsTConfig = + StatsTConfig { host :: !String + , port :: !Int + , flushInterval :: !Int + , prefix :: !ByteString + , suffix :: !ByteString + , defaultTags :: !Tags + } deriving (Eq, Read, Show, Generic, Typeable) + +defaultStatsTConfig :: StatsTConfig +defaultStatsTConfig = StatsTConfig { host = "127.0.0.1" + , port = 8125 + , flushInterval = 1000 + , prefix = "" + , suffix = "" + , defaultTags = [] + } + +data StatsTEnvironment = StatsTEnvironment (StatsTConfig, TMVar Socket, IORef StatsTState) + | NoStatsTEnvironment + deriving (Eq, Generic, Typeable) + +envConfig :: StatsTEnvironment -> StatsTConfig +envConfig NoStatsTEnvironment = error "called envConfig inside a runNoStatsT" +envConfig (StatsTEnvironment (a, _, _)) = a + +envSocket :: StatsTEnvironment -> TMVar Socket +envSocket NoStatsTEnvironment = error "called envSocket inside a runNoStatsT" +envSocket (StatsTEnvironment (_, b, _)) = b + +envState :: StatsTEnvironment -> IORef StatsTState +envState NoStatsTEnvironment = error "called envState inside a runNoStatsT" +envState (StatsTEnvironment (_, _, c)) = c + +type MetricMap = HashMap MetricStoreKey MetricStore + +metricMapLookup :: Metric m => m -> MetricMap -> Maybe MetricStore +metricMapLookup = HashMap.lookup . metricStoreKey + +metricMapInsert :: Metric m => m -> MetricStore -> MetricMap -> MetricMap +metricMapInsert = HashMap.insert . metricStoreKey + +data NonMetricEvent = HistogramEvent Histogram MetricStore + | SetEvent Set MetricStore + | ServiceCheckEvent ServiceCheck + | EventEvent Event + deriving (Eq, Show, Generic, Typeable) + +data StatsTState = + StatsTState { registeredMetrics :: HashMap MetricStoreKey MetricStore + , queuedLines :: BankersDequeue ByteString + } deriving (Eq, Read, Show, Generic, Typeable) + +mkStatsTEnv :: (MonadIO m, Monad m) => StatsTConfig -> TMVar Socket -> m StatsTEnvironment +mkStatsTEnv conf socket = liftIO $ + StatsTEnvironment . (conf,socket,) <$> newIORef (StatsTState HashMap.empty empty) diff --git a/src/Control/Monad/Stats/Util.hs b/src/Control/Monad/Stats/Util.hs new file mode 100644 --- /dev/null +++ b/src/Control/Monad/Stats/Util.hs @@ -0,0 +1,19 @@ +module Control.Monad.Stats.Util where + +import Control.Monad.IO.Class +import Data.Time.Clock.POSIX + +getMicrotime :: (MonadIO m) => m Int +getMicrotime = posixToMicros <$> liftIO getPOSIXTime + +posixToMicros :: POSIXTime -> Int +posixToMicros = round . (* 1000000.0) . toDouble + +getMillitime :: (MonadIO m) => m Int +getMillitime = posixToMillis <$> liftIO getPOSIXTime + +posixToMillis :: POSIXTime -> Int +posixToMillis = round . (* 1000.0) . toDouble + +toDouble :: (Real n) => n -> Double +toDouble = realToFrac diff --git a/statsdi.cabal b/statsdi.cabal new file mode 100644 --- /dev/null +++ b/statsdi.cabal @@ -0,0 +1,59 @@ +name: statsdi +version: 0.1.0.0 +synopsis: A lovely [Dog]StatsD implementation +description: An implementation of DogStatsD for collecting and pushing metrics +homepage: https://github.com/iostat/statsdi#readme +license: MIT +license-file: LICENSE +author: Ilya Ostrovskiy +maintainer: first-name@thenumber200-thewordproof.cc +copyright: 2017 Ilya Ostrovskiy +category: Metrics +build-type: Simple +extra-source-files: README.md +cabal-version: >=1.10 + +library + hs-source-dirs: src + exposed-modules: Control.Monad.Stats + , Control.Monad.Stats.Ethereal + , Control.Monad.Stats.TH + , Control.Monad.Stats.Types + , Control.Monad.Stats.Util + other-modules: Control.Monad.Stats.Monad + , Control.Monad.Stats.MTL + build-depends: base >= 4.7 && < 5 + , network + , random + , dequeue + , stm + , time + , bytestring + , ether + , transformers + , transformers-lift + , template-haskell >= 2.11.1.0 && < 2.12 + , unordered-containers >= 0.2.7.0 && < 0.3 + , hashable >= 1.2.4.0 && < 1.3 + default-language: Haskell2010 + +test-suite statsdi-test + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Spec.hs + other-modules: Harness + build-depends: base + , statsdi + , bytestring + , stm + , network + , time + , tasty + , tasty-hspec + , hspec + ghc-options: -threaded -rtsopts -with-rtsopts=-N + default-language: Haskell2010 + +source-repository head + type: git + location: https://github.com/iostat/statsdi diff --git a/test/Harness.hs b/test/Harness.hs new file mode 100644 --- /dev/null +++ b/test/Harness.hs @@ -0,0 +1,69 @@ +{-# LANGUAGE LambdaCase #-} +module Harness + ( runStatsTCapturingOutput + ) where + +import Control.Concurrent +import Control.Concurrent.STM +import Control.Monad +import Control.Monad.IO.Class +import Control.Monad.Stats +import Data.ByteString (ByteString) +import qualified Data.ByteString as ByteString +import Data.Time.Clock.POSIX +import Network.Socket hiding (recv) +import Network.Socket.ByteString (recv) + +import System.IO.Unsafe (unsafePerformIO) + +harnessLock :: TMVar Int +harnessLock = unsafePerformIO $ newTMVarIO 0 +{-# NOINLINE harnessLock #-} + +tQueueToList :: TQueue a -> IO [a] +tQueueToList = fmap reverse . loop [] + where loop read q = atomically (tryReadTQueue q) >>= \case + Nothing -> return read + Just x -> loop (x:read) q + +listenForUDP :: (MonadIO m) => StatsTConfig -> TMVar [ByteString] -> m ThreadId +listenForUDP cfg var = liftIO . withSocketsDo $ do + sock <- getAddrInfo opts' host' port' >>= \case + [] -> error $ "Unsupported address: " ++ host cfg ++ ":" ++ show (port cfg) + (a:_) -> do + socket' <- socket (addrFamily a) Datagram defaultProtocol + setSocketOption socket' ReuseAddr 1 + setSocketOption socket' ReusePort 1 + bind socket' (addrAddress a) + return socket' + queue <- newTQueueIO + forkFinally (udpLoop sock queue) . const $ do + close sock + atomically . putTMVar var =<< tQueueToList queue + + where udpLoop :: Socket -> TQueue ByteString -> IO () + udpLoop sock queue = + recv sock 10240 >>= atomically . writeTQueue queue >> udpLoop sock queue + + opts' = Just $ defaultHints { addrFlags = [AI_PASSIVE] } + host' = Just $ host cfg + port' = Just . show $ port cfg + +forkAndListen :: StatsTConfig -> IO (ThreadId, TMVar [ByteString]) +forkAndListen cfg = do + var <- newEmptyTMVarIO + tid <- listenForUDP cfg var + return (tid, var) + +runStatsTCapturingOutput :: (MonadIO m) => StatsTConfig -> Int -> StatsT m a -> m ([ByteString], a) +runStatsTCapturingOutput c lingerTime m = do + (tid, var) <- liftIO $ do + atomically $ takeTMVar harnessLock + forkAndListen c + ret <- runStatsT m c + liftIO $ do + threadDelay $ lingerTime * 1000 + killThread tid + val <- atomically (takeTMVar var) + atomically (putTMVar harnessLock 0) + return (val, ret) diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,174 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} + +import Control.Concurrent +import Control.Monad +import Control.Monad.IO.Class +import Control.Monad.Stats +import Data.ByteString (ByteString) +import qualified Data.ByteString as ByteString +import qualified Data.ByteString.Char8 as Char8 +import Data.Char (isNumber) +import Data.Proxy +import Data.Time.Clock.POSIX + +import Harness + +import Test.Hspec +import Test.Tasty +import Test.Tasty.Hspec +import Test.Tasty.Options (OptionDescription (..)) +import Test.Tasty.Runners (NumThreads (..)) + +import Debug.Trace (traceShowId) + +-- these tests effectively test the TH for sanity +-- otherwise the test suite wouldnt even compile +defineCounter "ctr.hello.world" [] +defineCounter "ctr.tagged" [("env","test")] +defineGauge "gau.testing.things" [] +defineTimer "time.test" [] +defineHistogram "hist.stuff.things" [] 1.0 +defineSet "set.of.people" [] +defineServiceCheck "svc.haskell" [] + +ourStatsTConfig :: StatsTConfig +ourStatsTConfig = defaultStatsTConfig { flushInterval = 250 } + +st :: (MonadIO m) => Int -> StatsT m a -> m ([ByteString], a) +st = st' ourStatsTConfig + +st' :: (MonadIO m) => StatsTConfig -> Int -> StatsT m a -> m ([ByteString], a) +st' = runStatsTCapturingOutput + +sleepMs :: MonadIO m => Int -> m () +sleepMs = liftIO . threadDelay . (1000 *) + +main :: IO () +main = do + putStrLn "" + let toRun = [sillyTests, noStatsTTest, counterTests, gaugeTests, timerTests] + forM toRun testSpecs >>= defaultMain . withTests + +withTests :: [[TestTree]] -> TestTree +withTests = testGroup "Statsdi" . concat + +sillyTests :: Spec +sillyTests = describe "The test harness" $ do + it "should run and capture something (1s linger)" $ do + (capture, _) <- st 1000 $ tick ctr_hello_world + capture `shouldSatisfy` (not . null) + + it "should run and capture with a delay before the tick (250ms flushInterval / 1s linger / 500ms delay)" $ do + (capture, _) <- st 1000 $ do + sleepMs 500 + tick ctr_hello_world + capture `shouldSatisfy` (not . null) + + it "should linger after the test runs and get everything (100ms linger / 5000ms+ test)" $ do + (capture, _) <- st 100 $ do + sleepMs 5000 + tick ctr_hello_world + capture `shouldSatisfy` (not . null) + +counterTests :: Spec +counterTests = describe "A Counter" $ do + it "should have a kind tag of |c" $ do + (capture, _) <- st 1000 $ setCounter 0 ctr_hello_world + capture `shouldSatisfy` (not . null) + capture `shouldSatisfy` (ByteString.isPrefixOf "ctr.hello.world:0|c" . head) + + it "should increment by one when calling `tick`" $ do + (capture, _) <- st 1000 $ tick ctr_hello_world + capture `shouldSatisfy` (not . null) + capture `shouldSatisfy` (ByteString.isPrefixOf "ctr.hello.world:1|c" . head) + + it "should send a multi-event with a zeroing-out before being set to a negative number" $ do + (capture, _) <- st 1000 $ setCounter (-20) ctr_hello_world + capture `shouldSatisfy` (not . null) + capture `shouldSatisfy` (ByteString.isPrefixOf "ctr.hello.world:0|c\nctr.hello.world:-20|c" . head) + +gaugeTests :: Spec +gaugeTests = describe "A Gauge" $ do + it "should have a kind tag of |g" $ do + (capture, _) <- st 1000 $ setGauge 0 gau_testing_things + capture `shouldSatisfy` (not . null) + capture `shouldSatisfy` (ByteString.isPrefixOf "gau.testing.things:0|g" . head) + + it "should send a multi-event with a zeroing-out before being set to a negative number" $ do + (capture, _) <- st 1000 $ setGauge (-20) gau_testing_things + capture `shouldSatisfy` (not . null) + capture `shouldSatisfy` (ByteString.isPrefixOf "gau.testing.things:0|g\ngau.testing.things:-20|g" . head) + +timerTests :: Spec +timerTests = describe "A Timer" $ do + it "should have a kind tag of |ms" $ do + (capture, _) <- st 1000 $ time 0 time_test + capture `shouldSatisfy` (not . null) + capture `shouldSatisfy` (ByteString.isPrefixOf "time.test:0|ms" . head) + + it "should report in milliseconds" $ do + (capture, _) <- st 1000 $ time 1.0 time_test + capture `shouldSatisfy` (not . null) + capture `shouldSatisfy` (ByteString.isPrefixOf "time.test:1000|ms" . head) + + it "should handle DiffTimes appropriately" $ do + (capture, _) <- st 1000 $ do + now <- liftIO getPOSIXTime + sleepMs 250 + then' <- liftIO getPOSIXTime + time (then' - now) time_test + + capture `shouldSatisfy` (not . null) + capture `shouldSatisfy` (isRoughlyMillis 250 10 . head) + + where isRoughlyMillis target leeway bs = abs (actual - target) <= leeway + where actual = read pluckedTime + pluckedTime = takeWhile isNumber nameStripped + nameStripped = drop (ByteString.length (timerName time_test) + 1) (Char8.unpack bs) + +noStatsTTest :: Spec +noStatsTTest = describe "runNoStatsT" $ do + it "should successfully run its inner monad without any funny hiccups" $ + runNoStatsT (return ()) >>= shouldBe () + + describe "should successfully run its inner monad even it performs Counter metrics" $ do + it "works with tick" $ + runNoStatsT (tick ctr_hello_world) >>= shouldBe () + it "works with tickBy" $ + runNoStatsT (tickBy 2 ctr_hello_world) >>= shouldBe () + + describe "should successfully run its inner monad even it performs Gauge metrics" $ + it "works with setGauge" $ + runNoStatsT (setGauge 20 gau_testing_things) >>= shouldBe () + + describe "should successfully run its inner monad even it performs Timer metrics" $ + it "works with time" $ do + ret <- runNoStatsT $ do + now <- liftIO getPOSIXTime + sleepMs 250 + then' <- liftIO getPOSIXTime + time (then' - now) time_test + ret `shouldBe` () + + describe "should successfully run its inner monad even it performs Histogram metrics" $ do + it "works with histoSample" $ + runNoStatsT (histoSample 3 hist_stuff_things) >>= shouldBe () + + it "works with multiple calls to histoSample" $ do + ret <- runNoStatsT $ do + histoSample 18 hist_stuff_things + histoSample 19 hist_stuff_things + ret `shouldBe` () + + describe "should successfully run its inner monad even it performs Set metrics" $ do + it "works with addSetMember" $ + runNoStatsT (addSetMember 12 set_of_people) >>= shouldBe () + + it "works with multiple calls to addSetMember" $ do + ret <- runNoStatsT $ do + addSetMember 12 set_of_people + addSetMember 24 set_of_people + ret `shouldBe` ()