live-stats-0.1.0.0: src/LiveStats.hs

module LiveStats
  ( StatState (..)
  , Stats
  , newStats
  , addCounter
  , readStats
  , incStats
  , obtainReport
  , reportStats
  , reportStatsForever
  ) where

import Relude
import Relude.Extra

import Data.Time.Clock.POSIX
import Control.Concurrent
import Numeric

newtype Stats = Stats (IORef (Map Text StatState))

data StatState = StatState
  { statStart  :: !Double
  -- ^ Origin start of run
  , statChkPnt :: !Double
  -- ^ Last time checkpoint
  , statGen    :: !Int
  -- ^ Hits since last checkpoint
  , statTotal  :: !Int
  -- ^ Total hits since original start
  }

newStats :: (MonadIO m) => m Stats
newStats = Stats <$> newIORef mempty

-- | Each named counter needs to be added (initialized) before it can be used.
addCounter :: (MonadIO m) => Stats -> Text -> m ()
addCounter (Stats s) nm = do
  now <- getTime
  atomicModifyIORef' s $ \m -> (insert nm (StatState now now 0 0) m, ())

readStats :: (MonadIO m) => Stats -> m (Map Text StatState)
readStats (Stats s) = readIORef s

-- | Increment given counter by n.
incStats :: (MonadIO m) => Stats -> Text -> Int -> m ()
incStats (Stats stats) nm n = atomicModifyIORef' stats $ \m -> (alter f nm m, ())
  where
    f = \case
      Nothing -> error "Counter should be added before usage"
      Just old@StatState {statGen, statTotal} ->
        let !x  = Just $! old { statGen = statGen + n, statTotal = statTotal + n }
        in x

-- | Produce stats messages - one for each registered counter.
obtainReport :: (MonadIO m) => Stats -> m [Text]
obtainReport (Stats stats) = do
  now <- liftIO getTime
  olds <- atomicModifyIORef' stats $ \m -> do
    let
      olds = toPairs m
      reset (nm, old) = (nm, old { statGen = 0, statChkPnt = now })
      news = fromList $ reset <$> olds
    (news, olds)
  pure $ go now <$> olds
  where
    go now (nm, StatState {statGen, statTotal, statStart, statChkPnt}) = do
      let
        chk = fromIntegral statGen / (now - statChkPnt)
        tot = fromIntegral statTotal / (now - statStart)
      toText $ "[" <> toString nm <> "]" <>
        " Processed " <> showInt statGen
        " items at " <> showFFloat (Just 2) chk
        " items/sec for a total of " <> showInt statTotal
        " items at " <> showFFloat (Just 2) tot " items/sec."

-- | Obtain report and call function for each line.
reportStats
  :: (MonadIO m)
  => (Text -> m ())
  -- ^ How to render each line
  -> Stats
  -> m ()
reportStats f stats = obtainReport stats >>= mapM_ f

-- | Report stats every n seconds. Blocks forever.
reportStatsForever :: (MonadIO m) => (Text -> m ()) -> Stats -> Int -> m ()
reportStatsForever f stats n = forever $ do
  reportStats f stats
  liftIO $ threadDelay $ n*1000*1000

getTime :: (MonadIO m) => m Double
getTime = realToFrac <$> liftIO getPOSIXTime