fix(observability): extract upstream traceparent, re-parent server span, and inject traceparent on egress

Signed-off-by: Guangya Liu <gyliu513@gmail.com>

Author: gyliu513

pkg/handlers/request.go

@@ -24,6 +24,8 @@ import (
 	"time"
 
 	eppb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
 	envoy "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/envoy"
@@ -46,15 +48,32 @@ func (s *Server) HandleRequestHeaders(ctx context.Context, reqCtx *RequestContex
 		}
 	}
 
+	// Inject the active trace context into the egress headers so the next
+	// filter in the chain is parented to this span rather than the upstream one.
+	traceCarrier := propagation.MapCarrier{}
+	otel.GetTextMapPropagator().Inject(ctx, traceCarrier)
+	for key, value := range traceCarrier {
+		reqCtx.Request.SetHeader(key, value)
+	}
+
 	if !headers.GetEndOfStream() {
 		log.FromContext(ctx).V(logutil.VERBOSE).Info("captured request headers, deferring response until body arrives...")
 		return nil
 	}
 	// EndOfStream means no body is expected, return HeadersResponse immediately
+	headersResponse := &eppb.HeadersResponse{}
+	if len(reqCtx.Request.MutatedHeaders()) > 0 || len(reqCtx.Request.RemovedHeaders()) > 0 {
+		headersResponse.Response = &eppb.CommonResponse{
+			HeaderMutation: &eppb.HeaderMutation{
+				SetHeaders:    envoy.GenerateHeadersMutation(reqCtx.Request.MutatedHeaders()),
+				RemoveHeaders: reqCtx.Request.RemovedHeaders(),
+			},
+		}
+	}
 	return []*eppb.ProcessingResponse{
 		{
 			Response: &eppb.ProcessingResponse_RequestHeaders{
-				RequestHeaders: &eppb.HeadersResponse{},
+				RequestHeaders: headersResponse,
 			},
 		},
 	}

pkg/handlers/request_test.go

@@ -26,6 +26,7 @@ import (
 	basepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"github.com/google/go-cmp/cmp"
+	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/protobuf/testing/protocmp"
 	metricsutils "k8s.io/component-base/metrics/testutil"
 	crmetrics "sigs.k8s.io/controller-runtime/pkg/metrics"
@@ -170,6 +171,85 @@ func TestHandleRequestHeaders(t *testing.T) {
 	}
 }
 
