{-# LANGUAGE UndecidableInstances #-}

module Hasura.Tracing.Monad
  ( TraceT (..),
    runTraceT,
    ignoreTraceT,
  )
where

import Control.Lens
import Control.Monad.Catch (MonadCatch, MonadMask, MonadThrow)
import Control.Monad.Morph
import Control.Monad.Trans.Control
import Data.IORef
import Hasura.Prelude
import Hasura.RQL.Types.Session (UserInfoM (..))
import Hasura.Server.Types (MonadGetPolicies (..))
import Hasura.Tracing.Class
import Hasura.Tracing.Context
import Hasura.Tracing.Reporter
import Hasura.Tracing.Sampling

--------------------------------------------------------------------------------
-- TraceT

-- | TraceT is the standard implementation of 'MonadTrace'. Via a 'Reader', it
-- keeps track of the default policy and reporter to use thoughout the stack, as
-- well as the current trace.
newtype TraceT m a = TraceT (ReaderT (Reporter, Maybe TraceEnv) m a)
  deriving
    ( (forall a b. (a -> b) -> TraceT m a -> TraceT m b)
-> (forall a b. a -> TraceT m b -> TraceT m a)
-> Functor (TraceT m)
forall a b. a -> TraceT m b -> TraceT m a
forall a b. (a -> b) -> TraceT m a -> TraceT m b
forall (m :: * -> *) a b.
Functor m =>
a -> TraceT m b -> TraceT m a
forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> TraceT m a -> TraceT m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> TraceT m a -> TraceT m b
fmap :: forall a b. (a -> b) -> TraceT m a -> TraceT m b
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> TraceT m b -> TraceT m a
<$ :: forall a b. a -> TraceT m b -> TraceT m a
Functor,
      Functor (TraceT m)
Functor (TraceT m)
-> (forall a. a -> TraceT m a)
-> (forall a b. TraceT m (a -> b) -> TraceT m a -> TraceT m b)
-> (forall a b c.
    (a -> b -> c) -> TraceT m a -> TraceT m b -> TraceT m c)
-> (forall a b. TraceT m a -> TraceT m b -> TraceT m b)
-> (forall a b. TraceT m a -> TraceT m b -> TraceT m a)
-> Applicative (TraceT m)
forall a. a -> TraceT m a
forall a b. TraceT m a -> TraceT m b -> TraceT m a
forall a b. TraceT m a -> TraceT m b -> TraceT m b
forall a b. TraceT m (a -> b) -> TraceT m a -> TraceT m b
forall a b c.
(a -> b -> c) -> TraceT m a -> TraceT m b -> TraceT m c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
forall {m :: * -> *}. Applicative m => Functor (TraceT m)
forall (m :: * -> *) a. Applicative m => a -> TraceT m a
forall (m :: * -> *) a b.
Applicative m =>
TraceT m a -> TraceT m b -> TraceT m a
forall (m :: * -> *) a b.
Applicative m =>
TraceT m a -> TraceT m b -> TraceT m b
forall (m :: * -> *) a b.
Applicative m =>
TraceT m (a -> b) -> TraceT m a -> TraceT m b
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> c) -> TraceT m a -> TraceT m b -> TraceT m c
$cpure :: forall (m :: * -> *) a. Applicative m => a -> TraceT m a
pure :: forall a. a -> TraceT m a
$c<*> :: forall (m :: * -> *) a b.
Applicative m =>
TraceT m (a -> b) -> TraceT m a -> TraceT m b
<*> :: forall a b. TraceT m (a -> b) -> TraceT m a -> TraceT m b
$cliftA2 :: forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> c) -> TraceT m a -> TraceT m b -> TraceT m c
liftA2 :: forall a b c.
(a -> b -> c) -> TraceT m a -> TraceT m b -> TraceT m c
$c*> :: forall (m :: * -> *) a b.
Applicative m =>
TraceT m a -> TraceT m b -> TraceT m b
*> :: forall a b. TraceT m a -> TraceT m b -> TraceT m b
$c<* :: forall (m :: * -> *) a b.
Applicative m =>
TraceT m a -> TraceT m b -> TraceT m a
<* :: forall a b. TraceT m a -> TraceT m b -> TraceT m a
Applicative,
      Applicative (TraceT m)
