{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}
module System.Remote.Monitoring.Prometheus
  ( toPrometheusRegistry
  , registerEKGStore
  , AdapterOptions(..)
  , labels
  , namespace
  , samplingFrequency
  , defaultOptions
  ) where

import           Control.Concurrent (forkIO, threadDelay)
import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.State.Strict
import qualified Data.HashMap.Strict as HMap
import qualified Data.Map.Strict as Map
import           Data.Monoid
import Lens.Micro.TH
import qualified Data.Text as T
import qualified System.Metrics as EKG
import qualified System.Metrics.Prometheus.Metric.Counter as Counter
import qualified System.Metrics.Prometheus.Metric.Gauge as Gauge
import qualified System.Metrics.Prometheus.MetricId as Prometheus
import qualified System.Metrics.Prometheus.Registry as Prometheus
import           System.Metrics.Prometheus.RegistryT (RegistryT(..))

--------------------------------------------------------------------------------
data AdapterOptions = AdapterOptions {
    AdapterOptions -> Labels
_labels :: Prometheus.Labels
  , AdapterOptions -> Maybe Text
_namespace :: Maybe T.Text
  , AdapterOptions -> Int
_samplingFrequency :: !Int
    -- ^ How often update the registry (in seconds).
  }

makeLenses ''AdapterOptions

--------------------------------------------------------------------------------
data Metric =
    C Counter.Counter
  | G Gauge.Gauge

type MetricsMap = Map.Map Prometheus.Name Metric

--------------------------------------------------------------------------------
defaultOptions :: Prometheus.Labels -> AdapterOptions
defaultOptions :: Labels -> AdapterOptions
defaultOptions Labels
l = Labels -> Maybe Text -> Int -> AdapterOptions
AdapterOptions Labels
l Maybe Text
forall a. Maybe a
Nothing Int
15

--------------------------------------------------------------------------------
registerEKGStore :: MonadIO m => EKG.Store -> AdapterOptions -> RegistryT m ()
registerEKGStore :: forall (m :: * -> *).
MonadIO m =>
Store -> AdapterOptions -> RegistryT m ()
registerEKGStore Store
store AdapterOptions
opts = StateT Registry m () -> RegistryT m ()
forall (m :: * -> *) a. StateT Registry m a -> RegistryT m a
RegistryT (StateT Registry m () -> RegistryT m ())
-> StateT Registry m () -> RegistryT m ()
forall a b. (a -> b) -> a -> b
$ (Registry -> m ((), Registry)) -> StateT Registry m ()
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
StateT ((Registry -> m ((), Registry)) -> StateT Registry m ())
-> (Registry -> m ((), Registry)) -> StateT Registry m ()
forall a b. (a -> b) -> a -> b
$ \Registry
_ -> do
  (Registry
r, MetricsMap
mmap) <- IO (Registry, MetricsMap) -> m (Registry, MetricsMap)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Registry, MetricsMap) -> m (Registry, MetricsMap))
-> IO (Registry, MetricsMap) -> m (Registry, MetricsMap)
forall a b. (a -> b) -> a -> b
$ Store -> AdapterOptions -> IO (Registry, MetricsMap)
toPrometheusRegistry' Store
store AdapterOptions
opts
  IO ThreadId -> m ThreadId
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ThreadId -> m ThreadId) -> IO ThreadId -> m ThreadId
forall a b. (a -> b) -> a -> b
$ IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ do
    let loop :: IO a
loop = IO a -> IO a
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ do
                 Int -> IO ()
threadDelay (AdapterOptions -> Int
_samplingFrequency AdapterOptions
opts Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
10Int -> Integer -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^Integer
6)
                 Store -> AdapterOptions -> MetricsMap -> IO ()
updateMetrics Store
store AdapterOptions
opts MetricsMap
mmap
                 IO a
loop
    IO ()
forall {a}. IO a
loop
  ((), Registry) -> m ((), Registry)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((), Registry
r)

--------------------------------------------------------------------------------
toPrometheusRegistry' :: EKG.Store -> AdapterOptions -> IO (Prometheus.Registry, MetricsMap)
toPrometheusRegistry' :: Store -> AdapterOptions -> IO (Registry, MetricsMap)
toPrometheusRegistry' Store
store AdapterOptions
opts = do
  let registry :: Registry
