{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NoFieldSelectors #-}

module System.Metrics.StatsD
  ( StatCounter,
    StatGauge,
    StatTiming,
    StatSet,
    Stats,
    StatConfig (..),
    newStatCounter,
    newStatGauge,
    newStatTiming,
    newStatSet,
    incrementCounter,
    setGauge,
    incrementGauge,
    decrementGauge,
    addTiming,
    newSetElement,
    withStats,
    defStatConfig,
    parseReport,
  )
where

import Control.Monad (MonadPlus (..))
import Data.ByteString (ByteString)
import Data.ByteString.Char8 qualified as C
import Data.HashSet qualified as HashSet
import System.Metrics.StatsD.Internal
  ( Key,
    MetricData (..),
    Report (..),
    Sampling,
    StatConfig (..),
    StatCounter (..),
    StatGauge (..),
    StatSet (..),
    StatTiming (..),
    Stats,
    Value (..),
    newMetric,
    newStats,
    processSample,
    statsLoop,
    validateKey,
  )
import Text.Read (readMaybe)
import UnliftIO (MonadIO, MonadUnliftIO)
import UnliftIO.Async (link, withAsync)

defStatConfig :: StatConfig
defStatConfig :: StatConfig
defStatConfig =
  StatConfig
    { $sel:reportStats:StatConfig :: Bool
reportStats = Bool
True,
      $sel:reportSamples:StatConfig :: Bool
reportSamples = Bool
True,
      $sel:namespace:StatConfig :: String
namespace = String
"",
      $sel:statsPrefix:StatConfig :: String
statsPrefix = String
"stats",
      $sel:prefixCounter:StatConfig :: String
prefixCounter = String
"counters",
      $sel:prefixTimer:StatConfig :: String
prefixTimer = String
"timers",
      $sel:prefixGauge:StatConfig :: String
prefixGauge = String
"gauges",
      $sel:prefixSet:StatConfig :: String
prefixSet = String
"sets",
      $sel:server:StatConfig :: String
server = String
"127.0.0.1",
      $sel:port:StatConfig :: Int
port = Int
8125,
      $sel:flushInterval:StatConfig :: Int
flushInterval = Int
1000,
      $sel:timingPercentiles:StatConfig :: [Int]
timingPercentiles = [Int
90, Int
95],
      $sel:newline:StatConfig :: Bool
newline = Bool
False
    }

newStatCounter ::
  (MonadIO m) => Stats -> Key -> Sampling -> m StatCounter
newStatCounter :: forall (m :: * -> *).
MonadIO m =>
Stats -> String -> Int -> m StatCounter
newStatCounter Stats
stats String
key Int
sampling = do
  Stats -> String -> MetricData -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> String -> MetricData -> m ()
newMetric Stats
stats String
key (Int -> MetricData
CounterData Int
0)
  StatCounter -> m StatCounter
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (StatCounter -> m StatCounter) -> StatCounter -> m StatCounter
forall a b. (a -> b) -> a -> b
$ Stats -> String -> Int -> StatCounter
StatCounter Stats
stats String
key Int
sampling

newStatGauge ::
  (MonadIO m) => Stats -> Key -> Int -> m StatGauge
newStatGauge :: forall (m :: * -> *).
MonadIO m =>
Stats -> String -> Int -> m StatGauge
newStatGauge Stats
stats String
key Int
ini = do
  Stats -> String -> MetricData -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> String -> MetricData -> m ()
newMetric Stats
stats String
key (Int -> MetricData
GaugeData Int
ini)
  StatGauge -> m StatGauge
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (StatGauge -> m StatGauge) -> StatGauge -> m StatGauge
forall a b. (a -> b) -> a -> b
$ Stats -> String -> StatGauge
StatGauge Stats
stats String
key

newStatTiming :: (MonadIO m) => Stats -> Key -> Int -> m StatTiming
newStatTiming :: forall (m :: * -> *).
MonadIO m =>
Stats -> String -> Int -> m StatTiming
newStatTiming Stats
stats String
key Int
sampling = do
  Stats -> String -> MetricData -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> String -> MetricData -> m ()
newMetric Stats
stats String
key ([Int] -> MetricData
TimingData [])
  StatTiming -> m StatTiming
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (StatTiming -> m StatTiming) -> StatTiming -> m StatTiming
forall a b. (a -> b) -> a -> b
$ Stats -> String -> Int -> StatTiming
StatTiming Stats
stats String
key Int
sampling

newStatSet :: (MonadIO m) => Stats -> Key -> m StatSet
newStatSet :: forall (m :: * -> *). MonadIO m => Stats -> String -> m StatSet
newStatSet Stats
stats String
key = do
  Stats -> String -> MetricData -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> String -> MetricData -> m ()
newMetric Stats
stats String
key (HashSet String -> MetricData
SetData HashSet String
forall a. HashSet a
HashSet.empty)
  StatSet -> m StatSet
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (StatSet -> m StatSet) -> StatSet -> m StatSet
forall a b. (a -> b) -> a -> b
$ Stats -> String -> StatSet
StatSet Stats
stats String
key

incrementCounter :: (MonadIO m) => StatCounter -> Int -> m ()
incrementCounter :: forall (m :: * -> *). MonadIO m => StatCounter -> Int -> m ()
incrementCounter StatCounter {Int
String
Stats
stats :: Stats
key :: String
sampling :: Int
$sel:stats:StatCounter :: StatCounter -> Stats
$sel:key:StatCounter :: StatCounter -> String
$sel:sampling:StatCounter :: StatCounter -> Int
..} =
  Stats -> Int -> String -> Value -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> Int -> String -> Value -> m ()
processSample Stats
stats Int
sampling String
key (Value -> m ()) -> (Int -> Value) -> Int -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Value
Counter

setGauge :: (MonadIO m) => StatGauge -> Int -> m ()
setGauge :: forall (m :: * -> *). MonadIO m => StatGauge -> Int -> m ()
setGauge StatGauge {String
Stats
stats :: Stats
key :: String
$sel:stats:StatGauge :: StatGauge -> Stats
$sel:key:StatGauge :: StatGauge -> String
..} Int
i =
  Stats -> Int -> String -> Value -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> Int -> String -> Value -> m ()
processSample Stats
stats Int
1 String
key (Int -> Bool -> Value
Gauge Int
i Bool
False)

incrementGauge :: (MonadIO m) => StatGauge -> Int -> m ()
incrementGauge :: forall (m :: * -> *). MonadIO m => StatGauge -> Int -> m ()
incrementGauge StatGauge {String
Stats
$sel:stats:StatGauge :: StatGauge -> Stats
$sel:key:StatGauge :: StatGauge -> String
stats :: Stats
key :: String
..} Int
i =
  Stats -> Int -> String -> Value -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> Int -> String -> Value -> m ()
processSample Stats
stats Int
1 String
key (Int -> Bool -> Value
Gauge Int
i Bool
True)

decrementGauge :: (MonadIO m) => StatGauge -> Int -> m ()
decrementGauge :: forall (m :: * -> *). MonadIO m => StatGauge -> Int -> m ()
decrementGauge StatGauge
x Int
i = StatGauge -> Int -> m ()
forall (m :: * -> *). MonadIO m => StatGauge -> Int -> m ()
incrementGauge StatGauge
x (Int -> Int
forall a. Num a => a -> a
negate Int
i)

addTiming :: (MonadIO m) => StatTiming -> Int -> m ()
addTiming :: forall (m :: * -> *). MonadIO m => StatTiming -> Int -> m ()
addTiming StatTiming {Int
String
Stats
stats :: Stats
key :: String
sampling :: Int
$sel:stats:StatTiming :: StatTiming -> Stats
$sel:key:StatTiming :: StatTiming -> String
$sel:sampling:StatTiming :: StatTiming -> Int
..} =
  Stats -> Int -> String -> Value -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> Int -> String -> Value -> m ()
processSample Stats
stats Int
sampling String
key (Value -> m ()) -> (Int -> Value) -> Int -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Value
Timing

newSetElement :: (MonadIO m) => StatSet -> String -> m ()
newSetElement :: forall (m :: * -> *). MonadIO m => StatSet -> String -> m ()
newSetElement StatSet {String
Stats
stats :: Stats
key :: String
$sel:stats:StatSet :: StatSet -> Stats
$sel:key:StatSet :: StatSet -> String
..} =
  Stats -> Int -> String -> Value -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> Int -> String -> Value -> m ()
processSample Stats
stats Int
1 String
key (Value -> m ()) -> (String -> Value) -> String -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Value
Set

withStats :: (MonadUnliftIO m) => StatConfig -> (Stats -> m a) -> m a
withStats :: forall (m :: * -> *) a.
MonadUnliftIO m =>
StatConfig -> (Stats -> m a) -> m a
withStats StatConfig
cfg Stats -> m a
go = do
  Stats
stats <- StatConfig -> m Stats
forall (m :: * -> *). MonadIO m => StatConfig -> m Stats
newStats StatConfig
cfg
  if StatConfig
cfg.reportStats
    then m () -> (Async () -> m a) -> m a
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
withAsync (Stats -> m ()
forall (m :: * -> *). MonadIO m => Stats -> m ()
statsLoop Stats
stats) (\Async ()
a -> Async () -> m ()
forall (m :: * -> *) a. MonadIO m => Async a -> m ()
link Async ()
a m () -> m a -> m a
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Stats -> m a
go Stats
stats)
    else Stats -> m a
go Stats
stats

parseReport :: (MonadPlus m) => ByteString -> m Report
parseReport :: forall (m :: * -> *). MonadPlus m => ByteString -> m Report
parseReport ByteString
bs =
  case Char -> ByteString -> [ByteString]
C.split Char
'|' ByteString
bs of
    [ByteString
kv, ByteString
t] -> do
      (String
k, Value
v) <- ByteString -> ByteString -> m (String, Value)
forall {m :: * -> *} {a}.
(MonadPlus m, Eq a, IsString a) =>
ByteString -> a -> m (String, Value)
parseKeyValue ByteString
kv ByteString
t
      Report -> m Report
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Report -> m Report) -> Report -> m Report
forall a b. (a -> b) -> a -> b
$ String -> Value -> Double -> Report
Report String
k Value
v Double
1
    [ByteString
kv, ByteString
t, ByteString
r] -> do
      (String
k, Value
v) <- ByteString -> ByteString -> m (String, Value)
forall {m :: * -> *} {a}.
(MonadPlus m, Eq a, IsString a) =>
ByteString -> a -> m (String, Value)
parseKeyValue ByteString
kv ByteString
t
      Double