Applicative (TraceT m)
-> (forall a b. TraceT m a -> (a -> TraceT m b) -> TraceT m b)
-> (forall a b. TraceT m a -> TraceT m b -> TraceT m b)
-> (forall a. a -> TraceT m a)
-> Monad (TraceT m)
forall a. a -> TraceT m a
forall a b. TraceT m a -> TraceT m b -> TraceT m b
forall a b. TraceT m a -> (a -> TraceT m b) -> TraceT m b
forall {m :: * -> *}. Monad m => Applicative (TraceT m)
forall (m :: * -> *) a. Monad m => a -> TraceT m a
forall (m :: * -> *) a b.
Monad m =>
TraceT m a -> TraceT m b -> TraceT m b
forall (m :: * -> *) a b.
Monad m =>
TraceT m a -> (a -> TraceT m b) -> TraceT m b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall (m :: * -> *) a b.
Monad m =>
TraceT m a -> (a -> TraceT m b) -> TraceT m b
>>= :: forall a b. TraceT m a -> (a -> TraceT m b) -> TraceT m b
$c>> :: forall (m :: * -> *) a b.
Monad m =>
TraceT m a -> TraceT m b -> TraceT m b
>> :: forall a b. TraceT m a -> TraceT m b -> TraceT m b
$creturn :: forall (m :: * -> *) a. Monad m => a -> TraceT m a
return :: forall a. a -> TraceT m a
Monad,
      Monad (TraceT m)
Monad (TraceT m)
-> (forall a. IO a -> TraceT m a) -> MonadIO (TraceT m)
forall a. IO a -> TraceT m a
forall (m :: * -> *).
Monad m -> (forall a. IO a -> m a) -> MonadIO m
forall {m :: * -> *}. MonadIO m => Monad (TraceT m)
forall (m :: * -> *) a. MonadIO m => IO a -> TraceT m a
$cliftIO :: forall (m :: * -> *) a. MonadIO m => IO a -> TraceT m a
liftIO :: forall a. IO a -> TraceT m a
MonadIO,
      Monad (TraceT m)
Monad (TraceT m)
-> (forall a. (a -> TraceT m a) -> TraceT m a)
-> MonadFix (TraceT m)
forall a. (a -> TraceT m a) -> TraceT m a
forall (m :: * -> *).
Monad m -> (forall a. (a -> m a) -> m a) -> MonadFix m
forall {m :: * -> *}. MonadFix m => Monad (TraceT m)
forall (m :: * -> *) a.
MonadFix m =>
(a -> TraceT m a) -> TraceT m a
$cmfix :: forall (m :: * -> *) a.
MonadFix m =>
(a -> TraceT m a) -> TraceT m a
mfix :: forall a. (a -> TraceT m a) -> TraceT m a
MonadFix,
      MonadCatch (TraceT m)
MonadCatch (TraceT m)
-> (forall b.
    ((forall a. TraceT m a -> TraceT m a) -> TraceT m b) -> TraceT m b)
-> (forall b.
    ((forall a. TraceT m a -> TraceT m a) -> TraceT m b) -> TraceT m b)
-> (forall a b c.
    TraceT m a
    -> (a -> ExitCase b -> TraceT m c)
    -> (a -> TraceT m b)
    -> TraceT m (b, c))
