{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module OpenTelemetry.Utils.Exceptions (inSpanM, inSpanM', inSpanM'') where

import Control.Monad (forM_)
import Control.Monad.Catch (MonadMask, SomeException)
import qualified Control.Monad.Catch as MonadMask
import Control.Monad.IO.Class (MonadIO)
import Data.Text (Text)
import qualified Data.Text as T
import GHC.Exception (SrcLoc (..), getCallStack)
import GHC.Stack (CallStack, callStack)
import GHC.Stack.Types (HasCallStack)
import OpenTelemetry.Context (insertSpan, lookupSpan, removeSpan)
import OpenTelemetry.Context.ThreadLocal (adjustContext)
import qualified OpenTelemetry.Context.ThreadLocal as TraceCore.SpanContext
import qualified OpenTelemetry.Trace as Trace
import OpenTelemetry.Trace.Core (ToAttribute (..), endSpan, recordException, setStatus, whenSpanIsRecording)
import qualified OpenTelemetry.Trace.Core as TraceCore


bracketError' :: (MonadMask m) => m a -> (Maybe SomeException -> a -> m b) -> (a -> m c) -> m c
bracketError' :: forall (m :: * -> *) a b c.
MonadMask m =>
m a -> (Maybe SomeException -> a -> m b) -> (a -> m c) -> m c
bracketError' m a
before Maybe SomeException -> a -> m b
after a -> m c
thing = forall (m :: * -> *) b.
MonadMask m =>
((forall a. m a -> m a) -> m b) -> m b
MonadMask.mask forall a b. (a -> b) -> a -> b
$ \forall a. m a -> m a
restore -> do
  a
x <- m a
before
  Either SomeException c
res1 <- forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MonadMask.try forall a b. (a -> b) -> a -> b
$ forall a. m a -> m a
restore forall a b. (a -> b) -> a -> b
$ a -> m c
thing a
x
  case Either SomeException c
res1 of
    Left (SomeException
e1 :: SomeException) -> do
      -- explicitly ignore exceptions from after. We know that
      -- no async exceptions were thrown there, so therefore
      -- the stronger exception must come from thing
      --
      -- https://github.com/fpco/safe-exceptions/issues/2
      Either SomeException b
_ :: Either SomeException b <-
        forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MonadMask.try forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadMask m => m a -> m a
MonadMask.uninterruptibleMask_ forall a b. (a -> b) -> a -> b
$ Maybe SomeException -> a -> m b
after (forall a. a -> Maybe a
Just SomeException
e1) a
x
      forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
MonadMask.throwM SomeException
e1
    Right c
y -> do
      forall (m :: * -> *) a. MonadMask m => m a -> m a
MonadMask.uninterruptibleMask_ forall a b. (a -> b) -> a -> b
$ Maybe SomeException -> a -> m b
after forall a. Maybe a
Nothing a
x
      forall (m :: * -> *) a. Monad m => a -> m a
return c
y


-- | The simplest function for annotating code with trace information.
inSpanM
  :: (MonadIO m, MonadMask m, HasCallStack)
  => Trace.Tracer
  -> Text
  -- ^ The name of the span. This may be updated later via 'updateName'
  -> Trace.SpanArguments
  -- ^ Additional options for creating the span, such as 'SpanKind',
  -- span links, starting attributes, etc.
  -> m a
  -- ^ The action to perform. 'inSpan' will record the time spent on the
  -- action without forcing strict evaluation of the result. Any uncaught
  -- exceptions will be recorded and rethrown.
  -> m a
inSpanM :: forall (m :: * -> *) a.
(MonadIO m, MonadMask m, HasCallStack) =>
Tracer -> Text -> SpanArguments -> m a -> m a
inSpanM Tracer
t Text
n SpanArguments
args m a
m = forall (m :: * -> *) a.
(MonadMask m, HasCallStack, MonadIO m) =>
Tracer
-> CallStack -> Text -> SpanArguments -> (Span -> m a) -> m a
inSpanM'' Tracer
t HasCallStack => CallStack
callStack Text
n SpanArguments
args (forall a b. a -> b -> a
const m a
m)


inSpanM'
  :: (MonadIO m, MonadMask m, HasCallStack)
  => Trace.Tracer
  -> Text
  -- ^ The name of the span. This may be updated later via 'updateName'
  -> Trace.SpanArguments
  -> (Trace.Span -> m a)
  -> m a
inSpanM' :: forall (m :: * -> *) a.
(MonadIO m, MonadMask m, HasCallStack) =>
Tracer -> Text -> SpanArguments -> (Span -> m a) -> m a
inSpanM' Tracer
t = forall (m :: * -> *) a.
(MonadMask m, HasCallStack, MonadIO m) =>
Tracer
-> CallStack -> Text -> SpanArguments -> (Span -> m a) -> m a
inSpanM'' Tracer
t HasCallStack => CallStack
callStack


inSpanM''
  :: (MonadMask m, HasCallStack, MonadIO m)
  => Trace.Tracer
  -> CallStack
  -- ^ Record the location of the span in the codebase using the provided
  -- callstack for source location info.
  -> Text
  -- ^ The name of the span. This may be updated later via 'updateName'
  -> Trace.SpanArguments
  -> (Trace.Span -> m a)
  -> m a
inSpanM'' :: forall (m :: * -> *) a.
(MonadMask m, HasCallStack, MonadIO m) =>
Tracer
-> CallStack -> Text -> SpanArguments -> (Span -> m a) -> m a
inSpanM'' Tracer
t CallStack
cs Text
n SpanArguments
args Span -> m a
f = forall (m :: * -> *) a b c.
MonadMask m =>
m a -> (Maybe SomeException -> a -> m b) -> (a -> m c) -> m c
bracketError' m (Maybe Span, Span)
before forall {m :: * -> *} {t :: * -> *}.
(Foldable t, MonadIO m) =>
t SomeException -> (Maybe Span, Span) -> m ()
after (Span -> m a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd)
  where
    before :: m (Maybe Span, Span)
before = do
      Context
ctx <- forall (m :: * -> *). MonadIO m => m Context
TraceCore.SpanContext.getContext
      Span
s <- forall (m :: * -> *).
MonadIO m =>
Tracer -> Context -> Text -> SpanArguments -> m Span
TraceCore.createSpanWithoutCallStack Tracer
t Context
ctx Text
n SpanArguments
args
      forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext (Span -> Context -> Context
insertSpan Span
s)
      forall (m :: * -> *). MonadIO m => Span -> m () -> m ()
whenSpanIsRecording Span
s forall a b. (a -> b) -> a -> b
$ do
        case CallStack -> [([Char], SrcLoc)]
getCallStack CallStack
cs of
          [] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          ([Char]
fn, SrcLoc
loc) : [([Char], SrcLoc)]
_ -> do
            forall (m :: * -> *).
MonadIO m =>
Span -> HashMap Text Attribute -> m ()
TraceCore.addAttributes
              Span
s
              [ (Text
"code.function", forall a. ToAttribute a => a -> Attribute
toAttribute forall a b. (a -> b) -> a -> b
$ [Char] -> Text
T.pack [Char]
fn)
              , (Text
"code.namespace", forall a. ToAttribute a => a -> Attribute
toAttribute forall a b. (a -> b) -> a -> b
$ [Char] -> Text
T.pack forall a b. (a -> b) -> a -> b
$ SrcLoc -> [Char]
srcLocModule SrcLoc
loc)
              , (Text
"code.filepath", forall a. ToAttribute a => a -> Attribute
toAttribute forall a b. (a -> b) -> a -> b
$ [Char] -> Text
T.pack forall a b. (a -> b) -> a -> b
$ SrcLoc -> [Char]
srcLocFile SrcLoc
loc)
              , (Text
"code.lineno", forall a. ToAttribute a => a -> Attribute
toAttribute forall a b. (a -> b) -> a -> b
$ SrcLoc -> Int
srcLocStartLine SrcLoc
loc)
              , (Text
"code.package", forall a. ToAttribute a => a -> Attribute
toAttribute forall a b. (a -> b) -> a -> b
$ [Char] -> Text
T.pack forall a b. (a -> b) -> a -> b
$ SrcLoc -> [Char]
srcLocPackage SrcLoc
loc)
              ]
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Context -> Maybe Span
lookupSpan Context
ctx, Span
s)

    after :: t SomeException -> (Maybe Span, Span) -> m ()
after t SomeException
e (Maybe Span
parent, Span
s) = do
      forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ t SomeException
e forall a b. (a -> b) -> a -> b
$ \(MonadMask.SomeException e
inner) -> do
        forall (m :: * -> *). MonadIO m => Span -> SpanStatus -> m ()
setStatus Span
s forall a b. (a -> b) -> a -> b
$ Text -> SpanStatus
Trace.Error forall a b. (a -> b) -> a -> b
$ [Char] -> Text
T.pack forall a b. (a -> b) -> a -> b
$ forall e. Exception e => e -> [Char]
MonadMask.displayException e
inner
        forall (m :: * -> *) e.
(MonadIO m, Exception e) =>
Span -> HashMap Text Attribute -> Maybe Timestamp -> e -> m ()
recordException Span
s [(Text
"exception.escaped", forall a. ToAttribute a => a -> Attribute
toAttribute Bool
True)] forall a. Maybe a
Nothing e
inner
      forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
endSpan Span
s forall a. Maybe a
Nothing
      forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext forall a b. (a -> b) -> a -> b
$ \Context
ctx ->
        forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Context -> Context
removeSpan Context
ctx) (Span -> Context -> Context
`insertSpan` Context
ctx) Maybe Span
parent