context/tracing/xray.go

// Package tracing parses distributed tracing headers. The X-Ray format is
// handled here; the OpenTelemetry format lives in a sibling file because
// they parse differently enough that sharing code would hurt clarity.
//
// See mercemay.top/src/lambdalog/context/tracing/.
package tracing

import (
	"context"
	"os"
	"strings"
)

// XRayTraceEnv is the environment variable Lambda populates with the
// X-Amzn-Trace-Id header.
const XRayTraceEnv = "_X_AMZN_TRACE_ID"

// XRayHeader is the raw header value split into its semicolon-separated
// components. Each component keeps the original case of its left-hand
// side; callers should not rely on Header ordering.
type XRayHeader struct {
	Root     string
	Parent   string
	Sampled  string
	Lineage  string
	Extras   map[string]string
}

type xrayKey struct{}

var xrayCtxKey xrayKey

// ParseXRay parses the raw header string. Unknown keys accumulate in Extras
// so the caller does not lose information.
func ParseXRay(raw string) XRayHeader {
	h := XRayHeader{Extras: map[string]string{}}
	for _, part := range strings.Split(raw, ";") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		k, v, ok := strings.Cut(part, "=")
		if !ok {
			continue
		}
		switch strings.ToLower(k) {
		case "root":
			h.Root = v
		case "parent":
			h.Parent = v
		case "sampled":
			h.Sampled = v
		case "lineage":
			h.Lineage = v
		default:
			h.Extras[k] = v
		}
	}
	return h
}

// FromEnv parses the _X_AMZN_TRACE_ID environment variable set by the
// Lambda runtime. Callers outside Lambda should use ParseXRay directly.
func FromEnv() XRayHeader {
	return ParseXRay(os.Getenv(XRayTraceEnv))
}

// WithXRay attaches h to ctx under this package's key.
func WithXRay(ctx context.Context, h XRayHeader) context.Context {
	if ctx == nil {
		ctx = context.Background()
	}
	return context.WithValue(ctx, xrayCtxKey, h)
}

// XRayFromContext returns the header attached by WithXRay.
func XRayFromContext(ctx context.Context) (XRayHeader, bool) {
	if ctx == nil {
		return XRayHeader{}, false
	}
	h, ok := ctx.Value(xrayCtxKey).(XRayHeader)
	return h, ok
}

// IsSampled reports whether the header indicates the trace is sampled. The
// X-Ray convention is the literal string "1".
func (h XRayHeader) IsSampled() bool {
	return h.Sampled == "1"
}

// SpanID returns the parent span id if present, otherwise the root. This
// matches the W3C trace-context notion of "current span".
func (h XRayHeader) SpanID() string {
	if h.Parent != "" {
		return h.Parent
	}
	return h.Root
}

// LogFields returns structured fields for inclusion in a log record. The
// keys match the X-Ray conventions so the fields can be joined with X-Ray
// segment data downstream.
func (h XRayHeader) LogFields() []Field {
	if h.Root == "" {
		return nil
	}
	fs := []Field{{Key: "xray_root", Value: h.Root}}
	if h.Parent != "" {
		fs = append(fs, Field{Key: "xray_parent", Value: h.Parent})
	}
	if h.Sampled != "" {
		fs = append(fs, Field{Key: "xray_sampled", Value: h.Sampled})
	}
	return fs
}

// Field is a local copy of encoder.Field to avoid an import cycle.
type Field struct {
	Key   string
	Value any
}