registry = Registry
Prometheus.new
  Sample
samples <- Store -> IO Sample
EKG.sampleAll Store
store
  ((Registry, MetricsMap)
 -> (Text, Value) -> IO (Registry, MetricsMap))
-> (Registry, MetricsMap)
-> [(Text, Value)]
-> IO (Registry, MetricsMap)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (AdapterOptions
-> (Registry, MetricsMap)
-> (Text, Value)
-> IO (Registry, MetricsMap)
mkMetric AdapterOptions
opts) (Registry
registry, MetricsMap
forall k a. Map k a
Map.empty) (Sample -> [(Text, Value)]
forall k v. HashMap k v -> [(k, v)]
HMap.toList Sample
samples)

--------------------------------------------------------------------------------
toPrometheusRegistry :: EKG.Store -> AdapterOptions -> IO Prometheus.Registry
toPrometheusRegistry :: Store -> AdapterOptions -> IO Registry
toPrometheusRegistry Store
store AdapterOptions
opts = (Registry, MetricsMap) -> Registry
forall a b. (a, b) -> a
fst ((Registry, MetricsMap) -> Registry)
-> IO (Registry, MetricsMap) -> IO Registry
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Store -> AdapterOptions -> IO (Registry, MetricsMap)
toPrometheusRegistry' Store
store AdapterOptions
opts

--------------------------------------------------------------------------------
mkMetric :: AdapterOptions -> (Prometheus.Registry, MetricsMap) -> (T.Text, EKG.Value) -> IO (Prometheus.Registry, MetricsMap)
mkMetric :: AdapterOptions
-> (Registry, MetricsMap)
-> (Text, Value)
-> IO (Registry, MetricsMap)
mkMetric AdapterOptions{Int
Maybe Text
Labels
_labels :: AdapterOptions -> Labels
_namespace :: AdapterOptions -> Maybe Text
_samplingFrequency :: AdapterOptions -> Int
_labels :: Labels
_namespace :: Maybe Text
_samplingFrequency :: Int
..} (Registry
oldRegistry, MetricsMap
mmap) (Text
key, Value
value) = do
  let k :: Name
k = Maybe Text -> Text -> Name
mkKey Maybe Text
_namespace Text
key
  case Value
value of
   EKG.Counter Int64
c -> do
     (Counter
counter, Registry
newRegistry) <- Name -> Labels -> Registry -> IO (Counter, Registry)
Prometheus.registerCounter Name
k Labels
_labels Registry
oldRegistry
     Int -> Counter -> IO ()
Counter.add (Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
c) Counter
counter
     (Registry, MetricsMap) -> IO (Registry, MetricsMap)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Registry, MetricsMap) -> IO (Registry, MetricsMap))
-> (Registry, MetricsMap) -> IO (Registry, MetricsMap)
forall a b. (a -> b) -> a -> b
$! (Registry
newRegistry, Name -> Metric -> MetricsMap -> MetricsMap
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Name
k (Counter -> Metric
C Counter
counter) (MetricsMap -> MetricsMap) -> MetricsMap -> MetricsMap
forall a b. (a -> b) -> a -> b
$! MetricsMap
mmap)
   EKG.Gauge Int64
g   -> do
     (Gauge
gauge, Registry
newRegistry) <- Name -> Labels -> Registry -> IO (Gauge, Registry)
Prometheus.registerGauge Name
k Labels
_labels Registry
oldRegistry
     Double -> Gauge -> IO ()
Gauge.set (Int64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
g) Gauge
gauge
     (Registry, MetricsMap) -> IO (Registry, MetricsMap)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Registry, MetricsMap) -> IO (Registry, MetricsMap))
-> (Registry, MetricsMap) -> IO (Registry, MetricsMap)
forall a b. (a -> b) -> a -> b
$! (Registry
newRegistry, Name -> Metric -> MetricsMap -> MetricsMap
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Name
k (Gauge -> Metric
G Gauge
gauge) (MetricsMap -> MetricsMap) -> MetricsMap -> MetricsMap
forall a b. (a -> b) -> a -> b
$! MetricsMap
mmap)
   EKG.Label Text
