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

Code before:
  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 Byte String to 16 character hex string
encodeBstoHex :: BS.ByteString -> BS.ByteString
encodeBstoHex randBString = Hex.encode $ randBString 

-- | Create a Random 16 character hex string
randomBS :: IO BS.ByteString
randomBS = do
  randNum <- liftIO (Rand.randomIO :: IO Word64)
  let randomByteString =  (BL.toStrict (BIN.encode  randNum))
  return randomByteString

-- | Inject the trace context as a set of HTTP headers.
injectHttpContext :: TraceContext -> [HTTP.Header]
injectHttpContext TraceContext{..} = 
  case tcCurrentParent of
    Nothing -> [ ("X-B3-TraceId", encodeBstoHex tcCurrentTrace)
               , ("X-B3-SpanId", encodeBstoHex tcCurrentSpan)
               ]
    Just parentID ->  [ ("X-B3-TraceId", encodeBstoHex tcCurrentTrace)
                      , ("X-B3-SpanId", encodeBstoHex tcCurrentSpan)
                      , ("X-B3-ParentSpanId", encodeBstoHex parentID)
                      ]

-- | 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.

Review comment (suggestion):
Should `encodeBstoHex` return back `Text` instead of ByteString?

Because I think Text type makes more sense, since having ByteString might suggest that we might have characters outside the UTF-8 range.

Code after:
  
  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 
-- | Hex.Decode returns two tuples: (properly decoded data, string starts at the first invalid base16 sequence)
hexToWord64 :: Text -> Maybe Word64
hexToWord64 randText = do
  let (decoded, leftovers) = Hex.decode $ txtToBs randText
      decodedWord64 = Bin.decode $ BL.fromStrict decoded
  guard (BS.null leftovers)
  pure decodedWord64
        

-- | 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.