Tempo.hs

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

module Effectful.OpenTelemetry.Exporter.Grafana.Tempo where

import Arbitrary
import Data.Aeson
import Data.Aeson.Key qualified as Key
import Data.Aeson.KeyMap qualified as KeyMap
import Data.Aeson.Types (Parser, typeMismatch)
import Data.ByteString (ByteString)
import Data.ByteString.Base16 qualified as Base16
import Data.ByteString.Base64 qualified as Base64
import Data.Foldable (for_)
import Data.Functor.Syntax ((<&&>))
import Data.Maybe (listToMaybe)
import Data.Sequence qualified as Seq
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text
import Data.Tuple.Extra (thd3)
import Effectful
import Effectful.Concurrent (Concurrent)
import Effectful.Environment (Environment)
import Effectful.Exception (catchSync, throwIO)
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.Protocol (Resource, Scope)
import Effectful.OpenTelemetry.Protocol.Attributes (Attributes)
import Effectful.OpenTelemetry.Protocol.Resource qualified as Resource
import Effectful.OpenTelemetry.Protocol.Scope qualified as Scope
import Effectful.OpenTelemetry.Protocol.Transport (Protocol (..))
import Effectful.OpenTelemetry.Timestamp (Timestamp (..))
import Effectful.OpenTelemetry.Tracing
    ( Tracing
    , addEvent
    , currentContext
    , inOkSpan
    , inSpan
    , withContext
    )
import Effectful.OpenTelemetry.Tracing.Span (Span (..))
import Effectful.OpenTelemetry.Tracing.Span.Context (Context (..))
import Effectful.OpenTelemetry.Tracing.Span.Context qualified as Span.Context
import Effectful.OpenTelemetry.Tracing.Span.Event (Event (..))
import Effectful.OpenTelemetry.Tracing.Span.ID qualified as Span.ID
import Effectful.OpenTelemetry.Tracing.Span.Kind (Kind (..))
import Effectful.OpenTelemetry.Tracing.Span.Kind qualified as Kind
import Effectful.OpenTelemetry.Tracing.Span.Status (Code, Status (..))
import Effectful.OpenTelemetry.Tracing.Span.Status qualified as Status
import Effectful.OpenTelemetry.Tracing.Trace.ID qualified as Trace.ID
import Effectful.OpenTelemetry.Tracing.Trace.State qualified as Trace.State
import Effectful.QuickCheck (arbitrary, generate)
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

instance FromJSON (Resource.Items Span) where
    parseJSON = withObject "batch" \batch -> do
        resource <- batch .:? "resource" .!= mempty
        scopeItems <- batch .:? "scopeSpans" .!= []
        pure Resource.Items{resource, scopeItems}

instance FromJSON Resource.Resource where
    parseJSON = withObject "resource" $ fmap Resource.Resource . attributesAt "attributes"

instance FromJSON (Scope.Items Span) where
    parseJSON = withObject "scopeSpans" \ss -> do
        scope <- ss .:? "scope" .!= emptyScope
        items <- ss .:? "spans" .!= []
        pure Scope.Items{scope, items}
      where
        emptyScope = Scope.Scope{name = "", version = "", attributes = mempty}

instance FromJSON Scope.Scope where
    parseJSON = withObject "scope" \sc -> do
        name <- sc .:? "name" .!= ""
        version <- sc .:? "version" .!= ""
        attributes <- attributesAt "attributes" sc
        pure Scope.Scope{..}

instance FromJSON Span where
    parseJSON = withObject "Span" \o -> do
        context <- do
            traceId <- parseId Trace.ID.fromBytes =<< o .: "traceId"
            spanId <- parseId Span.ID.fromBytes =<< o .: "spanId"
            traceFlags <- o .:? "flags" .!= mempty
            traceState <- maybe mempty Trace.State.fromText <$> o .:? "traceState"
            pure Context{..}
        name <- o .:? "name" .!= ""
        parentSpanId <- traverse (parseId Span.ID.fromBytes) =<< o .:? "parentSpanId"
        kind <- o .: "kind"
        startTime <- parseTimestamp =<< o .: "startTimeUnixNano"
        endTime <- mapM parseTimestamp =<< o .:? "endTimeUnixNano"
        attributes <- attributesAt "attributes" o
        events <- Seq.fromList <$> o .:? "events" .!= []
        status <- o .:? "status" .!= Status.unset
        pure Span{..}
      where
        parseId :: (ByteString -> Either String a) -> Text.Text -> Parser a
        parseId fromBytes t@(Text.encodeUtf8 -> bs) =
            case (fromBytes =<< Base16.decode bs, fromBytes =<< Base64.decode bs) of
                (Right a, _) -> pure a
                (_, Right a) -> pure a
                (Left hexErr, Left b64Err) ->
                    fail $ "neither valid hex (" <> hexErr <> ") nor base64 (" <> b64Err <> "): " <> Text.unpack t

instance FromJSON Event where
    parseJSON = withObject "event" \e -> do
        name <- e .: "name"
        time <- parseTimestamp =<< e .: "timeUnixNano"
        attributes <- attributesAt "attributes" e
        pure Event{..}

instance FromJSON Status where
    parseJSON = withObject "status" \st -> do
        message <- st .:? "message" .!= ""
        code <- st .:? "code" .!= Status.Unset
        pure Status{..}

