TracingSpec.hs

{- HLINT ignore "Use head" -}
{-# LANGUAGE OverloadedLists #-}

module Effectful.OpenTelemetry.TracingSpec where

import Arbitrary (Marker (..))
import Control.Monad (forM_)
import Data.Either (isLeft)
import Data.Functor ((<&>))
import Data.Text qualified as Text
import Effectful
import Effectful.Concurrent (runConcurrent)
import Effectful.Concurrent.Async (concurrently_, replicateConcurrently_)
import Effectful.Error.Static (runErrorNoCallStack, throwError)
import Effectful.Exception (AssertionFailed (..), throwIO, try)
import Effectful.HUnit (HUnit)
import Effectful.Hspec
import Effectful.OpenTelemetry.Tracing
import Effectful.OpenTelemetry.Tracing.Span.Context qualified as Span.Context
import Effectful.OpenTelemetry.Tracing.Span.Event (Event (..))
import Effectful.OpenTelemetry.Tracing.Span.Kind qualified as Span.Kind
import Effectful.OpenTelemetry.Tracing.Span.Status qualified as Span.Status
import Util
import Prelude

spec :: (IOE :> es, HUnit :> es, Hspec :> es) => Eff es ()
spec = runConcurrent . parallel $ do
    describe "inSpan" do
        prop "exports a span" \(Marker name) spanKind a -> do
            s <- only . snd =<< (runInMemoryTracing . inSpan name spanKind a . pure) ()
            s.name `shouldBe` name
            s.kind `shouldBe` spanKind
            s.attributes `shouldBe` a
            s.endTime `shouldSatisfy` (>= Just s.startTime)

        it "no parent, unique ids" do
            (_, spans) <- runInMemoryTracing do
                inSpan "a" Span.Kind.Internal mempty $ pure ()
                inSpan "b" Span.Kind.Internal mempty $ pure ()
                inSpan "c" Span.Kind.Internal mempty $ pure ()
            mapM_ ((`shouldBe` Nothing) . (.parentSpanId)) spans
            (spans <&> (.context.spanId)) `shouldSatisfy` allUnique

        describe "nesting" do
            it "parent-child chain" do
                (_, spans) <- runInMemoryTracing do
                    inSpan "root" Span.Kind.Internal mempty
                        . inSpan "middle" Span.Kind.Internal mempty
                        . inSpan "leaf" Span.Kind.Internal mempty
                        $ pure ()
                length spans `shouldBe` 3
                let leaf = spans !! 0
                    middle = spans !! 1
                    root = spans !! 2
                leaf.parentSpanId `shouldBe` Just middle.context.spanId
                middle.parentSpanId `shouldBe` Just root.context.spanId
                root.parentSpanId `shouldBe` Nothing

            it "shared parent" do
                (_, spans) <- runInMemoryTracing do
                    inSpan "parent" Span.Kind.Internal mempty do
                        inSpan "child-a" Span.Kind.Internal mempty $ pure ()
                        inSpan "child-b" Span.Kind.Internal mempty $ pure ()
                length spans `shouldBe` 3
                let childA = spans !! 0
                    childB = spans !! 1
                    parent = spans !! 2
                childA.parentSpanId `shouldBe` Just parent.context.spanId
                childB.parentSpanId `shouldBe` Just parent.context.spanId

            it "restores parent context" do
                (ctx, _) <- runInMemoryTracing do
                    inSpan "outer" Span.Kind.Internal mempty do
                        inSpan "inner" Span.Kind.Internal mempty $ pure ()
                        currentContext
                ctx `shouldSatisfy` \case
                    Just _ -> True
                    Nothing -> False

        describe "concurrency" do
            it "no loss, unique ids" do
                let n = 500 :: Int
                (_, spans) <-
                    runInMemoryTracing
                        . runConcurrent
                        . replicateConcurrently_ n
                        . inSpan "parallel" Span.Kind.Internal mempty
                        $ pure ()
                length spans `shouldBe` n
                (spans <&> (.context.spanId)) `shouldSatisfy` allUnique

            it "concurrent nesting" do
                (_, spans) <-
                    runInMemoryTracing
                        . runConcurrent
                        . inSpan "root" Span.Kind.Internal mempty
                        $ concurrently_
                            (inSpan "branch-a" Span.Kind.Internal mempty $ pure ())
                            (inSpan "branch-b" Span.Kind.Internal mempty $ pure ())
                length spans `shouldBe` 3
                case filter (("root" ==) . (.name)) spans of
                    [root] -> do
                        let branches = filter ((/= "root") . (.name)) spans
                        mapM_ ((`shouldBe` Just root.context.spanId) . (.parentSpanId)) branches
                    _ -> expectationFailure "expected exactly one root span"

        describe "exception handling" do
            prop "Error status on exception" \(Marker msg) -> do
                (result, spans) <-
                    runInMemoryTracing
                        . try @AssertionFailed
                        . inSpan @_ @() "item" Span.Kind.Internal mempty
                        . throwIO
                        $ AssertionFailed (Text.unpack msg)
                result `shouldSatisfy` isLeft
                length spans `shouldBe` 1
                forM_ spans \s -> do
                    s.status.code `shouldBe` Span.Status.Error
                    s.status.message `shouldSatisfy` Text.isPrefixOf msg
                    length s.events `shouldBe` 1
                    fmap (.name) s.events `shouldBe` ["exception"]

        describe "events" do
            prop "records events in order" \(Marker n1) a1 (Marker n2) a2 -> do
                s <-
                    only . snd =<< runInMemoryTracing do
                        inSpan "test-span" Span.Kind.Internal mempty do
                            addEventNow n1 a1
                            addEventNow n2 a2
                            pure ()
                fmap (\e -> (e.name, e.attributes)) s.events `shouldBe` [(n1, a1), (n2, a2)]

            it "events go to innermost span" do
                (_, spans) <- runInMemoryTracing do
                    inSpan "parent" Span.Kind.Internal mempty
                        . inSpan "child" Span.Kind.Internal mempty
                        $ do
                            addEventNow "event" mempty
                            pure ()
                length spans `shouldBe` 2
                let child = spans !! 0
                fmap (.name) child.events `shouldBe` ["event"]

    describe "inErrorSpan" do
        prop "exports a span" \(Marker name) spanKind a -> do
            s <-
                only . snd
                    =<< ( runInMemoryTracing
                            . runErrorNoCallStack @AssertionFailed
                            . inErrorSpan @AssertionFailed @_ @_ name spanKind a
                            . pure
                        )
                        ()
            s.name `shouldBe` name
            s.kind `shouldBe` spanKind
            s.attributes `shouldBe` a
            s.endTime `shouldSatisfy` (>= Just s.startTime)

        it "no parent, unique ids" do
            (_, spans) <- runInMemoryTracing . runErrorNoCallStack @AssertionFailed $ do
                inErrorSpan @AssertionFailed @_ @() "a" Span.Kind.Internal mempty $ pure ()
                inErrorSpan @AssertionFailed @_ @() "b" Span.Kind.Internal mempty $ pure ()
                inErrorSpan @AssertionFailed @_ @() "c" Span.Kind.Internal mempty $ pure ()
            mapM_ ((`shouldBe` Nothing) . (.parentSpanId)) spans
            (spans <&> (.context.spanId)) `shouldSatisfy` allUnique

        describe "nesting" do
            it "parent-child chain" do
                (_, spans) <- runInMemoryTracing . runErrorNoCallStack @AssertionFailed $ do
                    inErrorSpan @AssertionFailed @_ @() "root" Span.Kind.Internal mempty
                        . inErrorSpan @AssertionFailed @_ @() "middle" Span.Kind.Internal mempty
                        . inErrorSpan @AssertionFailed @_ @() "leaf" Span.Kind.Internal mempty
                        $ pure ()
                length spans `shouldBe` 3
                let leaf = spans !! 0
                    middle = spans !! 1
                    root = spans !! 2
                leaf.parentSpanId `shouldBe` Just middle.context.spanId
                middle.parentSpanId `shouldBe` Just root.context.spanId
                root.parentSpanId `shouldBe` Nothing

            it "shared parent" do
                (_, spans) <- runInMemoryTracing . runErrorNoCallStack @AssertionFailed $ do
                    inErrorSpan @AssertionFailed @_ @() "parent" Span.Kind.Internal mempty do
                        inErrorSpan @AssertionFailed @_ @() "child-a" Span.Kind.Internal mempty $ pure ()
                        inErrorSpan @AssertionFailed @_ @() "child-b" Span.Kind.Internal mempty $ pure ()
                length spans `shouldBe` 3
                let childA = spans !! 0
                    childB = spans !! 1
                    parent = spans !! 2
                childA.parentSpanId `shouldBe` Just parent.context.spanId
                childB.parentSpanId `shouldBe` Just parent.context.spanId

            it "restores parent context" do
                (ctx, _) <- runInMemoryTracing . runErrorNoCallStack @AssertionFailed $ do
                    inErrorSpan @AssertionFailed @_ @_ "outer" Span.Kind.Internal mempty do
                        inErrorSpan @AssertionFailed @_ @() "inner" Span.Kind.Internal mempty $ pure ()
                        currentContext
                ctx `shouldSatisfy` \case
                    Right (Just _) -> True
                    _ -> False

        describe "concurrency" do
            it "no loss, unique ids" do
                let n = 500 :: Int
                (_, spans) <-
                    runInMemoryTracing
                        . runErrorNoCallStack @AssertionFailed
                        . runConcurrent
                        . replicateConcurrently_ n
                        . inErrorSpan @AssertionFailed @_ @() "parallel" Span.Kind.Internal mempty
                        $ pure ()
                length spans `shouldBe` n
                (spans <&> (.context.spanId)) `shouldSatisfy` allUnique

            it "concurrent nesting" do
                (_, spans) <-
                    runInMemoryTracing
                        . runErrorNoCallStack @AssertionFailed
                        . runConcurrent
                        . inErrorSpan @AssertionFailed @_ @() "root" Span.Kind.Internal mempty
                        $ concurrently_
                            (inSpan "branch-a" Span.Kind.Internal mempty $ pure ())
                            (inSpan "branch-b" Span.Kind.Internal mempty $ pure ())
                length spans `shouldBe` 3
                case filter (("root" ==) . (.name)) spans of
                    [root] -> do
                        let branches = filter ((/= "root") . (.name)) spans
                        mapM_ ((`shouldBe` Just root.context.spanId) . (.parentSpanId)) branches
                    _ -> expectationFailure "expected exactly one root span"

        describe "exception handling" do
            prop "Error status on exception" \(Marker msg) -> do
                (result, spans) <-
                    runInMemoryTracing
                        . runErrorNoCallStack @AssertionFailed
                        . inErrorSpan
                            @AssertionFailed
                            @_
                            @()
                            "item"
                            Span.Kind.Internal
                            mempty
                        . throwError
                        $ AssertionFailed (Text.unpack msg)
                result `shouldSatisfy` isLeft
                length spans `shouldBe` 1
                forM_ spans \s -> do
                    s.status.code `shouldBe` Span.Status.Error
                    s.status.message `shouldBe` msg
                    length s.events `shouldBe` 1
                    fmap (.name) s.events `shouldBe` ["exception"]

            it "Ok status without exception" do
                (_, spans) <-
                    runInMemoryTracing
                        . runErrorNoCallStack @AssertionFailed
                        . inErrorSpan
                            @AssertionFailed
                            @_
                            @()
                            "item"
                            Span.Kind.Internal
                            mempty
                        $ pure ()
                length spans `shouldBe` 1
                forM_ (spans <&> (.status)) \s -> do
                    s.code `shouldBe` Span.Status.Unset
                    s.message `shouldBe` mempty

        describe "events" do
            prop "records events in order" \(Marker n1) a1 (Marker n2) a2 -> do
                s <-
                    only . snd =<< runInMemoryTracing do
                        runErrorNoCallStack @AssertionFailed
                            . inErrorSpan @AssertionFailed @_ @() "test-span" Span.Kind.Internal mempty
                            $ do
                                addEventNow n1 a1
                                addEventNow n2 a2
                                pure ()
                fmap (\e -> (e.name, e.attributes)) s.events `shouldBe` [(n1, a1), (n2, a2)]

            it "events go to innermost span" do
                (_, spans) <-
                    runInMemoryTracing
                        . runErrorNoCallStack @AssertionFailed
                        . inErrorSpan @AssertionFailed @_ @() "parent" Span.Kind.Internal mempty
                        . inErrorSpan @AssertionFailed @_ @() "child" Span.Kind.Internal mempty
                        $ do
                            addEventNow "event" mempty
                            pure ()
                length spans `shouldBe` 2
                let child = spans !! 0
                fmap (.name) child.events `shouldBe` ["event"]

    describe "inOkSpan" do
        prop "exports a span" \(Marker name) spanKind a -> do
            s <- only . snd =<< (runInMemoryTracing . inOkSpan name spanKind a . pure) ()
            s.name `shouldBe` name
            s.kind `shouldBe` spanKind
            s.attributes `shouldBe` a
            s.endTime `shouldSatisfy` (>= Just s.startTime)

        it "no parent, unique ids" do
            (_, spans) <- runInMemoryTracing do
                inOkSpan "a" Span.Kind.Internal mempty $ pure ()
                inOkSpan "b" Span.Kind.Internal mempty $ pure ()
                inOkSpan "c" Span.Kind.Internal mempty $ pure ()
            mapM_ ((`shouldBe` Nothing) . (.parentSpanId)) spans
            (spans <&> (.context.spanId)) `shouldSatisfy` allUnique

        describe "nesting" do
            it "parent-child chain" do
                (_, spans) <- runInMemoryTracing do
                    inOkSpan "root" Span.Kind.Internal mempty
                        . inOkSpan "middle" Span.Kind.Internal mempty
                        . inOkSpan "leaf" Span.Kind.Internal mempty
                        $ pure ()
                length spans `shouldBe` 3
                let leaf = spans !! 0
                    middle = spans !! 1
                    root = spans !! 2
                leaf.parentSpanId `shouldBe` Just middle.context.spanId
                middle.parentSpanId `shouldBe` Just root.context.spanId
                root.parentSpanId `shouldBe` Nothing

            it "shared parent" do
                (_, spans) <- runInMemoryTracing do
                    inOkSpan "parent" Span.Kind.Internal mempty do
                        inOkSpan "child-a" Span.Kind.Internal mempty $ pure ()
                        inOkSpan "child-b" Span.Kind.Internal mempty $ pure ()
                length spans `shouldBe` 3
                let childA = spans !! 0
                    childB = spans !! 1
                    parent = spans !! 2
                childA.parentSpanId `shouldBe` Just parent.context.spanId
                childB.parentSpanId `shouldBe` Just parent.context.spanId

            it "restores parent context" do
                (ctx, _) <- runInMemoryTracing do
                    inOkSpan "outer" Span.Kind.Internal mempty do
                        inOkSpan "inner" Span.Kind.Internal mempty $ pure ()
                        currentContext
                ctx `shouldSatisfy` \case
                    Just _ -> True
                    Nothing -> False

        describe "concurrency" do
            it "no loss, unique ids" do
                let n = 500 :: Int
                (_, spans) <-
                    runInMemoryTracing
                        . runConcurrent
                        . replicateConcurrently_ n
                        . inOkSpan "parallel" Span.Kind.Internal mempty
                        $ pure ()
                length spans `shouldBe` n
                (spans <&> (.context.spanId)) `shouldSatisfy` allUnique

            it "concurrent nesting" do
                (_, spans) <-
                    runInMemoryTracing
                        . runConcurrent
                        . inOkSpan "root" Span.Kind.Internal mempty
                        $ concurrently_
                            (inOkSpan "branch-a" Span.Kind.Internal mempty $ pure ())
                            (inOkSpan "branch-b" Span.Kind.Internal mempty $ pure ())
                length spans `shouldBe` 3
                case filter (("root" ==) . (.name)) spans of
                    [root] -> do
                        let branches = filter ((/= "root") . (.name)) spans
                        mapM_ ((`shouldBe` Just root.context.spanId) . (.parentSpanId)) branches
                    _ -> expectationFailure "expected exactly one root span"

        describe "exception handling" do
            it "no span on exception" do
                (_, spans) <-
                    runInMemoryTracing
                        . try @AssertionFailed
                        . inOkSpan "item" Span.Kind.Internal mempty
                        . throwIO
                        $ AssertionFailed "bla"
                length spans `shouldBe` 0

            it "Ok status without exception" do
                (_, spans) <-
                    runInMemoryTracing . inOkSpan "item" Span.Kind.Internal mempty $
                        pure ()
                length spans `shouldBe` 1
                forM_ (spans <&> (.status)) \s -> do
                    s.code `shouldBe` Span.Status.Ok
                    s.message `shouldBe` mempty

        describe "events" do
            prop "records events in order" \(Marker n1) a1 (Marker n2) a2 -> do
                s <-
                    only . snd =<< runInMemoryTracing do
                        inOkSpan "test-span" Span.Kind.Internal mempty do
                            addEventNow n1 a1
                            addEventNow n2 a2
                            pure ()
                fmap (\e -> (e.name, e.attributes)) s.events `shouldBe` [(n1, a1), (n2, a2)]

            it "events go to innermost span" do
                (_, spans) <- runInMemoryTracing do
                    inSpan "parent" Span.Kind.Internal mempty
                        . inOkSpan "child" Span.Kind.Internal mempty
                        $ do
                            addEventNow "event" mempty
                            pure ()
                length spans `shouldBe` 2
                let child = spans !! 0
                fmap (.name) child.events `shouldBe` ["event"]