module Network.ABCI.Server.Middleware.Metrics
    ( defaultBuckets
    , mkMetricsMiddleware
    ) where

import           Control.Monad                                 (forM_)
import           Control.Monad.IO.Class                        (MonadIO, liftIO)
import qualified Data.IORef                                    as Ref
import qualified Data.Map.Strict                               as Map
import           Data.String.Conversions                       (cs)
import           Data.Time                                     (diffUTCTime,
                                                                getCurrentTime)
import           Network.ABCI.Server.App                       (App (..),
                                                                MessageType (..),
                                                                Middleware,
                                                                demoteRequestType,
                                                                msgTypeKey)
import qualified System.Metrics.Prometheus.Concurrent.Registry as Registry
import qualified System.Metrics.Prometheus.Metric.Counter      as Counter
import qualified System.Metrics.Prometheus.Metric.Histogram    as Histogram
import qualified System.Metrics.Prometheus.MetricId            as MetricId

---------------------------------------------------------------------------
-- mkMetrics
---------------------------------------------------------------------------
-- | Metrics logger middleware for ABCI server already within the KatipContext.
-- Great for `App m` with a `KatipContext` instance.

mkMetricsMiddleware
  :: MonadIO m
  => [Histogram.UpperBound]
  -> Registry.Registry
  -> IO (Middleware m)
mkMetricsMiddleware :: [UpperBound] -> Registry -> IO (Middleware m)
mkMetricsMiddleware buckets :: [UpperBound]
buckets registry :: Registry
registry = do
  Config{..} <- [UpperBound] -> Registry -> IO Config
makeConfig [UpperBound]
buckets Registry
registry
  Middleware m -> IO (Middleware m)
forall (m :: * -> *) a. Monad m => a -> m a
return (Middleware m -> IO (Middleware m))
-> Middleware m -> IO (Middleware m)
forall a b. (a -> b) -> a -> b
$ \(App app :: forall (t :: MessageType). Request t -> m (Response t)
app) -> (forall (t :: MessageType). Request t -> m (Response t)) -> App m
forall (m :: * -> *).
(forall (t :: MessageType). Request t -> m (Response t)) -> App m
App ((forall (t :: MessageType). Request t -> m (Response t)) -> App m)
-> (forall (t :: MessageType). Request t -> m (Response t))
-> App m
forall a b. (a -> b) -> a -> b
$ \ req :: Request t
req -> do
       UTCTime
startTime <- IO UTCTime -> m UTCTime
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UTCTime
getCurrentTime
       Response t
res <- Request t -> m (Response t)
forall (t :: MessageType). Request t -> m (Response t)
app Request t
req
       UTCTime
endTime <- IO UTCTime -> m UTCTime
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UTCTime
getCurrentTime
       let msgType :: MessageType
msgType = Request t -> MessageType
forall (t :: MessageType). Request t -> MessageType
demoteRequestType Request t
req
           duration :: UpperBound
duration = NominalDiffTime -> UpperBound
forall a b. (Real a, Fractional b) => a -> b
realToFrac (NominalDiffTime -> UpperBound) -> NominalDiffTime -> UpperBound
forall a b. (a -> b) -> a -> b
$ UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
endTime UTCTime
startTime
       IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
         IORef (Map MessageType Counter) -> MessageType -> IO ()
incRequestCounter IORef (Map MessageType Counter)
cfgCounterMap MessageType
msgType
         IORef (Map MessageType Histogram)
-> MessageType -> UpperBound -> IO ()
addToHistogram IORef (Map MessageType Histogram)
cfgHistogramMap MessageType
msgType UpperBound
duration
       Response t -> m (Response t)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Response t
res
  where

    incRequestCounter :: IORef (Map MessageType Counter) -> MessageType -> IO ()
incRequestCounter counterMapRef :: IORef (Map MessageType Counter)
counterMapRef msgType :: MessageType
msgType = do
      Counter
counter <- do
        Map MessageType Counter
