module OpenTelemetry.Contrib.SpanTraversals (
  alterSpansUpwards,
  IterationInstruction (..),
) where

import Control.Monad.IO.Class
import Data.IORef
import OpenTelemetry.Internal.Trace.Types


data IterationInstruction a = Continue a | Halt


{- | Alter traces upwards from the provided span to the highest available mutable span.

The callback receives the 'ImmutableSpan' (for reading cold fields like parent)
and the current 'SpanHot' (for reading\/modifying mutable fields). It returns an
'IterationInstruction' and the (possibly modified) 'SpanHot'.

Iteration continues upward until a non-mutable span is reached, there are no
more parents, or the callback returns 'Halt'.
-}
alterSpansUpwards :: (MonadIO m) => Span -> st -> (st -> ImmutableSpan -> SpanHot -> (IterationInstruction st, SpanHot)) -> m st
alterSpansUpwards :: forall (m :: * -> *) st.
MonadIO m =>
Span
-> st
-> (st
    -> ImmutableSpan -> SpanHot -> (IterationInstruction st, SpanHot))
-> m st
alterSpansUpwards (Span ImmutableSpan
imm) st
st st
-> ImmutableSpan -> SpanHot -> (IterationInstruction st, SpanHot)
f = IO st -> m st
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO st -> m st) -> IO st -> m st
forall a b. (a -> b) -> a -> b
$ do
  step <- IORef SpanHot
-> (SpanHot -> (SpanHot, IterationInstruction st))
-> IO (IterationInstruction st)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (ImmutableSpan -> IORef SpanHot
spanHot ImmutableSpan
imm) ((SpanHot -> (SpanHot, IterationInstruction st))
 -> IO (IterationInstruction st))
-> (SpanHot -> (SpanHot, IterationInstruction st))
-> IO (IterationInstruction st)
forall a b. (a -> b) -> a -> b
$ \SpanHot
h ->
    let (IterationInstruction st
step, SpanHot
h') = st
-> ImmutableSpan -> SpanHot -> (IterationInstruction st, SpanHot)
f st
st ImmutableSpan
imm SpanHot
h in (SpanHot
h', IterationInstruction st
step)
  case step of
    Continue st
st' -> case ImmutableSpan -> Maybe Span
spanParent ImmutableSpan
imm of
      Maybe Span
Nothing -> st -> IO st
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return st
st'
      Just Span
s -> Span
-> st
-> (st
    -> ImmutableSpan -> SpanHot -> (IterationInstruction st, SpanHot))
-> IO st
forall (m :: * -> *) st.
MonadIO m =>
Span
-> st
-> (st
    -> ImmutableSpan -> SpanHot -> (IterationInstruction st, SpanHot))
-> m st
alterSpansUpwards Span
s st
st' st
-> ImmutableSpan -> SpanHot -> (IterationInstruction st, SpanHot)
f
    IterationInstruction st
Halt -> st -> IO st
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return st
st
alterSpansUpwards (FrozenSpan SpanContext
_) st
st st
-> ImmutableSpan -> SpanHot -> (IterationInstruction st, SpanHot)
_ = st -> m st
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return st
st
alterSpansUpwards (Dropped SpanContext
_) st
st st
-> ImmutableSpan -> SpanHot -> (IterationInstruction st, SpanHot)
_ = st -> m st
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return st
st