module OpenTelemetry.Contrib.CarryOns (
  alterCarryOns,
  withCarryOnProcessor,
) where

import Control.Monad.IO.Class
import qualified Data.HashMap.Strict as H
import Data.IORef (modifyIORef')
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import OpenTelemetry.Attributes
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 (H.HashMap Text Attribute)
carryOnKey :: Key (HashMap Text Attribute)
carryOnKey = IO (Key (HashMap Text Attribute)) -> Key (HashMap Text Attribute)
forall a. IO a -> a
unsafePerformIO (IO (Key (HashMap Text Attribute)) -> Key (HashMap Text Attribute))
-> IO (Key (HashMap Text Attribute))
-> Key (HashMap Text Attribute)
forall a b. (a -> b) -> a -> b
$ Text -> IO (Key (HashMap Text Attribute))
forall (m :: * -> *) a. MonadIO m => Text -> m (Key a)
newKey Text
"carryOn"
{-# NOINLINE carryOnKey #-}


alterCarryOns :: (MonadIO m) => (H.HashMap Text Attribute -> H.HashMap Text Attribute) -> m ()
alterCarryOns :: forall (m :: * -> *).
MonadIO m =>
(HashMap Text Attribute -> HashMap Text Attribute) -> m ()
alterCarryOns HashMap Text Attribute -> HashMap Text Attribute
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 (HashMap Text Attribute)
-> HashMap Text Attribute -> Context -> Context
forall a. Key a -> a -> Context -> Context
Context.insert Key (HashMap Text Attribute)
carryOnKey (HashMap Text Attribute -> HashMap Text Attribute
f (HashMap Text Attribute -> HashMap Text Attribute)
-> HashMap Text Attribute -> HashMap Text Attribute
forall a b. (a -> b) -> a -> b
$ HashMap Text Attribute
-> Maybe (HashMap Text Attribute) -> HashMap Text Attribute
forall a. a -> Maybe a -> a
fromMaybe HashMap Text Attribute
forall a. Monoid a => a
mempty (Maybe (HashMap Text Attribute) -> HashMap Text Attribute)
-> Maybe (HashMap Text Attribute) -> HashMap Text Attribute
forall a b. (a -> b) -> a -> b
$ Key (HashMap Text Attribute)
-> Context -> Maybe (HashMap Text Attribute)
forall a. Key a -> Context -> Maybe a
Context.lookup Key (HashMap Text Attribute)
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'.
-}
withCarryOnProcessor :: SpanProcessor -> SpanProcessor
withCarryOnProcessor :: SpanProcessor -> SpanProcessor
withCarryOnProcessor SpanProcessor
p =
  SpanProcessor
    { spanProcessorOnStart :: IORef ImmutableSpan -> Context -> IO ()
spanProcessorOnStart = SpanProcessor -> IORef ImmutableSpan -> Context -> IO ()
spanProcessorOnStart SpanProcessor
p
    , spanProcessorOnEnd :: IORef ImmutableSpan -> IO ()
spanProcessorOnEnd = \IORef ImmutableSpan
spanRef -> do
        Context
ctxt <- IO Context
forall (m :: * -> *). MonadIO m => m Context
getContext
        let carryOns :: HashMap Text Attribute
carryOns = HashMap Text Attribute
-> Maybe (HashMap Text Attribute) -> HashMap Text Attribute
forall a. a -> Maybe a -> a
fromMaybe HashMap Text Attribute
forall a. Monoid a => a
mempty (Maybe (HashMap Text Attribute) -> HashMap Text Attribute)
-> Maybe (HashMap Text Attribute) -> HashMap Text Attribute
forall a b. (a -> b) -> a -> b
$ Key (HashMap Text Attribute)
-> Context -> Maybe (HashMap Text Attribute)
forall a. Key a -> Context -> Maybe a
Context.lookup Key (HashMap Text Attribute)
carryOnKey Context
ctxt
        if HashMap Text Attribute -> Bool
forall k v. HashMap k v -> Bool
H.null HashMap Text Attribute
carryOns
          then () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          else do
            -- I doubt we need atomicity at this point. Hopefully people aren't trying to modify the same span after it has ended from multiple threads.
            IORef ImmutableSpan -> (ImmutableSpan -> ImmutableSpan) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef ImmutableSpan
spanRef ((ImmutableSpan -> ImmutableSpan) -> IO ())
-> (ImmutableSpan -> ImmutableSpan) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ImmutableSpan
is ->
              ImmutableSpan
is
                { spanAttributes =
                    OpenTelemetry.Attributes.addAttributes
                      (tracerProviderAttributeLimits $ tracerProvider $ spanTracer is)
                      (spanAttributes is)
                      carryOns
                }
        SpanProcessor -> IORef ImmutableSpan -> IO ()
spanProcessorOnEnd SpanProcessor
p IORef ImmutableSpan
spanRef
    , spanProcessorShutdown :: IO (Async ShutdownResult)
spanProcessorShutdown = SpanProcessor -> IO (Async ShutdownResult)
spanProcessorShutdown SpanProcessor
p
    , spanProcessorForceFlush :: IO ()
spanProcessorForceFlush = SpanProcessor -> IO ()
spanProcessorForceFlush SpanProcessor
p
    }