{-# language FlexibleContexts    #-}
{-# language OverloadedStrings   #-}
{-# language PolyKinds           #-}
{-# language RankNTypes          #-}
{-# language ScopedTypeVariables #-}
module Mu.Instrumentation.Prometheus (
  initPrometheus
, prometheus
, prometheusWai
) where

import           Control.Concurrent.MVar.Lifted
import           Control.Exception.Lifted
import           Control.Monad.Trans.Control
import           Data.Text                         (Text)
import           Mu.Rpc
import           Mu.Server
import           Network.Wai
import qualified Network.Wai.Middleware.Prometheus as Wai
import           Prometheus

-- Taken from https://github.com/higherkindness/mu-scala/blob/master/modules/metrics/prometheus/src/main/scala/higherkindness/mu/rpc/prometheus/PrometheusMetrics.scala

data MuMetrics
  = MuMetrics {
      MuMetrics -> Gauge
activeCalls      :: Gauge
    , MuMetrics -> Vector Label2 Counter
messagesSent     :: Vector Label2 Counter
    , MuMetrics -> Vector Label2 Counter
messagesReceived :: Vector Label2 Counter
    , MuMetrics -> Vector Label2 Histogram
callsTotal       :: Vector Label2 Histogram
    }

initPrometheus :: Text -> IO MuMetrics
initPrometheus :: Text -> IO MuMetrics
initPrometheus prefix :: Text
prefix =
  Gauge
-> Vector Label2 Counter
-> Vector Label2 Counter
-> Vector Label2 Histogram
-> MuMetrics
MuMetrics (Gauge
 -> Vector Label2 Counter
 -> Vector Label2 Counter
 -> Vector Label2 Histogram
 -> MuMetrics)
-> IO Gauge
-> IO
     (Vector Label2 Counter
      -> Vector Label2 Counter -> Vector Label2 Histogram -> MuMetrics)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Metric Gauge -> IO Gauge
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Info -> Metric Gauge
gauge (Info -> Metric Gauge) -> Info -> Metric Gauge
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Info
Info (Text
prefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> "_active_calls") "")
            IO
  (Vector Label2 Counter
   -> Vector Label2 Counter -> Vector Label2 Histogram -> MuMetrics)
-> IO (Vector Label2 Counter)
-> IO
     (Vector Label2 Counter -> Vector Label2 Histogram -> MuMetrics)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Metric (Vector Label2 Counter) -> IO (Vector Label2 Counter)
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Label2 -> Metric Counter -> Metric (Vector Label2 Counter)
forall l m. Label l => l -> Metric m -> Metric (Vector l m)
vector ("service", "method")
                         (Metric Counter -> Metric (Vector Label2 Counter))
-> Metric Counter -> Metric (Vector Label2 Counter)
forall a b. (a -> b) -> a -> b
$ Info -> Metric Counter
counter (Info -> Metric Counter) -> Info -> Metric Counter
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Info
Info (Text
prefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> "_messages_sent") "")
            IO (Vector Label2 Counter -> Vector Label2 Histogram -> MuMetrics)
-> IO (Vector Label2 Counter)
-> IO (Vector Label2 Histogram -> MuMetrics)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Metric (Vector Label2 Counter) -> IO (Vector Label2 Counter)
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Label2 -> Metric Counter -> Metric (Vector Label2 Counter)
forall l m. Label l => l -> Metric m -> Metric (Vector l m)
vector ("service", "method")
                         (Metric Counter -> Metric (Vector Label2 Counter))
-> Metric Counter -> Metric (Vector Label2 Counter)
forall a b. (a -> b) -> a -> b
$ Info -> Metric Counter
counter (Info -> Metric Counter) -> Info -> Metric Counter
forall a b. (a -> b) -> a -> b
$ Text -> Text -> Info
Info (Text
prefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> "_messages_received") "")
            IO (Vector Label2 Histogram -> MuMetrics)
-> IO (Vector Label2 Histogram) -> IO MuMetrics
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Metric (Vector Label2 Histogram) -> IO (Vector Label2 Histogram)
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register (Label2 -> Metric Histogram -> Metric (Vector Label2 Histogram)
forall l m. Label l => l -> Metric m -> Metric (Vector l m)
vector ("service", "method")
                         (Metric Histogram -> Metric (Vector Label2 Histogram))
-> Metric Histogram -> Metric (Vector Label2 Histogram)
forall a b. (a -> b) -> a -> b
$ Info -> [Bucket] -> Metric Histogram
histogram (Text -> Text -> Info
Info (Text
prefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> "_calls_total") "")
                                     [Bucket]
defaultBuckets)

prometheus :: (MonadBaseControl IO m, MonadMonitor m)
           => MuMetrics -> ServerT chn info p m topHs -> ServerT chn info p m topHs
prometheus :: MuMetrics
-> ServerT chn info p m topHs -> ServerT chn info p m topHs
prometheus m :: MuMetrics
m = (forall a. RpcInfo info -> m a -> m a)
-> ServerT chn info p m topHs -> ServerT chn info p m topHs
forall snm mnm anm (chn :: ServiceChain snm) info
       (p :: Package snm mnm anm (TypeRef snm)) (m :: * -> *)
       (topHs :: [[*]]).
(forall a. RpcInfo info -> m a -> m a)
-> ServerT chn info p m topHs -> ServerT chn info p m topHs
wrapServer (MuMetrics -> RpcInfo info -> m a -> m a
forall (m :: * -> *) a info.
(MonadBaseControl IO m, MonadMonitor m) =>
MuMetrics -> RpcInfo info -> m a -> m a
prometheusMetrics MuMetrics
m)

prometheusMetrics :: forall m a info. (MonadBaseControl IO m, MonadMonitor m)
                  => MuMetrics -> RpcInfo info -> m a -> m a
prometheusMetrics :: MuMetrics -> RpcInfo info -> m a -> m a
prometheusMetrics metrics :: MuMetrics
metrics NoRpcInfo run :: m a
run = do
  Gauge -> m ()
forall (m :: * -> *). MonadMonitor m => Gauge -> m ()
incGauge (MuMetrics -> Gauge
activeCalls MuMetrics
metrics)
  m a
run m a -> m () -> m a
forall (m :: * -> *) a b.
MonadBaseControl IO m =>
m a -> m b -> m a
`finally` Gauge -> m ()
forall (m :: * -> *). MonadMonitor m => Gauge -> m ()
decGauge (MuMetrics -> Gauge
activeCalls MuMetrics
metrics)
prometheusMetrics metrics :: MuMetrics
metrics (RpcInfo _pkg :: Package Text Text Text TyInfo
_pkg ss :: Service Text Text Text TyInfo
ss mm :: Maybe (Method Text Text Text TyInfo)
mm _ _) run :: m a
run = do
  let sname' :: Text
sname' = case Service Text Text Text TyInfo
ss of
                 Service sname :: Text
sname _ -> Text
sname
                 OneOf   sname :: Text
sname _ -> Text
sname
      mname' :: Text
mname' = case Maybe (Method Text Text Text TyInfo)
mm of
                 Just (Method mname :: Text
mname _ _) -> Text
mname
                 Nothing                 -> "<noname>"
  Gauge -> m ()
forall (m :: * -> *). MonadMonitor m => Gauge -> m ()
incGauge (MuMetrics -> Gauge
activeCalls MuMetrics
metrics)
  Vector Label2 Counter -> Label2 -> (Counter -> IO ()) -> m ()
forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> label -> (metric -> IO ()) -> m ()
withLabel (MuMetrics -> Vector Label2 Counter
messagesReceived MuMetrics
metrics) (Text
sname', Text
mname') Counter -> IO ()
forall (m :: * -> *). MonadMonitor m => Counter -> m ()
incCounter
  ( do -- We are forced to use a MVar because 'withLabel' only allows IO ()
       StM m a
r <- (RunInBase m IO -> IO (StM m a)) -> m (StM m a)
forall (b :: * -> *) (m :: * -> *) a.
MonadBaseControl b m =>
(RunInBase m b -> b a) -> m a
liftBaseWith ((RunInBase m IO -> IO (StM m a)) -> m (StM m a))
-> (RunInBase m IO -> IO (StM m a)) -> m (StM m a)
forall a b. (a -> b) -> a -> b
$ \runInIO :: RunInBase m IO
runInIO -> do
         result :: MVar (StM m a) <- IO (MVar (StM m a))
forall (m :: * -> *) a. MonadBase IO m => m (MVar a)
newEmptyMVar
         Vector Label2 Histogram -> Label2 -> (Histogram -> IO ()) -> IO ()
forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> label -> (metric -> IO ()) -> m ()
withLabel (MuMetrics -> Vector Label2 Histogram
callsTotal MuMetrics
metrics) (Text
sname', Text
mname') ((Histogram -> IO ()) -> IO ()) -> (Histogram -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \h :: Histogram
h ->
           Histogram
h Histogram -> IO () -> IO ()
forall metric (m :: * -> *) a.
(Observer metric, MonadIO m, MonadMonitor m) =>
metric -> m a -> m a
`observeDuration` (m a -> IO (StM m a)
RunInBase m IO
runInIO m a
run IO (StM m a) -> (StM m a -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVar (StM m a) -> StM m a -> IO ()
forall (m :: * -> *) a. MonadBase IO m => MVar a -> a -> m ()
putMVar MVar (StM m a)
result)
         MVar (StM m a) -> IO (StM m a)
forall (m :: * -> *) a. MonadBase IO m => MVar a -> m a
takeMVar MVar (StM m a)
result
       a
x <- StM m a -> m a
forall (b :: * -> *) (m :: * -> *) a.
MonadBaseControl b m =>
StM m a -> m a
restoreM StM m a
r
       Vector Label2 Counter -> Label2 -> (Counter -> IO ()) -> m ()
forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> label -> (metric -> IO ()) -> m ()
withLabel (MuMetrics -> Vector Label2 Counter
messagesSent MuMetrics
metrics) (Text
sname', Text
mname') Counter -> IO ()
forall (m :: * -> *). MonadMonitor m => Counter -> m ()
incCounter
       a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x )
  m a -> m () -> m a
forall (m :: * -> *) a b.
MonadBaseControl IO m =>
m a -> m b -> m a
`finally` Gauge -> m ()
forall (m :: * -> *). MonadMonitor m => Gauge -> m ()
decGauge (MuMetrics -> Gauge
activeCalls MuMetrics
metrics)

prometheusWai :: [Text] -> Middleware
prometheusWai :: [Text] -> Middleware
prometheusWai endpoint :: [Text]
endpoint
  = PrometheusSettings -> Middleware
Wai.prometheus ([Text] -> Bool -> Bool -> PrometheusSettings
Wai.PrometheusSettings [Text]
endpoint Bool
False Bool
False)