-- | Described metrics and compatibility aliases backed by the Prometheus client.
module Hydra.Logging.Monitoring.Registry where

import Hydra.Prelude hiding (toList)

import Data.ByteString.Builder (Builder, byteString, lazyByteString, toLazyByteString)
import Data.ByteString.Lazy qualified as LBS
import Data.Map.Strict qualified as Map
import Data.Text qualified as Text
import System.Metrics.Prometheus.Concurrent.Registry qualified as Prom
import System.Metrics.Prometheus.Encode.Text (encodeMetrics)
import System.Metrics.Prometheus.Metric (MetricSample (..))
import System.Metrics.Prometheus.Metric.Counter (Counter)
import System.Metrics.Prometheus.Metric.Gauge (Gauge, GaugeSample (..))
import System.Metrics.Prometheus.Metric.Histogram (Histogram, HistogramSample (..))
import System.Metrics.Prometheus.MetricId (Labels (..), MetricId (..), Name (..))
import System.Metrics.Prometheus.Registry (RegistrySample (..))

data Registry = Registry
  { Registry -> Registry
underlying :: Prom.Registry
  , Registry -> IORef (Map Text (Text, Text))
descriptions :: IORef (Map Text (Text, Text))
  , Registry -> IORef [(Text, Text, Double)]
aliases :: IORef [(Text, Text, Double)]
  }

new :: IO Registry
new :: IO Registry
new = Registry
-> IORef (Map Text (Text, Text))
-> IORef [(Text, Text, Double)]
-> Registry
Registry (Registry
 -> IORef (Map Text (Text, Text))
 -> IORef [(Text, Text, Double)]
 -> Registry)
-> IO Registry
-> IO
     (IORef (Map Text (Text, Text))
      -> IORef [(Text, Text, Double)] -> Registry)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Registry
Prom.new IO
  (IORef (Map Text (Text, Text))
   -> IORef [(Text, Text, Double)] -> Registry)
-> IO (IORef (Map Text (Text, Text)))
-> IO (IORef [(Text, Text, Double)] -> Registry)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map Text (Text, Text) -> IO (IORef (Map Text (Text, Text)))
forall (m :: * -> *) a. MonadIO m => a -> m (IORef a)
newIORef Map Text (Text, Text)
forall a. Monoid a => a
mempty IO (IORef [(Text, Text, Double)] -> Registry)
-> IO (IORef [(Text, Text, Double)]) -> IO Registry
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [(Text, Text, Double)] -> IO (IORef [(Text, Text, Double)])
forall (m :: * -> *) a. MonadIO m => a -> m (IORef a)
newIORef []

describe :: Registry -> Text -> Text -> Text -> IO ()
describe :: Registry -> Text -> Text -> Text -> IO ()
describe Registry{IORef (Map Text (Text, Text))
$sel:descriptions:Registry :: Registry -> IORef (Map Text (Text, Text))
descriptions :: IORef (Map Text (Text, Text))
descriptions} Text
name Text
kind Text
help =
  IORef (Map Text (Text, Text))
-> (Map Text (Text, Text) -> (Map Text (Text, Text), ())) -> IO ()
forall (m :: * -> *) a b.
MonadIO m =>
IORef a -> (a -> (a, b)) -> m b
atomicModifyIORef' IORef (Map Text (Text, Text))
descriptions ((Map Text (Text, Text) -> (Map Text (Text, Text), ())) -> IO ())
-> (Map Text (Text, Text) -> (Map Text (Text, Text), ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Map Text (Text, Text)
ds -> (Text
-> (Text, Text) -> Map Text (Text, Text) -> Map Text (Text, Text)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
name (Text
kind, Text
help) Map Text (Text, Text)
ds, ())

counter :: Registry -> Text -> Text -> [(Text, Text)] -> IO Counter
counter :: Registry -> Text -> Text -> [(Text, Text)] -> IO Counter
counter registry :: Registry
registry@Registry{Registry
$sel:underlying:Registry :: Registry -> Registry
underlying :: Registry
underlying} Text
name Text
help [(Text, Text)]
labels = do
  Registry -> Text -> Text -> Text -> IO ()
describe Registry
registry Text
name Text
"counter" Text
help
  Name -> Labels -> Registry -> IO Counter
Prom.registerCounter (Text -> Name
Name Text
name) (Map Text Text -> Labels
Labels (Map Text Text -> Labels) -> Map Text Text -> Labels
forall a b. (a -> b) -> a -> b
$ [(Text, Text)] -> Map Text Text
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Text, Text)]
labels) Registry
underlying

