-- | Defines the Tracing API.
--
-- The 'MonadTrace' class defines the "public API" of this component.
module Hasura.Tracing.Class
  ( MonadTrace (..),
    newTrace,
    newSpan,
  )
where

import Control.Monad.Morph
import Control.Monad.Trans.Maybe
import Hasura.Prelude
import Hasura.Tracing.Context
import Hasura.Tracing.Sampling
import Hasura.Tracing.TraceId

--------------------------------------------------------------------------------
-- MonadTrace

class (Monad m) => MonadTrace m where
  -- | Trace the execution of a block of code, attaching a human-readable
  -- name. This starts a new trace and its corresponding root span, to which
  -- subsequent spans will be attached.
  newTraceWith ::
    TraceContext ->
    SamplingPolicy ->
    Text ->
    m a ->
    m a

  -- | Starts a new span within the current trace. No-op if there's no current
  -- trace.
  --
  -- TODO: we could rewrite this to start a new trace if there isn't one, using
  -- the default reporter and policy? This would guarantee that no span is ever
  -- lost, but would also risk reporting undesired spans.
  newSpanWith ::
    SpanId ->
    Text ->
    m a ->
    m a

  -- | Ask for the current tracing context, so that we can provide it to any
  -- downstream services, e.g. in HTTP headers. Returns 'Nothing' if we're not
  -- currently tracing anything.
  currentContext :: m (Maybe TraceContext)

  -- | Log some arbitrary metadata to be attached to the current span, if any.
  attachMetadata :: TraceMetadata -> m ()

instance (MonadTrace m) => MonadTrace (ReaderT r m) where
  newTraceWith :: forall a.
TraceContext
-> SamplingPolicy -> Text -> ReaderT r m a -> ReaderT r m a
newTraceWith TraceContext
c SamplingPolicy
p Text
n = (m a -> m a) -> ReaderT r m a -> ReaderT r m a
forall (m :: * -> *) a (n :: * -> *) b r.
(m a -> n b) -> ReaderT r m a -> ReaderT r n b
mapReaderT (TraceContext -> SamplingPolicy -> Text -> m a -> m a
forall a. TraceContext -> SamplingPolicy -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
TraceContext -> SamplingPolicy -> Text -> m a -> m a
newTraceWith TraceContext
c SamplingPolicy
p Text
n)
  newSpanWith :: forall a. SpanId -> Text -> ReaderT r m a -> ReaderT r m a
newSpanWith SpanId
i Text
n = (m a -> m a) -> ReaderT r m a -> ReaderT r m a
forall (m :: * -> *) a (n :: * -> *) b r.
(m a -> n b) -> ReaderT r m a -> ReaderT r n b
mapReaderT (SpanId -> Text -> m a -> m a
forall a. SpanId -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
SpanId -> Text -> m a -> m a
newSpanWith SpanId
i Text
n)
  currentContext :: ReaderT r m (Maybe TraceContext)
currentContext = m (Maybe TraceContext) -> ReaderT r m (Maybe TraceContext)
forall (m :: * -> *) a. Monad m => m a -> ReaderT r m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m (Maybe TraceContext)
forall (m :: * -> *). MonadTrace m => m (Maybe TraceContext)
currentContext
  attachMetadata :: TraceMetadata -> ReaderT r m ()
attachMetadata = m () -> ReaderT r m ()
forall (m :: * -> *) a. Monad m => m a -> ReaderT r m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> ReaderT r m ())
-> (TraceMetadata -> m ()) -> TraceMetadata -> ReaderT r m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TraceMetadata -> m ()
forall (m :: * -> *). MonadTrace m => TraceMetadata -> m ()
attachMetadata

instance (MonadTrace m) => MonadTrace (StateT e m) where
  newTraceWith :: forall a.
