graphql-engine/server/src-lib/Hasura/Tracing.hs
Antoine Leblanc dd1192ca2c server: cleanup of language extensions [gardening]
GitOrigin-RevId: d862c724715cb8a4c2f37d2e0e525d12c46b18eb
2021-04-16 06:55:56 +00:00

287 lines
9.8 KiB
Haskell

{-# LANGUAGE UndecidableInstances #-}
module Hasura.Tracing
( MonadTrace(..)
, TraceT
, runTraceT
, runTraceTWith
, runTraceTWithReporter
, runTraceTInContext
, interpTraceT
, TraceContext(..)
, Reporter(..)
, noReporter
, HasReporter(..)
, TracingMetadata
, extractHttpContext
, tracedHttpRequest
, injectEventContext
, extractEventContext
) where
import Hasura.Prelude
import qualified Data.Aeson as J
import qualified Data.Aeson.Lens as JL
import qualified Data.Binary as Bin
import qualified Data.ByteString as BS
import qualified Data.ByteString.Base16 as Hex
import qualified Data.ByteString.Lazy as BL
import qualified Network.HTTP.Client.Extended as HTTP
import qualified Network.HTTP.Types.Header as HTTP
import qualified System.Random as Rand
import qualified Web.HttpApiData as HTTP
import Control.Lens ((^?))
import Control.Monad.Catch (MonadCatch, MonadMask, MonadThrow)
import Control.Monad.Morph
import Control.Monad.Trans.Control
import Control.Monad.Unique
import Data.String (fromString)
import Network.URI (URI)
-- | Any additional human-readable key-value pairs relevant
-- to the execution of a block of code.
type TracingMetadata = [(Text, Text)]
newtype Reporter = Reporter
{ 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 \_ _ -> fmap fst
-- | A type class for monads which support some
-- way to report execution traces.
class Monad m => HasReporter m where
-- | Get the current tracer
askReporter :: m Reporter
default askReporter :: m Reporter
askReporter = pure noReporter
instance HasReporter m => HasReporter (ReaderT r m) where
askReporter = lift askReporter
instance HasReporter m => HasReporter (ExceptT e m) where
askReporter = lift 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
{ tcCurrentTrace :: !Word64
, tcCurrentSpan :: !Word64
, tcCurrentParent :: !(Maybe Word64)
}
-- | The 'TraceT' monad transformer adds the ability to keep track of
-- the current trace context.
newtype TraceT m a = TraceT { unTraceT :: ReaderT (TraceContext, Reporter) (WriterT TracingMetadata m) a }
deriving (Functor, Applicative, Monad, MonadIO, MonadUnique, MonadMask, MonadCatch, MonadThrow)
instance MonadTrans TraceT where
lift = TraceT . lift . lift
instance MFunctor TraceT where
hoist f (TraceT rwma) = TraceT (hoist (hoist f) rwma)
deriving instance MonadBase b m => MonadBase b (TraceT m)
deriving instance MonadBaseControl b m => MonadBaseControl b (TraceT m)
instance MonadError e m => MonadError e (TraceT m) where
throwError = lift . throwError
catchError (TraceT m) f = TraceT (catchError m (unTraceT . f))
instance MonadReader r m => MonadReader r (TraceT m) where
ask = TraceT $ lift ask
local f m = TraceT $ mapReaderT (local f) (unTraceT m)
instance (HTTP.HasHttpManagerM m) => HTTP.HasHttpManagerM (TraceT m) where
askHttpManager = lift HTTP.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 name tma = do
rep <- askReporter
runTraceTWithReporter rep name tma
runTraceTWith :: MonadIO m => TraceContext -> Reporter -> Text -> TraceT m a -> m a
runTraceTWith ctx rep name tma =
runReporter rep ctx name
$ runWriterT
$ runReaderT (unTraceT tma) (ctx, 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 ctx name tma = do
rep <- askReporter
runTraceTWith ctx rep name tma
-- | Run an action in the 'TraceT' monad transformer in an
-- existing context.
runTraceTWithReporter :: MonadIO m => Reporter -> Text -> TraceT m a -> m a
runTraceTWithReporter rep name tma = do
ctx <- TraceContext
<$> liftIO Rand.randomIO
<*> liftIO Rand.randomIO
<*> pure Nothing
runTraceTWith ctx rep name 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.
interpTraceT
:: MonadTrace n
=> (m (a, TracingMetadata) -> n (b, TracingMetadata))
-> TraceT m a
-> n b
interpTraceT f (TraceT rwma) = do
ctx <- currentContext
rep <- currentReporter
(b, meta) <- f (runWriterT (runReaderT rwma (ctx, rep)))
attachMetadata meta
pure 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
trace name ma = TraceT . ReaderT $ \(ctx, rep) -> do
spanId <- liftIO (Rand.randomIO :: IO Word64)
let subCtx = ctx { tcCurrentSpan = spanId
, tcCurrentParent = Just (tcCurrentSpan ctx)
}
lift . runReporter rep subCtx name . runWriterT $ runReaderT (unTraceT ma) (subCtx, rep)
currentContext = TraceT (asks fst)
currentReporter = TraceT (asks snd)
attachMetadata = TraceT . tell
instance MonadTrace m => MonadTrace (ReaderT r m) where
trace = mapReaderT . trace
currentContext = lift currentContext
currentReporter = lift currentReporter
attachMetadata = lift . attachMetadata
instance MonadTrace m => MonadTrace (StateT e m) where
trace = mapStateT . trace
currentContext = lift currentContext
currentReporter = lift currentReporter
attachMetadata = lift . attachMetadata
instance MonadTrace m => MonadTrace (ExceptT e m) where
trace = mapExceptT . trace
currentContext = lift currentContext
currentReporter = lift currentReporter
attachMetadata = lift . attachMetadata
-- | Encode Word64 to 16 character hex string
word64ToHex :: Word64 -> Text
word64ToHex randNum = bsToTxt $ Hex.encode numInBytes
where numInBytes = BL.toStrict (Bin.encode randNum)
-- | Decode 16 character hex string to Word64
hexToWord64 :: Text -> Maybe Word64
hexToWord64 randText = do
case Hex.decode $ txtToBs randText of
Left _ -> Nothing
Right decoded -> Just $ Bin.decode $ BL.fromStrict decoded
-- | Inject the trace context as a set of HTTP headers.
injectHttpContext :: TraceContext -> [HTTP.Header]
injectHttpContext TraceContext{..} =
("X-B3-TraceId", txtToBs $ word64ToHex tcCurrentTrace)
: ("X-B3-SpanId", txtToBs $ word64ToHex tcCurrentSpan)
: [ ("X-B3-ParentSpanId", txtToBs $ word64ToHex parentID)
| parentID <- maybeToList 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 hdrs = do
freshSpanId <- liftIO Rand.randomIO
pure $ TraceContext
<$> (hexToWord64 =<< HTTP.parseHeaderMaybe =<< lookup "X-B3-TraceId" hdrs)
<*> pure freshSpanId
<*> pure (hexToWord64 =<< HTTP.parseHeaderMaybe =<< lookup "X-B3-SpanId" 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{..} =
J.object
[ "trace_id" J..= word64ToHex tcCurrentTrace
, "span_id" J..= word64ToHex tcCurrentSpan
]
-- | Extract a trace context from an event trigger payload.
extractEventContext :: J.Value -> IO (Maybe TraceContext)
extractEventContext e = do
freshSpanId <- liftIO Rand.randomIO
pure $ TraceContext
<$> (hexToWord64 =<< e ^? JL.key "trace_context" . JL.key "trace_id" . JL._String)
<*> pure freshSpanId
<*> pure (hexToWord64 =<< e ^? JL.key "trace_context" . JL.key "span_id" . JL._String)
-- | Perform HTTP request which supports Trace headers
tracedHttpRequest
:: MonadTrace m
=> HTTP.Request
-- ^ http request that needs to be made
-> (HTTP.Request -> m a)
-- ^ a function that takes the traced request and executes it
-> m a
tracedHttpRequest req f = do
let method = bsToTxt (HTTP.method req)
uri = show @URI (HTTP.getUri req)
trace (method <> " " <> fromString uri) do
let reqBytes = case HTTP.requestBody req of
HTTP.RequestBodyBS bs -> Just (fromIntegral (BS.length bs))
HTTP.RequestBodyLBS bs -> Just (BL.length bs)
HTTP.RequestBodyBuilder len _ -> Just len
HTTP.RequestBodyStream len _ -> Just len
_ -> Nothing
for_ reqBytes \b ->
attachMetadata [("request_body_bytes", fromString (show b))]
ctx <- currentContext
let req' = req { HTTP.requestHeaders =
injectHttpContext ctx <> HTTP.requestHeaders req
}
f req'