-- | This module provides "Network.Wai" middlware for exporting "Prometheus"
-- metrics and for instrumenting WAI applications.
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module Network.Wai.Middleware.Prometheus
  ( prometheus
  , PrometheusSettings(..)
  , Default.def
  , instrumentHandlerValue
  , instrumentHandlerValueWithFilter
  , instrumentHandlerValueWithHistogramAndFilter
  , ignoreRawResponses
  , instrumentApp
  , instrumentIO
  , observeSeconds
  , metricsApp
  ) where

import qualified Data.Default as Default
import Data.Maybe (fromMaybe)
import Data.Ratio ((%))
import Data.Text (Text)
import qualified Data.Text as T
import Data.Text.Encoding (decodeUtf8)
import qualified Network.HTTP.Types as HTTP
import qualified Network.Wai as Wai
import qualified Network.Wai.Internal as Wai (Response(ResponseRaw))
import qualified Prometheus as Prom
import System.Clock (Clock(..), TimeSpec, diffTimeSpec, getTime, toNanoSecs)


-- | Settings that control the behavior of the Prometheus middleware.
data PrometheusSettings = PrometheusSettings {
        PrometheusSettings -> [Text]
prometheusEndPoint             :: [T.Text]
        -- ^ The path that will be used for exporting metrics. The default value
        -- is ["metrics"] which corresponds to the path /metrics.
    ,   PrometheusSettings -> Bool
prometheusInstrumentApp        :: Bool
        -- ^ Whether the default instrumentation should be applied to the
        -- application. If this is set to false the application can still be
        -- instrumented using the 'instrumentApp' function. The default value is
        -- True.
    ,   PrometheusSettings -> Bool
prometheusInstrumentPrometheus :: Bool
        -- ^ Whether the default instrumentation should be applied to the
        -- middleware that serves the metrics endpoint. The default value is
        -- True.
    }

instance Default.Default PrometheusSettings where
    def :: PrometheusSettings
def = PrometheusSettings {
        prometheusEndPoint :: [Text]
prometheusEndPoint             = [Text
"metrics"]
    ,   prometheusInstrumentApp :: Bool
prometheusInstrumentApp        = Bool
True
    ,   prometheusInstrumentPrometheus :: Bool
prometheusInstrumentPrometheus = Bool
True
    }