counterMap <- IORef (Map MessageType Counter) -> IO (Map MessageType Counter)
forall a. IORef a -> IO a
Ref.readIORef IORef (Map MessageType Counter)
counterMapRef
        case MessageType -> Map MessageType Counter -> Maybe Counter
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup MessageType
msgType Map MessageType Counter
counterMap of
          Nothing -> [Char] -> IO Counter
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO Counter) -> [Char] -> IO Counter
forall a b. (a -> b) -> a -> b
$ "Impossible missing counter for " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> MessageType -> [Char]
msgTypeKey MessageType
msgType
          Just c :: Counter
c -> Counter -> IO Counter
forall (m :: * -> *) a. Monad m => a -> m a
return Counter
c
      Counter -> IO ()
Counter.inc Counter
counter

    addToHistogram :: IORef (Map MessageType Histogram)
-> MessageType -> UpperBound -> IO ()
addToHistogram histogramMapRef :: IORef (Map MessageType Histogram)
histogramMapRef msgType :: MessageType
msgType duration :: UpperBound
duration = do
      Histogram
histogram <- do
        Map MessageType Histogram
histMap <- IORef (Map MessageType Histogram) -> IO (Map MessageType Histogram)
forall a. IORef a -> IO a
Ref.readIORef IORef (Map MessageType Histogram)
histogramMapRef
        case MessageType -> Map MessageType Histogram -> Maybe Histogram
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup MessageType
msgType Map MessageType Histogram
histMap of
          Nothing -> [Char] -> IO Histogram
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO Histogram) -> [Char] -> IO Histogram
forall a b. (a -> b) -> a -> b
$ "Impossible missing histogram for " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> MessageType -> [Char]
msgTypeKey MessageType
msgType
          Just c :: Histogram
c -> Histogram -> IO Histogram
forall (m :: * -> *) a. Monad m => a -> m a
return Histogram
c
      UpperBound -> Histogram -> IO ()
Histogram.observe UpperBound
duration Histogram
histogram

data Config = Config
  { Config -> Registry
cfgRegistry         :: Registry.Registry
  , Config -> [UpperBound]
cfgHistogramBuckets :: [Histogram.UpperBound]
  , Config -> IORef (Map MessageType Counter)
cfgCounterMap       :: Ref.IORef (Map.Map MessageType Counter.Counter)
  , Config -> IORef (Map MessageType Histogram)
cfgHistogramMap     :: Ref.IORef (Map.Map MessageType Histogram.Histogram)
  }

makeConfig
  :: [Histogram.UpperBound]
  -> Registry.Registry
  -> IO Config
makeConfig :: [UpperBound] -> Registry -> IO Config
makeConfig bounds :: [UpperBound]
bounds registry :: Registry
registry = do
  IORef (Map MessageType Counter)
counterMap <- Map MessageType Counter -> IO (IORef (Map MessageType Counter))
forall a. a -> IO (IORef a)
Ref.newIORef Map MessageType Counter
forall k a. Map k a
Map.empty
  IORef (Map MessageType Histogram)
histMap <- Map MessageType Histogram -> IO (IORef (Map MessageType Histogram))
forall a. a -> IO (IORef a)
Ref.newIORef Map MessageType Histogram
forall k a. Map k a
Map.empty
  let cfg :: Config
cfg = Config :: Registry
-> [UpperBound]
-> IORef (Map MessageType Counter)
-> IORef (Map MessageType Histogram)
-> Config
Config
        { cfgRegistry :: Registry
cfgRegistry = Registry
registry
        , cfgHistogramBuckets :: [UpperBound]
cfgHistogramBuckets = [UpperBound]
bounds
        , cfgCounterMap :: IORef (Map MessageType Counter)
cfgCounterMap = IORef (Map MessageType Counter)
counterMap
        , cfgHistogramMap :: IORef (Map MessageType Histogram)
cfgHistogramMap = IORef (Map MessageType Histogram)
histMap
        }
  Config -> IO ()
registerMetrics Config
cfg
  Config -> IO Config
forall (m :: * -> *) a. Monad m => a -> m a
return Config
cfg

registerMetrics
  :: Config
  -> IO ()
registerMetrics :: Config -> IO ()
registerMetrics Config{..} = do
  [UpperBound]