+func TestHandleRequestHeaders_InjectsTraceContext(t *testing.T) {
+	withTraceContextPropagator(t)
+
+	traceID, err := trace.TraceIDFromHex(testTraceID)
+	if err != nil {
+		t.Fatalf("failed to parse trace ID: %v", err)
+	}
+	spanID, err := trace.SpanIDFromHex(testSpanID)
+	if err != nil {
+		t.Fatalf("failed to parse span ID: %v", err)
+	}
+	ctx := trace.ContextWithSpanContext(context.Background(), trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID:    traceID,
+		SpanID:     spanID,
+		TraceFlags: trace.FlagsSampled,
+	}))
+
+	wantTraceparent := "00-" + testTraceID + "-" + testSpanID + "-01"
+
+	t.Run("records traceparent as a header mutation", func(t *testing.T) {
+		server := newServerForTest(newTestProfiles())
+		reqCtx := &RequestContext{Request: requesthandling.NewInferenceRequest()}
+
+		resp := server.HandleRequestHeaders(ctx, reqCtx, &extProcPb.HttpHeaders{
+			Headers: &basepb.HeaderMap{Headers: []*basepb.HeaderValue{}},
+		})
+		if resp != nil {
+			t.Fatalf("expected deferred response, got %v", resp)
+		}
+		if got := reqCtx.Request.MutatedHeaders()["traceparent"]; got != wantTraceparent {
+			t.Errorf("mutated traceparent = %q, want %q", got, wantTraceparent)
+		}
+	})
+
+	t.Run("end of stream response carries the traceparent mutation", func(t *testing.T) {
+		server := newServerForTest(newTestProfiles())
+		reqCtx := &RequestContext{Request: requesthandling.NewInferenceRequest()}
+
+		resp := server.HandleRequestHeaders(ctx, reqCtx, &extProcPb.HttpHeaders{
+			Headers:     &basepb.HeaderMap{Headers: []*basepb.HeaderValue{}},
+			EndOfStream: true,
+		})
+		if len(resp) != 1 {
+			t.Fatalf("expected a single response, got %d", len(resp))
+		}
+		mutation := resp[0].GetRequestHeaders().GetResponse().GetHeaderMutation()
+		if mutation == nil {
+			t.Fatal("expected a header mutation in the end-of-stream response")
+		}
+		var got string
+		for _, h := range mutation.GetSetHeaders() {
+			if h.GetHeader().GetKey() == "traceparent" {
+				got = string(h.GetHeader().GetRawValue())
+			}
+		}
+		if got != wantTraceparent {
+			t.Errorf("traceparent in header mutation = %q, want %q", got, wantTraceparent)
+		}
+	})
+
+	t.Run("upstream traceparent matching current context is not re-set", func(t *testing.T) {
+		server := newServerForTest(newTestProfiles())
+		reqCtx := &RequestContext{Request: requesthandling.NewInferenceRequest()}
+
+		resp := server.HandleRequestHeaders(ctx, reqCtx, &extProcPb.HttpHeaders{
+			Headers: &basepb.HeaderMap{Headers: []*basepb.HeaderValue{
+				{Key: "traceparent", RawValue: []byte(wantTraceparent)},
+			}},
+			EndOfStream: true,
+		})
+		if len(resp) != 1 {
+			t.Fatalf("expected a single response, got %d", len(resp))
+		}
+		if mutation := resp[0].GetRequestHeaders().GetResponse().GetHeaderMutation(); mutation != nil {
+			t.Errorf("expected no header mutation, got %v", mutation)
+		}
+	})
+}
+
 // === Request Body Tests (built-in plugins) ===
 
 func TestHandleRequestBody_BuiltInPlugins(t *testing.T) {

pkg/handlers/server.go

@@ -20,11 +20,13 @@ import (
 	"context"
 	"errors"
 	"io"
+	"strings"
 	"time"
 
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -81,19 +83,37 @@ type RequestContext struct {
 	Response                    *requesthandling.InferenceResponse
 }
 
+// extractTraceContext returns ctx augmented with the upstream trace context
+// (traceparent, tracestate, baggage) propagated via the Envoy request headers.
+func extractTraceContext(ctx context.Context, headers *extProcPb.HttpHeaders) context.Context {
+	carrier := propagation.MapCarrier{}
+	if headers != nil && headers.Headers != nil {
+		for _, header := range headers.Headers.Headers {
+			carrier[strings.ToLower(header.Key)] = envoy.GetHeaderValue(header)
+		}
+	}
+	return otel.GetTextMapPropagator().Extract(ctx, carrier)
+}
+
 func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 	ctx := srv.Context()
 
-	// Start tracing span for the request
+	// The server span is started when the request headers arrive, so it can be
+	// parented to the upstream trace context they carry instead of starting an
+	// orphan root trace.
 	tracer := otel.Tracer(
 		"llm-d/inference-payload-processor/extproc",
 		trace.WithInstrumentationVersion(version.BuildRef),
 		trace.WithInstrumentationAttributes(
 			attribute.String("commit-sha", version.CommitSHA),
 		),
 	)
-	ctx, span := tracer.Start(ctx, "gateway.request", trace.WithSpanKind(trace.SpanKindServer))
-	defer span.End()
+	var span trace.Span
+	defer func() {
+		if span != nil {
+			span.End()
+		}
+	}()
 
 	logger := log.FromContext(ctx)
 	loggerVerbose := logger.V(logutil.VERBOSE)
@@ -129,6 +149,10 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 		var err error
 		switch v := req.Request.(type) {
 		case *extProcPb.ProcessingRequest_RequestHeaders:
+			if span == nil {
+				ctx, span = tracer.Start(extractTraceContext(ctx, v.RequestHeaders),
+					"gateway.request", trace.WithSpanKind(trace.SpanKindServer))
+			}
 			if requestId := envoy.ExtractHeaderValue(v, requestIdHeaderKey); len(requestId) > 0 {
 				logger = logger.WithValues(requestIdHeaderKey, requestId)
 				loggerVerbose = logger.V(logutil.VERBOSE)

pkg/handlers/server_test.go

@@ -25,6 +25,9 @@ import (
 	basepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"github.com/google/go-cmp/cmp"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/protobuf/testing/protocmp"
 
 	envoytest "github.com/llm-d/llm-d-inference-payload-processor/pkg/common/envoy/test"
@@ -40,6 +43,82 @@ import (
 
 const testProfileName = "default"
 
+const (
+	testTraceID = "0af7651916cd43dd8448eb211c80319c"
+	testSpanID  = "b7ad6b7169203331"
+)
+
+// withTraceContextPropagator installs the W3C trace context propagator for the
+// duration of the test and restores the previous global propagator afterwards.
+func withTraceContextPropagator(t *testing.T) {
+	t.Helper()
+	prev := otel.GetTextMapPropagator()
+	otel.SetTextMapPropagator(propagation.TraceContext{})
+	t.Cleanup(func() { otel.SetTextMapPropagator(prev) })
+}
+
+func TestExtractTraceContext(t *testing.T) {
+	withTraceContextPropagator(t)
+
+	tests := []struct {
+		name      string
+		headers   *extProcPb.HttpHeaders
+		wantValid bool
+	}{
+		{
+			name: "extracts upstream traceparent",
+			headers: &extProcPb.HttpHeaders{
+				Headers: &basepb.HeaderMap{
+					Headers: []*basepb.HeaderValue{
+						{Key: "traceparent", RawValue: []byte("00-" + testTraceID + "-" + testSpanID + "-01")},
+					},
+				},
+			},
+			wantValid: true,
+		},
+		{
+			name: "header key lookup is case-insensitive",
+			headers: &extProcPb.HttpHeaders{
+				Headers: &basepb.HeaderMap{
+					Headers: []*basepb.HeaderValue{
+						{Key: "Traceparent", RawValue: []byte("00-" + testTraceID + "-" + testSpanID + "-01")},
+					},
+				},
+			},
+			wantValid: true,
+		},
+		{
+			name:      "no traceparent leaves context untouched",
+			headers:   &extProcPb.HttpHeaders{Headers: &basepb.HeaderMap{}},
+			wantValid: false,
+		},
+		{
+			name:      "nil headers leaves context untouched",
+			headers:   nil,
+			wantValid: false,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := extractTraceContext(context.Background(), tc.headers)
+			sc := trace.SpanContextFromContext(ctx)
+			if sc.IsValid() != tc.wantValid {
+				t.Fatalf("extracted span context validity = %v, want %v", sc.IsValid(), tc.wantValid)
+			}
+			if !tc.wantValid {
+				return
+			}
+			if got := sc.TraceID().String(); got != testTraceID {
+				t.Errorf("trace ID = %s, want %s", got, testTraceID)
+			}
+			if got := sc.SpanID().String(); got != testSpanID {
+				t.Errorf("span ID = %s, want %s", got, testSpanID)
+			}
+		})
+	}
+}
+
 func TestHandleRequestBody(t *testing.T) {
 	ctx := logutil.NewTestLoggerIntoContext(context.Background())