gauge :: Registry -> Text -> Text -> [(Text, Text)] -> IO Gauge
gauge :: Registry -> Text -> Text -> [(Text, Text)] -> IO Gauge
gauge registry :: Registry
registry@Registry{Registry
$sel:underlying:Registry :: Registry -> Registry
underlying :: Registry
underlying} Text
name Text
help [(Text, Text)]
labels = do
  Registry -> Text -> Text -> Text -> IO ()
describe Registry
registry Text
name Text
"gauge" Text
help
  Name -> Labels -> Registry -> IO Gauge
Prom.registerGauge (Text -> Name
Name Text
name) (Map Text Text -> Labels
Labels (Map Text Text -> Labels) -> Map Text Text -> Labels
forall a b. (a -> b) -> a -> b
$ [(Text, Text)] -> Map Text Text
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Text, Text)]
labels) Registry
underlying

-- | Absolute, possibly fractional counters (e.g. RTS CPU seconds). The client's
-- incrementable Counter stores Int; sample these through a Gauge handle and
-- publish their counter type. Callers supply monotonically increasing values.
sampledCounter :: Registry -> Text -> Text -> IO Gauge
sampledCounter :: Registry -> Text -> Text -> IO Gauge
sampledCounter registry :: Registry
registry@Registry{Registry
$sel:underlying:Registry :: Registry -> Registry
underlying :: Registry
underlying} Text
name Text
help = do
  Registry -> Text -> Text -> Text -> IO ()
describe Registry
registry Text
name Text
"counter" Text
help
  Name -> Labels -> Registry -> IO Gauge
Prom.registerGauge (Text -> Name
Name Text
name) Labels
forall a. Monoid a => a
mempty Registry
underlying

histogram :: Registry -> Text -> Text -> [Double] -> [(Text, Text)] -> IO Histogram
histogram :: Registry
-> Text -> Text -> [Double] -> [(Text, Text)] -> IO Histogram
histogram registry :: Registry
registry@Registry{Registry
$sel:underlying:Registry :: Registry -> Registry
underlying :: Registry
underlying} Text
name Text
help [Double]
buckets [(Text, Text)]
labels = do
  Registry -> Text -> Text -> Text -> IO ()
describe Registry
registry Text
name Text
"histogram" Text
help
  Name -> Labels -> [Double] -> Registry -> IO Histogram
Prom.registerHistogram (Text -> Name
Name Text
name) (Map Text Text -> Labels
Labels (Map Text Text -> Labels) -> Map Text Text -> Labels
forall a b. (a -> b) -> a -> b
$ [(Text, Text)] -> Map Text Text
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Text, Text)]
labels) [Double]
buckets Registry
underlying

-- | Both names are rendered from exactly the same sample. Histogram aliases
-- scale the sum and bucket boundaries together; counts remain unchanged.
alias :: Registry -> Text -> Text -> Double -> IO ()
alias :: Registry -> Text -> Text -> Double -> IO ()
alias Registry{IORef [(Text, Text, Double)]
$sel:aliases:Registry :: Registry -> IORef [(Text, Text, Double)]
aliases :: IORef [(Text, Text, Double)]
aliases} Text
canonical Text
legacy Double
scale =
  IORef [(Text, Text, Double)]
-> ([(Text, Text, Double)] -> ([(Text, Text, Double)], ()))
-> IO ()
forall (m :: * -> *) a b.
MonadIO m =>
IORef a -> (a -> (a, b)) -> m b
atomicModifyIORef' IORef [(Text, Text, Double)]
aliases (([(Text, Text, Double)] -> ([(Text, Text, Double)], ())) -> IO ())
-> ([(Text, Text, Double)] -> ([(Text, Text, Double)], ()))
-> IO ()
forall a b. (a -> b) -> a -> b
$ \[(Text, Text, Double)]
as -> ((Text
canonical, Text
legacy, Double
scale) (Text, Text, Double)
-> [(Text, Text, Double)] -> [(Text, Text, Double)]
forall a. a -> [a] -> [a]
: [(Text, Text, Double)]
as, ())

