{-# LANGUAGE UndecidableInstances #-}

module Hasura.Tracing
  ( MonadTrace (..),
    TraceT,
    runTraceT,
    runTraceTWith,
    runTraceTWithReporter,
    runTraceTInContext,
    interpTraceT,
    TraceContext (..),
    Reporter (..),
    noReporter,
    HasReporter (..),
    TracingMetadata,
    extractHttpContext,
    tracedHttpRequest,
    injectEventContext,
    extractEventContext,
    word64ToHex,
  )
where

import Control.Lens (over, view, (^?))
import Control.Monad.Catch (MonadCatch, MonadMask, MonadThrow)
import Control.Monad.Morph
import Control.Monad.Trans.Control
import Data.Aeson qualified as J
import Data.Aeson.Lens qualified as JL
import Data.Binary qualified as Bin
import Data.ByteString.Base16 qualified as Hex
import Data.ByteString.Lazy qualified as BL
import Data.String (fromString)
import Hasura.Prelude
import Network.HTTP.Client.Manager (HasHttpManagerM (..))
import Network.HTTP.Client.Transformable qualified as HTTP
import System.Random qualified as Rand
import Web.HttpApiData qualified as HTTP

-- | Any additional human-readable key-value pairs relevant
-- to the execution of a block of code.
type TracingMetadata = [(Text, Text)]

newtype Reporter = Reporter
  { Reporter
-> forall (io :: * -> *) a.
   MonadIO io =>
   TraceContext -> Text -> io (a, TracingMetadata) -> io a
runReporter ::
      forall io a.
      MonadIO io =>
      TraceContext ->
      -- the current trace context
      Text ->
      -- human-readable name for this block of code
      io (a, TracingMetadata) ->
      -- the action whose execution we want to report, returning
      -- any metadata emitted
      io a
  }

noReporter :: Reporter
noReporter :: Reporter
noReporter = (forall (io :: * -> *) a.
 MonadIO io =>
 TraceContext -> Text -> io (a, TracingMetadata) -> io a)
-> Reporter
Reporter \TraceContext
_ Text
_ -> ((a, TracingMetadata) -> a) -> io (a, TracingMetadata) -> io a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, TracingMetadata) -> a
forall a b. (a, b) -> a
fst

-- | A type class for monads which support some way to report execution traces.
--
-- See @instance Tracing.HasReporter (AppM impl)@ in @HasuraPro.App@.
class Monad m => HasReporter m where
  -- | Get the current tracer
  askReporter :: m Reporter
  default askReporter :: m Reporter
  askReporter = Reporter -> m Reporter
forall (f :: * -> *) a. Applicative f => a -> f a
pure Reporter
noReporter

instance HasReporter m => HasReporter (ReaderT r m) where
  askReporter :: ReaderT r m Reporter
askReporter = m Reporter -> ReaderT r m Reporter
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m Reporter
forall (m :: * -> *). HasReporter m => m Reporter
askReporter

instance HasReporter m => HasReporter (ExceptT e m) where
  askReporter :: ExceptT e m Reporter
askReporter = m Reporter -> ExceptT e m Reporter
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m Reporter
forall (m :: * -> *). HasReporter m => m Reporter
askReporter

-- | A trace context records the current active trace,
-- the active span within that trace, and the span's parent,
-- unless the current span is the root.
data TraceContext = TraceContext
  { -- | TODO what is this exactly? The topmost span id?
    TraceContext -> Word64
tcCurrentTrace :: !Word64,
    TraceContext -> Word64
tcCurrentSpan :: !Word64,
    TraceContext -> Maybe Word64
tcCurrentParent :: !(Maybe Word64)
  }