-> MonadMask (TraceT m)
forall b.
((forall a. TraceT m a -> TraceT m a) -> TraceT m b) -> TraceT m b
forall a b c.
TraceT m a
-> (a -> ExitCase b -> TraceT m c)
-> (a -> TraceT m b)
-> TraceT m (b, c)
forall {m :: * -> *}. MonadMask m => MonadCatch (TraceT m)
forall (m :: * -> *) b.
MonadMask m =>
((forall a. TraceT m a -> TraceT m a) -> TraceT m b) -> TraceT m b
forall (m :: * -> *) a b c.
MonadMask m =>
TraceT m a
-> (a -> ExitCase b -> TraceT m c)
-> (a -> TraceT m b)
-> TraceT m (b, c)
forall (m :: * -> *).
MonadCatch m
-> (forall b. ((forall a. m a -> m a) -> m b) -> m b)
-> (forall b. ((forall a. m a -> m a) -> m b) -> m b)
-> (forall a b c.
    m a -> (a -> ExitCase b -> m c) -> (a -> m b) -> m (b, c))
-> MonadMask m
$cmask :: forall (m :: * -> *) b.
MonadMask m =>
((forall a. TraceT m a -> TraceT m a) -> TraceT m b) -> TraceT m b
mask :: forall b.
((forall a. TraceT m a -> TraceT m a) -> TraceT m b) -> TraceT m b
$cuninterruptibleMask :: forall (m :: * -> *) b.
MonadMask m =>
((forall a. TraceT m a -> TraceT m a) -> TraceT m b) -> TraceT m b
uninterruptibleMask :: forall b.
((forall a. TraceT m a -> TraceT m a) -> TraceT m b) -> TraceT m b
$cgeneralBracket :: forall (m :: * -> *) a b c.
MonadMask m =>
TraceT m a
-> (a -> ExitCase b -> TraceT m c)
-> (a -> TraceT m b)
-> TraceT m (b, c)
generalBracket :: forall a b c.
TraceT m a
-> (a -> ExitCase b -> TraceT m c)
-> (a -> TraceT m b)
-> TraceT m (b, c)
MonadMask,
      MonadThrow (TraceT m)
MonadThrow (TraceT m)
-> (forall e a.
    Exception e =>
    TraceT m a -> (e -> TraceT m a) -> TraceT m a)
-> MonadCatch (TraceT m)
forall e a.
Exception e =>
TraceT m a -> (e -> TraceT m a) -> TraceT m a
forall {m :: * -> *}. MonadCatch m => MonadThrow (TraceT m)
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
TraceT m a -> (e -> TraceT m a) -> TraceT m a
forall (m :: * -> *).
MonadThrow m
-> (forall e a. Exception e => m a -> (e -> m a) -> m a)
-> MonadCatch m
$ccatch :: forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
TraceT m a -> (e -> TraceT m a) -> TraceT m a
catch :: forall e a.
Exception e =>
TraceT m a -> (e -> TraceT m a) -> TraceT m a
MonadCatch,
      Monad (TraceT m)
Monad (TraceT m)
-> (forall e a. Exception e => e -> TraceT m a)
-> MonadThrow (TraceT m)
forall e a. Exception e => e -> TraceT m a
forall (m :: * -> *).
Monad m -> (forall e a. Exception e => e -> m a) -> MonadThrow m
forall {m :: * -> *}. MonadThrow m => Monad (TraceT m)
forall (m :: * -> *) e a.
(MonadThrow m, Exception e) =>
e -> TraceT m a
$cthrowM :: forall (m :: * -> *) e a.
(MonadThrow m, Exception e) =>
e -> TraceT m a
throwM :: forall e a. Exception e => e -> TraceT m a
MonadThrow,
      MonadState s,
      MonadError e,
      MonadBase b,
      MonadBaseControl b
    )

-- | Runs the 'TraceT' monad, by providing the default reporter. This does NOT
-- start a trace.
--
-- TODO: we could change this to always start a trace with a default name? This
-- would allow us to guarantee that there is always a current trace, but this
-- might not always be the correct behaviour: in practice, we would end up
-- generating one that spans the entire lifetime of the engine if 'runTraceT'
-- were to be used from 'main'.
runTraceT :: Reporter -> TraceT m a -> m a
runTraceT :: forall (m :: * -> *) a. Reporter -> TraceT m a -> m a
runTraceT Reporter
reporter (TraceT ReaderT (Reporter, Maybe TraceEnv) m a
m) = ReaderT (Reporter, Maybe TraceEnv) m a
-> (Reporter, Maybe TraceEnv) -> m a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (Reporter, Maybe TraceEnv) m a
m (Reporter
reporter, Maybe TraceEnv
forall a. Maybe a
Nothing)

