{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Prod.Prometheus (
    handlePrometheus,
    PrometheusApi,
    CORSAllowOrigin (..),
    PrometheusResult (..),
    initPrometheus,
    inc,
    obs,
    timeIt,
)
where

import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.ByteString.Lazy (ByteString)
import Data.Coerce (coerce)
import Data.Text (Text)
import Data.Time.Clock (diffUTCTime, getCurrentTime)
import qualified Network.HTTP.Media as M
import Prometheus (Counter, Vector, exportMetricsAsText, register)
import qualified Prometheus as Prometheus
import Prometheus.Metric.GHC (GHCMetrics, ghcMetrics)
import Servant
import Servant.Server (Handler, Server)

newtype PrometheusResult = PrometheusResult {PrometheusResult -> ByteString
toLBS :: ByteString}

instance MimeRender PlainText PrometheusResult where
    mimeRender :: Proxy PlainText -> PrometheusResult -> ByteString
mimeRender Proxy PlainText
_ = PrometheusResult -> ByteString
toLBS

type PrometheusApi =
    Summary "Prometheus metrics"
        :> "metrics"
        :> Get '[PlainText] (Headers '[Header "Access-Control-Allow-Origin" CORSAllowOrigin] PrometheusResult)

newtype CORSAllowOrigin = CORSAllowOrigin Text
    deriving (CORSAllowOrigin -> Text
CORSAllowOrigin -> ByteString
CORSAllowOrigin -> Builder
(CORSAllowOrigin -> Text)
-> (CORSAllowOrigin -> Builder)
-> (CORSAllowOrigin -> ByteString)
-> (CORSAllowOrigin -> Text)
-> (CORSAllowOrigin -> Builder)
-> ToHttpApiData CORSAllowOrigin
forall a.
(a -> Text)
-> (a -> Builder)
-> (a -> ByteString)
-> (a -> Text)
-> (a -> Builder)
-> ToHttpApiData a
$ctoUrlPiece :: CORSAllowOrigin -> Text
toUrlPiece :: CORSAllowOrigin -> Text
$ctoEncodedUrlPiece :: CORSAllowOrigin -> Builder
toEncodedUrlPiece :: CORSAllowOrigin -> Builder
$ctoHeader :: CORSAllowOrigin -> ByteString
toHeader :: CORSAllowOrigin -> ByteString
$ctoQueryParam :: CORSAllowOrigin -> Text
toQueryParam :: CORSAllowOrigin -> Text
$ctoEncodedQueryParam :: CORSAllowOrigin -> Builder
toEncodedQueryParam :: CORSAllowOrigin -> Builder
ToHttpApiData)

handlePrometheus :: CORSAllowOrigin -> Server PrometheusApi
handlePrometheus :: CORSAllowOrigin -> Server PrometheusApi
handlePrometheus CORSAllowOrigin
corsAllow = Handler
  (Headers
     '[Header "Access-Control-Allow-Origin" CORSAllowOrigin]
     PrometheusResult)
Server PrometheusApi
handleMetrics
  where
    handleMetrics :: Handler (Headers '[Header "Access-Control-Allow-Origin" CORSAllowOrigin] PrometheusResult)
    handleMetrics :: Handler
  (Headers
     '[Header "Access-Control-Allow-Origin" CORSAllowOrigin]
     PrometheusResult)
handleMetrics = do
        ByteString
metrics <- IO ByteString -> Handler ByteString
forall a. IO a -> Handler a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> Handler ByteString)
-> IO ByteString -> Handler ByteString
forall a b. (a -> b) -> a -> b
$ IO ByteString
forall (m :: * -> *). MonadIO m => m ByteString
exportMetricsAsText
        Headers
  '[Header "Access-Control-Allow-Origin" CORSAllowOrigin]
  PrometheusResult
-> Handler
     (Headers
        '[Header "Access-Control-Allow-Origin" CORSAllowOrigin]
        PrometheusResult)
forall a. a -> Handler a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Headers
   '[Header "Access-Control-Allow-Origin" CORSAllowOrigin]
   PrometheusResult
 -> Handler
      (Headers
         '[Header "Access-Control-Allow-Origin" CORSAllowOrigin]
         PrometheusResult))