-- | The 'TraceT' monad transformer adds the ability to keep track of
-- the current trace context.
newtype TraceT m a = TraceT {TraceT m a
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
unTraceT :: ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a}
  deriving (a -> TraceT m b -> TraceT m a
(a -> b) -> TraceT m a -> TraceT m b
(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
<$ :: a -> TraceT m b -> TraceT m a
$c<$ :: forall (m :: * -> *) a b.
Functor m =>
a -> TraceT m b -> TraceT m a
fmap :: (a -> b) -> TraceT m a -> TraceT m b
$cfmap :: forall (m :: * -> *) a b.
Functor m =>
(a -> b) -> TraceT m a -> TraceT m b
Functor, Functor (TraceT m)
a -> TraceT m a
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)
TraceT m a -> TraceT m b -> TraceT m b
TraceT m a -> TraceT m b -> TraceT m a
TraceT m (a -> b) -> TraceT m a -> TraceT m b
(a -> b -> c) -> TraceT m a -> TraceT m b -> TraceT m c
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
<* :: TraceT m a -> TraceT m b -> TraceT m a
$c<* :: forall (m :: * -> *) a b.
Applicative m =>
TraceT m a -> TraceT m b -> TraceT m a
*> :: TraceT m a -> TraceT m b -> TraceT m b
$c*> :: forall (m :: * -> *) a b.
Applicative m =>
TraceT m a -> TraceT m b -> TraceT m b
liftA2 :: (a -> b -> c) -> TraceT m a -> TraceT m b -> TraceT m c
$cliftA2 :: forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> c) -> TraceT m a -> TraceT m b -> TraceT m c
<*> :: TraceT m (a -> b) -> TraceT m a -> TraceT m b
$c<*> :: forall (m :: * -> *) a b.
Applicative m =>
TraceT m (a -> b) -> TraceT m a -> TraceT m b
pure :: a -> TraceT m a
$cpure :: forall (m :: * -> *) a. Applicative m => a -> TraceT m a
$cp1Applicative :: forall (m :: * -> *). Applicative m => Functor (TraceT m)
Applicative, Applicative (TraceT m)
a -> TraceT m a
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)
TraceT m a -> (a -> TraceT m b) -> TraceT m b
TraceT m a -> TraceT m b -> TraceT m b
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
return :: a -> TraceT m a
$creturn :: forall (m :: * -> *) a. Monad m => a -> TraceT m a
>> :: TraceT m a -> TraceT m b -> TraceT m b
$c>> :: forall (m :: * -> *) a b.
Monad m =>
TraceT m a -> TraceT m b -> TraceT m b
>>= :: TraceT m a -> (a -> TraceT m b) -> TraceT m b
$c>>= :: forall (m :: * -> *) a b.
Monad m =>
TraceT m a -> (a -> TraceT m b) -> TraceT m b
$cp1Monad :: forall (m :: * -> *). Monad m => Applicative (TraceT m)
Monad, Monad (TraceT m)
Monad (TraceT m)
-> (forall a. IO a -> TraceT m a) -> MonadIO (TraceT m)
IO a -> TraceT m a
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
liftIO :: IO a -> TraceT m a
$cliftIO :: forall (m :: * -> *) a. MonadIO m => IO a -> TraceT m a
$cp1MonadIO :: forall (m :: * -> *). MonadIO m => Monad (TraceT m)
MonadIO, Monad (TraceT m)
Monad (TraceT m)
-> (forall a. (a -> TraceT m a) -> TraceT m a)
-> MonadFix (TraceT m)
(a -> TraceT m a) -> TraceT m a
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
mfix :: (a -> TraceT m a) -> TraceT m a
$cmfix :: forall (m :: * -> *) a.
MonadFix m =>
(a -> TraceT m a) -> TraceT m a
$cp1MonadFix :: forall (m :: * -> *). MonadFix m => Monad (TraceT m)
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)
TraceT m a
-> (a -> ExitCase b -> TraceT m c)
-> (a -> TraceT m b)
-> TraceT m (b, c)
((forall a. TraceT m a -> TraceT m a) -> TraceT m b) -> TraceT m 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)
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
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)
generalBracket :: TraceT m a
-> (a -> ExitCase b -> TraceT m c)
-> (a -> TraceT m b)
-> TraceT m (b, c)
$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)
uninterruptibleMask :: ((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
mask :: ((forall a. TraceT m a -> TraceT m a) -> TraceT m b) -> TraceT m b
$cmask :: forall (m :: * -> *) b.
MonadMask m =>
((forall a. TraceT m a -> TraceT m a) -> TraceT m b) -> TraceT m b
$cp1MonadMask :: forall (m :: * -> *). MonadMask m => MonadCatch (TraceT m)
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)
TraceT m a -> (e -> TraceT m a) -> TraceT m a
forall e a.
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
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
catch :: TraceT m a -> (e -> TraceT m a) -> TraceT m a
$ccatch :: forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
TraceT m a -> (e -> TraceT m a) -> TraceT m a
$cp1MonadCatch :: forall (m :: * -> *). MonadCatch m => MonadThrow (TraceT m)
MonadCatch, Monad (TraceT m)
e -> TraceT m a
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
throwM :: e -> TraceT m a
$cthrowM :: forall (m :: * -> *) e a.
(MonadThrow m, Exception e) =>
e -> TraceT m a
$cp1MonadThrow :: forall (m :: * -> *). MonadThrow m => Monad (TraceT m)
MonadThrow, MonadBase b, MonadBaseControl b)

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

instance MFunctor TraceT where
  hoist :: (forall a. m a -> n a) -> TraceT m b -> TraceT n b
hoist forall a. m a -> n a
f (TraceT ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) b
rwma) = ReaderT (TraceContext, Reporter) (WriterT TracingMetadata n) b
-> TraceT n b
forall (m :: * -> *) a.
ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
TraceT ((forall a.
 WriterT TracingMetadata m a -> WriterT TracingMetadata n a)
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) b
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata n) b
forall k (t :: (* -> *) -> k -> *) (m :: * -> *) (n :: * -> *)
       (b :: k).
