{-# 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' :: 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 a. m a -> m a) -> m c) -> m c
forall (m :: * -> *) b.
MonadMask m =>
((forall a. m a -> m a) -> m b) -> m b
MonadMask.mask (((forall a. m a -> m a) -> m c) -> m c)
-> ((forall a. m a -> m a) -> m c) -> m c
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> m a
restore -> do
  a
x <- m a
before
  Either SomeException c
res1 <- m c -> m (Either SomeException c)
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MonadMask.try (m c -> m (Either SomeException c))
-> m c -> m (Either SomeException c)
forall a b. (a -> b) -> a -> b
$ m c -> m c
forall a. m a -> m a
restore (m c -> m c) -> m c -> m c
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 <-
        m b -> m (Either SomeException b)
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MonadMask.try (m b -> m (Either SomeException b))
-> m b -> m (Either SomeException b)
forall a b. (a -> b) -> a -> b
$ m b -> m b
forall (m :: * -> *) a. MonadMask m => m a -> m a
MonadMask.uninterruptibleMask_ (m b -> m b) -> m b -> m b
forall a b. (a -> b) -> a -> b
$ Maybe SomeException -> a -> m b
after (SomeException -> Maybe SomeException
forall a. a -> Maybe a
Just SomeException
e1) a
x
      SomeException -> m c
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
MonadMask.throwM SomeException
e1
    Right c
y -> do
      m b -> m b
forall (m :: * -> *) a. MonadMask m => m a -> m a
MonadMask.uninterruptibleMask_ (m b -> m b) -> m b -> m b
forall a b. (a -> b) -> a -> b
$ Maybe SomeException -> a -> m b
after Maybe SomeException
forall a. Maybe a
Nothing a
x
      c -> m c
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 ->
  -- | The name of the span. This may be updated later via 'updateName'
  Text ->
  -- | Additional options for creating the span, such as 'SpanKind',
  -- span links, starting attributes, etc.
  Trace.SpanArguments ->
  -- | 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 ->
  m a
inSpanM :: Tracer -> Text -> SpanArguments -> m a -> m a
inSpanM Tracer
t Text
n SpanArguments
args m a
m = Tracer
-> CallStack -> Text -> SpanArguments -> (Span -> m a) -> m a
forall (m :: * -> *) a.
(MonadMask m, HasCallStack, MonadIO m) =>
Tracer
-> CallStack -> Text -> SpanArguments -> (Span -> m a) -> m a
inSpanM'' Tracer
t CallStack
HasCallStack => CallStack
callStack Text
n SpanArguments
args (m a -> Span -> m a
forall a b. a -> b -> a
const m a
m)

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

inSpanM'' ::
  (MonadMask m, HasCallStack, MonadIO m) =>
  Trace.Tracer ->
  -- | Record the location of the span in the codebase using the provided
  -- callstack for source location info.
  CallStack ->
  -- | The name of the span. This may be updated later via 'updateName'
  Text ->
  Trace.SpanArguments ->
  (Trace.Span -> m a) ->
  m a
inSpanM'' :: Tracer
-> CallStack -> Text -> SpanArguments -> (Span -> m a) -> m a
inSpanM'' Tracer
t CallStack
cs Text
n SpanArguments
args Span -> m a
f = m (Maybe Span, Span)
-> (Maybe SomeException -> (Maybe Span, Span) -> m ())
-> ((Maybe Span, Span) -> m a)
-> m a
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 Maybe SomeException -> (Maybe Span, Span) -> m ()
forall (m :: * -> *) (t :: * -> *).
(Foldable t, MonadIO m) =>
t SomeException -> (Maybe Span, Span) -> m ()
after (Span -> m a
f (Span -> m a)
-> ((Maybe Span, Span) -> Span) -> (Maybe Span, Span) -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Maybe Span, Span) -> Span
forall a b. (a, b) -> b
snd)
  where
    before :: m (Maybe Span, Span)
before = do
      Context
ctx <- m Context
forall (m :: * -> *). MonadIO m => m Context
TraceCore.SpanContext.getContext
      Span