-> Headers
     '[Header "Access-Control-Allow-Origin" CORSAllowOrigin]
     PrometheusResult
-> Handler
     (Headers
        '[Header "Access-Control-Allow-Origin" CORSAllowOrigin]
        PrometheusResult)
forall a b. (a -> b) -> a -> b
$ CORSAllowOrigin
-> PrometheusResult
-> Headers
     '[Header "Access-Control-Allow-Origin" CORSAllowOrigin]
     PrometheusResult
forall (h :: Symbol) v orig new.
AddHeader h v orig new =>
v -> orig -> new
addHeader (CORSAllowOrigin -> CORSAllowOrigin
forall a b. Coercible a b => a -> b
coerce CORSAllowOrigin
corsAllow) (PrometheusResult
 -> Headers
      '[Header "Access-Control-Allow-Origin" CORSAllowOrigin]
      PrometheusResult)
-> PrometheusResult
-> Headers
     '[Header "Access-Control-Allow-Origin" CORSAllowOrigin]
     PrometheusResult
forall a b. (a -> b) -> a -> b
$ ByteString -> PrometheusResult
PrometheusResult ByteString
metrics

initPrometheus :: IO GHCMetrics
initPrometheus :: IO GHCMetrics
initPrometheus = Metric GHCMetrics -> IO GHCMetrics
forall (m :: * -> *) s. MonadIO m => Metric s -> m s
register Metric GHCMetrics
ghcMetrics

inc ::
    (MonadIO m) =>
    (a -> Vector Text Counter) ->
    Text ->
    a ->
    m ()
inc :: forall (m :: * -> *) a.
MonadIO m =>
(a -> Vector Text Counter) -> Text -> a -> m ()
inc a -> Vector Text Counter
f Text
s a
cnts =
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Vector Text Counter -> Text -> (Counter -> IO ()) -> IO ()
forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> label -> (metric -> IO ()) -> m ()
Prometheus.withLabel (a -> Vector Text Counter
f a
cnts) Text
s Counter -> IO ()
forall (m :: * -> *). MonadMonitor m => Counter -> m ()
Prometheus.incCounter

obs ::
    (MonadIO m) =>
    (a -> Prometheus.Summary) ->
    Double ->
    a ->
    m ()
obs :: forall (m :: * -> *) a.
MonadIO m =>
(a -> Summary) -> Double -> a -> m ()
obs a -> Summary
f Double
v a
cnts =
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Summary -> Double -> IO ()
forall metric (m :: * -> *).
(Observer metric, MonadMonitor m) =>
metric -> Double -> m ()
forall (m :: * -> *). MonadMonitor m => Summary -> Double -> m ()
Prometheus.observe (a -> Summary
f a
cnts) Double
v

timeIt ::
    (MonadIO m) =>
    (a -> Prometheus.Summary) ->
    a ->
    m b ->
    m b
timeIt :: forall (m :: * -> *) a b.
MonadIO m =>
(a -> Summary) -> a -> m b -> m b
timeIt a -> Summary
f a
cnts m b
action = do
    UTCTime
t0 <- IO UTCTime -> m UTCTime
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO UTCTime -> m UTCTime) -> IO UTCTime -> m UTCTime
forall a b. (a -> b) -> a -> b
$ IO UTCTime
getCurrentTime
    !b
ret <- m b
action
    UTCTime
t1 <- IO UTCTime -> m UTCTime
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO UTCTime -> m UTCTime) -> IO UTCTime -> m UTCTime
forall a b. (a -> b) -> a -> b
$ IO UTCTime
getCurrentTime
    (a -> Summary) -> Double -> a -> m ()
forall (m :: * -> *) a.
MonadIO m =>
(a -> Summary) -> Double -> a -> m ()
obs a -> Summary
f (Rational -> Double
forall a. Fractional a => Rational -> a
fromRational (Rational -> Double) -> Rational -> Double
forall a b. (a -> b) -> a -> b
$ NominalDiffTime -> Rational
forall a. Real a => a -> Rational
toRational (NominalDiffTime -> Rational) -> NominalDiffTime -> Rational
forall a b. (a -> b) -> a -> b
$ UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
t1 UTCTime
t0) a
cnts
    b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
ret