Make Gauge polymorphic

Author: brandonchinn178

prometheus.cabal

@@ -103,6 +103,7 @@ library
                  , Prometheus.V3.Name
                  , Prometheus.V3.Registry
                  , Prometheus.V3.Sample
+                 , Prometheus.V3.Utils.FromRealFrac
 
   build-depends: base            >= 4.9  && < 5
                , atomic-primops  >= 0.8  && < 0.9

src/Prometheus/V3/Metric/Gauge.hs

@@ -1,6 +1,10 @@
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE ImportQualifiedPost #-}
 {-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE NoFieldSelectors #-}
 -- Remove after inlining V2.Gauge
 {-# OPTIONS_GHC -Wno-orphans #-}
@@ -13,6 +17,7 @@ module Prometheus.V3.Metric.Gauge (
 
     -- * Methods
     IsGauge,
+    IsNumGauge,
     inc,
     add,
     dec,
@@ -36,22 +41,23 @@ import Prometheus.V3.Metric.Base
 import Prometheus.V3.Metric.Labelled (labels)
 import Prometheus.V3.Registry qualified as Registry
 import Prometheus.V3.Sample
+import Prometheus.V3.Utils.FromRealFrac (FromRealFrac, fromRealFrac)
 import System.Metrics.Prometheus.Metric.Gauge qualified as V2
 import UnliftIO (MonadUnliftIO)
 import UnliftIO.Exception (bracket_)
 
 
 -- | An alias for @'Registry.register' (new ...)@.
-register :: MetricName -> Description -> V2.Gauge
+register :: (Num a, ToSampleValue a) => MetricName -> Description -> V2.Gauge a
 register name description = Registry.register $ new name description
 {-# INLINE register #-}
 
 
-new :: MetricName -> Description -> Metric V2.Gauge
+new :: (Num a) => MetricName -> Description -> Metric (V2.Gauge a)
 new name description = newAt name description 0
 
 
-newAt :: MetricName -> Description -> Double -> Metric V2.Gauge
+newAt :: (Num a) => MetricName -> Description -> a -> Metric (V2.Gauge a)
 newAt name description val =
     Metric
         { name
@@ -63,74 +69,84 @@ newAt name description val =
         }
 
 
-instance IsMetric V2.Gauge where
+instance (SampleValueNum a) => IsMetric (V2.Gauge a) where
     getMetricType _ = MetricTypeGauge
     getMetricSamples gauge = do
         n <- sample gauge
         pure [defaultSample{value = toSampleValue n}]
 
 
 class IsGauge g where
-    getGauge :: g -> IO V2.Gauge
-instance IsGauge V2.Gauge where
+    type GaugeElem g
+    getGauge :: g -> IO (V2.Gauge (GaugeElem g))
+instance IsGauge (V2.Gauge a) where
+    type GaugeElem (V2.Gauge a) = a
     getGauge = pure
-instance IsGauge (IO V2.Gauge) where
+instance IsGauge (IO (V2.Gauge a)) where
+    type GaugeElem (IO (V2.Gauge a)) = a
     getGauge = id
 
 
-inc :: (IsGauge g, MonadIO m) => g -> m ()
+type IsNumGauge g a = (IsGauge g, a ~ GaugeElem g, Num a)
+
+
+inc :: (IsNumGauge g a, MonadIO m) => g -> m ()
 inc g = liftIO $ V2.inc =<< getGauge g
 {-# INLINE inc #-}
 
 
-add :: (IsGauge g, MonadIO m) => g -> Double -> m ()
+add :: (IsNumGauge g a, MonadIO m) => g -> a -> m ()
 add g x = liftIO $ V2.add x =<< getGauge g
 {-# INLINE add #-}
 
 
-dec :: (IsGauge g, MonadIO m) => g -> m ()
+dec :: (IsNumGauge g a, MonadIO m) => g -> m ()
 dec g = liftIO $ V2.dec =<< getGauge g
 {-# INLINE dec #-}
 
 
-sub :: (IsGauge g, MonadIO m) => g -> Double -> m ()
+sub :: (IsNumGauge g a, MonadIO m) => g -> a -> m ()
 sub g x = liftIO $ V2.sub x =<< getGauge g
 {-# INLINE sub #-}
 
 
-set :: (IsGauge g, MonadIO m) => g -> Double -> m ()
+set :: (IsNumGauge g a, MonadIO m) => g -> a -> m ()
 set g x = liftIO $ V2.set x =<< getGauge g
 {-# INLINE set #-}
 
 
-sample :: (IsGauge g, MonadIO m) => g -> m Double
+sample :: (IsNumGauge g a, MonadIO m) => g -> m a
 sample g = liftIO $ fmap V2.unGaugeSample . V2.sample =<< getGauge g
 {-# INLINE sample #-}
 
 
-modifyAndSample :: (IsGauge g, MonadIO m) => g -> (Double -> Double) -> m Double
+modifyAndSample :: (IsNumGauge g a, MonadIO m) => g -> (a -> a) -> m a
 modifyAndSample g f = liftIO $ fmap V2.unGaugeSample . V2.modifyAndSample f =<< getGauge g
 {-# INLINE modifyAndSample #-}
 
 
 {----- Helpers -----}
 
 -- | Set to the current number of seconds since the epoch.
-setToCurrentTime :: (IsGauge g, MonadIO m) => g -> m ()
+setToCurrentTime ::
+    (IsGauge g, FromRealFrac (GaugeElem g), MonadIO m) =>
+    g -> m ()
 setToCurrentTime g = do
     now <- liftIO getPOSIXTime
     g' <- liftIO $ getGauge g
-    set g' (realToFrac now)
+    set g' (fromRealFrac now)
 {-# INLINE setToCurrentTime #-}
 
 
 -- | Increment the gauge when the given action is running and decrement when it's exited.
-trackInProgress :: (IsGauge g, MonadUnliftIO m) => g -> m a -> m a
+trackInProgress :: (IsNumGauge g a, MonadUnliftIO m) => g -> m a -> m a
 trackInProgress g = bracket_ (inc g) (dec g)
 {-# INLINE trackInProgress #-}
 
 
 -- | Set the gauge to the duration of the given action.
-time :: (IsGauge g, MonadUnliftIO m) => g -> m a -> m a
-time g = withDuration (set g)
+time ::
+    (IsGauge g, FromRealFrac (GaugeElem g), MonadUnliftIO m) =>
+    g -> m a -> m a
+time g = withDuration (set g . fromRealFrac)
 {-# INLINE time #-}

src/Prometheus/V3/Sample.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE ConstraintKinds #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE NoFieldSelectors #-}
 
@@ -8,6 +9,7 @@ module Prometheus.V3.Sample (
     -- * SampleValue
     SampleValue (..),
     ToSampleValue (..),
+    SampleValueNum,
 ) where
 
 import Data.Int (Int64)
@@ -49,3 +51,6 @@ instance ToSampleValue Double where
     toSampleValue = SampleValueDouble
 instance ToSampleValue Float where
     toSampleValue = SampleValueDouble . realToFrac
+
+
+type SampleValueNum a = (Num a, ToSampleValue a)

src/Prometheus/V3/Utils/FromRealFrac.hs

@@ -0,0 +1,21 @@
+module Prometheus.V3.Utils.FromRealFrac (
+    FromRealFrac (..),
+) where
+
+import Data.Word (Word64)
+
+
+class (Num a) => FromRealFrac a where
+    fromRealFrac :: (RealFrac x) => x -> a
+
+
+instance FromRealFrac Int where
+    fromRealFrac = round
+instance FromRealFrac Word64 where
+    fromRealFrac = round
+
+
+instance FromRealFrac Double where
+    fromRealFrac = realToFrac
+instance FromRealFrac Float where
+    fromRealFrac = realToFrac

src/System/Metrics/Prometheus/Concurrent/Registry.hs

@@ -46,7 +46,7 @@ registerCounter name labels = flip modifyMVarMasked register . unRegistry
     register = fmap swap . R.registerCounter name labels
 
 
-registerGauge :: Name -> Labels -> Registry -> IO Gauge
+registerGauge :: Name -> Labels -> Registry -> IO (Gauge Double)
 registerGauge name labels = flip modifyMVarMasked register . unRegistry
   where
     register = fmap swap . R.registerGauge name labels

src/System/Metrics/Prometheus/Concurrent/RegistryT.hs

@@ -39,7 +39,7 @@ registerCounter :: MonadIO m => Name -> Labels -> RegistryT m Counter
 registerCounter n l = RegistryT ask >>= liftIO . R.registerCounter n l
 
 
-registerGauge :: MonadIO m => Name -> Labels -> RegistryT m Gauge
+registerGauge :: MonadIO m => Name -> Labels -> RegistryT m (Gauge Double)
 registerGauge n l = RegistryT ask >>= liftIO . R.registerGauge n l
 
 

src/System/Metrics/Prometheus/Encode/Text.hs

@@ -64,5 +64,5 @@ encodeCounter :: MetricId -> CounterSample -> Builder
 encodeCounter mid counter = encodeMetricId mid <> space <> encodeInt (unCounterSample counter)
 
 
-encodeGauge :: MetricId -> GaugeSample -> Builder
+encodeGauge :: MetricId -> GaugeSample Double -> Builder
 encodeGauge mid gauge = encodeMetricId mid <> space <> encodeDouble (unGaugeSample gauge)

src/System/Metrics/Prometheus/Metric.hs

@@ -14,21 +14,20 @@ import System.Metrics.Prometheus.Metric.Summary (SummarySample)
 
 data Metric
     = CounterMetric Counter
-    | GaugeMetric Gauge
-    | -- | Summary S.Summary
-      HistogramMetric Histogram
+    | GaugeMetric (Gauge Double) -- \| Summary S.Summary
+    | HistogramMetric Histogram
 
 
 data MetricSample
     = CounterMetricSample CounterSample
-    | GaugeMetricSample GaugeSample
+    | GaugeMetricSample (GaugeSample Double)
     | HistogramMetricSample HistogramSample
     | SummaryMetricSample SummarySample
 
 
 metricSample ::
     (CounterSample -> a) ->
-    (GaugeSample -> a) ->
+    (GaugeSample Double -> a) ->
     (HistogramSample -> a) ->
     (SummarySample -> a) ->
     MetricSample ->

src/System/Metrics/Prometheus/Metric/Gauge.hs

@@ -15,39 +15,39 @@ import Control.Applicative ((<$>))
 import Data.IORef (IORef, atomicModifyIORef', newIORef)
 
 
-newtype Gauge = Gauge {unGauge :: IORef Double}
-newtype GaugeSample = GaugeSample {unGaugeSample :: Double} deriving Show
+newtype Gauge a = Gauge {unGauge :: IORef a}
+newtype GaugeSample a = GaugeSample {unGaugeSample :: a} deriving (Show)
 
 
-new :: IO Gauge
+new :: (Num a) => IO (Gauge a)
 new = Gauge <$> newIORef 0
 
 
-modifyAndSample :: (Double -> Double) -> Gauge -> IO GaugeSample
+modifyAndSample :: (a -> a) -> Gauge a -> IO (GaugeSample a)
 modifyAndSample f = flip atomicModifyIORef' g . unGauge
   where
     g v = (f v, GaugeSample $ f v)
 
 
-add :: Double -> Gauge -> IO ()
+add :: (Num a) => a -> Gauge a -> IO ()
 add x g = modifyAndSample (+ x) g >> pure ()
 
 
-sub :: Double -> Gauge -> IO ()
+sub :: (Num a) => a -> Gauge a -> IO ()
 sub x g = modifyAndSample (subtract x) g >> pure ()
 
 
-inc :: Gauge -> IO ()
+inc :: (Num a) => Gauge a -> IO ()
 inc = add 1
 
 
-dec :: Gauge -> IO ()
+dec :: (Num a) => Gauge a -> IO ()
 dec = sub 1
 
 
-set :: Double -> Gauge -> IO ()
+set :: a -> Gauge a -> IO ()
 set x g = modifyAndSample (const x) g >> pure ()
 
 
-sample :: Gauge -> IO GaugeSample
+sample :: Gauge a -> IO (GaugeSample a)
 sample = modifyAndSample id

src/System/Metrics/Prometheus/Registry.hs

@@ -59,7 +59,7 @@ registerCounter name labels registry = do
     collision k _ _ = throw (KeyError k)
 
 
-registerGauge :: Name -> Labels -> Registry -> IO (Gauge, Registry)
+registerGauge :: Name -> Labels -> Registry -> IO (Gauge Double, Registry)
 registerGauge name labels registry = do
     gauge <- Gauge.new
     return (gauge, Registry $ Map.insertWithKey collision mid (GaugeMetric gauge) (unRegistry registry))