s <- Tracer -> Context -> Text -> SpanArguments -> m Span
forall (m :: * -> *).
MonadIO m =>
Tracer -> Context -> Text -> SpanArguments -> m Span
TraceCore.createSpanWithoutCallStack Tracer
t Context
ctx Text
n SpanArguments
args
      (Context -> Context) -> m ()
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext (Span -> Context -> Context
insertSpan Span
s)
      Span -> m () -> m ()
forall (m :: * -> *). MonadIO m => Span -> m () -> m ()
whenSpanIsRecording Span
s (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        case CallStack -> [([Char], SrcLoc)]
getCallStack CallStack
cs of
          [] -> () -> m ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          ([Char]
fn, SrcLoc
loc) : [([Char], SrcLoc)]
_ -> do
            Span -> [(Text, Attribute)] -> m ()
forall (m :: * -> *).
MonadIO m =>
Span -> [(Text, Attribute)] -> m ()
TraceCore.addAttributes
              Span
s
              [ (Text
"code.function", Text -> Attribute
forall a. ToAttribute a => a -> Attribute
toAttribute (Text -> Attribute) -> Text -> Attribute
forall a b. (a -> b) -> a -> b
$ [Char] -> Text
T.pack [Char]
fn),
                (Text
"code.namespace", Text -> Attribute
forall a. ToAttribute a => a -> Attribute
toAttribute (Text -> Attribute) -> Text -> Attribute
forall a b. (a -> b) -> a -> b
$ [Char] -> Text
T.pack ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ SrcLoc -> [Char]
srcLocModule SrcLoc
loc),
                (Text
"code.filepath", Text -> Attribute
forall a. ToAttribute a => a -> Attribute
toAttribute (Text -> Attribute) -> Text -> Attribute
forall a b. (a -> b) -> a -> b
$ [Char] -> Text
T.pack ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ SrcLoc -> [Char]
srcLocFile SrcLoc
loc),
                (Text
"code.lineno", Int -> Attribute
forall a. ToAttribute a => a -> Attribute
toAttribute (Int -> Attribute) -> Int -> Attribute
forall a b. (a -> b) -> a -> b
$ SrcLoc -> Int
srcLocStartLine SrcLoc
loc),
                (Text
"code.package", Text -> Attribute
forall a. ToAttribute a => a -> Attribute
toAttribute (Text -> Attribute) -> Text -> Attribute
forall a b. (a -> b) -> a -> b
$ [Char] -> Text
T.pack ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ SrcLoc -> [Char]
srcLocPackage SrcLoc
loc)
              ]
      (Maybe Span, Span) -> m (Maybe Span, Span)
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
      t SomeException -> (SomeException -> m ()) -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ t SomeException
e ((SomeException -> m ()) -> m ())
-> (SomeException -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \(MonadMask.SomeException e
inner) -> do
        Span -> SpanStatus -> m ()
forall (m :: * -> *). MonadIO m => Span -> SpanStatus -> m ()
setStatus Span
s (SpanStatus -> m ()) -> SpanStatus -> m ()
forall a b. (a -> b) -> a -> b
$ Text -> SpanStatus
Trace.Error (Text -> SpanStatus) -> Text -> SpanStatus
forall a b. (a -> b) -> a -> b
$ [Char] -> Text
T.pack ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ e -> [Char]
forall e. Exception e => e -> [Char]
MonadMask.displayException e
inner
        Span -> [(Text, Attribute)] -> Maybe Timestamp -> e -> m ()
forall (m :: * -> *) e.
(MonadIO m, Exception e) =>
Span -> [(Text, Attribute)] -> Maybe Timestamp -> e -> m ()
recordException Span
s [] Maybe Timestamp
forall a. Maybe a
Nothing e
inner
      Span -> Maybe Timestamp -> m ()
forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
endSpan Span
s Maybe Timestamp
forall a. Maybe a
Nothing
      (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
ctx ->
        Context -> (Span -> Context) -> Maybe Span -> Context
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Context -> Context
removeSpan Context
ctx) (Span -> Context -> Context
`insertSpan` Context
ctx) Maybe Span
parent