-- | Run the 'TraceT' monad, but without actually tracing anything: no report
-- will be emitted, even if calls to 'newTraceWith' force the trace to be
-- sampled.
ignoreTraceT :: TraceT m a -> m a
ignoreTraceT :: forall (m :: * -> *) a. TraceT m a -> m a
ignoreTraceT = Reporter -> TraceT m a -> m a
forall (m :: * -> *) a. Reporter -> TraceT m a -> m a
runTraceT Reporter
noReporter

instance MonadTrans TraceT where
  lift :: forall (m :: * -> *) a. Monad m => m a -> TraceT m a
lift = ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a
forall (m :: * -> *) a.
ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a
TraceT (ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a)
-> (m a -> ReaderT (Reporter, Maybe TraceEnv) m a)
-> m a
-> TraceT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> ReaderT (Reporter, Maybe TraceEnv) m a
forall (m :: * -> *) a.
Monad m =>
m a -> ReaderT (Reporter, Maybe TraceEnv) m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift

-- | Hides the fact that TraceT is a reader to the rest of the stack.
instance (MonadReader r m) => MonadReader r (TraceT m) where
  ask :: TraceT m r
ask = m r -> TraceT m r
forall (m :: * -> *) a. Monad m => m a -> TraceT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m r
forall r (m :: * -> *). MonadReader r m => m r
ask
  local :: forall a. (r -> r) -> TraceT m a -> TraceT m a
local r -> r
f (TraceT ReaderT (Reporter, Maybe TraceEnv) m a
m) = ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a
forall (m :: * -> *) a.
ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a
TraceT (ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a)
-> ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a
forall a b. (a -> b) -> a -> b
$ (m a -> m a)
-> ReaderT (Reporter, Maybe TraceEnv) m a
-> ReaderT (Reporter, Maybe TraceEnv) m a
forall (m :: * -> *) a (n :: * -> *) b r.
(m a -> n b) -> ReaderT r m a -> ReaderT r n b
mapReaderT ((r -> r) -> m a -> m a
forall a. (r -> r) -> m a -> m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local r -> r
f) ReaderT (Reporter, Maybe TraceEnv) m a
m

instance (MonadIO m, MonadBaseControl IO m) => MonadTrace (TraceT m) where
  newTraceWith :: forall a.
TraceContext -> SamplingPolicy -> Text -> TraceT m a -> TraceT m a
newTraceWith TraceContext
context SamplingPolicy
policy Text
name (TraceT ReaderT (Reporter, Maybe TraceEnv) m a
body) = ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a
forall (m :: * -> *) a.
ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a
TraceT do
    Reporter
reporter <- ((Reporter, Maybe TraceEnv) -> Reporter)
-> ReaderT (Reporter, Maybe TraceEnv) m Reporter
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (Reporter, Maybe TraceEnv) -> Reporter
forall a b. (a, b) -> a
fst
    SamplingDecision
samplingDecision <- SamplingState
-> SamplingPolicy
-> ReaderT (Reporter, Maybe TraceEnv) m SamplingDecision
forall (m :: * -> *).
MonadIO m =>
SamplingState -> SamplingPolicy -> m SamplingDecision
decideSampling (TraceContext -> SamplingState
tcSamplingState TraceContext
context) SamplingPolicy
policy
    IORef TraceMetadata
metadataRef <- IO (IORef TraceMetadata)
-> ReaderT (Reporter, Maybe TraceEnv) m (IORef TraceMetadata)
forall a. IO a -> ReaderT (Reporter, Maybe TraceEnv) m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef TraceMetadata)
 -> ReaderT (Reporter, Maybe TraceEnv) m (IORef TraceMetadata))