(MFunctor t, Monad m) =>
(forall a. m a -> n a) -> t m b -> t n b
hoist ((forall a. m a -> n a)
-> WriterT TracingMetadata m a -> WriterT TracingMetadata n a
forall k (t :: (* -> *) -> k -> *) (m :: * -> *) (n :: * -> *)
       (b :: k).
(MFunctor t, Monad m) =>
(forall a. m a -> n a) -> t m b -> t n b
hoist forall a. m a -> n a
f) ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) b
rwma)

instance MonadError e m => MonadError e (TraceT m) where
  throwError :: e -> TraceT m a
throwError = m a -> TraceT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m a -> TraceT m a) -> (e -> m a) -> e -> TraceT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> m a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError
  catchError :: TraceT m a -> (e -> TraceT m a) -> TraceT m a
catchError (TraceT ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
m) e -> TraceT m a
f = ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
forall (m :: * -> *) a.
ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
TraceT (ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> (e
    -> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a)
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
forall e (m :: * -> *) a.
MonadError e m =>
m a -> (e -> m a) -> m a
catchError ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
m (TraceT m a
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
forall (m :: * -> *) a.
TraceT m a
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
unTraceT (TraceT m a
 -> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a)
-> (e -> TraceT m a)
-> e
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> TraceT m a
f))

instance MonadReader r m => MonadReader r (TraceT m) where
  ask :: TraceT m r
ask = ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) r
-> TraceT m r
forall (m :: * -> *) a.
ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
TraceT (ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) r
 -> TraceT m r)
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) r
-> TraceT m r
forall a b. (a -> b) -> a -> b
$ WriterT TracingMetadata m r
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) r
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift WriterT TracingMetadata m r
forall r (m :: * -> *). MonadReader r m => m r
ask
  local :: (r -> r) -> TraceT m a -> TraceT m a
local r -> r
f TraceT m a
m = ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
forall (m :: * -> *) a.
ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
TraceT (ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
 -> TraceT m a)
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
forall a b. (a -> b) -> a -> b
$ (WriterT TracingMetadata m a -> WriterT TracingMetadata m a)
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
forall (m :: * -> *) a (n :: * -> *) b r.
(m a -> n b) -> ReaderT r m a -> ReaderT r n b
mapReaderT ((r -> r)
-> WriterT TracingMetadata m a -> WriterT TracingMetadata m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local r -> r
f) (TraceT m a
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
forall (m :: * -> *) a.
TraceT m a
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
unTraceT TraceT m a
m)

instance (HasHttpManagerM m) => HasHttpManagerM (TraceT m) where
  askHttpManager :: TraceT m Manager
askHttpManager = m Manager -> TraceT m Manager
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m Manager
forall (m :: * -> *). HasHttpManagerM m => m Manager
askHttpManager

-- | Run an action in the 'TraceT' monad transformer.
-- 'runTraceT' delimits a new trace with its root span, and the arguments
-- specify a name and metadata for that span.
runTraceT :: (HasReporter m, MonadIO m) => Text -> TraceT m a -> m a
runTraceT :: Text -> TraceT m a -> m a
runTraceT Text
name TraceT m a
tma = do
  Reporter
rep <- m Reporter
forall (m :: * -> *). HasReporter m => m Reporter
askReporter
  Reporter -> Text -> TraceT m a -> m a
forall (m :: * -> *) a.
MonadIO m =>
Reporter -> Text -> TraceT m a -> m a
runTraceTWithReporter Reporter
rep Text
name TraceT m a
tma

runTraceTWith :: MonadIO m => TraceContext -> Reporter -> Text -> TraceT m a -> m a
runTraceTWith :: TraceContext -> Reporter -> Text -> TraceT m a -> m a
runTraceTWith TraceContext
ctx Reporter
rep Text
name TraceT m a
tma =
  Reporter -> TraceContext -> Text -> m (a, TracingMetadata) -> m a
Reporter
-> forall (io :: * -> *) a.
   MonadIO io =>
   TraceContext -> Text -> io (a, TracingMetadata) -> io a
runReporter Reporter
rep TraceContext
ctx Text
name (m (a, TracingMetadata) -> m a) -> m (a, TracingMetadata) -> m a
forall a b. (a -> b) -> a -> b
$
    WriterT TracingMetadata m a -> m (a, TracingMetadata)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT TracingMetadata m a -> m (a, TracingMetadata))
-> WriterT TracingMetadata m a -> m (a, TracingMetadata)
forall a b. (a -> b) -> a -> b
$
      ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> (TraceContext, Reporter) -> WriterT TracingMetadata m a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (TraceT m a
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
forall (m :: * -> *) a.
TraceT m a
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
unTraceT TraceT m a
tma) (TraceContext
ctx, Reporter
rep)