TraceContext
-> SamplingPolicy -> Text -> StateT e m a -> StateT e m a
newTraceWith TraceContext
c SamplingPolicy
p Text
n = (m (a, e) -> m (a, e)) -> StateT e m a -> StateT e m a
forall (m :: * -> *) a s (n :: * -> *) b.
(m (a, s) -> n (b, s)) -> StateT s m a -> StateT s n b
mapStateT (TraceContext -> SamplingPolicy -> Text -> m (a, e) -> m (a, e)
forall a. TraceContext -> SamplingPolicy -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
TraceContext -> SamplingPolicy -> Text -> m a -> m a
newTraceWith TraceContext
c SamplingPolicy
p Text
n)
  newSpanWith :: forall a. SpanId -> Text -> StateT e m a -> StateT e m a
newSpanWith SpanId
i Text
n = (m (a, e) -> m (a, e)) -> StateT e m a -> StateT e m a
forall (m :: * -> *) a s (n :: * -> *) b.
(m (a, s) -> n (b, s)) -> StateT s m a -> StateT s n b
mapStateT (SpanId -> Text -> m (a, e) -> m (a, e)
forall a. SpanId -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
SpanId -> Text -> m a -> m a
newSpanWith SpanId
i Text
n)
  currentContext :: StateT e m (Maybe TraceContext)
currentContext = m (Maybe TraceContext) -> StateT e m (Maybe TraceContext)
forall (m :: * -> *) a. Monad m => m a -> StateT e m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m (Maybe TraceContext)
forall (m :: * -> *). MonadTrace m => m (Maybe TraceContext)
currentContext
  attachMetadata :: TraceMetadata -> StateT e m ()
attachMetadata = m () -> StateT e m ()
forall (m :: * -> *) a. Monad m => m a -> StateT e m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> StateT e m ())
-> (TraceMetadata -> m ()) -> TraceMetadata -> StateT e m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TraceMetadata -> m ()
forall (m :: * -> *). MonadTrace m => TraceMetadata -> m ()
attachMetadata

instance (MonadTrace m) => MonadTrace (ExceptT e m) where
  newTraceWith :: forall a.