-> IO (IORef TraceMetadata)
-> ReaderT (Reporter, Maybe TraceEnv) m (IORef TraceMetadata)
forall a b. (a -> b) -> a -> b
$ TraceMetadata -> IO (IORef TraceMetadata)
forall a. a -> IO (IORef a)
newIORef []
    let report :: ReaderT (Reporter, Maybe TraceEnv) m a
-> ReaderT (Reporter, Maybe TraceEnv) m a
report = case SamplingDecision
samplingDecision of
          SamplingDecision
SampleNever -> ReaderT (Reporter, Maybe TraceEnv) m a
-> ReaderT (Reporter, Maybe TraceEnv) m a
forall a. a -> a
id
          SamplingDecision
SampleAlways -> Reporter
-> forall (m :: * -> *) a.
   (MonadIO m, MonadBaseControl IO m) =>
   TraceContext -> Text -> IO TraceMetadata -> m a -> m a
runReporter Reporter
reporter TraceContext
context Text
name (IORef TraceMetadata -> IO TraceMetadata
forall a. IORef a -> IO a
readIORef IORef TraceMetadata
metadataRef)
        updatedContext :: TraceContext
updatedContext =
          TraceContext
context
            { tcSamplingState :: SamplingState
tcSamplingState = SamplingDecision -> SamplingState -> SamplingState
updateSamplingState SamplingDecision
samplingDecision (TraceContext -> SamplingState
tcSamplingState TraceContext
context)
            }
        traceEnv :: TraceEnv
traceEnv = TraceContext -> IORef TraceMetadata -> SamplingDecision -> TraceEnv
TraceEnv TraceContext
updatedContext IORef TraceMetadata
metadataRef SamplingDecision
samplingDecision
    ReaderT (Reporter, Maybe TraceEnv) m a
-> ReaderT (Reporter, Maybe TraceEnv) m a
report (ReaderT (Reporter, Maybe TraceEnv) m a
 -> ReaderT (Reporter, Maybe TraceEnv) m a)
-> ReaderT (Reporter, Maybe TraceEnv) m a
-> ReaderT (Reporter, Maybe TraceEnv) m a
forall a b. (a -> b) -> a -> b
$ ((Reporter, Maybe TraceEnv) -> (Reporter, Maybe TraceEnv))
-> ReaderT (Reporter, Maybe TraceEnv) m a
-> ReaderT (Reporter, Maybe TraceEnv) m a
forall a.
((Reporter, Maybe TraceEnv) -> (Reporter, Maybe TraceEnv))
-> ReaderT (Reporter, Maybe TraceEnv) m a
-> ReaderT (Reporter, Maybe TraceEnv) m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Maybe TraceEnv -> Identity (Maybe TraceEnv))
-> (Reporter, Maybe TraceEnv)
-> Identity (Reporter, Maybe TraceEnv)
forall s t a b. Field2 s t a b => Lens s t a b
Lens
  (Reporter, Maybe TraceEnv)
  (Reporter, Maybe TraceEnv)
  (Maybe TraceEnv)
  (Maybe TraceEnv)
_2 ((Maybe TraceEnv -> Identity (Maybe TraceEnv))
 -> (Reporter, Maybe TraceEnv)
 -> Identity (Reporter, Maybe TraceEnv))
-> Maybe TraceEnv
-> (Reporter, Maybe TraceEnv)
-> (Reporter, Maybe TraceEnv)
forall s t a b. ASetter s t a b -> b -> s -> t
.~ TraceEnv -> Maybe TraceEnv
forall a. a -> Maybe a
Just TraceEnv
traceEnv) ReaderT (Reporter, Maybe TraceEnv) m a
body

  newSpanWith :: forall a. SpanId -> Text -> TraceT m a -> TraceT m a
newSpanWith SpanId
spanId Text
name (TraceT ReaderT (Reporter, Maybe TraceEnv) m a
body) = ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a
forall (m :: * -> *) a.
ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a
TraceT do
    (Reporter
reporter, Maybe TraceEnv
traceEnv) <- ReaderT (Reporter, Maybe TraceEnv) m (Reporter, Maybe TraceEnv)
forall r (m :: * -> *). MonadReader r m => m r
ask
    case Maybe TraceEnv
