{-# LANGUAGE ExistentialQuantification #-}
module Prometheus.Registry (
    register
,   registerIO
,   unsafeRegister
,   unsafeRegisterIO
,   collectMetrics
,   unregisterAll
) where

import Prometheus.Metric

import Control.Applicative ((<$>))
import Control.Monad.IO.Class
import System.IO.Unsafe (unsafePerformIO)
import qualified Control.Concurrent.STM as STM


-- $setup
-- >>> :module +Prometheus
-- >>> unregisterAll

-- | A 'Registry' is a list of all registered metrics, currently represented by
-- their sampling functions.
type Registry = [IO [SampleGroup]]

{-# NOINLINE globalRegistry #-}
globalRegistry :: STM.TVar Registry
globalRegistry :: TVar Registry
globalRegistry = IO (TVar Registry) -> TVar Registry
forall a. IO a -> a
unsafePerformIO (IO (TVar Registry) -> TVar Registry)
-> IO (TVar Registry) -> TVar Registry
forall a b. (a -> b) -> a -> b
$ Registry -> IO (TVar Registry)
forall a. a -> IO (TVar a)
STM.newTVarIO []

-- | Registers a metric with the global metric registry.
register :: MonadIO m => Metric s -> m s
register :: Metric s -> m s
register (Metric IO (s, IO [SampleGroup])
mk) = IO s -> m s
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO s -> m s) -> IO s -> m s
forall a b. (a -> b) -> a -> b
$ do
    (s
metric, IO [SampleGroup]
sampleGroups) <- IO (s, IO [SampleGroup])
mk
    let addToRegistry :: Registry -> Registry
addToRegistry = (IO [SampleGroup]
sampleGroups IO [SampleGroup] -> Registry -> Registry
forall a. a -> [a] -> [a]
:)
    IO () -> IO ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ STM () -> IO ()
forall a. STM a -> IO a
STM.atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Registry -> (Registry -> Registry) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
STM.modifyTVar' TVar Registry
globalRegistry Registry -> Registry
addToRegistry
    s -> IO s
forall (m :: * -> *) a. Monad m => a -> m a
return s
metric

-- | Registers a metric with the global metric registry.
registerIO :: MonadIO m => m (Metric s) -> m s
registerIO :: m (Metric s) -> m s
registerIO m (Metric s)
metricGen = m (Metric s)
metricGen m (Metric s) -> (Metric s -> m s) -> m s
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Metric s -> m s
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register

-- | Registers a metric with the global metric registry.
--
-- __IMPORTANT__: This method should only be used to register metrics as top
-- level symbols, it should not be run from other pure code.
unsafeRegister :: Metric s -> s
unsafeRegister :: Metric s -> s
unsafeRegister = IO s -> s
forall a. IO a -> a
unsafePerformIO (IO s -> s) -> (Metric s -> IO s) -> Metric s -> s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Metric s -> IO s
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register

-- | Registers a metric with the global metric registry.
--
-- __IMPORTANT__: This method should only be used to register metrics as top
-- level symbols, it should not be run from other pure code.
--
-- For example,
--
-- >>> :{
--  {-# NOINLINE c #-}
--  let c = unsafeRegisterIO $ counter (Info "my_counter" "An example metric")
-- :}
-- ...
unsafeRegisterIO :: IO (Metric s) -> s
unsafeRegisterIO :: IO (Metric s) -> s
unsafeRegisterIO = IO s -> s
forall a. IO a -> a
unsafePerformIO (IO s -> s) -> (IO (Metric s) -> IO s) -> IO (Metric s) -> s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Metric s) -> IO s
forall (m :: * -> *) s. MonadIO m => m (Metric s) -> m s
registerIO

-- | Removes all currently registered metrics from the registry.
unregisterAll :: MonadIO m => m ()
unregisterAll :: m ()
unregisterAll = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ STM () -> IO ()
forall a. STM a -> IO a
STM.atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Registry -> Registry -> STM ()
forall a. TVar a -> a -> STM ()
STM.writeTVar TVar Registry
globalRegistry []

-- | Collect samples from all currently registered metrics. In typical use cases
-- there is no reason to use this function, instead you should use
-- `exportMetricsAsText` or a convenience library.
--
-- This function is likely only of interest if you wish to export metrics in
-- a non-supported format for use with another monitoring service.
collectMetrics :: MonadIO m => m [SampleGroup]
collectMetrics :: m [SampleGroup]
collectMetrics = IO [SampleGroup] -> m [SampleGroup]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO [SampleGroup] -> m [SampleGroup])
-> IO [SampleGroup] -> m [SampleGroup]
forall a b. (a -> b) -> a -> b
$ do
    Registry
registry <- STM Registry -> IO Registry
forall a. STM a -> IO a
STM.atomically (STM Registry -> IO Registry) -> STM Registry -> IO Registry
forall a b. (a -> b) -> a -> b
$ TVar Registry -> STM Registry
forall a. TVar a -> STM a
STM.readTVar TVar Registry
globalRegistry
    [[SampleGroup]] -> [SampleGroup]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[SampleGroup]] -> [SampleGroup])
-> IO [[SampleGroup]] -> IO [SampleGroup]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Registry -> IO [[SampleGroup]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence Registry
registry