Repository: hasura/graphql-engine
File: server/src-lib/Hasura/Tracing.hs (Haskell)

Code before:
  trace = mapExceptT . trace
  currentContext = lift currentContext
  currentReporter = lift currentReporter
  attachMetadata = lift . attachMetadata

-- | 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
    <$> (HTTP.parseHeaderMaybe =<< lookup "X-Hasura-TraceId" hdrs)
    <*> pure freshSpanId
    <*> pure (HTTP.parseHeaderMaybe =<< lookup "X-Hasura-SpanId" hdrs)

-- | 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 = trace (T.decodeUtf8 (HTTP.path req)) 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))]
  TraceContext{..} <- currentContext
  let tracingHeaders =
        [ ("X-Hasura-TraceId", fromString (show tcCurrentTrace))
        , ("X-Hasura-SpanId", fromString (show tcCurrentSpan))
        ]
      req' = req { HTTP.requestHeaders =
                     tracingHeaders <> HTTP.requestHeaders req
                 }
  f req'


Review comment (suggestion):
We should use `bsToTxt'` here, because `decodeUtf8` throws exceptions in the event of a decoding error. It's very unlikely, but we should guard against it.

Code after:
  attachMetadata = lift . attachMetadata

instance MonadTrace m => MonadTrace (ExceptT e m) where
  trace = mapExceptT . trace
  currentContext = lift currentContext
  currentReporter = lift currentReporter
  attachMetadata = lift . attachMetadata

-- | Inject the trace context as a set of HTTP headers.
injectHttpContext :: TraceContext -> [HTTP.Header]
injectHttpContext TraceContext{..} = 
  [ ("X-Hasura-TraceId", fromString (show tcCurrentTrace))
  , ("X-Hasura-SpanId", fromString (show tcCurrentSpan))
  ]
  
-- | 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
    <$> (HTTP.parseHeaderMaybe =<< lookup "X-Hasura-TraceId" hdrs)
    <*> pure freshSpanId
    <*> pure (HTTP.parseHeaderMaybe =<< lookup "X-Hasura-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 ctx =
  J.object
    [ "trace_id" J..= tcCurrentTrace ctx
    , "span_id"  J..= tcCurrentSpan ctx
    ]

-- | Extract a trace context from an event trigger payload.
extractEventContext :: J.Value -> IO (Maybe TraceContext)
extractEventContext e = do
  freshSpanId <- liftIO Rand.randomIO
  pure $ TraceContext
    <$> (e ^? JL.key "trace_context" . JL.key "trace_id" . JL._Integral)
    <*> pure freshSpanId
    <*> pure (e ^? JL.key "trace_context" . JL.key "span_id" . JL._Integral)

-- | Perform HTTP request which supports Trace headers
tracedHttpRequest 
  :: MonadTrace m 
  =>  HTTP.Request
  -- ^ http request that needs to be made
  -> (HTTP.Request -> m a)