{-# LANGUAGE DeriveDataTypeable #-}

module System.Metrics.Prometheus.Registry (
    Registry,
    RegistrySample (..),
    new,
    registerCounter,
    registerGauge,
    registerHistogram,
    listMetricIds,
    removeMetric,
    sample,
) where

import Control.Applicative ((<$>))
import Control.Exception (Exception, throw)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Typeable (Typeable)

import System.Metrics.Prometheus.Metric (
    Metric (..),
    MetricSample (..),
 )
import System.Metrics.Prometheus.Metric.Counter (Counter)
import qualified System.Metrics.Prometheus.Metric.Counter as Counter
import System.Metrics.Prometheus.Metric.Gauge (Gauge)
import qualified System.Metrics.Prometheus.Metric.Gauge as Gauge
import System.Metrics.Prometheus.Metric.Histogram (
    Histogram,
    UpperBound,
 )
import qualified System.Metrics.Prometheus.Metric.Histogram as Histogram
import System.Metrics.Prometheus.MetricId (
    Labels (..),
    MetricId (MetricId),
    Name (..),
 )


newtype Registry = Registry {Registry -> Map MetricId Metric
unRegistry :: Map MetricId Metric}
newtype RegistrySample = RegistrySample {RegistrySample -> Map MetricId MetricSample
unRegistrySample :: Map MetricId MetricSample}


newtype KeyError = KeyError MetricId deriving (Int -> KeyError -> ShowS
[KeyError] -> ShowS
KeyError -> String
(Int -> KeyError -> ShowS)
-> (KeyError -> String) -> ([KeyError] -> ShowS) -> Show KeyError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> KeyError -> ShowS
showsPrec :: Int -> KeyError -> ShowS
$cshow :: KeyError -> String
show :: KeyError -> String
$cshowList :: [KeyError] -> ShowS
showList :: [KeyError] -> ShowS
Show, Typeable)
instance Exception KeyError


new :: Registry
new :: Registry
new = Map MetricId Metric -> Registry
Registry Map MetricId Metric
forall k a. Map k a
Map.empty


registerCounter :: Name -> Labels -> Registry -> IO (Counter, Registry)
registerCounter :: Name -> Labels -> Registry -> IO (Counter, Registry)
registerCounter Name
name Labels
labels Registry
registry = do
    Counter
counter <- IO Counter
Counter.new
    (Counter, Registry) -> IO (Counter, Registry)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Counter
counter, Map MetricId Metric -> Registry
Registry (Map MetricId Metric -> Registry)
-> Map MetricId Metric -> Registry
forall a b. (a -> b) -> a -> b
$ (MetricId -> Metric -> Metric -> Metric)
-> MetricId -> Metric -> Map MetricId Metric -> Map MetricId Metric
forall k a.
Ord k =>
(k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWithKey MetricId -> Metric -> Metric -> Metric
forall {p} {p} {a}. MetricId -> p -> p -> a
collision MetricId
mid (Counter -> Metric
CounterMetric Counter
counter) (Registry -> Map MetricId Metric
unRegistry Registry
registry))
  where
    mid :: MetricId
mid = Name -> Labels -> MetricId
MetricId Name
name Labels
labels
    collision :: MetricId -> p -> p -> a
collision MetricId
k p
_ p
_ = KeyError -> a
forall a e. Exception e => e -> a
throw (MetricId -> KeyError
KeyError MetricId
k)


registerGauge :: Name -> Labels -> Registry -> IO (Gauge, Registry)
registerGauge :: Name -> Labels -> Registry -> IO (Gauge, Registry)
registerGauge Name
name Labels
labels Registry
registry = do
    Gauge
gauge <- IO Gauge
Gauge.new
    (Gauge, Registry) -> IO (Gauge, Registry)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Gauge
gauge, Map MetricId Metric -> Registry
Registry (Map MetricId Metric -> Registry)
-> Map MetricId Metric -> Registry
forall a b. (a -> b) -> a -> b
$ (MetricId -> Metric -> Metric -> Metric)
-> MetricId -> Metric -> Map MetricId Metric -> Map MetricId Metric
forall k a.
Ord k =>
(k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWithKey MetricId -> Metric -> Metric -> Metric
forall {p} {p} {a}. MetricId -> p -> p -> a
collision MetricId
mid (Gauge -> Metric
GaugeMetric Gauge
gauge) (Registry -> Map MetricId Metric
unRegistry Registry
registry))
  where
    mid :: MetricId
mid = Name -> Labels -> MetricId
MetricId Name
name Labels
labels
    collision :: MetricId -> p -> p -> a
collision MetricId
k p
_ p
_ = KeyError -> a
forall a e. Exception e => e -> a
throw (MetricId -> KeyError
KeyError MetricId
k)


registerHistogram :: Name -> Labels -> [UpperBound] -> Registry -> IO (Histogram, Registry)
registerHistogram :: Name
-> Labels -> [UpperBound] -> Registry -> IO (Histogram, Registry)
registerHistogram Name
name Labels
labels [UpperBound]
buckets Registry
registry = do
    Histogram
histogram <- [UpperBound] -> IO Histogram
Histogram.new [UpperBound]
buckets
    (Histogram, Registry) -> IO (Histogram, Registry)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Histogram
histogram, Map MetricId Metric -> Registry
Registry (Map MetricId Metric -> Registry)
-> Map MetricId Metric -> Registry
forall a b. (a -> b) -> a -> b
$ (MetricId -> Metric -> Metric -> Metric)
-> MetricId -> Metric -> Map MetricId Metric -> Map MetricId Metric
forall k a.
Ord k =>
(k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWithKey MetricId -> Metric -> Metric -> Metric
forall {p} {p} {a}. MetricId -> p -> p -> a
collision MetricId
mid (Histogram -> Metric
HistogramMetric Histogram
histogram) (Registry -> Map MetricId Metric
unRegistry Registry
registry))
  where
    mid :: MetricId
mid = Name -> Labels -> MetricId
MetricId Name
name Labels
labels
    collision :: MetricId -> p -> p -> a
collision MetricId
k p
_ p
_ = KeyError -> a
forall a e. Exception e => e -> a
throw (MetricId -> KeyError
KeyError MetricId
k)


removeMetric :: MetricId -> Registry -> Registry
removeMetric :: MetricId -> Registry -> Registry
removeMetric MetricId
i (Registry Map MetricId Metric
m) = Map MetricId Metric -> Registry
Registry (Map MetricId Metric -> Registry)
-> (Map MetricId Metric -> Map MetricId Metric)
-> Map MetricId Metric
-> Registry
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MetricId -> Map MetricId Metric -> Map MetricId Metric
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete MetricId
i (Map MetricId Metric -> Registry)
-> Map MetricId Metric -> Registry
forall a b. (a -> b) -> a -> b
$ Map MetricId Metric
m


listMetricIds :: Registry -> [MetricId]
listMetricIds :: Registry -> [MetricId]
listMetricIds = Map MetricId Metric -> [MetricId]
forall k a. Map k a -> [k]
Map.keys (Map MetricId Metric -> [MetricId])
-> (Registry -> Map MetricId Metric) -> Registry -> [MetricId]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Registry -> Map MetricId Metric
unRegistry


sample :: Registry -> IO RegistrySample
sample :: Registry -> IO RegistrySample
sample = (Map MetricId MetricSample -> RegistrySample)
-> IO (Map MetricId MetricSample) -> IO RegistrySample
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Map MetricId MetricSample -> RegistrySample
RegistrySample (IO (Map MetricId MetricSample) -> IO RegistrySample)
-> (Registry -> IO (Map MetricId MetricSample))
-> Registry
-> IO RegistrySample
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Metric -> IO MetricSample)
-> Map MetricId Metric -> IO (Map MetricId MetricSample)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Map MetricId a -> m (Map MetricId b)
mapM Metric -> IO MetricSample
sampleMetric (Map MetricId Metric -> IO (Map MetricId MetricSample))
-> (Registry -> Map MetricId Metric)
-> Registry
-> IO (Map MetricId MetricSample)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Registry -> Map MetricId Metric
unRegistry
  where
    sampleMetric :: Metric -> IO MetricSample
    sampleMetric :: Metric -> IO MetricSample
sampleMetric (CounterMetric Counter
count) = CounterSample -> MetricSample
CounterMetricSample (CounterSample -> MetricSample)
-> IO CounterSample -> IO MetricSample
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Counter -> IO CounterSample
Counter.sample Counter
count
    sampleMetric (GaugeMetric Gauge
gauge) = GaugeSample -> MetricSample
GaugeMetricSample (GaugeSample -> MetricSample) -> IO GaugeSample -> IO MetricSample
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gauge -> IO GaugeSample
Gauge.sample Gauge
gauge
    sampleMetric (HistogramMetric Histogram
histogram) = HistogramSample -> MetricSample
HistogramMetricSample (HistogramSample -> MetricSample)
-> IO HistogramSample -> IO MetricSample
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Histogram -> IO HistogramSample
Histogram.sample Histogram
histogram