sample :: Registry -> IO LBS.ByteString
sample :: Registry -> IO ByteString
sample Registry{Registry
$sel:underlying:Registry :: Registry -> Registry
underlying :: Registry
underlying, IORef (Map Text (Text, Text))
$sel:descriptions:Registry :: Registry -> IORef (Map Text (Text, Text))
descriptions :: IORef (Map Text (Text, Text))
descriptions, IORef [(Text, Text, Double)]
$sel:aliases:Registry :: Registry -> IORef [(Text, Text, Double)]
aliases :: IORef [(Text, Text, Double)]
aliases} = do
  RegistrySample Map MetricId MetricSample
values <- Registry -> IO RegistrySample
Prom.sample Registry
underlying
  Map Text (Text, Text)
ds <- IORef (Map Text (Text, Text)) -> IO (Map Text (Text, Text))
forall (m :: * -> *) a. MonadIO m => IORef a -> m a
readIORef IORef (Map Text (Text, Text))
descriptions
  [(Text, Text, Double)]
as <- IORef [(Text, Text, Double)] -> IO [(Text, Text, Double)]
forall (m :: * -> *) a. MonadIO m => IORef a -> m a
readIORef IORef [(Text, Text, Double)]
aliases
  let extra :: [(MetricId, MetricSample)]
extra =
        [ (Name -> Labels -> MetricId
MetricId (Text -> Name
Name Text
legacy) Labels
labels, Double -> MetricSample -> MetricSample
scaleSample Double
scale MetricSample
value)
        | (Text
canonical, Text
legacy, Double
scale) <- [(Text, Text, Double)]
as
        , (MetricId (Name Text
name) Labels
labels, MetricSample
value) <- Map MetricId MetricSample -> [(MetricId, MetricSample)]
forall k a. Map k a -> [(k, a)]
Map.toList Map MetricId MetricSample
values
        , Text
name Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
canonical
        ]
      allDescriptions :: Map Text (Text, Text)
allDescriptions = (Map Text (Text, Text)
 -> (Text, Text, Double) -> Map Text (Text, Text))
-> Map Text (Text, Text)
-> [(Text, Text, Double)]
-> Map Text (Text, Text)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Map Text (Text, Text)
acc (Text
canonical, Text
legacy, Double
_) -> Map Text (Text, Text)
-> ((Text, Text) -> Map Text (Text, Text))
-> Maybe (Text, Text)
-> Map Text (Text, Text)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Map Text (Text, Text)
acc (\(Text, Text)
d -> Text
-> (Text, Text) -> Map Text (Text, Text) -> Map Text (Text, Text)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
legacy (Text, Text)
d Map Text (Text, Text)
acc) (Text -> Map Text (Text, Text) -> Maybe (Text, Text)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
canonical Map Text (Text, Text)
ds)) Map Text (Text, Text)
ds [(Text, Text, Double)]
as
      groups :: Map Text [(MetricId, MetricSample)]
groups = ([(MetricId, MetricSample)]
 -> [(MetricId, MetricSample)] -> [(MetricId, MetricSample)])
-> [(Text, [(MetricId, MetricSample)])]
-> Map Text [(MetricId, MetricSample)]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith [(MetricId, MetricSample)]
-> [(MetricId, MetricSample)] -> [(MetricId, MetricSample)]
forall a. Semigroup a => a -> a -> a
(<>) [(Name -> Text
unName Name
name, [(MetricId
mid, MetricSample
value)]) | (mid :: MetricId
mid@MetricId{Name
name :: Name
name :: MetricId -> Name
name}, MetricSample
value) <- Map MetricId MetricSample -> [(MetricId, MetricSample)]
forall k a. Map k a -> [(k, a)]
Map.toList Map MetricId MetricSample
values [(MetricId, MetricSample)]
-> [(MetricId, MetricSample)] -> [(MetricId, MetricSample)]
forall a. Semigroup a => a -> a -> a
<> [(MetricId, MetricSample)]
extra]
  ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString)
