{-# 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,
  )
where

import Control.Monad (when)
import Data.ByteString.Char8 qualified as C
import Data.HashSet qualified as HashSet
import System.Metrics.StatsD.Internal
  ( MetricData (..),
    StatConfig (..),
    StatCounter (..),
    StatGauge (..),
    StatSet (..),
    StatTiming (..),
    Stats,
    Value (..),
    connectStatsD,
    newMetric,
    newMetrics,
    newStats,
    processSample,
    statsLoop,
  )
import UnliftIO (MonadIO, MonadUnliftIO, throwIO)
import UnliftIO.Async (link, withAsync)

type Key = String

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:prefixStats:StatConfig :: String
prefixStats = 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:statsdServer:StatConfig :: String
statsdServer = String
"127.0.0.1",
      $sel:statsdPort:StatConfig :: Int
statsdPort = Int
8125,
      $sel:flushInterval:StatConfig :: Int
flushInterval = Int
1000,
      $sel:timingPercentiles:StatConfig :: [Int]
timingPercentiles = [Int
90, Int
95],
      $sel:appendNewline:StatConfig :: Bool
appendNewline = Bool
False
    }

newStatCounter ::
  (MonadIO m) => Stats -> Key -> Int -> m StatCounter
newStatCounter :: forall (m :: * -> *).
MonadIO m =>
Stats -> String -> Int -> m StatCounter
newStatCounter Stats
stats String
key Int
sampling = do
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
sampling Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
    IOError -> m ()
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO (IOError -> m ()) -> IOError -> m ()
forall a b. (a -> b) -> a -> b
$
      String -> IOError
userError String
"Counter sampling rate should not be negative"
  Stats -> ByteString -> MetricData -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> ByteString -> MetricData -> m ()
newMetric Stats
stats (String -> ByteString
C.pack 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 -> ByteString -> Int -> StatCounter
StatCounter Stats
stats (String -> ByteString
C.pack 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 -> ByteString -> MetricData -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> ByteString -> MetricData -> m ()
newMetric Stats
stats (String -> ByteString
C.pack 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 -> ByteString -> StatGauge
StatGauge Stats
stats (String -> ByteString
C.pack 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
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
sampling Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
    IOError -> m ()
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO (IOError -> m ()) -> IOError -> m ()
forall a b. (a -> b) -> a -> b
$
      String -> IOError
userError String
"Timing sampling rate should not be negative"
  Stats -> ByteString -> MetricData -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> ByteString -> MetricData -> m ()
newMetric Stats
stats (String -> ByteString
C.pack 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 -> ByteString -> Int -> StatTiming
StatTiming Stats
stats (String -> ByteString
C.pack 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 -> ByteString -> MetricData -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> ByteString -> MetricData -> m ()
newMetric Stats
stats (String -> ByteString
C.pack String
key) (HashSet ByteString -> MetricData
SetData HashSet ByteString
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 -> ByteString -> StatSet
StatSet Stats
stats (String -> ByteString
C.pack String
key)

incrementCounter :: (MonadIO m) => StatCounter -> Int -> m ()
incrementCounter :: forall (m :: * -> *). MonadIO m => StatCounter -> Int -> m ()
incrementCounter StatCounter {Int
ByteString
Stats
stats :: Stats
key :: ByteString
sampling :: Int
$sel:stats:StatCounter :: StatCounter -> Stats
$sel:key:StatCounter :: StatCounter -> ByteString
$sel:sampling:StatCounter :: StatCounter -> Int
..} =
  Stats -> Int -> ByteString -> Value -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> Int -> ByteString -> Value -> m ()
processSample Stats
stats Int
sampling ByteString
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 {ByteString
Stats
stats :: Stats
key :: ByteString
$sel:stats:StatGauge :: StatGauge -> Stats
$sel:key:StatGauge :: StatGauge -> ByteString
..} Int
i =
  Stats -> Int -> ByteString -> Value -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> Int -> ByteString -> Value -> m ()
processSample Stats
stats Int
1 ByteString
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 {ByteString
Stats
$sel:stats:StatGauge :: StatGauge -> Stats
$sel:key:StatGauge :: StatGauge -> ByteString
stats :: Stats
key :: ByteString
..} Int
i =
  Stats -> Int -> ByteString -> Value -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> Int -> ByteString -> Value -> m ()
processSample Stats
stats Int
1 ByteString
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
ByteString
Stats
stats :: Stats
key :: ByteString
sampling :: Int
$sel:stats:StatTiming :: StatTiming -> Stats
$sel:key:StatTiming :: StatTiming -> ByteString
$sel:sampling:StatTiming :: StatTiming -> Int
..} =
  Stats -> Int -> ByteString -> Value -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> Int -> ByteString -> Value -> m ()
processSample Stats
stats Int
sampling ByteString
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 {ByteString
Stats
stats :: Stats
key :: ByteString
$sel:stats:StatSet :: StatSet -> Stats
$sel:key:StatSet :: StatSet -> ByteString
..} =
  Stats -> Int -> ByteString -> Value -> m ()
forall (m :: * -> *).
MonadIO m =>
Stats -> Int -> ByteString -> Value -> m ()
processSample Stats
stats Int
1 ByteString
key (Value -> m ()) -> (String -> Value) -> String -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Value
Set (ByteString -> Value) -> (String -> ByteString) -> String -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ByteString
C.pack

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
  TVar Metrics
metrics <- m (TVar Metrics)
forall (m :: * -> *). MonadIO m => m (TVar Metrics)
newMetrics
  Socket
socket <- String -> Int -> m Socket
forall (m :: * -> *). MonadIO m => String -> Int -> m Socket
connectStatsD StatConfig
cfg.statsdServer StatConfig
cfg.statsdPort
  let stats :: Stats
stats = StatConfig -> TVar Metrics -> Socket -> Stats
newStats StatConfig
cfg TVar Metrics
metrics Socket
socket
  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