{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeApplications #-}

module Instrument.Counter
  ( Counter,
    newCounter,
    readCounter,
    resetCounter,
    add,
    increment,
  )
where

-------------------------------------------------------------------------------
import Data.IORef
import qualified Data.Atomics.Counter as A
-------------------------------------------------------------------------------

data Counter = Counter {Counter -> AtomicCounter
_atomicCounter :: A.AtomicCounter, Counter -> IORef Int
_lastResetValue :: IORef Int}

-------------------------------------------------------------------------------
newCounter :: IO Counter
newCounter :: IO Counter
newCounter = AtomicCounter -> IORef Int -> Counter
Counter forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO AtomicCounter
A.newCounter Int
0 forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. a -> IO (IORef a)
newIORef Int
0

-------------------------------------------------------------------------------

-- | Reads current counter value
readCounter :: Counter -> IO Integer
readCounter :: Counter -> IO Integer
readCounter (Counter AtomicCounter
i IORef Int
lastReset) = Int -> Int -> Integer
calculateDelta forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. IORef a -> IO a
readIORef IORef Int
lastReset forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> AtomicCounter -> IO Int
A.readCounter AtomicCounter
i

-- | Our counters are represented as Int (with machine word size) and increment
-- | only, with no possibility to reset. Since we report deltas, we need facility
-- | to reset counter: we do that by storing last reported value in IORef.
-- | When application will run for long-enough counters will inevitably overflow.
-- | In such case our delta would be negative – this does not make sense for increment only.
-- | To overcome this issue we apply heuristic: when current value of counter
-- | is lower last value the counter has been reset, rollover happened and we need to
-- | take it into account. Otherwise, we can simply subtract values.
calculateDelta :: Int -> Int -> Integer
calculateDelta :: Int -> Int -> Integer
calculateDelta Int
lastReset Int
current | Int
current forall a. Ord a => a -> a -> Bool
< Int
lastReset =
  forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Bounded a => a
maxBound @Int forall a. Num a => a -> a -> a
- Int
lastReset) forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
current forall a. Num a => a -> a -> a
- forall a. Bounded a => a
minBound @Int) forall a. Num a => a -> a -> a
+ Integer
1
calculateDelta Int
lastReset Int
current = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
current forall a. Num a => a -> a -> a
- forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
lastReset

-------------------------------------------------------------------------------

-- | Reset the counter while reading it
resetCounter :: Counter -> IO Integer
resetCounter :: Counter -> IO Integer
resetCounter (Counter AtomicCounter
i IORef Int
lastReset) = do
  Int
ctrValue <- AtomicCounter -> IO Int
A.readCounter AtomicCounter
i
  Int
oldLast <- forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
lastReset forall a b. (a -> b) -> a -> b
$ \Int
oldLast -> (Int
ctrValue, Int
oldLast)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Int -> Int -> Integer
calculateDelta Int
oldLast Int
ctrValue

-------------------------------------------------------------------------------
increment :: Counter -> IO ()
increment :: Counter -> IO ()
increment = Int -> Counter -> IO ()
add Int
1

-------------------------------------------------------------------------------
add :: Int -> Counter -> IO ()
add :: Int -> Counter -> IO ()
add Int
x (Counter AtomicCounter
i IORef Int
_) = Int -> AtomicCounter -> IO ()
A.incrCounter_ Int
x AtomicCounter
i