-> ([(Text, [(MetricId, MetricSample)])] -> ByteString)
-> [(Text, [(MetricId, MetricSample)])]
-> IO ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> ByteString
toLazyByteString (Builder -> ByteString)
-> ([(Text, [(MetricId, MetricSample)])] -> Builder)
-> [(Text, [(MetricId, MetricSample)])]
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Text, [(MetricId, MetricSample)]) -> Builder)
-> [(Text, [(MetricId, MetricSample)])] -> Builder
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Map Text (Text, Text)
-> (Text, [(MetricId, MetricSample)]) -> Builder
render Map Text (Text, Text)
allDescriptions) ([(Text, [(MetricId, MetricSample)])] -> IO ByteString)
-> [(Text, [(MetricId, MetricSample)])] -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Map Text [(MetricId, MetricSample)]
-> [(Text, [(MetricId, MetricSample)])]
forall k a. Map k a -> [(k, a)]
Map.toList Map Text [(MetricId, MetricSample)]
groups
 where
  render :: Map Text (Text, Text) -> (Text, [(MetricId, MetricSample)]) -> Builder
  render :: Map Text (Text, Text)
-> (Text, [(MetricId, MetricSample)]) -> Builder
render Map Text (Text, Text)
ds (Text
name, [(MetricId, MetricSample)]
values) =
    case Text -> Map Text (Text, Text) -> Maybe (Text, Text)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
name Map Text (Text, Text)
ds of
      Maybe (Text, Text)
Nothing -> RegistrySample -> Builder
encodeMetrics (Map MetricId MetricSample -> RegistrySample
RegistrySample (Map MetricId MetricSample -> RegistrySample)
-> Map MetricId MetricSample -> RegistrySample
forall a b. (a -> b) -> a -> b
$ [(MetricId, MetricSample)] -> Map MetricId MetricSample
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(MetricId, MetricSample)]
values)
      Just (Text
kind, Text
help) ->
        ByteString -> Builder
byteString (Text -> ByteString
forall a b. ConvertUtf8 a b => a -> b
encodeUtf8 (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$ Text
"# HELP " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
escape Text
help Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n# TYPE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
name Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
kind Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n")
          -- The library supplies samples, label escaping and histogram encoding.
          -- Replace only its TYPE header, adding the description it cannot store.
          Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
lazyByteString (Int64 -> ByteString -> ByteString
LBS.drop Int64
1 (ByteString -> ByteString)
-> (Map MetricId MetricSample -> ByteString)
-> Map MetricId MetricSample
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word8 -> Bool) -> ByteString -> ByteString
LBS.dropWhile (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word8
10) (ByteString -> ByteString)
-> (Map MetricId MetricSample -> ByteString)
-> Map MetricId MetricSample
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> ByteString
toLazyByteString (Builder -> ByteString)
-> (Map MetricId MetricSample -> Builder)
-> Map MetricId MetricSample
-> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RegistrySample -> Builder
encodeMetrics (RegistrySample -> Builder)
-> (Map MetricId MetricSample -> RegistrySample)
-> Map MetricId MetricSample
-> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map MetricId MetricSample -> RegistrySample
RegistrySample (Map MetricId MetricSample -> ByteString)
-> Map MetricId MetricSample -> ByteString
forall a b. (a -> b) -> a -> b
$ [(MetricId, MetricSample)] -> Map MetricId MetricSample
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(MetricId, MetricSample)]
values)

  escape :: Text -> Text
  escape :: Text -> Text
escape = HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
Text.replace Text
"\n" Text
"\\n" (Text -> Text) -> (Text -> Text) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
Text.replace Text
"\\" Text
"\\\\"

scaleSample :: Double -> MetricSample -> MetricSample
scaleSample :: Double -> MetricSample -> MetricSample
scaleSample Double
scale = \case
  HistogramMetricSample HistogramSample{Buckets
histBuckets :: Buckets
histBuckets :: HistogramSample -> Buckets
histBuckets, Double
histSum :: Double
histSum :: HistogramSample -> Double
histSum, Int
histCount :: Int
histCount :: HistogramSample -> Int
histCount} ->
    HistogramSample -> MetricSample
HistogramMetricSample HistogramSample{histBuckets :: Buckets
histBuckets = (Double -> Double) -> Buckets -> Buckets
forall k1 k2 a. (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeysMonotonic (Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
scale) Buckets
histBuckets, histSum :: Double
histSum = Double
histSum Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
scale, Int
histCount :: Int
histCount :: Int
histCount}
  GaugeMetricSample (GaugeSample Double
value) -> GaugeSample -> MetricSample
GaugeMetricSample (Double -> GaugeSample
GaugeSample (Double -> GaugeSample) -> Double -> GaugeSample
forall a b. (a -> b) -> a -> b
$ Double
value Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
scale)
  MetricSample
value -> MetricSample
value