traceEnv of
      -- we are not currently in a trace: ignore this span
      Maybe TraceEnv
Nothing -> ReaderT (Reporter, Maybe TraceEnv) m a
body
      Just TraceEnv
env -> case TraceEnv -> SamplingDecision
teSamplingDecision TraceEnv
env of
        -- this trace is not sampled: ignore this span
        SamplingDecision
SampleNever -> ReaderT (Reporter, Maybe TraceEnv) m a
body
        SamplingDecision
SampleAlways -> do
          IORef TraceMetadata
metadataRef <- IO (IORef TraceMetadata)
-> ReaderT (Reporter, Maybe TraceEnv) m (IORef TraceMetadata)
forall a. IO a -> ReaderT (Reporter, Maybe TraceEnv) m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef TraceMetadata)
 -> ReaderT (Reporter, Maybe TraceEnv) m (IORef TraceMetadata))
-> IO (IORef TraceMetadata)
-> ReaderT (Reporter, Maybe TraceEnv) m (IORef TraceMetadata)
forall a b. (a -> b) -> a -> b
$ TraceMetadata -> IO (IORef TraceMetadata)
forall a. a -> IO (IORef a)
newIORef []
          let subContext :: TraceContext
subContext =
                (TraceEnv -> TraceContext
teTraceContext TraceEnv
env)
                  { tcCurrentSpan :: SpanId
tcCurrentSpan = SpanId
spanId,
                    tcCurrentParent :: Maybe SpanId
tcCurrentParent = SpanId -> Maybe SpanId
forall a. a -> Maybe a
Just (TraceContext -> SpanId
tcCurrentSpan (TraceContext -> SpanId) -> TraceContext -> SpanId
forall a b. (a -> b) -> a -> b
$ TraceEnv -> TraceContext
teTraceContext TraceEnv
env)
                  }
              subTraceEnv :: TraceEnv
subTraceEnv =
                TraceEnv
env
                  { teTraceContext :: TraceContext
teTraceContext = TraceContext
subContext,
                    teMetadataRef :: IORef TraceMetadata
teMetadataRef = IORef TraceMetadata
metadataRef
                  }
          Reporter
-> forall (m :: * -> *) a.
   (MonadIO m, MonadBaseControl IO m) =>
   TraceContext -> Text -> IO TraceMetadata -> m a -> m a
runReporter Reporter
reporter TraceContext
subContext Text
name (IORef TraceMetadata -> IO TraceMetadata
forall a. IORef a -> IO a
readIORef IORef TraceMetadata
metadataRef)
            (ReaderT (Reporter, Maybe TraceEnv) m a
 -> ReaderT (Reporter, Maybe TraceEnv) m a)
-> ReaderT (Reporter, Maybe TraceEnv) m a
-> ReaderT (Reporter, Maybe TraceEnv) m a
forall a b. (a -> b) -> a -> b
$ ((Reporter, Maybe TraceEnv) -> (Reporter, Maybe TraceEnv))
-> ReaderT (Reporter, Maybe TraceEnv) m a
-> ReaderT (Reporter, Maybe TraceEnv) m a
forall a.
((Reporter, Maybe TraceEnv) -> (Reporter, Maybe TraceEnv))
-> ReaderT (Reporter, Maybe TraceEnv) m a
-> ReaderT (Reporter, Maybe TraceEnv) m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((Maybe TraceEnv -> Identity (Maybe TraceEnv))
-> (Reporter, Maybe TraceEnv)
-> Identity (Reporter, Maybe TraceEnv)
forall s t a b. Field2 s t a b => Lens s t a b
Lens
  (Reporter, Maybe TraceEnv)
  (Reporter, Maybe TraceEnv)
  (Maybe TraceEnv)
  (Maybe TraceEnv)
_2 ((Maybe TraceEnv -> Identity (Maybe TraceEnv))
 -> (Reporter, Maybe TraceEnv)
 -> Identity (Reporter, Maybe TraceEnv))
-> Maybe TraceEnv
-> (Reporter, Maybe TraceEnv)
-> (Reporter, Maybe TraceEnv)
forall s t a b. ASetter s t a b -> b -> s -> t
.~ TraceEnv -> Maybe TraceEnv
forall a. a -> Maybe a
Just TraceEnv
subTraceEnv) ReaderT (Reporter, Maybe TraceEnv) m a
body

  currentContext :: TraceT m (Maybe TraceContext)
