{-# 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
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [KeyError] -> ShowS
$cshowList :: [KeyError] -> ShowS
show :: KeyError -> String
$cshow :: KeyError -> String
showsPrec :: Int -> KeyError -> ShowS
$cshowsPrec :: Int -> KeyError -> ShowS
Show, Typeable)
instance Exception KeyError


new :: Registry
new :: Registry
new = Map MetricId Metric -> Registry
Registry 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
    forall (m :: * -> *) a. Monad m => a -> m a
return (Counter
counter, Map MetricId Metric -> Registry
Registry forall a b. (a -> b) -> a -> b
$ forall k a.
Ord k =>
(k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWithKey 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
_ = 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
    forall (m :: * -> *) a. Monad m => a -> m a
return (Gauge
gauge, Map MetricId Metric -> Registry
Registry forall a b. (a -> b) -> a -> b
$ forall k a.
Ord k =>
(k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWithKey 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
_ = 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
    forall (m :: * -> *) a. Monad m => a -> m a
return (Histogram
histogram, Map MetricId Metric -> Registry
Registry forall a b. (a -> b) -> a -> b
$ forall k a.
Ord k =>
(k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWithKey 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
_ = 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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => k -> Map k a -> Map k a
Map.delete MetricId
i forall a b. (a -> b) -> a -> b
$ Map MetricId Metric
m


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


sample :: Registry -> IO RegistrySample
sample :: Registry -> IO RegistrySample
sample = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Map MetricId MetricSample -> RegistrySample
RegistrySample forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Metric -> IO MetricSample
sampleMetric 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 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 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Histogram -> IO HistogramSample
Histogram.sample Histogram
histogram