Mimir.hs

module Effectful.OpenTelemetry.Exporter.Grafana.Mimir where

import Arbitrary
import Control.Applicative ((<|>))
import Data.Aeson (Key, Value (..), decode, toJSON)
import Data.Aeson.Key qualified as Key
import Data.Aeson.KeyMap qualified as KeyMap
import Data.Char (isAlphaNum, isAscii)
import Data.Foldable (toList)
import Data.Functor ((<&>))
import Data.Maybe (listToMaybe)
import Data.Scientific (Scientific)
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Word (Word64)
import Effectful
import Effectful.Concurrent (Concurrent)
import Effectful.Environment (Environment)
import Effectful.HUnit (HUnit, assertFailure)
import Effectful.Hspec
import Effectful.HttpClient (httpLbs, parseRequest_, responseBody, runHttpClientTls)
import Effectful.OpenTelemetry.Exporter.Grafana.Polling (checkReady, pollOrFail)
import Effectful.OpenTelemetry.Logging (Logging)
import Effectful.OpenTelemetry.Metrics (Metrics)
import Effectful.OpenTelemetry.Metrics.Counter qualified as Counter
import Effectful.OpenTelemetry.Metrics.Gauge qualified as Gauge
import Effectful.OpenTelemetry.Metrics.Histogram qualified as Histogram
import Effectful.OpenTelemetry.Metrics.Measurement
    ( HistogramDataPoint (..)
    , Measurement (..)
    , NumberDataPoint (..)
    )
import Effectful.OpenTelemetry.Metrics.Metadata (Metadata (..))
import Effectful.OpenTelemetry.Metrics.Metadata qualified as Metadata
import Effectful.OpenTelemetry.Metrics.UpDownCounter qualified as UpDownCounter
import Effectful.OpenTelemetry.Protocol.Attributes (Attributes)
import Effectful.OpenTelemetry.Protocol.Attributes qualified as Attributes
import Effectful.OpenTelemetry.Protocol.Transport (Protocol (..))
import Effectful.OpenTelemetry.Tracing (Tracing)
import Effectful.QuickCheck (NonNegative (..), Small (..))
import Effectful.Retry (Retry)
import Effectful.Timeout (Timeout)
import GHC.Stack (HasCallStack)
import Network.HTTP.Client (responseStatus)
import Network.HTTP.Types.Status (statusCode)
import Network.URI (URI (..))
import Text.Read (readMaybe)
import Util
import Prelude

data NumberSample = NumberSample
    { value :: Scientific
    , attributes :: Attributes
    }
    deriving stock (Eq, Show)

numberSample :: Measurement -> Maybe NumberSample
numberSample m = numberDataPoint m <&> \NumberDataPoint{..} -> NumberSample{..}

extractNumberSample :: Value -> Maybe NumberSample
extractNumberSample v = do
    Object res <- Just v
    value <- extractInstantValue v
    Object metric <- KeyMap.lookup "metric" res
    pure
        NumberSample
            { value
            , attributes =
                Attributes.fromList
                    [ (k, val)
                    | (k, val) <- KeyMap.toList metric
                    , k `notElem` reservedLabels
                    ]
            }

data HistogramStats = HistogramStats
    { count :: Word64
    , total :: Maybe Scientific
    , buckets :: Int
    }
    deriving stock (Eq, Show)

histogramStats :: Measurement -> Maybe HistogramStats
histogramStats m =
    histogramDataPoint m <&> \dp ->
        HistogramStats
            { count = dp.count
            , total = dp.sum
            , buckets = length dp.explicitBounds + 1
            }

fetchInstant :: (IOE :> es) => (Value -> Maybe a) -> URI -> Text -> Eff es (Either String a)
fetchInstant extract baseUri metricName = do
    let url =
            show
                baseUri
                    { uriPath = "/prometheus/api/v1/query"
                    , uriQuery = "?query=" <> Text.unpack metricName
                    }
    runHttpClientTls $ do
        resp <- httpLbs $ parseRequest_ url
        let code = statusCode $ responseStatus resp
        pure $ case code of
            200 -> case decode (responseBody resp) of
                Just (Object obj)
                    | Just (Object d) <- KeyMap.lookup "data" obj
                    , Just (Array results) <- KeyMap.lookup "result" d ->
                        case listToMaybe (toList results) >>= extract of
                            Just v -> Right v
                            Nothing -> Left $ "Mimir: no samples for " <> Text.unpack metricName
                _ -> Left "Mimir response was not the expected shape"
            _ -> Left $ "Mimir returned HTTP " <> show code

fetchMetric :: (IOE :> es) => URI -> Text -> Eff es (Either String Scientific)
fetchMetric = fetchInstant extractInstantValue