currentContext = ReaderT (Reporter, Maybe TraceEnv) m (Maybe TraceContext)
-> TraceT m (Maybe TraceContext)
forall (m :: * -> *) a.
ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a
TraceT (ReaderT (Reporter, Maybe TraceEnv) m (Maybe TraceContext)
 -> TraceT m (Maybe TraceContext))
-> ReaderT (Reporter, Maybe TraceEnv) m (Maybe TraceContext)
-> TraceT m (Maybe TraceContext)
forall a b. (a -> b) -> a -> b
$ ((Reporter, Maybe TraceEnv) -> Maybe TraceContext)
-> ReaderT (Reporter, Maybe TraceEnv) m (Maybe TraceContext)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (((Reporter, Maybe TraceEnv) -> Maybe TraceContext)
 -> ReaderT (Reporter, Maybe TraceEnv) m (Maybe TraceContext))
-> ((Reporter, Maybe TraceEnv) -> Maybe TraceContext)
-> ReaderT (Reporter, Maybe TraceEnv) m (Maybe TraceContext)
forall a b. (a -> b) -> a -> b
$ (TraceEnv -> TraceContext) -> Maybe TraceEnv -> Maybe TraceContext
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TraceEnv -> TraceContext
teTraceContext (Maybe TraceEnv -> Maybe TraceContext)
-> ((Reporter, Maybe TraceEnv) -> Maybe TraceEnv)
-> (Reporter, Maybe TraceEnv)
-> Maybe TraceContext
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Reporter, Maybe TraceEnv) -> Maybe TraceEnv
forall a b. (a, b) -> b
snd

  attachMetadata :: TraceMetadata -> TraceT m ()
attachMetadata TraceMetadata
metadata = ReaderT (Reporter, Maybe TraceEnv) m () -> TraceT m ()
forall (m :: * -> *) a.
ReaderT (Reporter, Maybe TraceEnv) m a -> TraceT m a
TraceT do
    ((Reporter, Maybe TraceEnv) -> Maybe (IORef TraceMetadata))
-> ReaderT
     (Reporter, Maybe TraceEnv) m (Maybe (IORef TraceMetadata))
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((TraceEnv -> IORef TraceMetadata)
-> Maybe TraceEnv -> Maybe (IORef TraceMetadata)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap TraceEnv -> IORef TraceMetadata
teMetadataRef (Maybe TraceEnv -> Maybe (IORef TraceMetadata))
-> ((Reporter, Maybe TraceEnv) -> Maybe TraceEnv)
-> (Reporter, Maybe TraceEnv)
-> Maybe (IORef TraceMetadata)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Reporter, Maybe TraceEnv) -> Maybe TraceEnv
forall a b. (a, b) -> b
snd) ReaderT (Reporter, Maybe TraceEnv) m (Maybe (IORef TraceMetadata))
-> (Maybe (IORef TraceMetadata)
    -> ReaderT (Reporter, Maybe TraceEnv) m ())