-- | Run an action in the 'TraceT' monad transformer in an
-- existing context.
runTraceTInContext :: (MonadIO m, HasReporter m) => TraceContext -> Text -> TraceT m a -> m a
runTraceTInContext :: TraceContext -> Text -> TraceT m a -> m a
runTraceTInContext TraceContext
ctx Text
name TraceT m a
tma = do
  Reporter
rep <- m Reporter
forall (m :: * -> *). HasReporter m => m Reporter
askReporter
  TraceContext -> Reporter -> Text -> TraceT m a -> m a
forall (m :: * -> *) a.
MonadIO m =>
TraceContext -> Reporter -> Text -> TraceT m a -> m a
runTraceTWith TraceContext
ctx Reporter
rep Text
name TraceT m a
tma

-- | Run an action in the 'TraceT' monad transformer in an
-- existing context.
runTraceTWithReporter :: MonadIO m => Reporter -> Text -> TraceT m a -> m a
runTraceTWithReporter :: Reporter -> Text -> TraceT m a -> m a
runTraceTWithReporter Reporter
rep Text
name TraceT m a
tma = do
  TraceContext
ctx <-
    Word64 -> Word64 -> Maybe Word64 -> TraceContext
TraceContext
      (Word64 -> Word64 -> Maybe Word64 -> TraceContext)
-> m Word64 -> m (Word64 -> Maybe Word64 -> TraceContext)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Word64 -> m Word64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Word64
forall a (m :: * -> *). (Random a, MonadIO m) => m a
Rand.randomIO
      m (Word64 -> Maybe Word64 -> TraceContext)
-> m Word64 -> m (Maybe Word64 -> TraceContext)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> IO Word64 -> m Word64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Word64
forall a (m :: * -> *). (Random a, MonadIO m) => m a
Rand.randomIO
      m (Maybe Word64 -> TraceContext)
-> m (Maybe Word64) -> m TraceContext
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Word64 -> m (Maybe Word64)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Word64
forall a. Maybe a
Nothing
  TraceContext -> Reporter -> Text -> TraceT m a -> m a
forall (m :: * -> *) a.
MonadIO m =>
TraceContext -> Reporter -> Text -> TraceT m a -> m a
runTraceTWith TraceContext
ctx Reporter
rep Text
name TraceT m a
tma

-- | Monads which support tracing. 'TraceT' is the standard example.
class Monad m => MonadTrace m where
  -- | Trace the execution of a block of code, attaching a human-readable name.
  trace :: 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.
  currentContext :: m TraceContext

  -- | Ask for the current tracing reporter
  currentReporter :: m Reporter

  -- | Log some metadata to be attached to the current span
  attachMetadata :: TracingMetadata -> m ()

-- | Reinterpret a 'TraceT' action in another 'MonadTrace'.
-- This can be useful when you need to reorganize a monad transformer stack, for
-- example, to embed an action in some monadic computation, while preserving tracing
-- metadata and context.
--
-- For example, we use this function in various places in 'BackendExecute',
-- where we receive an action to execute in some concrete monad transformer stack.
-- See the various implementations of 'runQuery' for examples.
-- Ideally, the input computation's type would be sufficiently polymorphic that
-- we would not need to reorder monads inthe transformer stack. However, the monad
-- transformer stacks must be concrete, because their types are defined by
-- an associated type family 'ExecutionMonad'. Hence, we need to use this function
-- to peel off the outermost 'TraceT' constructor, and embed the computation in some
-- other 'MonadTrace'.
--
-- A second example is related to caching. The 'cacheLookup' function returns an
-- action in a concrete transformer stack, again because we are constrained by the
-- usage of a type class. We need to reinterpret the 'TraceT' component of this
-- concrete stack in some other abstract monad transformer stack, using this function.
--
-- Laws:
--
-- > interpTraceT id (hoist f (TraceT x)) = interpTraceT f (TraceT x)
interpTraceT ::
  MonadTrace n =>
  (m (a, TracingMetadata) -> n (b, TracingMetadata)) ->
  TraceT m a ->
  n b
interpTraceT :: (m (a, TracingMetadata) -> n (b, TracingMetadata))
-> TraceT m a -> n b
interpTraceT m (a, TracingMetadata) -> n (b, TracingMetadata)
f (TraceT ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
rwma) = do
  TraceContext
ctx <- n TraceContext
forall (m :: * -> *). MonadTrace m => m TraceContext
currentContext
  Reporter
rep <- n Reporter
forall (m :: * -> *). MonadTrace m => m Reporter
currentReporter
  (b
b, TracingMetadata
meta) <- m (a, TracingMetadata) -> n (b, TracingMetadata)
f (WriterT TracingMetadata m a -> m (a, TracingMetadata)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> (TraceContext, Reporter) -> WriterT TracingMetadata m a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
rwma (TraceContext
ctx, Reporter
rep)))
  TracingMetadata -> n ()
