Propagator.hs

module Effectful.OpenTelemetry.Tracing.Propagator
    ( -- * Propagator
      Propagator (..)

      -- * HTTP header propagators
    , w3cTraceContext
    )
where

import Control.Applicative ((<|>))
import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.List qualified as List
import Data.Text (Text)
import Effectful
import Effectful.OpenTelemetry.Tracing.Span.Context qualified as Span (Context)
import Effectful.OpenTelemetry.Tracing.Span.Context qualified as Span.Context
import Effectful.OpenTelemetry.Tracing.Trace.State qualified as Trace.State
import Network.HTTP.Types (Header, HeaderName, RequestHeaders, ResponseHeaders)
import Prelude

-- | Propagates context information across process boundaries and
-- defines the restrictions imposed by a specific transport.
--
-- [@es@] An arbitrary effect stack.
-- [@context@] A cross-cutting concern context, for example a 'Span.Context'.
-- [@inboundCarrier@] Medium used by the propagator to read values from.
-- [@outboundCarrier@] Medium used by the propagator to write values to.
--
-- See <https://opentelemetry.io/docs/specs/otel/context/api-propagators/ the OpenTelemetry spec>.
data Propagator es context inboundCarrier outboundCarrier = Propagator
    { propagatorNames :: [Text]
    -- ^ The header (or field) names this propagator is responsible for.
    , extract :: inboundCarrier -> context -> Eff es context
    -- ^ Read propagation fields out of the inbound carrier, returning a
    -- context derived from the one passed in. If nothing valid can be parsed,
    -- the original context is returned unchanged.
    , inject :: context -> outboundCarrier -> Eff es outboundCarrier
    -- ^ Write the context's propagation fields into the outbound carrier.
    }

-- | The <https://www.w3.org/TR/trace-context/ W3C Trace Context> propagator,
-- carrying the @traceparent@ HTTP header.
--
-- See <https://www.w3.org/TR/trace-context/#traceparent-header the traceparent header specification>.
w3cTraceContext :: Propagator es (Maybe Span.Context) RequestHeaders ResponseHeaders
w3cTraceContext =
    Propagator
        { propagatorNames = ["traceparent", "tracestate"]
        , extract = \headers context -> pure . (<|> context) $ do
            context' <- Span.Context.fromTraceparent =<< List.lookup hTraceparent headers
            pure
                context'
                    { Span.Context.traceState =
                        foldMap (Trace.State.fromByteString . snd)
                            . filter ((hTracestate ==) . fst)
                            $ headers
                    }
        , inject = \context headers ->
            pure $ case context of
                Nothing -> headers
                Just ctx ->
                    addHeader hTracestate (Trace.State.toByteString ctx.traceState)
                        . setHeader hTraceparent (Span.Context.toTraceparent ctx)
                        $ headers
        }
  where
    setHeader :: HeaderName -> ByteString -> [Header] -> [Header]
    setHeader name value = addHeader name value . filter ((name /=) . fst)

    addHeader :: HeaderName -> ByteString -> [Header] -> [Header]
    addHeader name value
        | ByteString.null value = id
        | otherwise = ((name, value) :)

hTraceparent, hTracestate :: HeaderName
hTraceparent = "traceparent"
hTracestate = "tracestate"