feat(telemetry): bridge configurable gateway trace headers

telemetry/telemetry_test.go

@@ -7,6 +7,7 @@ import (
 	"git.yuntuo.hk/k2safe/go-omni/logx"
 	"git.yuntuo.hk/k2safe/go-omni/telemetry"
 	"github.com/redis/go-redis/v9"
+	"go.opentelemetry.io/otel/propagation"
 	sdkmetric "go.opentelemetry.io/otel/sdk/metric"
 	"go.opentelemetry.io/otel/sdk/metric/metricdata"
 	"go.opentelemetry.io/otel/sdk/trace/tracetest"
@@ -115,3 +116,115 @@ func TestHTTPTracePropagationMetricsAndRedisNil(t *testing.T) {
 		t.Fatal("request metric", count)
 	}
 }
+
+func TestConfiguredHeaderPropagation(t *testing.T) {
+	exporter := tracetest.NewInMemoryExporter()
+	propagator, err := telemetry.NewHeaderPropagator(telemetry.HeaderPropagationConfig{
+		TraceIDHeader: "X-Trace-Id", SpanIDHeader: "X-Span-Id", RequestIDHeader: "X-Request-Id",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	runtime, err := telemetry.New(telemetry.Config{Service: "header-test", SampleRatio: 1, Exporter: exporter, Propagator: propagator})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer runtime.Close()
+	const traceID = "11111111111111111111111111111111"
+	const parentSpanID = "2222222222222222"
+	const requestID = "gateway-request-1"
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("X-Trace-Id") != traceID || r.Header.Get("X-Request-Id") != requestID {
+			t.Errorf("custom propagation = trace %q request %q", r.Header.Get("X-Trace-Id"), r.Header.Get("X-Request-Id"))
+		}
+		if r.Header.Get("X-Span-Id") == parentSpanID || len(r.Header.Get("X-Span-Id")) != 16 || !strings.HasPrefix(r.Header.Get("Traceparent"), "00-"+traceID+"-") {
+			t.Errorf("outbound parent propagation = %q", r.Header)
+		}
+		io.WriteString(w, "ok")
+	}))
+	defer upstream.Close()
+	app := omni.New(omni.WithLogger(false), omni.WithRequestIDHeader("X-Request-Id"))
+	defer app.Close()
+	app.Use(runtime.Middleware())
+	app.GET("/", func(c *httpx.Context) error {
+		if got := logx.TraceID(c.Context()); got != traceID {
+			t.Errorf("inbound trace correlation = %q", got)
+		}
+		if got := logx.RequestID(c.Context()); got != requestID {
+			t.Errorf("inbound request correlation = %q", got)
+		}
+		req, _ := http.NewRequestWithContext(c.Context(), http.MethodGet, upstream.URL, nil)
+		resp, err := (telemetry.Transport{Runtime: runtime}).RoundTrip(req)
+		if err != nil {
+			return err
+		}
+		defer resp.Body.Close()
+		_, _ = io.ReadAll(resp.Body)
+		return c.OK("ok")
+	})
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	req.Header.Set("X-Trace-Id", traceID)
+	req.Header.Set("X-Span-Id", parentSpanID)
+	req.Header.Set("X-Request-Id", requestID)
+	rec := httptest.NewRecorder()
+	app.ServeHTTP(rec, req)
+	if rec.Code != http.StatusOK || rec.Header().Get("X-Request-Id") != requestID {
+		t.Fatalf("response status/request ID = %d/%q", rec.Code, rec.Header().Get("X-Request-Id"))
+	}
+	if err := runtime.Flush(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	spans := exporter.GetSpans()
+	var server, client *tracetest.SpanStub
+	for _, span := range spans {
+		switch span.SpanKind {
+		case trace.SpanKindServer:
+			span := span
+			server = &span
+		case trace.SpanKindClient:
+			span := span
+			client = &span
+		}
+	}
+	if server == nil || client == nil {
+		t.Fatalf("missing spans: %d", len(spans))
+	}
+	if server.SpanContext.TraceID().String() != traceID || server.Parent.TraceID().String() != traceID || server.Parent.SpanID().String() != parentSpanID || !server.Parent.IsRemote() {
+		t.Fatalf("server parent = %s/%s remote=%v", server.SpanContext.TraceID(), server.Parent.SpanID(), server.Parent.IsRemote())
+	}
+	if client.SpanContext.TraceID().String() != traceID || client.Parent.SpanID() != server.SpanContext.SpanID() {
+		t.Fatalf("client linkage = trace %s parent %s", client.SpanContext.TraceID(), client.Parent.SpanID())
+	}
+}
+
+func TestConfiguredHeaderPropagationPrefersW3CAndRejectsInvalidSetup(t *testing.T) {
+	propagator, err := telemetry.NewHeaderPropagator(telemetry.HeaderPropagationConfig{TraceIDHeader: "x-trace-id", SpanIDHeader: "x-span-id"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	carrier := propagation.MapCarrier{
+		"traceparent": "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-01",
+		"X-Trace-Id":  "11111111111111111111111111111111",
+		"X-Span-Id":   "2222222222222222",
+	}
+	ctx := propagator.Extract(context.Background(), carrier)
+	parent := trace.SpanContextFromContext(ctx)
+	if parent.TraceID().String() != "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" || parent.SpanID().String() != "bbbbbbbbbbbbbbbb" {
+		t.Fatalf("W3C parent was not preferred: %s/%s", parent.TraceID(), parent.SpanID())
+	}
+	for _, config := range []telemetry.HeaderPropagationConfig{
+		{TraceIDHeader: "", SpanIDHeader: "X-Span-Id"},
+		{TraceIDHeader: "X-Trace-Id", SpanIDHeader: ""},
+		{TraceIDHeader: "X-Trace-Id", SpanIDHeader: "X-Trace-Id"},
+		{TraceIDHeader: "traceparent", SpanIDHeader: "X-Span-Id"},
+		{TraceIDHeader: "X Trace", SpanIDHeader: "X-Span-Id"},
+	} {
+		if _, err := telemetry.NewHeaderPropagator(config); err == nil {
+			t.Fatalf("invalid propagation config accepted: %+v", config)
+		}
+	}
+	invalid := propagation.MapCarrier{"X-Trace-Id": "bad", "X-Span-Id": "bad"}
+	if got := trace.SpanContextFromContext(propagator.Extract(context.Background(), invalid)); got.IsValid() {
+		t.Fatal("invalid legacy IDs became a parent")
+	}
+}