forall (m :: * -> *). MonadTrace m => TracingMetadata -> m ()
attachMetadata TracingMetadata
meta
  b -> n b
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
b

-- | If the underlying monad can report trace data, then 'TraceT' will
-- collect it and hand it off to that reporter.
instance MonadIO m => MonadTrace (TraceT m) where
  -- Note: this implementation is so awkward because we don't want to give the
  -- derived MonadReader/Writer instances to TraceT
  trace :: Text -> TraceT m a -> TraceT m a
trace Text
name TraceT m a
ma = ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
forall (m :: * -> *) a.
ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
TraceT (ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
 -> TraceT m a)
-> (((TraceContext, Reporter) -> WriterT TracingMetadata m a)
    -> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a)
-> ((TraceContext, Reporter) -> WriterT TracingMetadata m a)
-> TraceT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((TraceContext, Reporter) -> WriterT TracingMetadata m a)
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT (((TraceContext, Reporter) -> WriterT TracingMetadata m a)
 -> TraceT m a)
-> ((TraceContext, Reporter) -> WriterT TracingMetadata m a)
-> TraceT m a
forall a b. (a -> b) -> a -> b
$ \(TraceContext
ctx, Reporter
rep) -> do
    Word64
spanId <- IO Word64 -> WriterT TracingMetadata m Word64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Word64
forall a (m :: * -> *). (Random a, MonadIO m) => m a
Rand.randomIO :: IO Word64)
    let subCtx :: TraceContext
subCtx =
          TraceContext
ctx
            { tcCurrentSpan :: Word64
tcCurrentSpan = Word64
spanId,
              tcCurrentParent :: Maybe Word64
tcCurrentParent = Word64 -> Maybe Word64
forall a. a -> Maybe a
Just (TraceContext -> Word64
tcCurrentSpan TraceContext
ctx)
            }
    m a -> WriterT TracingMetadata m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m a -> WriterT TracingMetadata m a)
-> (WriterT TracingMetadata m a -> m a)
-> WriterT TracingMetadata m a
-> WriterT TracingMetadata m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Reporter -> TraceContext -> Text -> m (a, TracingMetadata) -> m a
Reporter
-> forall (io :: * -> *) a.
   MonadIO io =>
   TraceContext -> Text -> io (a, TracingMetadata) -> io a
runReporter Reporter
rep TraceContext
subCtx Text
name (m (a, TracingMetadata) -> m a)
-> (WriterT TracingMetadata m a -> m (a, TracingMetadata))
-> WriterT TracingMetadata m a
-> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WriterT TracingMetadata m a -> m (a, TracingMetadata)
forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT (WriterT TracingMetadata m a -> WriterT TracingMetadata m a)
-> WriterT TracingMetadata m a -> WriterT TracingMetadata m a
forall a b. (a -> b) -> a -> b
$ ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> (TraceContext, Reporter) -> WriterT TracingMetadata m a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (TraceT m a
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
forall (m :: * -> *) a.
TraceT m a
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
unTraceT TraceT m a
ma) (TraceContext
subCtx, Reporter
rep)

  currentContext :: TraceT m TraceContext
currentContext = ReaderT
  (TraceContext, Reporter) (WriterT TracingMetadata m) TraceContext
-> TraceT m TraceContext
forall (m :: * -> *) a.
ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
TraceT (((TraceContext, Reporter) -> TraceContext)
-> ReaderT
     (TraceContext, Reporter) (WriterT TracingMetadata m) TraceContext
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (TraceContext, Reporter) -> TraceContext
forall a b. (a, b) -> a
fst)

  currentReporter :: TraceT m Reporter
currentReporter = ReaderT
  (TraceContext, Reporter) (WriterT TracingMetadata m) Reporter
-> TraceT m Reporter
forall (m :: * -> *) a.
ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
TraceT (((TraceContext, Reporter) -> Reporter)
-> ReaderT
     (TraceContext, Reporter) (WriterT TracingMetadata m) Reporter
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (TraceContext, Reporter) -> Reporter
forall a b. (a, b) -> b
snd)

  attachMetadata :: TracingMetadata -> TraceT m ()
attachMetadata = ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) ()
-> TraceT m ()
forall (m :: * -> *) a.
ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a
-> TraceT m a
TraceT (ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) ()
 -> TraceT m ())
-> (TracingMetadata
    -> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) ())
-> TracingMetadata
-> TraceT m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TracingMetadata
-> ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell

instance MonadTrace m => MonadTrace (ReaderT r m) where
  trace :: Text -> ReaderT r m a -> ReaderT r m a
