Middleware.hs

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

module System.Metrics.Prometheus.Wai.Middleware
  ( registerWaiMetrics,
    WaiMetrics (..),
    instrumentWaiMiddleware,
    combineRegistrySample,
    MetricsEndpoint (..),
    defaultMetricsEndpoint,
    metricsEndpointMiddleware,
  )
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 (..))
import qualified System.Metrics.Prometheus.Registry as Prometheus (RegistrySample (..))

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 .. 100],
            [100, 200 .. 900],
            [1_000, 2_000 .. 9_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 =
  let isWebSocketsReq =
        lookup "upgrade" (Wai.requestHeaders request) == Just "websocket"
      shouldInstrument =
        -- Don't instrument WebSocket requests because they don't have a
        -- response but some libraries still pretend that it does and will
        -- give it a 500 status code.
        -- Moreover, the 'latency' will be 'how long the connection was open'
        -- which is also useless.
        not isWebSocketsReq
   in if shouldInstrument
        then 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
        else application request sendResponse

combineRegistrySample :: Prometheus.RegistrySample -> Prometheus.RegistrySample -> Prometheus.RegistrySample
combineRegistrySample (Prometheus.RegistrySample a) (Prometheus.RegistrySample b) =
  Prometheus.RegistrySample (M.union a b)

data MetricsEndpoint = MetricsEndpoint
  { -- | Path to serve metrics on, e.g. @"/metrics"@
    metricsEndpointPath :: ByteString,
    -- | Registry to sample from
    metricsEndpointRegistry :: Registry,
    -- | Extra samples to be collected 'at the last second', to be combined
    -- with the main registry samples
    metricsEndpointLastSecond :: Maybe (IO Prometheus.RegistrySample)
  }

defaultMetricsEndpoint :: Registry -> MetricsEndpoint
defaultMetricsEndpoint registry =
  MetricsEndpoint
    { metricsEndpointPath = "/metrics",
      metricsEndpointRegistry = registry,
      metricsEndpointLastSecond = Nothing
    }

metricsEndpointMiddleware :: MetricsEndpoint -> Wai.Middleware
metricsEndpointMiddleware MetricsEndpoint {..} application request sendResponse =
  if Request.rawPathInfo request == metricsEndpointPath
    then do
      let sampleRegistry = Registry.sample metricsEndpointRegistry
      finalSamples <- case metricsEndpointLastSecond of
        Nothing -> sampleRegistry
        Just lastSecondSample ->
          combineRegistrySample
            <$> sampleRegistry
            <*> lastSecondSample
      sendResponse
        $ Wai.responseBuilder
          HTTP.ok200
          [(HTTP.hContentType, "text/plain; version=0.0.4; charset=utf-8")]
        $ Prometheus.encodeMetrics finalSamples
    else application request sendResponse