Loki.hs

{-# OPTIONS_GHC -Wno-name-shadowing #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Effectful.OpenTelemetry.Exporter.Grafana.Loki where

import Arbitrary
import Control.Monad ((>=>))
import Data.Aeson
import Data.Aeson.KeyMap (KeyMap)
import Data.Aeson.KeyMap qualified as KeyMap
import Data.Aeson.Types (Parser)
import Data.Text (Text)
import Data.Text qualified as Text
import Effectful
import Effectful.Concurrent (Concurrent)
import Effectful.Environment (Environment)
import Effectful.HUnit (HUnit)
import Effectful.Hspec
import Effectful.HttpClient (httpLbs, parseRequest_, responseBody, runHttpClientTls)
import Effectful.OpenTelemetry.Exporter.Grafana.Polling (checkReady, pollOrFail)
import Effectful.OpenTelemetry.Logging (Logging, log)
import Effectful.OpenTelemetry.Logging.LogRecord (LogRecord (..))
import Effectful.OpenTelemetry.Logging.Severity (Severity)
import Effectful.OpenTelemetry.Metrics (Metrics)
import Effectful.OpenTelemetry.Protocol.Attributes (Attributes (..))
import Effectful.OpenTelemetry.Protocol.Transport (Protocol (..))
import Effectful.OpenTelemetry.Timestamp (Timestamp)
import Effectful.OpenTelemetry.Tracing (Tracing, inSpan)
import Effectful.OpenTelemetry.Tracing.Span.Context (Context (..))
import Effectful.OpenTelemetry.Tracing.Span.Kind qualified as Kind
import Effectful.Retry (Retry)
import Effectful.Timeout (Timeout)
import GHC.Generics (Generic)
import GHC.Stack (HasCallStack)
import Network.HTTP.Client (responseStatus)
import Network.HTTP.Types.Status (statusCode)
import Network.URI (URI (..), escapeURIString, isUnreserved)
import Text.Read (readMaybe)
import Util
import Prelude hiding (log)

-- HACK: Loki does not include the timestamp in the log records
data Stream = Stream
    { stream :: KeyMap Value
    , values :: [(Text, Text)]
    }
    deriving stock (Generic)
    deriving anyclass (FromJSON)

instance {-# OVERLAPPING #-} FromJSON [LogRecord] where
    parseJSON = parseJSON @[Stream] >=> fmap concat . traverse \Stream{..} -> traverse (parseEntry stream) values
      where
        parseEntry :: KeyMap Value -> (Text, Text) -> Parser LogRecord
        parseEntry labels (timeUnixNano, lineText) = do
            Object line <- either fail pure $ eitherDecodeStrictText lineText
            observedTimestamp <- parseJSON @Timestamp $ String timeUnixNano
            let o = line <> labels
            traceId <- o .:? "traceid"
            spanId <- o .:? "spanid"
            traceFlags <- o .:? "flags" .!= mempty
            let context = case (traceId, spanId) of
                    (Just traceId', Just spanId') ->
                        Just Context{traceId = traceId', spanId = spanId', traceFlags, traceState = mempty}
                    _ -> Nothing
            severity <- o .: "level"
            body <- o .:? "body" .!= Null
            attributes <- Attributes <$> (o .:? "attributes" .!= mempty)
            eventName <- o .:? "name"
            pure LogRecord{timestamp = observedTimestamp, ..}

instance FromJSON Severity where
    parseJSON = withText "Severity" $ maybe (fail "invalid severity") pure . readMaybe . Text.unpack

fetchLines :: (IOE :> es) => URI -> Text -> Eff es (Either String [LogRecord])
fetchLines baseUri needle = do
    let q = "{service_name=\"otel-effectful-test\"} |= `" <> Text.unpack needle <> "`"
        url =
            show
                baseUri
                    { uriPath = "/loki/api/v1/query_range"
                    , uriQuery = "?query=" <> escapeURIString isUnreserved q
                    }
    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 results <- KeyMap.lookup "result" d ->
                        case fromJSON results of
                            Success [] -> Left "Loki responded with no logs"
                            Success (records :: [LogRecord]) -> pure records
                            Error err -> Left $ "Loki response failed to parse: " <> err
                _ -> Left "Loki response was not the expected shape"
            _ -> Left $ "Loki returned HTTP " <> show code

spec
    :: (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, [LogRecord])
       )
    -> Eff es ()
spec (protocolLabel -> label) lokiUri runTest = describe "Loki" . parallel $ do
    prop "log line round-trip" \(Marker marker) (MarkerAttributes attrs) severity -> do
        checkReady "Loki" lokiUri
        let needle = label <> "-log-" <> marker
            spanName = needle <> "-span"
            eventName = "log.event." <> marker
        (_, capturedLogs) <-
            runTest . inSpan spanName Kind.Internal mempty $
                log severity (String needle) attrs (Just eventName)
        theLog <- only capturedLogs
        -- Alloy does not forward eventName
        let expected = theLog{eventName = Nothing}
        pollOrFail (fetchLines lokiUri needle) (`shouldBe` pure expected)