trace = (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 ((m a -> m a) -> ReaderT r m a -> ReaderT r m a)
-> (Text -> m a -> m a) -> Text -> ReaderT r m a -> ReaderT r m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> m a -> m a
forall (m :: * -> *) a. MonadTrace m => Text -> m a -> m a
trace
  currentContext :: ReaderT r m TraceContext
currentContext = m TraceContext -> ReaderT r m TraceContext
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m TraceContext
forall (m :: * -> *). MonadTrace m => m TraceContext
currentContext
  currentReporter :: ReaderT r m Reporter
currentReporter = m Reporter -> ReaderT r m Reporter
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m Reporter
forall (m :: * -> *). MonadTrace m => m Reporter
currentReporter
  attachMetadata :: TracingMetadata -> ReaderT r m ()
attachMetadata = m () -> ReaderT r m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> ReaderT r m ())
-> (TracingMetadata -> m ()) -> TracingMetadata -> ReaderT r m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TracingMetadata -> m ()
forall (m :: * -> *). MonadTrace m => TracingMetadata -> m ()
attachMetadata

instance MonadTrace m => MonadTrace (StateT e m) where
  trace :: Text -> StateT e m a -> StateT e m a
trace = (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 ((m (a, e) -> m (a, e)) -> StateT e m a -> StateT e m a)
-> (Text -> m (a, e) -> m (a, e))
-> Text
-> StateT e m a
-> StateT e m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> m (a, e) -> m (a, e)
forall (m :: * -> *) a. MonadTrace m => Text -> m a -> m a
trace
  currentContext :: StateT e m TraceContext
currentContext = m TraceContext -> StateT e m TraceContext
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m TraceContext
forall (m :: * -> *). MonadTrace m => m TraceContext
currentContext
  currentReporter :: StateT e m Reporter
currentReporter = m Reporter -> StateT e m Reporter
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m Reporter
forall (m :: * -> *). MonadTrace m => m Reporter
currentReporter
  attachMetadata :: TracingMetadata -> StateT e m ()
attachMetadata = m () -> StateT e m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> StateT e m ())
-> (TracingMetadata -> m ()) -> TracingMetadata -> StateT e m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TracingMetadata -> m ()
forall (m :: * -> *). MonadTrace m => TracingMetadata -> m ()
attachMetadata

instance MonadTrace m => MonadTrace (ExceptT e m) where
  trace :: Text -> ExceptT e m a -> ExceptT e m a
trace = (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 ((m (Either e a) -> m (Either e a))
 -> ExceptT e m a -> ExceptT e m a)
-> (Text -> m (Either e a) -> m (Either e a))
-> Text
-> ExceptT e m a
-> ExceptT e m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> m (Either e a) -> m (Either e a)
forall (m :: * -> *) a. MonadTrace m => Text -> m a -> m a
trace
  currentContext :: ExceptT e m TraceContext
currentContext = m TraceContext -> ExceptT e m TraceContext
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m TraceContext
forall (m :: * -> *). MonadTrace m => m TraceContext
currentContext
  currentReporter :: ExceptT e m Reporter
currentReporter = m Reporter -> ExceptT e m Reporter
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift m Reporter
forall (m :: * -> *). MonadTrace m => m Reporter
currentReporter
  attachMetadata :: TracingMetadata -> ExceptT e m ()
attachMetadata = m () -> ExceptT e m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> ExceptT e m ())
-> (TracingMetadata -> m ()) -> TracingMetadata -> ExceptT e m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TracingMetadata -> m ()
forall (m :: * -> *). MonadTrace m => TracingMetadata -> m ()
attachMetadata

-- | Encode Word64 to 16 character hex string
word64ToHex :: Word64 -> Text
word64ToHex :: Word64 -> Text
word64ToHex Word64
randNum = ByteString -> Text
bsToTxt (ByteString -> Text) -> ByteString -> Text
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
Hex.encode ByteString
numInBytes
  where
    numInBytes :: ByteString
numInBytes = ByteString -> ByteString
BL.toStrict (Word64 -> ByteString
forall a. Binary a => a -> ByteString
Bin.encode Word64
randNum)

-- | Decode 16 character hex string to Word64
hexToWord64 :: Text -> Maybe Word64
hexToWord64 :: Text -> Maybe Word64
hexToWord64 Text
randText = do
  case ByteString -> Either String ByteString
Hex.decode (ByteString -> Either String ByteString)
-> ByteString -> Either String ByteString
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
txtToBs Text
randText of
    Left String
_ -> Maybe Word64
forall a. Maybe a
Nothing
    Right ByteString
decoded -> Word64 -> Maybe Word64
forall a. a -> Maybe a
Just (Word64 -> Maybe Word64) -> Word64 -> Maybe Word64
forall a b. (a -> b) -> a -> b
$ ByteString -> Word64
forall a. Binary a => ByteString -> a
Bin.decode (ByteString -> Word64) -> ByteString -> Word64
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BL.fromStrict ByteString
decoded

