module OpenTelemetry.AWSXRay.Propagator
  ( awsXRayContextPropagator
  , awsXRayContextPropagatorOnError
  ) where

import Prelude

import Control.Error.Util (note)
import Network.HTTP.Types.Header (HeaderName, RequestHeaders, ResponseHeaders)
import OpenTelemetry.AWSXRay.TraceInfo
import OpenTelemetry.Context
  (Context, insertBaggage, insertSpan, lookupBaggage, lookupSpan)
import OpenTelemetry.Propagator
import OpenTelemetry.Trace.Core (getSpanContext, wrapSpanContext)

awsXRayContextPropagator :: Propagator Context RequestHeaders ResponseHeaders
awsXRayContextPropagator :: Propagator Context RequestHeaders RequestHeaders
awsXRayContextPropagator = (RequestHeaders -> String -> IO ())
-> Propagator Context RequestHeaders RequestHeaders
awsXRayContextPropagatorOnError forall a b. (a -> b) -> a -> b
$ \RequestHeaders
_ String
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

awsXRayContextPropagatorOnError
  :: (RequestHeaders -> String -> IO ())
  -- ^ Called on failure to find or parse an @X-Amzn-Trace-Id@ header
  -> Propagator Context RequestHeaders ResponseHeaders
awsXRayContextPropagatorOnError :: (RequestHeaders -> String -> IO ())
-> Propagator Context RequestHeaders RequestHeaders
awsXRayContextPropagatorOnError RequestHeaders -> String -> IO ()
onErr = Propagator
  { propagatorNames :: [Text]
propagatorNames = [Text
"awsxray trace context"]
  , extractor :: RequestHeaders -> Context -> IO Context
extractor = \RequestHeaders
hs Context
c -> do
    case ByteString -> Either String TraceInfo
fromXRayHeader forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall a b. a -> Maybe b -> Either a b
note String
"not found" (forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup HeaderName
hAmznTraceId RequestHeaders
hs) of
      Left String
err -> Context
c forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ RequestHeaders -> String -> IO ()
onErr RequestHeaders
hs String
err
      Right TraceInfo {Maybe Baggage
SpanContext
baggage :: TraceInfo -> Maybe Baggage
spanContext :: TraceInfo -> SpanContext
baggage :: Maybe Baggage
spanContext :: SpanContext
..} -> do
        let wrapped :: Span
wrapped = SpanContext -> Span
wrapSpanContext SpanContext
spanContext
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. a -> a
id Baggage -> Context -> Context
insertBaggage Maybe Baggage
baggage forall a b. (a -> b) -> a -> b
$ Span -> Context -> Context
insertSpan Span
wrapped Context
c
  , injector :: Context -> RequestHeaders -> IO RequestHeaders
injector = \Context
c RequestHeaders
hs -> case Context -> Maybe Span
lookupSpan Context
c of
    Maybe Span
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure RequestHeaders
hs
    Just Span
sp -> do
      TraceInfo
info <- SpanContext -> Maybe Baggage -> TraceInfo
TraceInfo forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). MonadIO m => Span -> m SpanContext
getSpanContext Span
sp forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Context -> Maybe Baggage
lookupBaggage Context
c)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (HeaderName
hAmznTraceId, TraceInfo -> ByteString
toXRayHeader TraceInfo
info) forall a. a -> [a] -> [a]
: RequestHeaders
hs
  }

hAmznTraceId :: HeaderName
hAmznTraceId :: HeaderName
hAmznTraceId = HeaderName
"X-Amzn-Trace-Id"