fix: use internal id for routerProcessorsPerReqID map (#1344)

**Description**
The fix ensures that even if multiple requests arrive with identical
x-request-id headers, they will have unique internal IDs in the
`routerProcessorsPerReqID` map, preventing processor overwrites and
incorrect cleanup.

**Related Issues/PRs (if applicable)**
Fixes #1340 


**Special notes for reviewers (if applicable)**
- Created internal request ID header constant: Added internalReqIDHeader
= "x-ai-eg-internal-req-id"
  - Modified request ID generation logic:
- Router filter: Creates unique internal request ID by appending UUID
suffix: `originalReqID + "-" + uuid.NewString()`
- Upstream filter: Uses the internal request ID passed from router
filter via header

---------

Signed-off-by: Dan Sun <[email protected]>
Co-authored-by: achoo30 <[email protected]>

Author: yuzisun

internal/extproc/server.go

@@ -20,6 +20,7 @@ import (
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	typev3 "github.com/envoyproxy/go-control-plane/envoy/type/v3"
 	"github.com/google/cel-go/cel"
+	"github.com/google/uuid"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/status"
@@ -45,6 +46,7 @@ type Server struct {
 	processorFactories            map[string]ProcessorFactory
 	routerProcessorsPerReqID      map[string]Processor
 	routerProcessorsPerReqIDMutex sync.RWMutex
+	uuidFn                        func() string
 }
 
 // NewServer creates a new external processor server.
@@ -54,6 +56,7 @@ func NewServer(logger *slog.Logger, tracing tracing.Tracing) (*Server, error) {
 		tracing:                  tracing,
 		processorFactories:       make(map[string]ProcessorFactory),
 		routerProcessorsPerReqID: make(map[string]Processor),
+		uuidFn:                   uuid.NewString,
 	}
 	return srv, nil
 }
@@ -131,6 +134,10 @@ func (s *Server) processorForPath(requestHeaders map[string]string, isUpstreamFi
 // This is used in the upstream filter level to determine the original path of the request on retry.
 const originalPathHeader = "x-ai-eg-original-path"
 
+// internalReqIDHeader is the header used to pass the unique internal request ID to the upstream filter.
+// This ensures that the upstream filter uses the same unique ID as the router filter to avoid race conditions.
+const internalReqIDHeader = "x-ai-eg-internal-req-id"
+
 // Process implements [extprocv3.ExternalProcessorServer].
 func (s *Server) Process(stream extprocv3.ExternalProcessor_ProcessServer) error {
 	s.logger.Debug("handling a new stream", slog.Any("config_uuid", s.config.uuid))
@@ -145,13 +152,14 @@ func (s *Server) Process(stream extprocv3.ExternalProcessor_ProcessServer) error
 	// to pass the request through without any processing as there would be nothing to process from AI Gateway's perspective.
 	var p Processor = passThroughProcessor{}
 	var isUpstreamFilter bool
-	var reqID string
+	var internalReqID string
+	var originalReqID string
 	var logger *slog.Logger
 	defer func() {
 		if !isUpstreamFilter {
 			s.routerProcessorsPerReqIDMutex.Lock()
 			defer s.routerProcessorsPerReqIDMutex.Unlock()
-			delete(s.routerProcessorsPerReqID, reqID)
+			delete(s.routerProcessorsPerReqID, internalReqID)
 		}
 	}()
 
@@ -177,9 +185,21 @@ func (s *Server) Process(stream extprocv3.ExternalProcessor_ProcessServer) error
 		// request, and the processor will be instantiated only once.
 		if headers := req.GetRequestHeaders().GetHeaders(); headers != nil {
 			headersMap := headersToMap(headers)
-			reqID = headersMap["x-request-id"]
+			originalReqID = headersMap["x-request-id"]
 			// Assume that when attributes are set, this stream is for the upstream filter level.
 			isUpstreamFilter = req.GetAttributes() != nil
+
+			if isUpstreamFilter {
+				// For upstream filter, use the internal request ID passed from the router filter
+				internalReqID = headersMap[internalReqIDHeader]
+				if internalReqID == "" {
+					return status.Errorf(codes.Internal, "missing internal request ID header from router filter")
+				}
+			} else {
+				// For router filter, create a unique internal request ID to avoid race conditions
+				// with duplicate x-request-id values by appending a UUID suffix to the original request ID
+				internalReqID = originalReqID + "-" + s.uuidFn()
+			}
 			p, err = s.processorForPath(headersMap, isUpstreamFilter)
 			if err != nil {
 				if errors.Is(err, errNoProcessor) {
@@ -200,22 +220,22 @@ func (s *Server) Process(stream extprocv3.ExternalProcessor_ProcessServer) error
 			}
 			_, isEndpoinPicker := headersMap[internalapi.EndpointPickerHeaderKey]
 			if isUpstreamFilter {
-				if err = s.setBackend(ctx, p, reqID, isEndpoinPicker, req); err != nil {
+				if err = s.setBackend(ctx, p, internalReqID, isEndpoinPicker, req); err != nil {
 					s.logger.Error("error processing request message", slog.String("error", err.Error()))
 					return status.Errorf(codes.Unknown, "error processing request message: %v", err)
 				}
 			} else {
 				s.routerProcessorsPerReqIDMutex.Lock()
-				s.routerProcessorsPerReqID[reqID] = p
+				s.routerProcessorsPerReqID[internalReqID] = p
 				s.routerProcessorsPerReqIDMutex.Unlock()
 			}
 		}
 		if logger == nil {
-			logger = s.logger.With("request_id", reqID, "is_upstream_filter", isUpstreamFilter)
+			logger = s.logger.With("request_id", originalReqID, "is_upstream_filter", isUpstreamFilter)
 		}
 
 		// At this point, p is guaranteed to be a valid processor either from the concrete processor or the passThroughProcessor.
-		resp, err := s.processMsg(ctx, logger, p, req)
+		resp, err := s.processMsg(ctx, logger, p, req, internalReqID, isUpstreamFilter)
 		if err != nil {
 			s.logger.Error("error processing request message", slog.String("error", err.Error()))
 			return status.Errorf(codes.Unknown, "error processing request message: %v", err)
@@ -227,7 +247,7 @@ func (s *Server) Process(stream extprocv3.ExternalProcessor_ProcessServer) error
 	}
 }
 
-func (s *Server) processMsg(ctx context.Context, l *slog.Logger, p Processor, req *extprocv3.ProcessingRequest) (*extprocv3.ProcessingResponse, error) {
+func (s *Server) processMsg(ctx context.Context, l *slog.Logger, p Processor, req *extprocv3.ProcessingRequest, internalReqID string, isUpstreamFilter bool) (*extprocv3.ProcessingResponse, error) {
 	switch value := req.Request.(type) {
 	case *extprocv3.ProcessingRequest_RequestHeaders:
 		requestHdrs := req.GetRequestHeaders().Headers
@@ -240,6 +260,35 @@ func (s *Server) processMsg(ctx context.Context, l *slog.Logger, p Processor, re
 		if err != nil {
 			return nil, fmt.Errorf("cannot process request headers: %w", err)
 		}
+
+		// For router filter, inject the internal request ID header so upstream filter can use it
+		if !isUpstreamFilter && resp != nil {
+			if requestHeaders, ok := resp.Response.(*extprocv3.ProcessingResponse_RequestHeaders); ok {
+				// Ensure we have header mutation to add the internal request ID
+				if requestHeaders.RequestHeaders == nil {
+					requestHeaders.RequestHeaders = &extprocv3.HeadersResponse{}
+				}
+				if requestHeaders.RequestHeaders.Response == nil {
+					requestHeaders.RequestHeaders.Response = &extprocv3.CommonResponse{}
+				}
+				if requestHeaders.RequestHeaders.Response.HeaderMutation == nil {
+					requestHeaders.RequestHeaders.Response.HeaderMutation = &extprocv3.HeaderMutation{}
+				}
+
+				// Add the internal request ID header
+				internalReqIDHeaderValue := &corev3.HeaderValueOption{
+					Header: &corev3.HeaderValue{
+						Key:      internalReqIDHeader,
+						RawValue: []byte(internalReqID),
+					},
+				}
+				requestHeaders.RequestHeaders.Response.HeaderMutation.SetHeaders = append(
+					requestHeaders.RequestHeaders.Response.HeaderMutation.SetHeaders,
+					internalReqIDHeaderValue,
+				)
+			}
+		}
+
 		l.Debug("request headers processed", slog.Any("response", resp))
 		return resp, nil
 	case *extprocv3.ProcessingRequest_RequestBody:
@@ -279,7 +328,7 @@ func (s *Server) processMsg(ctx context.Context, l *slog.Logger, p Processor, re
 
 // setBackend retrieves the backend from the request attributes and sets it in the processor. This is only called
 // if the processor is an upstream filter.
-func (s *Server) setBackend(ctx context.Context, p Processor, reqID string, isEndpointPicker bool, req *extprocv3.ProcessingRequest) error {
+func (s *Server) setBackend(ctx context.Context, p Processor, internalReqID string, isEndpointPicker bool, req *extprocv3.ProcessingRequest) error {
 	attributes := req.GetAttributes()["envoy.filters.http.ext_proc"]
 	if attributes == nil || len(attributes.Fields) == 0 { // coverage-ignore
 		return status.Error(codes.Internal, "missing attributes in request")
@@ -317,10 +366,10 @@ func (s *Server) setBackend(ctx context.Context, p Processor, reqID string, isEn
 
 	s.routerProcessorsPerReqIDMutex.RLock()
 	defer s.routerProcessorsPerReqIDMutex.RUnlock()
-	routerProcessor, ok := s.routerProcessorsPerReqID[reqID]
+	routerProcessor, ok := s.routerProcessorsPerReqID[internalReqID]
 	if !ok {
 		return status.Errorf(codes.Internal, "no router processor found, request_id=%s, backend=%s",
-			reqID, backendName.GetStringValue())
+			internalReqID, backendName.GetStringValue())
 	}
 
 	if err := p.SetBackend(ctx, backend.b, backend.handler, routerProcessor); err != nil {

internal/extproc/server_test.go

@@ -120,7 +120,7 @@ func TestServer_List(t *testing.T) {
 func TestServer_processMsg(t *testing.T) {
 	t.Run("unknown request type", func(t *testing.T) {
 		s, p := requireNewServerWithMockProcessor(t)
-		_, err := s.processMsg(t.Context(), slog.Default(), p, &extprocv3.ProcessingRequest{})
+		_, err := s.processMsg(t.Context(), slog.Default(), p, &extprocv3.ProcessingRequest{}, "test-req-id", false)
 		require.ErrorContains(t, err, "unknown request type")
 	})
 	t.Run("request headers", func(t *testing.T) {
@@ -134,7 +134,7 @@ func TestServer_processMsg(t *testing.T) {
 		req := &extprocv3.ProcessingRequest{
 			Request: &extprocv3.ProcessingRequest_RequestHeaders{RequestHeaders: &extprocv3.HttpHeaders{Headers: hm}},
 		}
-		resp, err := s.processMsg(t.Context(), slog.Default(), p, req)
+		resp, err := s.processMsg(t.Context(), slog.Default(), p, req, "test-req-id", false)
 		require.NoError(t, err)
 		require.NotNil(t, resp)
 		require.Equal(t, expResponse, resp)
@@ -150,7 +150,7 @@ func TestServer_processMsg(t *testing.T) {
 		req := &extprocv3.ProcessingRequest{
 			Request: &extprocv3.ProcessingRequest_RequestBody{RequestBody: reqBody},
 		}
-		resp, err := s.processMsg(t.Context(), slog.Default(), p, req)
+		resp, err := s.processMsg(t.Context(), slog.Default(), p, req, "test-req-id", false)
 		require.NoError(t, err)
 		require.NotNil(t, resp)
 		require.Equal(t, expResponse, resp)
@@ -166,7 +166,7 @@ func TestServer_processMsg(t *testing.T) {
 		req := &extprocv3.ProcessingRequest{
 			Request: &extprocv3.ProcessingRequest_ResponseHeaders{ResponseHeaders: &extprocv3.HttpHeaders{Headers: hm}},
 		}
-		resp, err := s.processMsg(t.Context(), slog.Default(), p, req)
+		resp, err := s.processMsg(t.Context(), slog.Default(), p, req, "test-req-id", false)
 		require.NoError(t, err)
 		require.NotNil(t, resp)
 		require.Equal(t, expResponse, resp)
@@ -182,7 +182,7 @@ func TestServer_processMsg(t *testing.T) {
 		req := &extprocv3.ProcessingRequest{
 			Request: &extprocv3.ProcessingRequest_ResponseHeaders{ResponseHeaders: &extprocv3.HttpHeaders{Headers: hm}},
 		}
-		resp, err := s.processMsg(t.Context(), slog.Default(), p, req)
+		resp, err := s.processMsg(t.Context(), slog.Default(), p, req, "test-req-id", false)
 		require.NoError(t, err)
 		require.NotNil(t, resp)
 		require.Equal(t, expResponse, resp)
@@ -198,7 +198,7 @@ func TestServer_processMsg(t *testing.T) {
 		req := &extprocv3.ProcessingRequest{
 			Request: &extprocv3.ProcessingRequest_ResponseBody{ResponseBody: reqBody},
 		}
-		resp, err := s.processMsg(t.Context(), slog.Default(), p, req)
+		resp, err := s.processMsg(t.Context(), slog.Default(), p, req, "test-req-id", false)
 		require.NoError(t, err)
 		require.NotNil(t, resp)
 		require.Equal(t, expResponse, resp)
@@ -236,7 +236,7 @@ func TestServer_Process(t *testing.T) {
 	t.Run("upstream filter", func(t *testing.T) {
 		s, p := requireNewServerWithMockProcessor(t)
 
-		hm := &corev3.HeaderMap{Headers: []*corev3.HeaderValue{{Key: originalPathHeader, Value: "/"}, {Key: "foo", Value: "bar"}}}
+		hm := &corev3.HeaderMap{Headers: []*corev3.HeaderValue{{Key: originalPathHeader, Value: "/"}, {Key: internalReqIDHeader, Value: "test-req-id-123"}, {Key: "foo", Value: "bar"}}}
 		p.t = t
 		p.expHeaderMap = hm
 		req := &extprocv3.ProcessingRequest{
@@ -351,7 +351,7 @@ func TestServer_ProcessorSelection(t *testing.T) {
 	s.Register("/two", func(*processorConfig, map[string]string, *slog.Logger, tracing.Tracing, bool) (Processor, error) {
 		return &mockProcessor{
 			t:                     t,
-			expHeaderMap:          &corev3.HeaderMap{Headers: []*corev3.HeaderValue{{Key: ":path", Value: "/two"}}},
+			expHeaderMap:          &corev3.HeaderMap{Headers: []*corev3.HeaderValue{{Key: ":path", Value: "/two"}, {Key: "x-request-id", Value: "original-req-id"}}},
 			retProcessingResponse: &extprocv3.ProcessingResponse{Response: &extprocv3.ProcessingResponse_RequestHeaders{}},
 		}, nil
 	})
@@ -388,11 +388,36 @@ func TestServer_ProcessorSelection(t *testing.T) {
 		req := &extprocv3.ProcessingRequest{
 			Request: &extprocv3.ProcessingRequest_RequestHeaders{
 				RequestHeaders: &extprocv3.HttpHeaders{
-					Headers: &corev3.HeaderMap{Headers: []*corev3.HeaderValue{{Key: ":path", Value: "/two"}}},
+					Headers: &corev3.HeaderMap{Headers: []*corev3.HeaderValue{
+						{Key: ":path", Value: "/two"},
+						{Key: "x-request-id", Value: "original-req-id"},
+					}},
 				},
 			},
 		}
-		expResponse := &extprocv3.ProcessingResponse{Response: &extprocv3.ProcessingResponse_RequestHeaders{}}
+		expResponse := &extprocv3.ProcessingResponse{
+			Response: &extprocv3.ProcessingResponse_RequestHeaders{
+				RequestHeaders: &extprocv3.HeadersResponse{
+					Response: &extprocv3.CommonResponse{
+						HeaderMutation: &extprocv3.HeaderMutation{
+							SetHeaders: []*corev3.HeaderValueOption{
+								{
+									Header: &corev3.HeaderValue{
+										Key:      internalReqIDHeader,
+										RawValue: []byte("original-req-id-test-internal-req-id"),
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		}
+
+		s.uuidFn = func() string {
+			return "test-internal-req-id"
+		}
+
 		ms := &mockExternalProcessingStream{t: t, ctx: ctx, retRecv: req, expResponseOnSend: expResponse}
 
 		err = s.Process(ms)