TraceContext
-> SamplingPolicy -> Text -> ExceptT e m a -> ExceptT e m a
newTraceWith TraceContext
c SamplingPolicy
p Text
n = (m (Either e a) -> m (Either e a))
-> ExceptT e m a -> ExceptT e m a
forall (m :: * -> *) e a (n :: * -> *) e' b.
(m (Either e a) -> n (Either e' b))
-> ExceptT e m a -> ExceptT e' n b
mapExceptT (TraceContext
-> SamplingPolicy -> Text -> m (Either e a) -> m (Either e a)
forall a. TraceContext -> SamplingPolicy -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
TraceContext -> SamplingPolicy -> Text -> m a -> m a
newTraceWith TraceContext
c SamplingPolicy
p Text
n)
  newSpanWith :: forall a. SpanId -> Text -> ExceptT e m a -> ExceptT e m a
newSpanWith SpanId
i Text
n = (m (Either e a) -> m (Either e a))
-> ExceptT e m a -> ExceptT e m a
forall (m :: * -> *) e a (n :: * -> *) e' b.
(m (Either e a) -> n (Either e' b))
-> ExceptT e m a -> ExceptT e' n b
mapExceptT (SpanId -> Text -> m (Either e a) -> m (Either e a)
forall a. SpanId -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
SpanId -> Text -> m a -> m a
newSpanWith SpanId
i Text
n)
  currentContext :: ExceptT e m (Maybe TraceContext)
currentContext = m (Maybe TraceContext) -> ExceptT e m (Maybe TraceContext)
forall (m :: * -> *) a. Monad m => m a -> ExceptT e m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m (Maybe TraceContext)
forall (m :: * -> *). MonadTrace m => m (Maybe TraceContext)
currentContext
  attachMetadata :: TraceMetadata -> ExceptT e m ()
attachMetadata = m () -> ExceptT e m ()
forall (m :: * -> *) a. Monad m => m a -> ExceptT e m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> ExceptT e m ())
-> (TraceMetadata -> m ()) -> TraceMetadata -> ExceptT e m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TraceMetadata -> m ()
forall (m :: * -> *). MonadTrace m => TraceMetadata -> m ()
attachMetadata

instance (MonadTrace m) => MonadTrace (MaybeT m) where
  newTraceWith :: forall a.
TraceContext -> SamplingPolicy -> Text -> MaybeT m a -> MaybeT m a
newTraceWith TraceContext
c SamplingPolicy
p Text
n = (m (Maybe a) -> m (Maybe a)) -> MaybeT m a -> MaybeT m a
forall (m :: * -> *) a (n :: * -> *) b.
(m (Maybe a) -> n (Maybe b)) -> MaybeT m a -> MaybeT n b
mapMaybeT (TraceContext
-> SamplingPolicy -> Text -> m (Maybe a) -> m (Maybe a)
forall a. TraceContext -> SamplingPolicy -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
TraceContext -> SamplingPolicy -> Text -> m a -> m a
newTraceWith TraceContext
c SamplingPolicy
p Text
n)
  newSpanWith :: forall a. SpanId -> Text -> MaybeT m a -> MaybeT m a
newSpanWith SpanId
i Text
n = (m (Maybe a) -> m (Maybe a)) -> MaybeT m a -> MaybeT m a
forall (m :: * -> *) a (n :: * -> *) b.
(m (Maybe a) -> n (Maybe b)) -> MaybeT m a -> MaybeT n b
mapMaybeT (SpanId -> Text -> m (Maybe a) -> m (Maybe a)
forall a. SpanId -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
SpanId -> Text -> m a -> m a
newSpanWith SpanId
i Text
n)
  currentContext :: MaybeT m (Maybe TraceContext)
currentContext = m (Maybe TraceContext) -> MaybeT m (Maybe TraceContext)
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m (Maybe TraceContext)
forall (m :: * -> *). MonadTrace m => m (Maybe TraceContext)
currentContext
  attachMetadata :: TraceMetadata -> MaybeT m ()
attachMetadata = m () -> MaybeT m ()
forall (m :: * -> *) a. Monad m => m a -> MaybeT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> MaybeT m ())
-> (TraceMetadata -> m ()) -> TraceMetadata -> MaybeT m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TraceMetadata -> m ()
forall (m :: * -> *). MonadTrace m => TraceMetadata -> m ()
attachMetadata

--------------------------------------------------------------------------------
-- Trace helpers

-- | Create a new trace using a randomly-generated context.
newTrace :: (MonadIO m, MonadTrace m) => SamplingPolicy -> Text -> m a -> m a
newTrace :: forall (m :: * -> *) a.
(MonadIO m, MonadTrace m) =>
SamplingPolicy -> Text -> m a -> m a
newTrace SamplingPolicy
policy Text
name m a
body = do
  TraceId
traceId <- m TraceId
forall (m :: * -> *). MonadIO m => m TraceId
randomTraceId
  SpanId
spanId <- m SpanId
forall (m :: * -> *). MonadIO m => m SpanId
randomSpanId
  let context :: TraceContext
context = TraceId -> SpanId -> Maybe SpanId -> SamplingState -> TraceContext
TraceContext TraceId
traceId SpanId
spanId Maybe SpanId
forall a. Maybe a
Nothing SamplingState
SamplingDefer
  TraceContext -> SamplingPolicy -> Text -> m a -> m a
forall a. TraceContext -> SamplingPolicy -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
TraceContext -> SamplingPolicy -> Text -> m a -> m a
newTraceWith TraceContext
context SamplingPolicy
policy Text
name m a
body

-- | Create a new span with a randomly-generated id.
newSpan :: (MonadIO m, MonadTrace m) => Text -> m a -> m a
newSpan :: forall (m :: * -> *) a.
(MonadIO m, MonadTrace m) =>
Text -> m a -> m a
newSpan Text
name m a
body = do
  SpanId
spanId <- m SpanId
forall (m :: * -> *). MonadIO m => m SpanId
randomSpanId
  SpanId -> Text -> m a -> m a
forall a. SpanId -> Text -> m a -> m a
forall (m :: * -> *) a.
MonadTrace m =>
SpanId -> Text -> m a -> m a
newSpanWith SpanId
spanId Text
name m a
body