-- | This module contains a collection of utility functions we use with tracing
-- throughout the codebase, but that are not a core part of the library. If we
-- were to move tracing to a separate library, those functions should be kept
-- here in the core engine code.
module Hasura.Tracing.Utils
  ( traceHTTPRequest,
    attachSourceConfigAttributes,
  )
where

import Control.Lens
import Data.String
import Data.Text.Extended (toTxt)
import Hasura.Prelude
import Hasura.RQL.Types.SourceConfiguration (HasSourceConfiguration (..))
import Hasura.Tracing.Class
import Hasura.Tracing.Context
import Hasura.Tracing.Sampling
import Hasura.Tracing.TraceId
import Network.HTTP.Client.Transformable qualified as HTTP

-- | Wrap the execution of an HTTP request in a span in the current
-- trace. Despite its name, this function does not start a new trace, and the
-- span will therefore not be recorded if the surrounding context isn't traced
-- (see 'spanWith').
--
-- Additionally, this function adds metadata regarding the request to the
-- created span, and injects the trace context into the HTTP header.
traceHTTPRequest ::
  (MonadIO m, 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
traceHTTPRequest :: forall (m :: * -> *) a.
(MonadIO m, MonadTrace m) =>
Request -> (Request -> m a) -> m a
traceHTTPRequest 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.
(MonadIO m, MonadTrace m) =>
Text -> m a -> m a
newSpan (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
    TraceMetadata -> m ()
forall (m :: * -> *). MonadTrace m => TraceMetadata -> 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))]
    [Header]
headers <- (Maybe TraceContext -> [Header])
-> m (Maybe TraceContext) -> m [Header]
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ([Header]
-> (TraceContext -> [Header]) -> Maybe TraceContext -> [Header]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] TraceContext -> [Header]
toHeaders) m (Maybe TraceContext)
forall (m :: * -> *). MonadTrace m => m (Maybe 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 ([Header]
headers [Header] -> [Header] -> [Header]
forall a. Semigroup a => a -> a -> a
<>) Request
req
  where
    toHeaders :: TraceContext -> [HTTP.Header]
    toHeaders :: TraceContext -> [Header]
toHeaders TraceContext {Maybe SpanId
SamplingState
SpanId
TraceId
tcCurrentTrace :: TraceId
tcCurrentSpan :: SpanId
tcCurrentParent :: Maybe SpanId
tcSamplingState :: SamplingState
tcCurrentTrace :: TraceContext -> TraceId
tcCurrentSpan :: TraceContext -> SpanId
tcCurrentParent :: TraceContext -> Maybe SpanId
tcSamplingState :: TraceContext -> SamplingState
..} =
      [Maybe Header] -> [Header]
forall a. [Maybe a] -> [a]
forall (f :: * -> *) a. Filterable f => f (Maybe a) -> f a
catMaybes
        [ Header -> Maybe Header
forall a. a -> Maybe a
Just (HeaderName
"X-B3-TraceId", TraceId -> ByteString
traceIdToHex TraceId
tcCurrentTrace),
          Header -> Maybe Header
forall a. a -> Maybe a
Just (HeaderName
"X-B3-SpanId", SpanId -> ByteString
spanIdToHex SpanId
tcCurrentSpan),
          (HeaderName
"X-B3-ParentSpanId",) (ByteString -> Header)
-> (SpanId -> ByteString) -> SpanId -> Header
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SpanId -> ByteString
spanIdToHex (SpanId -> Header) -> Maybe SpanId -> Maybe Header
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe SpanId
tcCurrentParent,
          (HeaderName
"X-B3-Sampled",) (ByteString -> Header) -> Maybe ByteString -> Maybe Header
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SamplingState -> Maybe ByteString
forall s. IsString s => SamplingState -> Maybe s
samplingStateToHeader SamplingState
tcSamplingState
        ]

attachSourceConfigAttributes :: forall b m. (HasSourceConfiguration b, MonadTrace m) => SourceConfig b -> m ()
attachSourceConfigAttributes :: forall (b :: BackendType) (m :: * -> *).
(HasSourceConfiguration b, MonadTrace m) =>
SourceConfig b -> m ()
attachSourceConfigAttributes SourceConfig b
sourceConfig = do
  let backendSourceKind :: BackendSourceKind b
backendSourceKind = forall (b :: BackendType).
HasSourceConfiguration b =>
SourceConfig b -> BackendSourceKind b
sourceConfigBackendSourceKind @b SourceConfig b
sourceConfig
  TraceMetadata -> m ()
forall (m :: * -> *). MonadTrace m => TraceMetadata -> m ()
attachMetadata [(Text
"source.kind", BackendSourceKind b -> Text
forall a. ToTxt a => a -> Text
toTxt (BackendSourceKind b -> Text) -> BackendSourceKind b -> Text
forall a b. (a -> b) -> a -> b
$ BackendSourceKind b
backendSourceKind)]