-- | Inject the trace context as a set of HTTP headers.
injectHttpContext :: TraceContext -> [HTTP.Header]
injectHttpContext :: TraceContext -> [Header]
injectHttpContext TraceContext {Maybe Word64
Word64
tcCurrentParent :: Maybe Word64
tcCurrentSpan :: Word64
tcCurrentTrace :: Word64
tcCurrentParent :: TraceContext -> Maybe Word64
tcCurrentSpan :: TraceContext -> Word64
tcCurrentTrace :: TraceContext -> Word64
..} =
  (HeaderName
"X-B3-TraceId", Text -> ByteString
txtToBs (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$ Word64 -> Text
word64ToHex Word64
tcCurrentTrace) Header -> [Header] -> [Header]
forall a. a -> [a] -> [a]
:
  (HeaderName
"X-B3-SpanId", Text -> ByteString
txtToBs (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$ Word64 -> Text
word64ToHex Word64
tcCurrentSpan) Header -> [Header] -> [Header]
forall a. a -> [a] -> [a]
:
    [ (HeaderName
"X-B3-ParentSpanId", Text -> ByteString
txtToBs (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$ Word64 -> Text
word64ToHex Word64
parentID)
      | Word64
parentID <- Maybe Word64 -> [Word64]
forall a. Maybe a -> [a]
maybeToList Maybe Word64
tcCurrentParent
    ]

-- | Extract the trace and parent span headers from a HTTP request
-- and create a new 'TraceContext'. The new context will contain
-- a fresh span ID, and the provided span ID will be assigned as
-- the immediate parent span.
extractHttpContext :: [HTTP.Header] -> IO (Maybe TraceContext)
extractHttpContext :: [Header] -> IO (Maybe TraceContext)
extractHttpContext [Header]
hdrs = do
  Word64
freshSpanId <- IO Word64 -> IO Word64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Word64
forall a (m :: * -> *). (Random a, MonadIO m) => m a
Rand.randomIO
  Maybe TraceContext -> IO (Maybe TraceContext)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe TraceContext -> IO (Maybe TraceContext))
-> Maybe TraceContext -> IO (Maybe TraceContext)
forall a b. (a -> b) -> a -> b
$
    Word64 -> Word64 -> Maybe Word64 -> TraceContext
TraceContext
      (Word64 -> Word64 -> Maybe Word64 -> TraceContext)
-> Maybe Word64 -> Maybe (Word64 -> Maybe Word64 -> TraceContext)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> Maybe Word64
hexToWord64 (Text -> Maybe Word64) -> Maybe Text -> Maybe Word64
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ByteString -> Maybe Text
forall a. FromHttpApiData a => ByteString -> Maybe a
HTTP.parseHeaderMaybe (ByteString -> Maybe Text) -> Maybe ByteString -> Maybe Text
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HeaderName -> [Header] -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup HeaderName
"X-B3-TraceId" [Header]
hdrs)
      Maybe (Word64 -> Maybe Word64 -> TraceContext)
-> Maybe Word64 -> Maybe (Maybe Word64 -> TraceContext)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Word64 -> Maybe Word64
forall (f :: * -> *) a. Applicative f => a -> f a
pure Word64
freshSpanId
      Maybe (Maybe Word64 -> TraceContext)
-> Maybe (Maybe Word64) -> Maybe TraceContext
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Word64 -> Maybe (Maybe Word64)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> Maybe Word64
hexToWord64 (Text -> Maybe Word64) -> Maybe Text -> Maybe Word64
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ByteString -> Maybe Text
forall a. FromHttpApiData a => ByteString -> Maybe a
HTTP.parseHeaderMaybe (ByteString -> Maybe Text) -> Maybe ByteString -> Maybe Text
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HeaderName -> [Header] -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup HeaderName
"X-B3-SpanId" [Header]
hdrs)

-- | Inject the trace context as a JSON value, appropriate for
-- storing in (e.g.) an event trigger payload.
injectEventContext :: TraceContext -> J.Value
injectEventContext :: TraceContext -> Value
injectEventContext TraceContext {Maybe Word64
Word64
tcCurrentParent :: Maybe Word64
tcCurrentSpan :: Word64
tcCurrentTrace :: Word64
tcCurrentParent :: TraceContext -> Maybe Word64
tcCurrentSpan :: TraceContext -> Word64
tcCurrentTrace :: TraceContext -> Word64
..} =
  [Pair] -> Value
J.object
    [ Key
"trace_id" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
J..= Word64 -> Text
word64ToHex Word64
tcCurrentTrace,
      Key
"span_id" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
J..= Word64 -> Text
word64ToHex Word64
tcCurrentSpan
    ]

-- | Extract a trace context from an event trigger payload.
extractEventContext :: J.Value -> IO (Maybe TraceContext)
extractEventContext :: Value -> IO (Maybe TraceContext)
extractEventContext Value
e = do
  Word64
freshSpanId <- IO Word64 -> IO Word64
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO Word64
forall a (m :: * -> *). (Random a, MonadIO m) => m a
Rand.randomIO
  Maybe TraceContext -> IO (Maybe TraceContext)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe TraceContext -> IO (Maybe TraceContext))