{-# NOINLINE requestLatency #-}
requestLatency :: Prom.Vector Prom.Label3 Prom.Histogram
requestLatency :: Vector Label3 Histogram
requestLatency = Metric (Vector Label3 Histogram) -> Vector Label3 Histogram
forall s. Metric s -> s
Prom.unsafeRegister (Metric (Vector Label3 Histogram) -> Vector Label3 Histogram)
-> Metric (Vector Label3 Histogram) -> Vector Label3 Histogram
forall a b. (a -> b) -> a -> b
$ Label3 -> Metric Histogram -> Metric (Vector Label3 Histogram)
forall l m. Label l => l -> Metric m -> Metric (Vector l m)
Prom.vector (Text
"handler", Text
"method", Text
"status_code")
                                     (Metric Histogram -> Metric (Vector Label3 Histogram))
-> Metric Histogram -> Metric (Vector Label3 Histogram)
forall a b. (a -> b) -> a -> b
$ Info -> [Double] -> Metric Histogram
Prom.histogram Info
info [Double]
Prom.defaultBuckets
    where info :: Info
info = Text -> Text -> Info
Prom.Info Text
"http_request_duration_seconds"
                           Text
"The HTTP request latencies in seconds."

-- | This function is used to populate the @handler@ label of all Prometheus metrics recorded by this library.
--
-- If you use this function you will likely want to override the default value
-- of 'prometheusInstrumentApp' to be false so that your app does not get double
-- instrumented.
--
-- WARNING: If you have 'ResponseRaw' values in your API, consider using
-- @instrumentHandlerValueWithFilter ignoreRawResponses@ instead.
instrumentHandlerValue ::
     (Wai.Request -> Text) -- ^ The function used to derive the "handler" value in Prometheus
  -> Wai.Application -- ^ The app to instrument
  -> Wai.Application -- ^ The instrumented app
instrumentHandlerValue :: (Request -> Text) -> Application -> Application
instrumentHandlerValue = (Response -> Maybe Response)
-> (Request -> Text) -> Application -> Application
instrumentHandlerValueWithFilter Response -> Maybe Response
forall a. a -> Maybe a
Just

-- | A more flexible variant of 'instrumentHandlerValue'.  The filter can change some
-- responses, or drop others entirely.
instrumentHandlerValueWithFilter ::
     (Wai.Response -> Maybe Wai.Response) -- ^ Response filter
  -> (Wai.Request -> Text) -- ^ The function used to derive the "handler" value in Prometheus
  -> Wai.Application -- ^ The app to instrument
  -> Wai.Application -- ^ The instrumented app
instrumentHandlerValueWithFilter :: (Response -> Maybe Response)
-> (Request -> Text) -> Application -> Application
instrumentHandlerValueWithFilter =
  Vector Label3 Histogram
-> (Response -> Maybe Response)
-> (Request -> Text)
-> Application
-> Application
instrumentHandlerValueWithHistogramAndFilter Vector Label3 Histogram
requestLatency

instrumentHandlerValueWithHistogramAndFilter ::
     Prom.Vector Prom.Label3 Prom.Histogram
  -> (Wai.Response -> Maybe Wai.Response) -- ^ Response filter
  -> (Wai.Request -> Text) -- ^ The function used to derive the "handler" value in Prometheus
  -> Wai.Application -- ^ The app to instrument
  -> Wai.Application -- ^ The instrumented app
instrumentHandlerValueWithHistogramAndFilter :: Vector Label3 Histogram
-> (Response -> Maybe Response)
-> (Request -> Text)
-> Application
-> Application
instrumentHandlerValueWithHistogramAndFilter Vector Label3 Histogram
histogram Response -> Maybe Response
resFilter Request -> Text
f Application
app Request
req Response -> IO ResponseReceived
respond = do
  TimeSpec
start <- Clock -> IO TimeSpec
getTime Clock
Monotonic
  Application
app Request
req ((Response -> IO ResponseReceived) -> IO ResponseReceived)
-> (Response -> IO ResponseReceived) -> IO ResponseReceived
forall a b. (a -> b) -> a -> b
$ \Response
res -> do
    case Response -> Maybe Response
resFilter Response
res of
      Maybe Response
Nothing -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
      Just Response
res' -> do
        TimeSpec
end <- Clock -> IO TimeSpec
getTime Clock
Monotonic
        let method :: Maybe Text
method = Text -> Maybe Text
forall a. a -> Maybe a
Just (Text -> Maybe Text) -> Text -> Maybe Text
forall a b. (a -> b) -> a -> b
$ Method -> Text
decodeUtf8 (Request -> Method
Wai.requestMethod Request
req)
        let status :: Maybe Text
status = Text -> Maybe Text
forall a. a -> Maybe a
Just (Text -> Maybe Text) -> Text -> Maybe Text
forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show (Status -> Int
HTTP.statusCode (Response -> Status
Wai.responseStatus Response
res')))
        Vector Label3 Histogram
-> Text
-> Maybe Text
-> Maybe Text
-> TimeSpec
-> TimeSpec
-> IO ()
observeSecondsWithHistogram Vector Label3 Histogram
histogram (Request -> Text
f Request
req) Maybe Text
method Maybe Text
status TimeSpec
start TimeSpec
end
    Response -> IO ResponseReceived
respond Response
res

-- | 'Wai.ResponseRaw' values have two parts: an action that can be executed to construct a
-- 'Wai.Response', and a pure "backup" 'Wai.Response' in case the computation fails.  Since
-- the pure selectors like 'Wai.responseStatus' are pure and it makes no sense for them to
-- call the action, they just go to the backup response and pull the value from that:
--
-- @
-- responseStatus (ResponseRaw ...
--   (ResponseBuilder (Status 500 "blargh") ... ...))
-- == Status {statusCode = 500, statusMessage = "blargh"}
-- @
--
-- This is often not what you want.  For example, if you have an end-point for establishing
-- websocket connections that has a backup response with status 5xx, every websocket
-- connection request, whether successful or not, will register as an internal server error.
--
-- This helper therefore filters out all raw requests so they won't create any metrics.  Use
-- together with 'instrumentHandlerValueWithFilter'.
ignoreRawResponses :: Wai.Response -> Maybe Wai.Response
ignoreRawResponses :: Response -> Maybe Response
ignoreRawResponses (Wai.ResponseRaw {}) = Maybe Response
forall a. Maybe a
Nothing
ignoreRawResponses Response
res = Response -> Maybe Response
forall a. a -> Maybe a
Just Response
res

-- | Instrument a WAI app with the default WAI metrics.
--
-- If you use this function you will likely want to override the default value
-- of 'prometheusInstrumentApp' to be false so that your app does not get double
-- instrumented.
instrumentApp ::
     Text -- ^ The label used to identify this app
  -> Wai.Application -- ^ The app to instrument
  -> Wai.Application -- ^ The instrumented app
instrumentApp :: Text -> Application -> Application
instrumentApp Text
handler Application
app Request
req Response -> IO ResponseReceived
respond =
  (Request -> Text) -> Application -> Application
instrumentHandlerValue (Text -> Request -> Text
forall a b. a -> b -> a
const Text
handler) Application
app Request
req Response -> IO ResponseReceived
respond

-- | Instrument an IO action with timing metrics. This function can be used if
-- you would like to get more fine grained metrics, for instance this can be
-- used to instrument individual end points.
--
-- If you use this function you will likely want to override the default value
-- of 'prometheusInstrumentApp' to be false so that your app does not get double
-- instrumented.
instrumentIO :: Text    -- ^ The label used to identify this IO operation
             -> IO a    -- ^ The IO action to instrument
             -> IO a    -- ^ The instrumented app
instrumentIO :: forall a. Text -> IO a -> IO a
instrumentIO Text
label IO a
io = do
    TimeSpec
start  <- Clock -> IO TimeSpec
getTime Clock
Monotonic
    a
result <- IO a
io
    TimeSpec
end    <- Clock -> IO TimeSpec
getTime Clock
Monotonic
    Text -> Maybe Text -> Maybe Text -> TimeSpec -> TimeSpec -> IO ()
observeSeconds Text
label Maybe Text
forall a. Maybe a
Nothing Maybe Text
forall a. Maybe a
Nothing TimeSpec
start TimeSpec
end
    a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
result

-- | Record an event to the middleware metric.
observeSeconds :: Text         -- ^ handler label
               -> Maybe Text   -- ^ method
               -> Maybe Text   -- ^ status
               -> TimeSpec     -- ^ start time
               -> TimeSpec     -- ^ end time
               -> IO ()
observeSeconds :: Text -> Maybe Text -> Maybe Text -> TimeSpec -> TimeSpec -> IO ()
observeSeconds = do
  Vector Label3 Histogram
-> Text
-> Maybe Text
-> Maybe Text
-> TimeSpec
-> TimeSpec
-> IO ()
observeSecondsWithHistogram Vector Label3 Histogram
requestLatency

-- | Record an event to the middleware metric.
observeSecondsWithHistogram :: Prom.Vector Prom.Label3 Prom.Histogram
                            -> Text         -- ^ handler label
                            -> Maybe Text   -- ^ method
                            -> Maybe Text   -- ^ status
                            -> TimeSpec     -- ^ start time
                            -> TimeSpec     -- ^ end time
                            -> IO ()
observeSecondsWithHistogram :: Vector Label3 Histogram
-> Text
-> Maybe Text
-> Maybe Text
-> TimeSpec
-> TimeSpec
-> IO ()
observeSecondsWithHistogram Vector Label3 Histogram
histograms Text
handler Maybe Text
method Maybe Text
status TimeSpec
start TimeSpec
end = do
    let latency :: Double
        latency :: Double
latency = Ratio Integer -> Double
forall a. Fractional a => Ratio Integer -> a
fromRational (Ratio Integer -> Double) -> Ratio Integer -> Double
forall a b. (a -> b) -> a -> b
$ Ratio Integer -> Ratio Integer
forall a. Real a => a -> Ratio Integer
toRational (TimeSpec -> Integer
toNanoSecs (TimeSpec
end TimeSpec -> TimeSpec -> TimeSpec
`diffTimeSpec` TimeSpec
start) Integer -> Integer -> Ratio Integer
forall a. Integral a => a -> a -> Ratio a
% Integer
1000000000)
    Vector Label3 Histogram -> Label3 -> (Histogram -> IO ()) -> IO ()
forall label (m :: * -> *) metric.
(Label label, MonadMonitor m) =>
Vector label metric -> label -> (metric -> IO ()) -> m ()
Prom.withLabel Vector Label3 Histogram
histograms
                   (Text
handler, Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
"" Maybe Text
method, Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
"" Maybe Text
status)
                   ((Histogram -> Double -> IO ()) -> Double -> Histogram -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip Histogram -> Double -> IO ()
forall metric (m :: * -> *).
(Observer metric, MonadMonitor m) =>
metric -> Double -> m ()
forall (m :: * -> *). MonadMonitor m => Histogram -> Double -> m ()
Prom.observe Double
latency)

-- | Expose Prometheus metrics and instrument an application with some basic
-- metrics (e.g. request latency).
prometheus :: PrometheusSettings -> Wai.Middleware
prometheus :: PrometheusSettings -> Application -> Application
prometheus PrometheusSettings{Bool
[Text]
prometheusEndPoint :: PrometheusSettings -> [Text]
prometheusInstrumentApp :: PrometheusSettings -> Bool
prometheusInstrumentPrometheus :: PrometheusSettings -> Bool
prometheusEndPoint :: [Text]
prometheusInstrumentApp :: Bool
prometheusInstrumentPrometheus :: Bool
..} Application
app Request
req Response -> IO ResponseReceived
respond =
    if     Request -> Method
Wai.requestMethod Request
req Method -> Method -> Bool
forall a. Eq a => a -> a -> Bool
== Method
HTTP.methodGet
        Bool -> Bool -> Bool
&& Request -> [Text]
Wai.pathInfo Request
req [Text] -> [Text] -> Bool
forall a. Eq a => a -> a -> Bool
== [Text]
prometheusEndPoint
        -- XXX: Should probably be "metrics" rather than "prometheus", since
        -- "prometheus" can be confused with actual prometheus.
    then
      if Bool
prometheusInstrumentPrometheus
        then Text -> Application -> Application
instrumentApp Text
"prometheus" (((Response -> IO ResponseReceived) -> IO ResponseReceived)
-> Application
forall a b. a -> b -> a
const (Response -> IO ResponseReceived) -> IO ResponseReceived
respondWithMetrics) Request
req Response -> IO ResponseReceived
respond
        else (Response -> IO ResponseReceived) -> IO ResponseReceived
respondWithMetrics Response -> IO ResponseReceived
respond
    else
      if Bool
prometheusInstrumentApp
        then Text -> Application -> Application
instrumentApp Text
"app" Application
app Request
req Response -> IO ResponseReceived
respond
        else Application
app Request
req Response -> IO ResponseReceived
respond


-- | WAI Application that serves the Prometheus metrics page regardless of
-- what the request is.
metricsApp :: Wai.Application
metricsApp :: Application
metricsApp = ((Response -> IO ResponseReceived) -> IO ResponseReceived)
-> Application
forall a b. a -> b -> a
const (Response -> IO ResponseReceived) -> IO ResponseReceived
respondWithMetrics

respondWithMetrics :: (Wai.Response -> IO Wai.ResponseReceived)
                   -> IO Wai.ResponseReceived
respondWithMetrics :: (Response -> IO ResponseReceived) -> IO ResponseReceived
respondWithMetrics Response -> IO ResponseReceived
respond = do
    ByteString
metrics <- IO ByteString
forall (m :: * -> *). MonadIO m => m ByteString
Prom.exportMetricsAsText
    Response -> IO ResponseReceived
respond (Response -> IO ResponseReceived)
-> Response -> IO ResponseReceived
forall a b. (a -> b) -> a -> b
$ Status -> ResponseHeaders -> ByteString -> Response
Wai.responseLBS Status
HTTP.status200 ResponseHeaders
headers ByteString
metrics
    where
        headers :: ResponseHeaders
headers = [(HeaderName
HTTP.hContentType, Method
"text/plain; version=0.0.4")]