instance FromJSON Code where
    parseJSON (Number n) = case round n :: Int of
        0 -> pure Status.Unset
        1 -> pure Status.Ok
        2 -> pure Status.Error
        n' -> fail $ "unknown status code: " <> show n'
    parseJSON (String s) = case s of
        "STATUS_CODE_UNSET" -> pure Status.Unset
        "STATUS_CODE_OK" -> pure Status.Ok
        "STATUS_CODE_ERROR" -> pure Status.Error
        _ -> fail $ "unknown status code: " <> Text.unpack s
    parseJSON v = typeMismatch "Status.Code" v

instance FromJSON Kind where
    parseJSON (Number n) = case round n :: Int of
        1 -> pure Internal
        2 -> pure Server
        3 -> pure Client
        4 -> pure Producer
        5 -> pure Consumer
        n' -> fail $ "unknown span kind: " <> show n'
    parseJSON (String s) = case s of
        "SPAN_KIND_INTERNAL" -> pure Internal
        "SPAN_KIND_SERVER" -> pure Server
        "SPAN_KIND_CLIENT" -> pure Client
        "SPAN_KIND_PRODUCER" -> pure Producer
        "SPAN_KIND_CONSUMER" -> pure Consumer
        _ -> fail $ "unknown span kind: " <> Text.unpack s
    parseJSON v = typeMismatch "Kind" v

parseTimestamp :: Value -> Parser Timestamp
parseTimestamp (Number n) = pure . Timestamp $ round n
parseTimestamp (String t) = maybe (fail "invalid timestamp") (pure . Timestamp) . readMaybe $ Text.unpack t
parseTimestamp v = typeMismatch "Timestamp" v

attributesAt :: Key.Key -> Object -> Parser Attributes
attributesAt k o = o .:? k .!= mempty

fetchTrace :: (IOE :> es) => URI -> Text -> Eff es (Either String [(Resource, Scope, Span)])
fetchTrace baseUri traceIdHex = do
    let url = show baseUri{uriPath = "/api/traces/" <> Text.unpack traceIdHex}
    runHttpClientTls $ do
        resp <- httpLbs $ parseRequest_ url
        let code = statusCode $ responseStatus resp
        pure $ case code of
            200 -> case decode (responseBody resp) of
                Just (Object o) -> case KeyMap.lookup "batches" o of
                    Just batches -> case fromJSON batches of
                        Success (batchItems :: [Resource.Items Span]) -> case flattenBatches batchItems of
                            [] -> Left "Tempo returned trace with no spans"
                            ss -> Right ss
                        Error err -> Left $ "Tempo response failed to parse: " <> err
                    Nothing -> Left "Tempo response had no batches field"
                _ -> Left "Tempo response was not a JSON object"
            _ -> Left $ "Tempo returned HTTP " <> show code

flattenBatches :: [Resource.Items a] -> [(Resource, Scope, a)]
flattenBatches batchItems =
    [ (resource, scope, s)
    | Resource.Items{resource, scopeItems} <- batchItems
    , Scope.Items{scope, items} <- scopeItems
    , s <- items
    ]

spec
    :: forall es
     . (HasCallStack, IOE :> es, HUnit :> es, Hspec :> es, Retry :> es, Timeout :> es)
    => Protocol
    -> URI
    -> Resource
    -> Scope
    -> ( forall a
          . Eff '[Metrics, Logging, Tracing, Environment, Timeout, Retry, Concurrent, IOE] a
         -> Eff es (a, [Span])
       )
    -> Eff es ()
spec (protocolLabel -> label) tempoUri resource scope runTest = describe "Tempo" . parallel $ do
    Marker suffix <- liftIO $ generate arbitrary
    let nameFor :: Text -> Text
        nameFor kind = label <> "-" <> kind <> "-" <> suffix
        roundTrip act check = do
            checkReady "Tempo" tempoUri
            ((), captured) <- runTest act
            firstSpan <- maybe (assertFailure "no in-memory spans") pure $ listToMaybe captured
            pollOrFail (fetchTrace tempoUri $ Trace.ID.toHex firstSpan.context.traceId) (check captured)

    prop "span round-trip" \(Marker marker) spanKind (MarkerAttributes attrs) traceState events -> do
        let innerSpanName = label <> "-trace-" <> marker
            outerSpanName = innerSpanName <> "-outer"
        roundTrip
            ( inSpan outerSpanName Kind.Internal mempty do
                ctx <- currentContext <&&> \ctx -> ctx{Span.Context.traceState}
                maybe id withContext ctx . inSpan innerSpanName spanKind attrs $
                    mapM_ @[] addEvent events
            )
            \captured fetched -> do
                captured `shouldBe` (thd3 <$> fetched)
                for_ fetched \(tempoResource, tempoScope, _) -> do
                    tempoScope `shouldBe` scope
                    tempoResource `shouldBe` resource

    it "Ok status round-trip" $
        roundTrip (inOkSpan (nameFor "ok") Kind.Internal mempty $ pure ()) \captured fetched ->
            (thd3 <$> fetched) `shouldBe` captured

    it "Error status round-trip" $
        roundTrip
            ( (inSpan (nameFor "error") Kind.Internal mempty . throwIO . userError $ "boom-" <> Text.unpack suffix)
                `catchSync` const (pure ())
            )
            \captured fetched -> (thd3 <$> fetched) `shouldBe` captured