Middleware.hs

{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module System.Metrics.Prometheus.Wai.Middleware
  ( registerWaiMetrics,
    instrumentWaiMiddleware,
    metricsEndpointMiddleware,
    metricsEndpointAtMiddleware,
  )
where

import Control.Monad (forM)
import Data.ByteString (ByteString)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import GHC.Clock (getMonotonicTimeNSec)
import qualified Network.HTTP.Types as HTTP
import Network.Wai as Wai (Middleware)
import qualified Network.Wai as Request
import qualified Network.Wai as Wai
import System.Metrics.Prometheus.Concurrent.Registry (Registry)
import qualified System.Metrics.Prometheus.Concurrent.Registry as Prometheus
import qualified System.Metrics.Prometheus.Concurrent.Registry as Registry
import qualified System.Metrics.Prometheus.Encode.Text as Prometheus
import qualified System.Metrics.Prometheus.Metric.Counter as Counter (inc)
import qualified System.Metrics.Prometheus.Metric.Counter as Prometheus (Counter)
import qualified System.Metrics.Prometheus.Metric.Histogram as Histogram (observe)
import qualified System.Metrics.Prometheus.Metric.Histogram as Prometheus (Histogram)
import qualified System.Metrics.Prometheus.MetricId as Labels
import qualified System.Metrics.Prometheus.MetricId as Prometheus (Labels (..))

-- TODO ignore websocket requests when gathering info

data WaiMetrics = WaiMetrics
  { waiMetricsStatusCode :: !(Map Int Prometheus.Counter),
    waiMetricsDuration :: !Prometheus.Histogram
  }

-- | Register the Wai metrics with the given labels at the given registry.
registerWaiMetrics :: Prometheus.Labels -> Registry -> IO WaiMetrics
registerWaiMetrics labels registry = do
  -- Status code counters
  -- Based on the codes defined at
  -- https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Status
  let codes =
        [100 .. 103]
          <> [200 .. 208]
          <> [226]
          <> [300 .. 304]
          <> [307, 308]
          <> [400 .. 418]
          <> [421 .. 426]
          <> [428, 429, 431, 451]
          <> [500 .. 508]
          <> [510, 511]
  let labelsForCode code = Labels.addLabel "http_response_code" (T.pack (show code)) labels
  waiMetricsStatusCode <- fmap M.fromList $ forM codes $ \code -> do
    counterForCode <- Prometheus.registerCounter "http_requests_total" (labelsForCode code) registry
    pure (code, counterForCode)

  -- Duration histogram
  let durationBounds =
        concat
          [ [1, 2, 3, 5],
            [10, 20, 30, 40, 50],
            [100, 200, 300, 400, 500],
            [1_000, 2_000, 3_000, 4_000, 5_000],
            [10_000]
          ]
  waiMetricsDuration <- Prometheus.registerHistogram "http_request_duration_milliseconds" labels durationBounds registry
  pure WaiMetrics {..}

-- | Record the given Wai metrics in a middleware.
instrumentWaiMiddleware :: WaiMetrics -> Wai.Middleware
instrumentWaiMiddleware WaiMetrics {..} application request sendResponse = do
  begin <- getMonotonicTimeNSec
  application request $ \response -> do
    end <- getMonotonicTimeNSec

    -- Count the status code
    mapM_ Counter.inc (M.lookup (HTTP.statusCode (Wai.responseStatus response)) waiMetricsStatusCode)

    -- Count the application response duration
    let nanos = end - begin
        millis = fromIntegral nanos / 1_000_000
    Histogram.observe millis waiMetricsDuration

    sendResponse response

-- | Add a metrics endpoint using the given industry.
metricsEndpointMiddleware :: Registry -> Wai.Middleware
metricsEndpointMiddleware = metricsEndpointAtMiddleware "/metrics"

metricsEndpointAtMiddleware :: ByteString -> Registry -> Wai.Middleware
metricsEndpointAtMiddleware path registry application request sendResponse =
  if Request.rawPathInfo request == path
    then do
      s <- Registry.sample registry
      sendResponse
        $ Wai.responseBuilder
          HTTP.ok200
          [(HTTP.hContentType, "text/plain; version=0.0.4; charset=utf-8")]
        $ Prometheus.encodeMetrics s
    else application request sendResponse