x <- ByteString -> m Double
forall {m :: * -> *} {a}.
(MonadPlus m, Read a) =>
ByteString -> m a
parseRate ByteString
r
      Report -> m Report
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Report -> m Report) -> Report -> m Report
forall a b. (a -> b) -> a -> b
$ String -> Value -> Double -> Report
Report String
k Value
v Double
x
    [ByteString]
_ -> m Report
forall a. m a
forall (m :: * -> *) a. MonadPlus m => m a
mzero
  where
    parseRead :: (MonadPlus m, Read a) => String -> m a
    parseRead :: forall (m :: * -> *) a. (MonadPlus m, Read a) => String -> m a
parseRead = m a -> (a -> m a) -> Maybe a -> m a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m a
forall a. m a
forall (m :: * -> *) a. MonadPlus m => m a
mzero a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe a -> m a) -> (String -> Maybe a) -> String -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Maybe a
forall a. Read a => String -> Maybe a
readMaybe
    parseKeyValue :: ByteString -> a -> m (String, Value)
parseKeyValue ByteString
kv a
t = do
      case Char -> ByteString -> [ByteString]
C.split Char
':' ByteString
kv of
        [ByteString
k, ByteString
v] -> do
          String
key <- ByteString -> m String
forall {m :: * -> *}. MonadPlus m => ByteString -> m String
parseKey ByteString
k
          Value