-> Registry -> IORef (Map MessageType Histogram) -> IO ()
registerHistograms [UpperBound]
cfgHistogramBuckets Registry
cfgRegistry IORef (Map MessageType Histogram)
cfgHistogramMap
  Registry -> IORef (Map MessageType Counter) -> IO ()
registerCounters Registry
cfgRegistry IORef (Map MessageType Counter)
cfgCounterMap
  where

    registerHistograms
      :: [Histogram.UpperBound]
      -> Registry.Registry
      -> Ref.IORef (Map.Map MessageType Histogram.Histogram)
      -> IO ()
    registerHistograms :: [UpperBound]
-> Registry -> IORef (Map MessageType Histogram) -> IO ()
registerHistograms buckets :: [UpperBound]
buckets registry :: Registry
registry histRef :: IORef (Map MessageType Histogram)
histRef =
      let histName :: Name
histName = "abci_request_duration_seconds"
      in [MessageType] -> (MessageType -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [MessageType
MTEcho .. MessageType
MTCommit] ((MessageType -> IO ()) -> IO ())
-> (MessageType -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \messageType :: MessageType
messageType  -> do
           let labels :: Labels
labels = Map Text Text -> Labels
MetricId.Labels (Map Text Text -> Labels)
-> ([(Text, Text)] -> Map Text Text) -> [(Text, Text)] -> Labels
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Text, Text)] -> Map Text Text
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Text, Text)] -> Labels) -> [(Text, Text)] -> Labels
forall a b. (a -> b) -> a -> b
$
                 [ ("message_type", [Char] -> Text
forall a b. ConvertibleStrings a b => a -> b
cs ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ MessageType -> [Char]
msgTypeKey MessageType
messageType)
                 ]
           Histogram
hist <- Name -> Labels -> [UpperBound] -> Registry -> IO Histogram
Registry.registerHistogram Name
histName Labels
labels [UpperBound]
buckets Registry
registry
           IORef (Map MessageType Histogram)
-> (Map MessageType Histogram -> Map MessageType Histogram)
-> IO ()
forall a. IORef a -> (a -> a) -> IO ()
Ref.modifyIORef' IORef (Map MessageType Histogram)
histRef (MessageType
-> Histogram
-> Map MessageType Histogram
-> Map MessageType Histogram
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert MessageType
messageType Histogram
hist)


    registerCounters
      :: Registry.Registry
      -> Ref.IORef (Map.Map MessageType Counter.Counter)
      -> IO ()
    registerCounters :: Registry -> IORef (Map MessageType Counter) -> IO ()
registerCounters registry :: Registry
registry counterRef :: IORef (Map MessageType Counter)
counterRef =
      let counterName :: Name
counterName = "abci_request_total"
      in [MessageType] -> (MessageType -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [MessageType
MTEcho .. MessageType
MTCommit] ((MessageType -> IO ()) -> IO ())
-> (MessageType -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \messageType :: MessageType
messageType  -> do
           let labels :: Labels
labels = Map Text Text -> Labels
MetricId.Labels (Map Text Text -> Labels)
-> ([(Text, Text)] -> Map Text Text) -> [(Text, Text)] -> Labels
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Text, Text)] -> Map Text Text
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Text, Text)] -> Labels) -> [(Text, Text)] -> Labels
forall a b. (a -> b) -> a -> b
$
                 [ ("message_type", [Char] -> Text
forall a b. ConvertibleStrings a b => a -> b
cs ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ MessageType -> [Char]
msgTypeKey MessageType
messageType)
                 ]
           Counter
counter <- Name -> Labels -> Registry -> IO Counter
Registry.registerCounter Name
counterName Labels
labels Registry
registry
           IORef (Map MessageType Counter)
-> (Map MessageType Counter -> Map MessageType Counter) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
Ref.modifyIORef' IORef (Map MessageType Counter)
counterRef (MessageType
-> Counter -> Map MessageType Counter -> Map MessageType Counter
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert MessageType
messageType Counter
counter)

-- buckets with upper bounds [0.005, 0.01, 0.015 ... 5.0]
-- measured in seconds
defaultBuckets :: [Histogram.UpperBound]
defaultBuckets :: [UpperBound]
defaultBuckets = [0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0, 10.0]