_   -> (Registry, MetricsMap) -> IO (Registry, MetricsMap)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Registry, MetricsMap) -> IO (Registry, MetricsMap))
-> (Registry, MetricsMap) -> IO (Registry, MetricsMap)
forall a b. (a -> b) -> a -> b
$! (Registry
oldRegistry, MetricsMap
mmap)
   EKG.Distribution Stats
_ -> (Registry, MetricsMap) -> IO (Registry, MetricsMap)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Registry, MetricsMap) -> IO (Registry, MetricsMap))
-> (Registry, MetricsMap) -> IO (Registry, MetricsMap)
forall a b. (a -> b) -> a -> b
$! (Registry
oldRegistry, MetricsMap
mmap)

--------------------------------------------------------------------------------
updateMetrics :: EKG.Store -> AdapterOptions -> MetricsMap -> IO ()
updateMetrics :: Store -> AdapterOptions -> MetricsMap -> IO ()
updateMetrics Store
store AdapterOptions
opts MetricsMap
mmap = do
  Sample
samples <- Store -> IO Sample
EKG.sampleAll Store
store
  () -> MetricsMap -> ()
forall a b. a -> b -> a
const () (MetricsMap -> ()) -> IO MetricsMap -> IO ()
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (MetricsMap -> (Text, Value) -> IO MetricsMap)
-> MetricsMap -> [(Text, Value)] -> IO MetricsMap
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (AdapterOptions -> MetricsMap -> (Text, Value) -> IO MetricsMap
updateMetric AdapterOptions
opts) MetricsMap
mmap (Sample -> [(Text, Value)]
forall k v. HashMap k v -> [(k, v)]
HMap.toList Sample
samples)


--------------------------------------------------------------------------------
mkKey :: Maybe T.Text -> T.Text -> Prometheus.Name
mkKey :: Maybe Text -> Text -> Name
mkKey Maybe Text
mbNs Text
k =
  Text -> Name
Prometheus.Name (Text -> Name) -> Text -> Name
forall a b. (a -> b) -> a -> b
$ (Text -> (Text -> Text) -> Maybe Text -> Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Text
forall a. Monoid a => a
mempty (\Text
x -> Text
x Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_") Maybe Text
mbNs) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
T.replace Text
"." Text
"_" Text
k

--------------------------------------------------------------------------------
updateMetric :: AdapterOptions -> MetricsMap -> (T.Text, EKG.Value) -> IO MetricsMap
updateMetric :: AdapterOptions -> MetricsMap -> (Text, Value) -> IO MetricsMap
updateMetric AdapterOptions{Int
Maybe Text
Labels
_labels :: AdapterOptions -> Labels
_namespace :: AdapterOptions -> Maybe Text
_samplingFrequency :: AdapterOptions -> Int
_labels :: Labels
_namespace :: Maybe Text
_samplingFrequency :: Int
..} MetricsMap
mmap (Text
key, Value
value) = do
  let k :: Name
k = Maybe Text -> Text -> Name
mkKey Maybe Text
_namespace Text
key
  case (Metric -> Value -> (Metric, Value))
-> Maybe Metric -> Maybe Value -> Maybe (Metric, Value)
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,) (Name -> MetricsMap -> Maybe Metric
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Name
k MetricsMap
mmap) (Value -> Maybe Value
forall a. a -> Maybe a
Just Value
value) of
    Just (C Counter
counter, EKG.Counter Int64
c)  -> do
      (Counter.CounterSample Int
oldCounterValue) <- Counter -> IO CounterSample
Counter.sample Counter
counter
      let slack :: Int64
slack = Int64
c Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
oldCounterValue
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int64
slack Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Int64
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> Counter -> IO ()
Counter.add (Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
slack) Counter
counter
      MetricsMap -> IO MetricsMap
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MetricsMap -> IO MetricsMap) -> MetricsMap -> IO MetricsMap
forall a b. (a -> b) -> a -> b
$! MetricsMap
mmap
    Just (G Gauge
gauge,   EKG.Gauge Int64
g) -> do
      Double -> Gauge -> IO ()
Gauge.set (Int64 -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
g) Gauge
gauge
      MetricsMap -> IO MetricsMap
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MetricsMap -> IO MetricsMap) -> MetricsMap -> IO MetricsMap
forall a b. (a -> b) -> a -> b
$! MetricsMap
mmap
    Maybe (Metric, Value)
_ -> MetricsMap -> IO MetricsMap
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return MetricsMap
mmap