Make Counter polymorphic

Author: brandonchinn178

src/Prometheus/V3/Metric/Counter.hs

@@ -1,19 +1,26 @@
 {-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE ImportQualifiedPost #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UndecidableInstances #-}
 {-# LANGUAGE NoFieldSelectors #-}
 -- Remove after inlining V2.Counter
 {-# OPTIONS_GHC -Wno-orphans #-}
 
 module Prometheus.V3.Metric.Counter (
     V2.Counter,
+    V2.Counter',
     register,
     new,
 
     -- * Methods
     IsCounter,
+    IsNumCounter,
     inc,
     add,
     reset,
@@ -40,12 +47,16 @@ import UnliftIO.Exception (Exception, withException)
 
 
 -- | An alias for @'Registry.register' (new ...)@.
-register :: MetricName -> Description -> V2.Counter
+register ::
+    (Num a, Ord a, ToSampleValue a, V2.HasCounterBackend a) =>
+    MetricName -> Description -> (V2.Counter' a)
 register name description = Registry.register $ new name description
 {-# INLINE register #-}
 
 
-new :: MetricName -> Description -> Metric V2.Counter
+new ::
+    (Num a, V2.HasCounterBackend a) =>
+    MetricName -> Description -> Metric (V2.Counter' a)
 new name description =
     Metric
         { name
@@ -54,42 +65,54 @@ new name description =
         }
 
 
-instance IsMetric V2.Counter where
+instance (SampleValueNum a, Ord a, V2.HasCounterBackend a) => IsMetric (V2.Counter' a) where
     getMetricType _ = MetricTypeCounter
     getMetricSamples counter = do
         n <- sample counter
         pure [defaultSample{value = toSampleValue n}]
 
 
 class IsCounter c where
-    getCounter :: c -> IO V2.Counter
-instance IsCounter V2.Counter where
+    type CounterElem c
+    getCounter :: c -> IO (V2.Counter' (CounterElem c))
+instance IsCounter (V2.Counter' a) where
+    type CounterElem (V2.Counter' a) = a
     getCounter = pure
-instance IsCounter (IO V2.Counter) where
+instance IsCounter (IO (V2.Counter' a)) where
+    type CounterElem (IO (V2.Counter' a)) = a
     getCounter = id
 
 
-inc :: (IsCounter c, MonadIO m) => c -> m ()
+type IsNumCounter c a =
+    ( IsCounter c
+    , a ~ CounterElem c
+    , V2.HasCounterBackend a
+    , Num a
+    , Ord a
+    )
+
+
+inc :: (IsNumCounter c a, MonadIO m) => c -> m ()
 inc c = liftIO $ V2.inc =<< getCounter c
 {-# INLINE inc #-}
 
 
-add :: (IsCounter c, MonadIO m) => c -> Int -> m ()
+add :: (IsNumCounter c a, MonadIO m) => c -> a -> m ()
 add c n = liftIO $ V2.add n =<< getCounter c
 {-# INLINE add #-}
 
 
-reset :: (IsCounter c, MonadIO m) => c -> m ()
+reset :: (IsNumCounter c a, MonadIO m) => c -> m ()
 reset c = liftIO $ V2.set 0 =<< getCounter c
 {-# INLINE reset #-}
 
 
-sample :: (IsCounter c, MonadIO m) => c -> m Int
+sample :: (IsNumCounter c a, MonadIO m) => c -> m a
 sample c = liftIO $ fmap V2.unCounterSample . V2.sample =<< getCounter c
 {-# INLINE sample #-}
 
 
-addAndSample :: (IsCounter c, MonadIO m) => c -> Int -> m Int
+addAndSample :: (IsNumCounter c a, MonadIO m) => c -> a -> m a
 addAndSample c n = liftIO $ fmap V2.unCounterSample . V2.addAndSample n =<< getCounter c
 {-# INLINE addAndSample #-}
 
@@ -99,7 +122,7 @@ addAndSample c n = liftIO $ fmap V2.unCounterSample . V2.addAndSample n =<< getC
 -- | Count the number of times a particular exception is thrown.
 countExceptions ::
     forall e c m a.
-    (IsCounter c, Exception e, MonadUnliftIO m) =>
+    (IsNumCounter c a, Exception e, MonadUnliftIO m) =>
     c -> m a -> m a
 countExceptions c = countExceptionsWhere c (\(_ :: e) -> True)
 {-# INLINE countExceptions #-}
@@ -109,7 +132,7 @@ countExceptions c = countExceptionsWhere c (\(_ :: e) -> True)
 -- include the exception in the count.
 countExceptionsWhere ::
     forall e c m a.
-    (IsCounter c, Exception e, MonadUnliftIO m) =>
+    (IsNumCounter c a, Exception e, MonadUnliftIO m) =>
     c -> (e -> Bool) -> m a -> m a
 countExceptionsWhere c f io = io `withException` \e -> when (f e) (inc c)
 {-# INLINE countExceptionsWhere #-}

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

@@ -60,7 +60,7 @@ encodeMetric (mid, sample) =
     encodeSummary = undefined
 
 
-encodeCounter :: MetricId -> CounterSample -> Builder
+encodeCounter :: MetricId -> CounterSample Int -> Builder
 encodeCounter mid counter = encodeMetricId mid <> space <> encodeInt (unCounterSample counter)
 
 

src/System/Metrics/Prometheus/Metric.hs

@@ -19,14 +19,14 @@ data Metric
 
 
 data MetricSample
-    = CounterMetricSample CounterSample
+    = CounterMetricSample (CounterSample Int)
     | GaugeMetricSample (GaugeSample Double)
     | HistogramMetricSample HistogramSample
     | SummaryMetricSample SummarySample
 
 
 metricSample ::
-    (CounterSample -> a) ->
+    (CounterSample Int -> a) ->
     (GaugeSample Double -> a) ->
     (HistogramSample -> a) ->
     (SummarySample -> a) ->

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

@@ -1,42 +1,100 @@
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+
 module System.Metrics.Prometheus.Metric.Counter (
     Counter,
+    Counter',
     CounterSample (..),
     new,
     add,
     inc,
     sample,
     addAndSample,
     set,
+
+    -- * Backends
+    HasCounterBackend,
+    CounterBackend,
+    IsCounterBackend,
 ) where
 
 import Control.Applicative ((<$>))
 import Control.Monad (when)
-import Data.Atomics.Counter (AtomicCounter, incrCounter, newCounter, writeCounter)
+import Data.Atomics.Counter (AtomicCounter)
+import qualified Data.Atomics.Counter as Atomic
+import Data.IORef (IORef, atomicModifyIORef', newIORef, writeIORef)
+
+
+-- | The default Counter tracking an Int.
+type Counter = Counter' Int
+
+
+-- | A Counter that works for any type, but with worse performance.
+newtype Counter' a = Counter {unCounter :: CounterBackend a}
+
+
+newtype CounterSample a = CounterSample {unCounterSample :: a} deriving (Show)
+
+
+type family CounterBackend a where
+    CounterBackend Int = AtomicCounter
+    CounterBackend a = IORef a
+
+
+class IsCounterBackend a where
+    type CounterElem a
+    newCounter :: CounterElem a -> IO a
+    incrCounter :: CounterElem a -> a -> IO (CounterElem a)
+    writeCounter :: a -> CounterElem a -> IO ()
+instance IsCounterBackend AtomicCounter where
+    type CounterElem AtomicCounter = Int
+    newCounter = Atomic.newCounter
+    incrCounter = Atomic.incrCounter
+    writeCounter = Atomic.writeCounter
+instance (Num a) => IsCounterBackend (IORef a) where
+    type CounterElem (IORef a) = a
+    newCounter = newIORef
+    incrCounter x ref = atomicModifyIORef' ref $ \a -> (a + x, a + x)
+    writeCounter = writeIORef
 
 
-newtype Counter = Counter {unCounter :: AtomicCounter}
-newtype CounterSample = CounterSample {unCounterSample :: Int} deriving Show
+type HasCounterBackend a =
+    ( IsCounterBackend (CounterBackend a)
+    , a ~ CounterElem (CounterBackend a)
+    )
 
 
-new :: IO Counter
+new ::
+    (Num a, HasCounterBackend a) =>
+    IO (Counter' a)
 new = Counter <$> newCounter 0
 
 
-addAndSample :: Int -> Counter -> IO CounterSample
+addAndSample ::
+    (Num a, Ord a, HasCounterBackend a) =>
+    a -> Counter' a -> IO (CounterSample a)
 addAndSample by
     | by >= 0 = fmap CounterSample . incrCounter by . unCounter
     | otherwise = error "must be >= 0"
 
 
-add :: Int -> Counter -> IO ()
+add ::
+    (Num a, Ord a, HasCounterBackend a) =>
+    a -> Counter' a -> IO ()
 add by c = addAndSample by c >> pure ()
 
 
-inc :: Counter -> IO ()
+inc ::
+    (Num a, Ord a, HasCounterBackend a) =>
+    Counter' a -> IO ()
 inc = add 1
 
 
-sample :: Counter -> IO CounterSample
+sample ::
+    (Num a, Ord a, HasCounterBackend a) =>
+    Counter' a -> IO (CounterSample a)
 sample = addAndSample 0
 
 
@@ -46,7 +104,7 @@ sample = addAndSample 0
 --
 -- WARNING: For multiple writers, the most recent one wins, which may not
 -- preserve the increasing property. If you have stronger requirements than this,
--- please check with the maintainers. 
+-- please check with the maintainers.
 -- See <https://github.com/bitnomial/prometheus/pull/44> for discussion.
-set :: Int -> Counter -> IO ()
+set :: (HasCounterBackend a) => a -> Counter' a -> IO ()
 set i (Counter c) = writeCounter c i