{- |
Module      : OpenTelemetry.Contrib.CarryOns
Description : Carry-on attributes that propagate from parent to child spans.
Stability   : experimental

Carry-ons are extra attributes merged into spans completed within a thread's
context, so values can flow down a trace without attaching them to every span.
-}
module OpenTelemetry.Contrib.CarryOns (
  alterCarryOns,
  withCarryOnProcessor,
) where

import Control.Monad.IO.Class
import qualified Data.HashMap.Strict as H
import Data.IORef (atomicModifyIORef')
import Data.Maybe (fromMaybe)
import qualified OpenTelemetry.Attributes as Attributes
import OpenTelemetry.Attributes.Map (AttributeMap)
import OpenTelemetry.Context
import qualified OpenTelemetry.Context as Context
import OpenTelemetry.Context.ThreadLocal
import OpenTelemetry.Internal.Trace.Types
import System.IO.Unsafe (unsafePerformIO)


carryOnKey :: Key AttributeMap
carryOnKey :: Key AttributeMap
carryOnKey = IO (Key AttributeMap) -> Key AttributeMap
forall a. IO a -> a
unsafePerformIO (IO (Key AttributeMap) -> Key AttributeMap)
-> IO (Key AttributeMap) -> Key AttributeMap
forall a b. (a -> b) -> a -> b
$ Text -> IO (Key AttributeMap)
forall (m :: * -> *) a. MonadIO m => Text -> m (Key a)
newKey Text
"carryOn"
{-# NOINLINE carryOnKey #-}


-- | @since 0.4.0.0
alterCarryOns :: (MonadIO m) => (AttributeMap -> AttributeMap) -> m ()
alterCarryOns :: forall (m :: * -> *).
MonadIO m =>
(AttributeMap -> AttributeMap) -> m ()
alterCarryOns AttributeMap -> AttributeMap
f = (Context -> Context) -> m ()
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext ((Context -> Context) -> m ()) -> (Context -> Context) -> m ()
forall a b. (a -> b) -> a -> b
$ \Context
ctxt ->
  Key AttributeMap -> AttributeMap -> Context -> Context
forall a. Key a -> a -> Context -> Context
Context.insert Key AttributeMap
carryOnKey (AttributeMap -> AttributeMap
f (AttributeMap -> AttributeMap) -> AttributeMap -> AttributeMap
forall a b. (a -> b) -> a -> b
$ AttributeMap -> Maybe AttributeMap -> AttributeMap
forall a. a -> Maybe a -> a
fromMaybe AttributeMap
forall a. Monoid a => a
mempty (Maybe AttributeMap -> AttributeMap)
-> Maybe AttributeMap -> AttributeMap
forall a b. (a -> b) -> a -> b
$ Key AttributeMap -> Context -> Maybe AttributeMap
forall a. Key a -> Context -> Maybe a
Context.lookup Key AttributeMap
carryOnKey Context
ctxt) Context
ctxt


{- |
"Carry ons" are extra attributes that are added to every span that is completed for within a thread's context.
This helps us propagate attributes across a trace without having to manually add them to every span.

Be cautious about adding too many additional attributes via carry ons. The attributes are added to every span,
and will be discarded if the span has attributes that exceed the configured attribute limits for the configured
'TracerProvider'.

 @since 0.4.0.0
-}
withCarryOnProcessor :: SpanProcessor -> SpanProcessor
withCarryOnProcessor :: SpanProcessor -> SpanProcessor
withCarryOnProcessor SpanProcessor
p =
  SpanProcessor
    { spanProcessorOnStart :: ImmutableSpan -> Context -> IO ()
spanProcessorOnStart = SpanProcessor -> ImmutableSpan -> Context -> IO ()
spanProcessorOnStart SpanProcessor
p
    , spanProcessorOnEnd :: ImmutableSpan -> IO ()
spanProcessorOnEnd = \ImmutableSpan
imm -> do
        ctxt <- IO Context
forall (m :: * -> *). MonadIO m => m Context
getContext
        let carryOns = AttributeMap -> Maybe AttributeMap -> AttributeMap
forall a. a -> Maybe a -> a
fromMaybe AttributeMap
forall a. Monoid a => a
mempty (Maybe AttributeMap -> AttributeMap)
-> Maybe AttributeMap -> AttributeMap
forall a b. (a -> b) -> a -> b
$ Key AttributeMap -> Context -> Maybe AttributeMap
forall a. Key a -> Context -> Maybe a
Context.lookup Key AttributeMap
carryOnKey Context
ctxt
        if H.null carryOns
          then pure ()
          else do
            atomicModifyIORef' (spanHot imm) $ \SpanHot
h ->
              ( SpanHot
h
                  { hotAttributes =
                      Attributes.addAttributes
                        (tracerProviderAttributeLimits $ tracerProvider $ spanTracer imm)
                        (hotAttributes h)
                        carryOns
                  }
              , ()
              )
        spanProcessorOnEnd p imm
    , spanProcessorShutdown :: IO ShutdownResult
spanProcessorShutdown = SpanProcessor -> IO ShutdownResult
spanProcessorShutdown SpanProcessor
p
    , spanProcessorForceFlush :: IO FlushResult
spanProcessorForceFlush = SpanProcessor -> IO FlushResult
spanProcessorForceFlush SpanProcessor
p
    }