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

Code before:
-- | 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{..} = 
  case tcCurrentParent of
    Nothing -> [ ("X-B3-TraceId", txtToBs $ word64ToHex tcCurrentTrace)
               , ("X-B3-SpanId", txtToBs $ word64ToHex tcCurrentSpan)
               ]
    Just parentID ->  [ ("X-B3-TraceId", txtToBs $ word64ToHex tcCurrentTrace)
                      , ("X-B3-SpanId", txtToBs $ word64ToHex tcCurrentSpan)
                      , ("X-B3-ParentSpanId", txtToBs $ word64ToHex 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.
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..= tcCurrentTrace 
    , "span_id"  J..= tcCurrentSpan 
    ]
  
-- | Extract a trace context from an event trigger payload.
extractEventContext :: J.Value -> IO (Maybe TraceContext)

Review comment (suggestion):
```suggestion
  ("X-B3-TraceId", txtToBs $ word64ToHex tcCurrentTrace)
  : ("X-B3-SpanId", txtToBs $ word64ToHex tcCurrentSpan)
  : [ ("X-B3-ParentSpanId", txtToBs $ word64ToHex parentID)
    | parentID <- maybeToList tcCurrentParent
    ]
```

Code after:
-- | 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.
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..= tcCurrentTrace 
    , "span_id"  J..= 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