Context.hs

module Effectful.OpenTelemetry.Tracing.Span.Context where

import Control.Monad ((>=>))
import Data.Aeson qualified as Aeson
import Data.ByteString (ByteString)
import Data.ByteString.Builder qualified as Builder
import Data.ByteString.Builder.Extra qualified as Builder
import Data.ByteString.Lazy qualified as LazyByteString
import Data.Either.Extra (eitherToMaybe)
import Data.List qualified as List
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text
import Effectful.OpenTelemetry.Tracing.Span.ID qualified as Span (ID)
import Effectful.OpenTelemetry.Tracing.Span.ID qualified as Span.ID
import Effectful.OpenTelemetry.Tracing.Trace.Flags qualified as Trace (Flags)
import Effectful.OpenTelemetry.Tracing.Trace.Flags qualified as Trace.Flags
import Effectful.OpenTelemetry.Tracing.Trace.ID qualified as Trace (ID)
import Effectful.OpenTelemetry.Tracing.Trace.ID qualified as Trace.ID
import Effectful.OpenTelemetry.Tracing.Trace.State qualified as Trace (State)
import GHC.Generics (Generic)
import Prelude

-- | A propagation mechanism that carries execution-scoped values across API boundaries.
--
-- See the <https://opentelemetry.io/docs/concepts/signals/traces/#span-context OpenTelemetry spec>.
data Context = Context
    { traceId :: Trace.ID
    , spanId :: Span.ID
    , traceFlags :: Trace.Flags
    , traceState :: Trace.State
    }
    deriving stock (Generic, Eq, Show)

-- | Create a new 'Context' with an optional parent and a new span 'Span.ID'.
-- If a parent is given, the new 'Context' inherits the parent's trace 'Trace.ID', 'Trace.Flags',
-- and 'Trace.State'.
-- Otherwise it gets a fresh trace 'Trace.ID', the 'Trace.Flags.sampled' and 'Trace.Flags.random'
-- trace 'Trace.Flags', and empty trace 'Trace.State'.
new :: Maybe Context -> IO Context
new parent = do
    traceId <- maybe Trace.ID.new (pure . traceId) parent
    spanId <- Span.ID.new
    pure
        Context
            { traceFlags =
                let flags =
                        maybe
                            mempty
                                { Trace.Flags.sampled = True
                                , Trace.Flags.random = True
                                }
                            traceFlags
                            parent
                 in flags{Trace.Flags.remote = Trace.Flags.IsNotRemote}
            , traceState = maybe mempty traceState parent
            , ..
            }

-- | Decode a 'Context' with an empty 'Trace.State' from a @traceparent@ header.
fromTraceparent :: ByteString -> Maybe Context
fromTraceparent =
    eitherToMaybe . Text.decodeUtf8' >=> \t ->
        case Aeson.String <$> Text.splitOn "-" t of
            [ "00"
                , Aeson.fromJSON -> Aeson.Success traceId
                , Aeson.fromJSON -> Aeson.Success spanId
                , Aeson.fromJSON -> Aeson.Success traceFlags
                ] ->
                    Just
                        Context
                            { traceState = mempty
                            , traceFlags = traceFlags{Trace.Flags.remote = Trace.Flags.IsRemote}
                            , ..
                            }
            _ -> Nothing

-- | Encode a 'Context' as a @traceparent@ header.
-- Drops the 'traceState'.
toTraceparent :: Context -> ByteString
toTraceparent Context{..} =
    LazyByteString.toStrict
        . Builder.toLazyByteStringWith (Builder.untrimmedStrategy 64 64) mempty
        . mconcat
        . List.intersperse "-"
        $ [ "00"
          , Builder.byteStringHex $ Trace.ID.toBytes traceId
          , Builder.byteStringHex $ Span.ID.toBytes spanId
          , Builder.word8HexFixed $ Trace.Flags.toWord8 traceFlags
          ]