-> ReaderT (Reporter, Maybe TraceEnv) m ()
forall a b.
ReaderT (Reporter, Maybe TraceEnv) m a
-> (a -> ReaderT (Reporter, Maybe TraceEnv) m b)
-> ReaderT (Reporter, Maybe TraceEnv) m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Maybe (IORef TraceMetadata)
Nothing -> () -> ReaderT (Reporter, Maybe TraceEnv) m ()
forall a. a -> ReaderT (Reporter, Maybe TraceEnv) m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Just IORef TraceMetadata
ref -> IO () -> ReaderT (Reporter, Maybe TraceEnv) m ()
forall a. IO a -> ReaderT (Reporter, Maybe TraceEnv) m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ReaderT (Reporter, Maybe TraceEnv) m ())
-> IO () -> ReaderT (Reporter, Maybe TraceEnv) m ()
forall a b. (a -> b) -> a -> b
$ IORef TraceMetadata -> (TraceMetadata -> TraceMetadata) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef TraceMetadata
ref (TraceMetadata
metadata TraceMetadata -> TraceMetadata -> TraceMetadata
forall a. [a] -> [a] -> [a]
++)

instance (UserInfoM m) => UserInfoM (TraceT m) where
  askUserInfo :: TraceT m UserInfo
askUserInfo = m UserInfo -> TraceT m UserInfo
forall (m :: * -> *) a. Monad m => m a -> TraceT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m UserInfo
forall (m :: * -> *). UserInfoM m => m UserInfo
askUserInfo

instance (MonadGetPolicies m) => MonadGetPolicies (TraceT m) where
  runGetApiTimeLimit :: TraceT m (Maybe MaxTime)
runGetApiTimeLimit = m (Maybe MaxTime) -> TraceT m (Maybe MaxTime)
forall (m :: * -> *) a. Monad m => m a -> TraceT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m (Maybe MaxTime)
forall (m :: * -> *). MonadGetPolicies m => m (Maybe MaxTime)
runGetApiTimeLimit
  runGetPrometheusMetricsGranularity :: TraceT m (IO GranularPrometheusMetricsState)
runGetPrometheusMetricsGranularity = m (IO GranularPrometheusMetricsState)
-> TraceT m (IO GranularPrometheusMetricsState)
forall (m :: * -> *) a. Monad m => m a -> TraceT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m (IO GranularPrometheusMetricsState)
forall (m :: * -> *).
MonadGetPolicies m =>
m (IO GranularPrometheusMetricsState)
runGetPrometheusMetricsGranularity

--------------------------------------------------------------------------------
-- Internal

-- | Information about the current trace and span.
data TraceEnv = TraceEnv
  { TraceEnv -> TraceContext
teTraceContext :: TraceContext,
    TraceEnv -> IORef TraceMetadata
teMetadataRef :: IORef TraceMetadata,
    TraceEnv -> SamplingDecision
teSamplingDecision :: SamplingDecision
  }

-- Helper for consistently deciding whether or not to sample a trace based on
-- trace context and sampling policy.
decideSampling :: (MonadIO m) => SamplingState -> SamplingPolicy -> m SamplingDecision
decideSampling :: forall (m :: * -> *).
MonadIO m =>
SamplingState -> SamplingPolicy -> m SamplingDecision
decideSampling SamplingState
samplingState SamplingPolicy
samplingPolicy =
  case SamplingState
samplingState of
    SamplingState
SamplingDefer -> SamplingPolicy -> m SamplingDecision
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO SamplingPolicy
samplingPolicy
    SamplingState
SamplingDeny -> SamplingDecision -> m SamplingDecision
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SamplingDecision
SampleNever
    SamplingState
SamplingAccept -> SamplingDecision -> m SamplingDecision
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SamplingDecision
SampleAlways

-- Helper for consistently updating the sampling state when a sampling decision
-- is made.
updateSamplingState :: SamplingDecision -> SamplingState -> SamplingState
updateSamplingState :: SamplingDecision -> SamplingState -> SamplingState
updateSamplingState SamplingDecision
samplingDecision = \case
  SamplingState
SamplingDefer ->
    case SamplingDecision
samplingDecision of
      SamplingDecision
SampleNever -> SamplingState
SamplingDefer
      SamplingDecision
SampleAlways -> SamplingState
SamplingAccept
  SamplingState
SamplingDeny -> SamplingState
SamplingDeny
  SamplingState
SamplingAccept -> SamplingState
SamplingAccept