-> Maybe TraceContext -> IO (Maybe TraceContext)
forall a b. (a -> b) -> a -> b
$
    Word64 -> Word64 -> Maybe Word64 -> TraceContext
TraceContext
      (Word64 -> Word64 -> Maybe Word64 -> TraceContext)
-> Maybe Word64 -> Maybe (Word64 -> Maybe Word64 -> TraceContext)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> Maybe Word64
hexToWord64 (Text -> Maybe Word64) -> Maybe Text -> Maybe Word64
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Value
e Value -> Getting (First Text) Value Text -> Maybe Text
forall s a. s -> Getting (First a) s a -> Maybe a
^? Text -> Traversal' Value Value
forall t. AsValue t => Text -> Traversal' t Value
JL.key Text
"trace_context" ((Value -> Const (First Text) Value)
 -> Value -> Const (First Text) Value)
-> Getting (First Text) Value Text
-> Getting (First Text) Value Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Traversal' Value Value
forall t. AsValue t => Text -> Traversal' t Value
JL.key Text
"trace_id" ((Value -> Const (First Text) Value)
 -> Value -> Const (First Text) Value)
-> Getting (First Text) Value Text
-> Getting (First Text) Value Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (First Text) Value Text
forall t. AsPrimitive t => Prism' t Text
JL._String)
      Maybe (Word64 -> Maybe Word64 -> TraceContext)
-> Maybe Word64 -> Maybe (Maybe Word64 -> TraceContext)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Word64 -> Maybe Word64
forall (f :: * -> *) a. Applicative f => a -> f a
pure Word64
freshSpanId
      Maybe (Maybe Word64 -> TraceContext)
-> Maybe (Maybe Word64) -> Maybe TraceContext
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Word64 -> Maybe (Maybe Word64)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> Maybe Word64
hexToWord64 (Text -> Maybe Word64) -> Maybe Text -> Maybe Word64
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Value
e Value -> Getting (First Text) Value Text -> Maybe Text
forall s a. s -> Getting (First a) s a -> Maybe a
^? Text -> Traversal' Value Value
forall t. AsValue t => Text -> Traversal' t Value
JL.key Text
"trace_context" ((Value -> Const (First Text) Value)
 -> Value -> Const (First Text) Value)
-> Getting (First Text) Value Text
-> Getting (First Text) Value Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Traversal' Value Value
forall t. AsValue t => Text -> Traversal' t Value
JL.key Text
"span_id" ((Value -> Const (First Text) Value)
 -> Value -> Const (First Text) Value)
-> Getting (First Text) Value Text
-> Getting (First Text) Value Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Getting (First Text) Value Text
forall t. AsPrimitive t => Prism' t Text
JL._String)

-- | Perform HTTP request which supports Trace headers using a
-- HTTP.Request value
--
-- TODO REFACTOR:
--   - inline 'HTTP.performRequest' so that we can be sure a trace is always logged
--   - Inline 'try' here since we always use that at call sites
tracedHttpRequest ::
  MonadTrace m =>
  -- | http request that needs to be made
  HTTP.Request ->
  -- | a function that takes the traced request and executes it
  (HTTP.Request -> m a) ->
  m a
tracedHttpRequest :: Request -> (Request -> m a) -> m a
tracedHttpRequest Request
req Request -> m a
f = do
  let method :: Text
method = ByteString -> Text
bsToTxt (Getting ByteString Request ByteString -> Request -> ByteString
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting ByteString Request ByteString
Lens' Request ByteString
HTTP.method Request
req)
      uri :: Text
uri = Getting Text Request Text -> Request -> Text
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting Text Request Text
Lens' Request Text
HTTP.url Request
req
  Text -> m a -> m a
forall (m :: * -> *) a. MonadTrace m => Text -> m a -> m a
trace (Text
method Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
uri) do
    let reqBytes :: Int64
reqBytes = Request -> Int64
HTTP.getReqSize Request
req
    TracingMetadata -> m ()
forall (m :: * -> *). MonadTrace m => TracingMetadata -> m ()
attachMetadata [(Text
"request_body_bytes", String -> Text
forall a. IsString a => String -> a
fromString (Int64 -> String
forall a. Show a => a -> String
show Int64
reqBytes))]
    TraceContext
ctx <- m TraceContext
forall (m :: * -> *). MonadTrace m => m TraceContext
currentContext
    Request -> m a
f (Request -> m a) -> Request -> m a
forall a b. (a -> b) -> a -> b
$ ASetter Request Request [Header] [Header]
-> ([Header] -> [Header]) -> Request -> Request
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter Request Request [Header] [Header]
Lens' Request [Header]
HTTP.headers (TraceContext -> [Header]
injectHttpContext TraceContext
ctx [Header] -> [Header] -> [Header]
forall a. Semigroup a => a -> a -> a
<>) Request
req