fetchNumberSample :: (IOE :> es) => URI -> Text -> Eff es (Either String NumberSample)
fetchNumberSample = fetchInstant extractNumberSample

extractInstantValue :: Value -> Maybe Scientific
extractInstantValue v = do
    Object res <- Just v
    Array a <- KeyMap.lookup "value" res
    case toList a of
        [_, String txt] -> readMaybe $ Text.unpack txt
        _ -> Nothing

reservedLabels :: [Key]
reservedLabels = ["__name__", "job", "instance"]

hasDescriptionAndUnit :: Metadata -> Measurement -> Bool
hasDescriptionAndUnit metadata m = m.description == metadata.description && m.unit == metadata.unit

hasMetadata :: Metadata -> Measurement -> Bool
hasMetadata metadata m =
    hasDescriptionAndUnit metadata m
        && (numberAttributes <|> histogramAttributes) == Just metadata.attributes
  where
    numberAttributes = numberDataPoint m <&> (.attributes)
    histogramAttributes = histogramDataPoint m <&> (.attributes)

sanitiseName :: Text -> Text
sanitiseName = Text.map \c -> if isAsciiAlphaNum c || c == '_' then c else '_'
  where
    isAsciiAlphaNum c = isAscii c && isAlphaNum c

spec
    :: forall es
     . (HasCallStack, IOE :> es, HUnit :> es, Hspec :> es, Retry :> es, Timeout :> es)
    => Protocol
    -> URI
    -> ( forall a
          . Eff '[Metrics, Logging, Tracing, Environment, Timeout, Retry, Concurrent, IOE] a
         -> Eff es (a, [Measurement])
       )
    -> Eff es ()
spec (protocolLabel -> label) mimirUri runTest = describe "Mimir" . parallel $ do
    let instrumentName kind marker = sanitiseName $ label <> "_" <> kind <> "_" <> marker
        numberRoundTrip metadata selector act = do
            checkReady "Mimir" mimirUri
            m <- only . snd =<< runTest act
            m `shouldSatisfy` hasDescriptionAndUnit metadata
            sample <- maybe (assertFailure "measurement had no numeric value") pure $ numberSample m
            pollOrFail (fetchNumberSample mimirUri selector) (`shouldBe` sample)

    prop "counter round-trip" \(Marker marker) metadata (NonNegative (Small value)) ->
        let name = instrumentName "counter" marker
         in numberRoundTrip metadata (name <> "_total") do
                c <- Counter.new name metadata
                Counter.add c value

    prop "up-down counter round-trip" \(Marker marker) metadata (Small value) ->
        let name = instrumentName "updown" marker
         in numberRoundTrip metadata name do
                c <- UpDownCounter.new name metadata
                UpDownCounter.add c value

    prop "gauge round-trip" \(Marker marker) metadata (NonNegative (Small value)) ->
        let name = instrumentName "gauge" marker
         in numberRoundTrip metadata name do
                g <- Gauge.new name metadata
                Gauge.set g value

    prop "instrument attributes as labels" \(Marker marker) (Marker labelValue) metadata (NonNegative (Small increment)) ->
        let name = instrumentName "labelled" marker
            labelKey = "attr_kind" :: Text
            metadata' = metadata{Metadata.attributes = Attributes.fromList [(Key.fromText labelKey, toJSON labelValue)]}
            selector = name <> "_total{" <> labelKey <> "=\"" <> labelValue <> "\"}"
         in numberRoundTrip metadata' selector do
                c <- Counter.new name metadata'
                Counter.add c increment

    prop "histogram round-trip" \(Marker marker) metadata (HistogramSamples samples) (HistogramBounds bounds) -> do
        checkReady "Mimir" mimirUri
        let base = sanitiseName $ label <> "_hist_" <> marker
            countSeries = base <> "_count"
            sumSeries = base <> "_sum"
        m <-
            only . snd =<< runTest do
                h <- Histogram.newWithBounds base bounds metadata
                mapM_ (Histogram.record h) samples
        m `shouldSatisfy` hasMetadata metadata
        stats <- maybe (assertFailure "measurement had no histogram data") pure $ histogramStats m
        total <- maybe (assertFailure "measurement had no sum") pure stats.total
        pollOrFail (fetchMetric mimirUri countSeries) (`shouldBe` fromIntegral stats.count)
        pollOrFail (fetchMetric mimirUri sumSeries) \fetched -> abs (fetched - total) `shouldSatisfy` (< 0.01)
        pollOrFail
            (fetchMetric mimirUri $ "count(" <> base <> "_bucket)")
            (`shouldBe` fromIntegral stats.buckets)
        pollOrFail
            (fetchMetric mimirUri $ "max(" <> base <> "_bucket)")
            (`shouldBe` fromIntegral stats.count)