Middleware.hs

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

module System.Metrics.Prometheus.Wai.Middleware
  ( registerWaiMetrics,
    WaiMetrics (..),
    instrumentWaiMiddleware,
    combineRegistrySample,
    MetricsEndpoint (..),
    defaultMetricsEndpoint,
    withLastSecondSamples,
    setSampleCache,
    metricsEndpointMiddleware,
  )
where

import Control.Concurrent.MVar
import Control.Monad
import Control.Monad.IO.Class
import Data.ByteString (ByteString)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import Data.Word
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,
    -- | Samples to collect.
    --
    -- Typically obtained via 'Registry.sample' but you can add others that you collect at the last second.
    metricsEndpointSample :: IO Prometheus.RegistrySample,
    -- | Cache duration in nanoseconds for the metrics samples.
    --
    -- You can set this to use multiple prometheus fetchers for redundancy
    -- without resampling every time.
    --
    -- Nothing means no caching
    metricsEndpointCacheDurationNanos :: Maybe Word64
  }

-- | Create a default 'MetricsEndpoint' serving at @"/metrics"@ with samples
-- from the given registry.
defaultMetricsEndpoint :: Registry -> MetricsEndpoint
defaultMetricsEndpoint registry =
  MetricsEndpoint
    { metricsEndpointPath = "/metrics",
      metricsEndpointSample = Registry.sample registry,
      metricsEndpointCacheDurationNanos = Nothing
    }

-- | Combine the given last second samples with the existing samples of the
-- endpoint.
withLastSecondSamples :: IO Prometheus.RegistrySample -> MetricsEndpoint -> MetricsEndpoint
withLastSecondSamples lastSecondSamples endpoint =
  endpoint
    { metricsEndpointSample =
        combineRegistrySample
          <$> metricsEndpointSample endpoint
          <*> lastSecondSamples
    }

setSampleCache :: Word64 -> MetricsEndpoint -> MetricsEndpoint
setSampleCache cacheDurationNanos endpoint = endpoint {metricsEndpointCacheDurationNanos = Just cacheDurationNanos}

metricsEndpointMiddleware :: (MonadIO m) => MetricsEndpoint -> m Wai.Middleware
metricsEndpointMiddleware MetricsEndpoint {..} =
  -- This function is almost always called in something like LoggingT IO, so we
  -- might as well lift already.
  liftIO $ do
    mCache <- forM metricsEndpointCacheDurationNanos $ \cacheDurationNanos -> do
      begin <- getMonotonicTimeNSec
      firstSample <- metricsEndpointSample
      var <- newMVar (begin, firstSample)
      pure (cacheDurationNanos, var)

    let sampleWithCache :: IO Prometheus.RegistrySample
        sampleWithCache = case mCache of
          Nothing -> metricsEndpointSample
          Just (cacheDurationNanos, var) -> do
            modifyMVar var $ \lastTup@(lastSampleTime, lastSample) -> do
              now <- getMonotonicTimeNSec
              if now - lastSampleTime < cacheDurationNanos
                then pure (lastTup, lastSample)
                else do
                  newSample <- metricsEndpointSample
                  let combinedSample = combineRegistrySample lastSample newSample
                  let newTup = (now, combinedSample)
                  pure (newTup, combinedSample)

    pure $ \application request sendResponse ->
      if Request.rawPathInfo request == metricsEndpointPath
        then do
          samples <- sampleWithCache
          sendResponse $
            Wai.responseBuilder HTTP.ok200 [(HTTP.hContentType, "text/plain; version=0.0.4; charset=utf-8")] $
              Prometheus.encodeMetrics samples
        else application request sendResponse