TraceContext.hs

{- | W3C trace-context helpers that read and write @traceparent@ /
@tracestate@ header strings inside Kiroku event metadata.

The on-the-wire JSON shape inside the event's @metadata@ JSONB column is:

> {
>   "traceparent": "00-<32-hex traceId>-<16-hex spanId>-<2-hex flags>",
>   "tracestate":  "<vendor entries, optional>"
> }

Other keys in @metadata@ are preserved by 'injectTraceContext'.
-}
module Kiroku.Otel.TraceContext (
    injectTraceContext,
    extractTraceContext,
) where

import Control.Lens ((&), (?~))
import Data.Aeson (Value (..))
import Data.Aeson.Key qualified as Key
import Data.Aeson.KeyMap (KeyMap)
import Data.Aeson.KeyMap qualified as KM
import Data.ByteString (ByteString)
import Data.Generics.Labels ()
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import GHC.IO (unsafePerformIO)
import Kiroku.Store.Types (EventData (..), RecordedEvent (..))
import OpenTelemetry.Propagator.W3CTraceContext (decodeSpanContext, encodeSpanContext)
import OpenTelemetry.Trace.Core (SpanContext, wrapSpanContext)

{- | Encode a 'SpanContext' as W3C @traceparent@ / @tracestate@ strings and
merge them into the @metadata@ JSON object of an 'EventData'. Existing
keys in @metadata@ are preserved; existing @traceparent@ / @tracestate@
keys (if any) are overwritten — the W3C spec mandates exactly one of
each value per propagation.

If the input @metadata@ is 'Nothing', or is a non-object JSON value,
the helper starts from an empty object before merging.

This function is pure: it uses 'unsafePerformIO' to call
'encodeSpanContext' (OpenTelemetry 1.0 type:
@Span -> IO (ByteString, ByteString)@), which is observably pure on the
frozen span returned by 'wrapSpanContext' (no shared mutable state, no exceptions).
The \"unsafe\" annotation is mandatory to bridge the propagator's
@IO@-typed encoder to the pure interface this module exposes.
-}
{-# NOINLINE injectTraceContext #-}
injectTraceContext :: SpanContext -> EventData -> EventData
injectTraceContext sc ed@EventData{metadata = oldMeta} =
    ed & (#metadata ?~ Object (mergeTraceContext sc oldMeta))

{- | Pull a 'SpanContext' back out of a 'RecordedEvent'\'s @metadata@.
Returns 'Nothing' when @metadata@ is absent, is not a JSON object,
lacks a @traceparent@ key, or contains a @traceparent@ value that fails
W3C parsing. Never throws.
-}
extractTraceContext :: RecordedEvent -> Maybe SpanContext
extractTraceContext re = case re of
    RecordedEvent{metadata = Just (Object o)} -> decodeFromObject o
    _ -> Nothing

-- ---------------------------------------------------------------------------
-- Internal helpers
-- ---------------------------------------------------------------------------

mergeTraceContext :: SpanContext -> Maybe Value -> KeyMap Value
mergeTraceContext sc mMeta =
    let (tp, ts) = unsafePerformIO (encodeSpanContext (wrapSpanContext sc))
        tpText = TE.decodeUtf8 tp
        tsText = TE.decodeUtf8 ts
        existing = case mMeta of
            Just (Object o) -> o
            _ -> KM.empty
        withTp = KM.insert (Key.fromText (T.pack "traceparent")) (String tpText) existing
     in if T.null tsText
            then withTp
            else KM.insert (Key.fromText (T.pack "tracestate")) (String tsText) withTp

decodeFromObject :: KeyMap Value -> Maybe SpanContext
decodeFromObject o = do
    String tpText <- KM.lookup (Key.fromText (T.pack "traceparent")) o
    let tpBs :: ByteString
        tpBs = TE.encodeUtf8 tpText
        tsBs = case KM.lookup (Key.fromText (T.pack "tracestate")) o of
            Just (String tsText) -> Just (TE.encodeUtf8 tsText)
            _ -> Nothing
    decodeSpanContext (Just tpBs) tsBs