Add tracing entry span with W3C propagation to EPP handler

Signed-off-by: sallyom <somalley@redhat.com>

Author: sallyom

pkg/epp/handlers/request.go

@@ -17,11 +17,14 @@ limitations under the License.
 package handlers
 
 import (
+	"context"
 	"strconv"
 	"time"
 
 	configPb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
 	"google.golang.org/protobuf/types/known/structpb"
 
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/common"
@@ -37,7 +40,7 @@ const (
 	defaultFairnessID = "default-flow"
 )
 
-func (s *StreamingServer) HandleRequestHeaders(reqCtx *RequestContext, req *extProcPb.ProcessingRequest_RequestHeaders) error {
+func (s *StreamingServer) HandleRequestHeaders(ctx context.Context, reqCtx *RequestContext, req *extProcPb.ProcessingRequest_RequestHeaders) error {
 	reqCtx.RequestReceivedTimestamp = time.Now()
 
 	// an EoS in the request headers means this request has no body or trailers.
@@ -52,7 +55,7 @@ func (s *StreamingServer) HandleRequestHeaders(reqCtx *RequestContext, req *extP
 		}
 		reqCtx.TargetEndpoint = pod.GetIPAddress() + ":" + pod.GetPort()
 		reqCtx.RequestSize = 0
-		reqCtx.reqHeaderResp = s.generateRequestHeaderResponse(reqCtx)
+		reqCtx.reqHeaderResp = s.generateRequestHeaderResponse(ctx, reqCtx)
 		return nil
 	}
 
@@ -91,7 +94,7 @@ func (s *StreamingServer) generateRequestBodyResponses(requestBodyBytes []byte)
 	return responses
 }
 
-func (s *StreamingServer) generateRequestHeaderResponse(reqCtx *RequestContext) *extProcPb.ProcessingResponse {
+func (s *StreamingServer) generateRequestHeaderResponse(ctx context.Context, reqCtx *RequestContext) *extProcPb.ProcessingResponse {
 	// The Endpoint Picker supports two approaches to communicating the target endpoint, as a request header
 	// and as an unstructure ext-proc response metadata key/value pair. This enables different integration
 	// options for gateway providers.
@@ -101,7 +104,7 @@ func (s *StreamingServer) generateRequestHeaderResponse(reqCtx *RequestContext)
 				Response: &extProcPb.CommonResponse{
 					ClearRouteCache: true,
 					HeaderMutation: &extProcPb.HeaderMutation{
-						SetHeaders: s.generateHeaders(reqCtx),
+						SetHeaders: s.generateHeaders(ctx, reqCtx),
 					},
 				},
 			},
@@ -110,7 +113,7 @@ func (s *StreamingServer) generateRequestHeaderResponse(reqCtx *RequestContext)
 	}
 }
 
-func (s *StreamingServer) generateHeaders(reqCtx *RequestContext) []*configPb.HeaderValueOption {
+func (s *StreamingServer) generateHeaders(ctx context.Context, reqCtx *RequestContext) []*configPb.HeaderValueOption {
 	// can likely refactor these two bespoke headers to be updated in PostDispatch, to centralize logic.
 	headers := []*configPb.HeaderValueOption{
 		{
@@ -131,6 +134,19 @@ func (s *StreamingServer) generateHeaders(reqCtx *RequestContext) []*configPb.He
 		})
 	}
 
+	// Inject trace context headers for propagation to downstream services
+	traceHeaders := make(map[string]string)
+	propagator := otel.GetTextMapPropagator()
+	propagator.Inject(ctx, propagation.MapCarrier(traceHeaders))
+	for key, value := range traceHeaders {
+		headers = append(headers, &configPb.HeaderValueOption{
+			Header: &configPb.HeaderValue{
+				Key:      key,
+				RawValue: []byte(value),
+			},
+		})
+	}
+
 	// Include any non-system-owned headers.
 	for key, value := range reqCtx.Request.Headers {
 		if request.IsSystemOwnedHeader(key) {

pkg/epp/handlers/request_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package handlers
 
 import (
+	"context"
 	"testing"
 
 	configPb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
@@ -64,7 +65,7 @@ func TestHandleRequestHeaders(t *testing.T) {
 				},
 			}
 
-			err := server.HandleRequestHeaders(reqCtx, req)
+			err := server.HandleRequestHeaders(context.Background(), reqCtx, req)
 			assert.NoError(t, err, "HandleRequestHeaders should not return an error")
 
 			assert.Equal(t, tc.wantFairnessID, reqCtx.FairnessID, "FairnessID should match expected value")
@@ -93,7 +94,7 @@ func TestGenerateHeaders_Sanitization(t *testing.T) {
 		},
 	}
 
-	results := server.generateHeaders(reqCtx)
+	results := server.generateHeaders(context.Background(), reqCtx)
 
 	gotHeaders := make(map[string]string)
 	for _, h := range results {

pkg/epp/handlers/server.go

@@ -27,6 +27,8 @@ import (
 	envoyTypePb "github.com/envoyproxy/go-control-plane/envoy/type/v3"
 	"github.com/go-logr/logr"
 	"github.com/google/uuid"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/trace"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
@@ -126,6 +128,12 @@ const (
 
 func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 	ctx := srv.Context()
+
+	// Start tracing span for the request
+	tracer := otel.Tracer("gateway-api-inference-extension")
+	ctx, span := tracer.Start(ctx, "gateway.request", trace.WithSpanKind(trace.SpanKindServer))
+	defer span.End()
+
 	logger := log.FromContext(ctx)
 	loggerTrace := logger.V(logutil.TRACE)
 	loggerTrace.Info("Processing")
@@ -204,7 +212,7 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 			loggerTrace = logger.V(logutil.TRACE)
 			ctx = log.IntoContext(ctx, logger)
 
-			err = s.HandleRequestHeaders(reqCtx, v)
+			err = s.HandleRequestHeaders(ctx, reqCtx, v)
 		case *extProcPb.ProcessingRequest_RequestBody:
 			loggerTrace.Info("Incoming body chunk", "EoS", v.RequestBody.EndOfStream)
 			// In the stream case, we can receive multiple request bodies.
@@ -240,7 +248,7 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 					break
 				}
 				reqCtx.RequestSize = len(requestBodyBytes)
-				reqCtx.reqHeaderResp = s.generateRequestHeaderResponse(reqCtx)
+				reqCtx.reqHeaderResp = s.generateRequestHeaderResponse(ctx, reqCtx)
 				reqCtx.reqBodyResp = s.generateRequestBodyResponses(requestBodyBytes)
 
 				metrics.RecordRequestCounter(reqCtx.IncomingModelName, reqCtx.TargetModelName)