value <- ByteString -> a -> m Value
forall {a} {f :: * -> *}.
(Eq a, IsString a, MonadPlus f) =>
ByteString -> a -> f Value
parseValue ByteString
v a
t
          (String, Value) -> m (String, Value)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (String
key, Value
value)
        [ByteString]
_ -> m (String, Value)
forall a. m a
forall (m :: * -> *) a. MonadPlus m => m a
mzero
    parseKey :: ByteString -> m String
parseKey ByteString
k =
      let s :: String
s = ByteString -> String
C.unpack ByteString
k
       in if String -> Bool
validateKey String
s
            then String -> m String
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return String
s
            else m String
forall a. m a
forall (m :: * -> *) a. MonadPlus m => m a
mzero
    parseValue :: ByteString -> a -> f Value
parseValue ByteString
v a
t =
      let s :: String
s = ByteString -> String
C.unpack ByteString
v
       in case a
t of
            a
"c" -> Int -> Value
Counter (Int -> Value) -> f Int -> f Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> f Int
forall (m :: * -> *) a. (MonadPlus m, Read a) => String -> m a
parseRead String
s
            a
"g" ->
              case String
s of
                Char
'+' : String
n -> Int -> Bool -> Value
Gauge (Int -> Bool -> Value) -> f Int -> f (Bool -> Value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> f Int
forall (m :: * -> *) a. (MonadPlus m, Read a) => String -> m a
parseRead String
n f (Bool -> Value) -> f Bool -> f Value
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> f Bool
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
                Char
'-' : String
_ -> Int -> Bool -> Value
Gauge (Int -> Bool -> Value) -> f Int -> f (Bool -> Value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> f Int
forall (m :: * -> *) a. (MonadPlus m, Read a) => String -> m a
parseRead String
s f (Bool -> Value) -> f Bool -> f Value
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> f Bool
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
                String
_ -> Int -> Bool -> Value
Gauge (Int -> Bool -> Value) -> f Int -> f (Bool -> Value)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> f Int
forall (m :: * -> *) a. (MonadPlus m, Read a) => String -> m a
parseRead String
s f (Bool -> Value) -> f Bool -> f Value
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> f Bool
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
            a
"s" -> Value -> f Value
forall a. a -> f a
forall (m :: * -> *) a. Monad m => a -> m a
return (Value -> f Value) -> Value -> f Value
forall a b. (a -> b) -> a -> b
$ String -> Value
Set String
s
            a
"ms" -> Int -> Value
Timing (Int -> Value) -> f Int -> f Value
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> f Int
forall (m :: * -> *) a. (MonadPlus m, Read a) => String -> m a
parseRead String
s
            a
_ -> f Value
forall a. f a
forall (m :: * -> *) a. MonadPlus m => m a
mzero
    parseRate :: ByteString -> m a
parseRate ByteString
r = case ByteString -> String
C.unpack ByteString
r of
      Char
'@' : String
s -> String -> m a
forall (m :: * -> *) a. (MonadPlus m, Read a) => String -> m a
parseRead String
s
      String
_ -> m a
forall a. m a
forall (m :: * -> *) a. MonadPlus m => m a
mzero