perf: Reduce allocations on the response-body streaming path

Signed-off-by: Alberto Perdomo <aperdomo@redhat.com>

Author: albertoperdomo2

pkg/common/envoy/metadata.go

@@ -21,11 +21,12 @@ import (
 )
 
 func ExtractMetadataValues(req *extProcPb.ProcessingRequest) map[string]any {
-	metadata := make(map[string]any)
-	if req != nil && req.MetadataContext != nil && req.MetadataContext.FilterMetadata != nil {
-		for key, val := range req.MetadataContext.FilterMetadata {
-			metadata[key] = val.AsMap()
-		}
+	if req == nil || req.MetadataContext == nil || len(req.MetadataContext.FilterMetadata) == 0 {
+		return nil
+	}
+	metadata := make(map[string]any, len(req.MetadataContext.FilterMetadata))
+	for key, val := range req.MetadataContext.FilterMetadata {
+		metadata[key] = val.AsMap()
 	}
 	return metadata
 }

pkg/common/envoy/metadata_test.go

@@ -39,12 +39,35 @@ func TestExtractMetadataValues(t *testing.T) {
 
 	tests := []struct {
 		name     string
-		metadata map[string]*structpb.Struct
+		req      *extProcPb.ProcessingRequest
 		expected map[string]any
 	}{
 		{
-			name:     "Exact match",
-			metadata: makeFilterMetadata(),
+			name:     "Nil request",
+			req:      nil,
+			expected: nil,
+		},
+		{
+			name:     "Nil MetadataContext",
+			req:      &extProcPb.ProcessingRequest{},
+			expected: nil,
+		},
+		{
+			name: "Empty FilterMetadata",
+			req: &extProcPb.ProcessingRequest{
+				MetadataContext: &corev3.Metadata{
+					FilterMetadata: map[string]*structpb.Struct{},
+				},
+			},
+			expected: nil,
+		},
+		{
+			name: "Populated metadata",
+			req: &extProcPb.ProcessingRequest{
+				MetadataContext: &corev3.Metadata{
+					FilterMetadata: makeFilterMetadata(),
+				},
+			},
 			expected: map[string]any{
 				"key-1": map[string]any{
 					"hello":      "world",
@@ -56,13 +79,7 @@ func TestExtractMetadataValues(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			req := &extProcPb.ProcessingRequest{
-				MetadataContext: &corev3.Metadata{
-					FilterMetadata: tt.metadata,
-				},
-			}
-
-			result := ExtractMetadataValues(req)
+			result := ExtractMetadataValues(tt.req)
 			if diff := cmp.Diff(result, tt.expected); diff != "" {
 				t.Errorf("ExtractMetadataValues() unexpected response (-want +got): %v", diff)
 			}

pkg/epp/handlers/server.go

@@ -406,10 +406,9 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 			return status.Errorf(codes.Unknown, "cannot receive stream request: %v", recvErr)
 		}
 
-		reqCtx.Request.Metadata = envoy.ExtractMetadataValues(req)
-
 		switch v := req.Request.(type) {
 		case *extProcPb.ProcessingRequest_RequestHeaders:
+			reqCtx.Request.Metadata = envoy.ExtractMetadataValues(req)
 			requestID := envoy.ExtractHeaderValue(v, reqcommon.RequestIDHeaderKey)
 			// request ID is a must for maintaining a state per request in plugins that hold internal state and use PluginState.
 			// if request id was not supplied as a header, we generate it ourselves.
@@ -439,6 +438,7 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 			// Message is buffered, we can read and decode.
 			if v.RequestBody.EndOfStream {
 				loggerTrace.Info("decoding")
+				reqCtx.Request.Metadata = envoy.ExtractMetadataValues(req)
 				reqCtx.Request.RawBody = make([]byte, buf.Len())
 				copy(reqCtx.Request.RawBody, buf.Bytes())
 
@@ -494,6 +494,7 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 		case *extProcPb.ProcessingRequest_RequestTrailers:
 			// This is currently unused.
 		case *extProcPb.ProcessingRequest_ResponseHeaders:
+			reqCtx.Request.Metadata = envoy.ExtractMetadataValues(req)
 			respHeadersReceivedAt := time.Now()
 			for _, header := range v.ResponseHeaders.Headers.GetHeaders() {
 				value := string(header.RawValue)

pkg/epp/requestcontrol/director.go

@@ -521,9 +521,14 @@ func (d *Director) HandleResponseHeader(ctx context.Context, reqCtx *handlers.Re
 // to the ext_proc response sent back to Envoy.
 func (d *Director) HandleResponseBody(ctx context.Context, reqCtx *handlers.RequestContext, endOfStream bool) *handlers.RequestContext {
 	logger := log.FromContext(ctx).WithValues("stage", "bodyChunk")
-	logger.V(logutil.TRACE).Info("Entering HandleResponseBodyChunk")
+	loggerTrace := logger.V(logutil.TRACE)
+	if loggerTrace.Enabled() {
+		loggerTrace.Info("Entering HandleResponseBodyChunk")
+	}
 	if len(d.requestControlPlugins.responseStreamingPlugins) == 0 {
-		logger.V(logutil.TRACE).Info("Exiting HandleResponseBodyChunk")
+		if loggerTrace.Enabled() {
+			loggerTrace.Info("Exiting HandleResponseBodyChunk")
+		}
 		return reqCtx
 	}
 
@@ -536,7 +541,6 @@ func (d *Director) HandleResponseBody(ctx context.Context, reqCtx *handlers.Requ
 		EndOfStream:   endOfStream,
 		Usage:         reqCtx.Usage,
 	}
-	requestID := reqCtx.Request.Headers[reqcommon.RequestIDHeaderKey]
 
 	if endOfStream {
 		// Drain the async queue: close the channel and wait for the goroutine to finish
@@ -558,10 +562,13 @@ func (d *Director) HandleResponseBody(ctx context.Context, reqCtx *handlers.Requ
 		}
 		q := d.loadOrCreateResponseBodyQueue(reqCtx)
 		if !q.enqueue(work) {
-			logger.V(logutil.DEBUG).Info("Skipping response body chunk because the async queue is closed", "requestID", requestID)
+			logger.V(logutil.DEBUG).Info("Skipping response body chunk because the async queue is closed",
+				"requestID", reqCtx.Request.Headers[reqcommon.RequestIDHeaderKey])
 		}
 	}
-	logger.V(logutil.TRACE).Info("Exiting HandleResponseBodyChunk")
+	if loggerTrace.Enabled() {
+		loggerTrace.Info("Exiting HandleResponseBodyChunk")
+	}
 	return reqCtx
 }
 
@@ -592,11 +599,12 @@ func (d *Director) runPreRequestPlugins(ctx context.Context, request *fwksched.I
 	schedulingResult *fwksched.SchedulingResult) {
 	loggerDebug := log.FromContext(ctx).V(logutil.DEBUG)
 	for _, plugin := range d.requestControlPlugins.preRequestPlugins {
-		loggerDebug.Info("Running PreRequest plugin", "plugin", plugin.TypedName())
+		tn := plugin.TypedName()
+		loggerDebug.Info("Running PreRequest plugin", "plugin", tn)
 		before := time.Now()
 		plugin.PreRequest(ctx, request, schedulingResult)
-		metrics.RecordPluginProcessingLatency(fwkrc.PreRequestExtensionPoint, plugin.TypedName().Type, plugin.TypedName().Name, time.Since(before))
-		loggerDebug.Info("Completed running PreRequest plugin successfully", "plugin", plugin.TypedName())
+		metrics.RecordPluginProcessingLatency(fwkrc.PreRequestExtensionPoint, tn.Type, tn.Name, time.Since(before))
+		loggerDebug.Info("Completed running PreRequest plugin successfully", "plugin", tn)
 	}
 }
 
@@ -606,13 +614,14 @@ func (d *Director) runRequestHeaderProcessors(ctx context.Context, request *fwks
 	}
 	loggerDebug := log.FromContext(ctx).V(logutil.DEBUG)
 	for _, plugin := range d.requestControlPlugins.requestHeaderPlugins {
-		loggerDebug.Info("Running RequestHeaderProcessor plugin", "plugin", plugin.TypedName())
+		tn := plugin.TypedName()
+		loggerDebug.Info("Running RequestHeaderProcessor plugin", "plugin", tn)
 		before := time.Now()
 		if err := plugin.RequestHeader(ctx, request); err != nil {
 			return err
 		}
-		metrics.RecordPluginProcessingLatency(fwkrc.RequestHeaderExtensionPoint, plugin.TypedName().Type, plugin.TypedName().Name, time.Since(before))
-		loggerDebug.Info("Completed running RequestHeaderProcessor plugin successfully", "plugin", plugin.TypedName())
+		metrics.RecordPluginProcessingLatency(fwkrc.RequestHeaderExtensionPoint, tn.Type, tn.Name, time.Since(before))
+		loggerDebug.Info("Completed running RequestHeaderProcessor plugin successfully", "plugin", tn)
 	}
 	return nil
 }
@@ -637,38 +646,45 @@ func (d *Director) runAdmissionPlugins(ctx context.Context,
 	request *fwksched.InferenceRequest, endpoints []fwksched.Endpoint) error {
 	loggerDebug := log.FromContext(ctx).V(logutil.DEBUG)
 	for _, plugin := range d.requestControlPlugins.admissionPlugins {
-		loggerDebug.Info("Running Admit plugin", "plugin", plugin.TypedName())
+		tn := plugin.TypedName()
+		loggerDebug.Info("Running Admit plugin", "plugin", tn)
 		before := time.Now()
 		denyReason := plugin.Admit(ctx, request, endpoints)
-		metrics.RecordPluginProcessingLatency(fwkrc.AdmissionExtensionPoint, plugin.TypedName().Type, plugin.TypedName().Name, time.Since(before))
+		metrics.RecordPluginProcessingLatency(fwkrc.AdmissionExtensionPoint, tn.Type, tn.Name, time.Since(before))
 		if denyReason != nil {
-			loggerDebug.Info("Admit plugin denied the request", "plugin", plugin.TypedName(), "reason", denyReason.Error())
+			loggerDebug.Info("Admit plugin denied the request", "plugin", tn, "reason", denyReason.Error())
 			return denyReason
 		}
-		loggerDebug.Info("Completed running Admit plugin successfully", "plugin", plugin.TypedName())
+		loggerDebug.Info("Completed running Admit plugin successfully", "plugin", tn)
 	}
 	return nil
 }
 
 func (d *Director) runResponseHeaderPlugins(ctx context.Context, request *fwksched.InferenceRequest, response *fwkrc.Response, targetEndpoint *fwkdl.EndpointMetadata) {
 	loggerDebug := log.FromContext(ctx).V(logutil.DEBUG)
 	for _, plugin := range d.requestControlPlugins.responseReceivedPlugins {
-		loggerDebug.Info("Running ResponseReceived plugin", "plugin", plugin.TypedName())
+		tn := plugin.TypedName()
+		loggerDebug.Info("Running ResponseReceived plugin", "plugin", tn)
 		before := time.Now()
 		plugin.ResponseHeader(ctx, request, response, targetEndpoint)
-		metrics.RecordPluginProcessingLatency(fwkrc.ResponseReceivedExtensionPoint, plugin.TypedName().Type, plugin.TypedName().Name, time.Since(before))
-		loggerDebug.Info("Completed running ResponseReceived plugin successfully", "plugin", plugin.TypedName())
+		metrics.RecordPluginProcessingLatency(fwkrc.ResponseReceivedExtensionPoint, tn.Type, tn.Name, time.Since(before))
+		loggerDebug.Info("Completed running ResponseReceived plugin successfully", "plugin", tn)
 	}
 }
 
 func (d *Director) runResponseBodyPlugins(ctx context.Context, request *fwksched.InferenceRequest, response *fwkrc.Response, targetEndpoint *fwkdl.EndpointMetadata) {
 	loggerTrace := log.FromContext(ctx).V(logutil.TRACE)
 	for _, plugin := range d.requestControlPlugins.responseStreamingPlugins {
-		loggerTrace.Info("Running ResponseStreaming plugin", "plugin", plugin.TypedName())
+		tn := plugin.TypedName()
+		if loggerTrace.Enabled() {
+			loggerTrace.Info("Running ResponseStreaming plugin", "plugin", tn)
+		}
 		before := time.Now()
 		plugin.ResponseBody(ctx, request, response, targetEndpoint)
-		metrics.RecordPluginProcessingLatency(fwkrc.ResponseStreamingExtensionPoint, plugin.TypedName().Type, plugin.TypedName().Name, time.Since(before))
-		loggerTrace.Info("Completed running ResponseStreaming plugin successfully", "plugin", plugin.TypedName())
+		metrics.RecordPluginProcessingLatency(fwkrc.ResponseStreamingExtensionPoint, tn.Type, tn.Name, time.Since(before))
+		if loggerTrace.Enabled() {
+			loggerTrace.Info("Completed running ResponseStreaming plugin successfully", "plugin", tn)
+		}
 	}
 }
 

pkg/epp/requestcontrol/director_bench_test.go

@@ -0,0 +1,125 @@
+/*
+Copyright 2026 The llm-d Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package requestcontrol
+
+import (
+	"context"
+	"testing"
+
+	"github.com/go-logr/logr"
+	"k8s.io/apimachinery/pkg/types"
+	"sigs.k8s.io/controller-runtime/pkg/log"
+
+	reqcommon "github.com/llm-d/llm-d-router/pkg/common/request"
+	fwkdl "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/datalayer"
+	fwkrh "github.com/llm-d/llm-d-router/pkg/epp/framework/interface/requesthandling"
+	"github.com/llm-d/llm-d-router/pkg/epp/handlers"
+)
+
+const (
+	benchChunksPerOp = 10
+	// maxAllocsPerResponse is the allocation ceiling for one complete
+	// response (benchChunksPerOp intermediate chunks + 1 end-of-stream).
+	// Bump this intentionally when a change adds justified allocations.
+	maxAllocsPerResponse = 60
+)
+
+func TestHandleResponseBodyAllocs(t *testing.T) {
+	plugin := newTestResponseStreaming("alloc-plugin")
+	director := NewDirectorWithConfig(nil, &mockScheduler{}, nil, nil,
+		NewConfig().WithResponseStreamingPlugins(plugin))
+
+	ctx := log.IntoContext(context.Background(), logr.Discard())
+
+	avg := testing.AllocsPerRun(100, func() {
+		reqCtx := &handlers.RequestContext{
+			Request: &handlers.Request{
+				Headers: map[string]string{
+					reqcommon.RequestIDHeaderKey: "alloc-request",
+				},
+			},
+			Response: &handlers.Response{
+				Headers: map[string]string{},
+			},
+			TargetPod: &fwkdl.EndpointMetadata{
+				NamespacedName: types.NamespacedName{Namespace: "ns", Name: "pod"},
+			},
+			Usage: fwkrh.Usage{},
+		}
+		for chunk := 0; chunk < benchChunksPerOp; chunk++ {
+			director.HandleResponseBody(ctx, reqCtx, false)
+		}
+		director.HandleResponseBody(ctx, reqCtx, true)
+
+		plugin.mu.Lock()
+		plugin.respsOnStreaming = plugin.respsOnStreaming[:0]
+		plugin.targetPodsOnStreaming = plugin.targetPodsOnStreaming[:0]
+		plugin.mu.Unlock()
+	})
+	if avg > maxAllocsPerResponse {
+		t.Errorf("HandleResponseBody allocations regressed: got %.0f, ceiling %d", avg, maxAllocsPerResponse)
+	}
+}
+
+// BenchmarkHandleResponseBody measures the per-response cost of
+// Director.HandleResponseBody with one streaming plugin registered.
+// Each operation simulates 10 intermediate chunks followed by one
+// end-of-stream chunk.
+//
+// Run:
+//
+//	go test -run='^$' -bench=BenchmarkHandleResponseBody -benchmem -count=10 \
+//	    ./pkg/epp/requestcontrol/ | tee bench.out
+//	benchstat bench.out
+func BenchmarkHandleResponseBody(b *testing.B) {
+	plugin := newTestResponseStreaming("bench-plugin")
+	director := NewDirectorWithConfig(nil, &mockScheduler{}, nil, nil,
+		NewConfig().WithResponseStreamingPlugins(plugin))
+
+	ctx := log.IntoContext(context.Background(), logr.Discard())
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		reqCtx := &handlers.RequestContext{
+			Request: &handlers.Request{
+				Headers: map[string]string{
+					reqcommon.RequestIDHeaderKey: "bench-request",
+				},
+			},
+			Response: &handlers.Response{
+				Headers: map[string]string{},
+			},
+			TargetPod: &fwkdl.EndpointMetadata{
+				NamespacedName: types.NamespacedName{Namespace: "ns", Name: "pod"},
+			},
+			Usage: fwkrh.Usage{},
+		}
+
+		for chunk := 0; chunk < benchChunksPerOp; chunk++ {
+			director.HandleResponseBody(ctx, reqCtx, false)
+		}
+		// Wait for async queue to drain before the final synchronous chunk.
+		director.HandleResponseBody(ctx, reqCtx, true)
+
+		// Reset plugin state for the next iteration.
+		plugin.mu.Lock()
+		plugin.respsOnStreaming = plugin.respsOnStreaming[:0]
+		plugin.targetPodsOnStreaming = plugin.targetPodsOnStreaming[:0]
+		plugin.mu.Unlock()
+	}
+}