Support full duplex streaming

Author: rramkumar1

cmd/body-based-routing/main.go

@@ -42,6 +42,8 @@ var (
 		"grpcHealthPort",
 		9003,
 		"The port used for gRPC liveness and readiness probes")
+	streaming = flag.Bool(
+		"streaming", false, "Enables streaming support for Envoy full-duplex streaming mode")
 	logVerbosity = flag.Int("v", logging.DEFAULT, "number for the log level verbosity")
 
 	setupLog = ctrl.Log.WithName("setup")
@@ -82,7 +84,7 @@ func run() error {
 	ctx := ctrl.SetupSignalHandler()
 
 	// Setup runner.
-	serverRunner := &runserver.ExtProcServerRunner{GrpcPort: *grpcPort}
+	serverRunner := runserver.NewDefaultExtProcServerRunner(*streaming)
 
 	// Register health server.
 	if err := registerHealthServer(mgr, ctrl.Log.WithName("health"), *grpcHealthPort); err != nil {

pkg/body-based-routing/handlers/request.go

@@ -18,7 +18,6 @@ package handlers
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 
 	basepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
@@ -28,14 +27,9 @@ import (
 )
 
 // HandleRequestBody handles request bodies.
-func (s *Server) HandleRequestBody(ctx context.Context, body *eppb.HttpBody) (*eppb.ProcessingResponse, error) {
+func (s *Server) HandleRequestBody(ctx context.Context, data map[string]any) (*eppb.ProcessingResponse, error) {
 	logger := log.FromContext(ctx)
 
-	var data map[string]any
-	if err := json.Unmarshal(body.GetBody(), &data); err != nil {
-		return nil, err
-	}
-
 	modelVal, ok := data["model"]
 	if !ok {
 		logger.V(logutil.DEFAULT).Info("Request body does not contain model parameter")
@@ -56,6 +50,27 @@ func (s *Server) HandleRequestBody(ctx context.Context, body *eppb.HttpBody) (*e
 		}, fmt.Errorf("the model parameter value %v is not a string", modelVal)
 	}
 
+	if s.streaming {
+		return &eppb.ProcessingResponse{
+			Response: &eppb.ProcessingResponse_RequestHeaders{
+				RequestHeaders: &eppb.HeadersResponse{
+					Response: &eppb.CommonResponse{
+						ClearRouteCache: true,
+						HeaderMutation: &eppb.HeaderMutation{
+							SetHeaders: []*basepb.HeaderValueOption{
+								{
+									Header: &basepb.HeaderValue{
+										Key:      "X-Gateway-Model-Name",
+										RawValue: []byte(modelStr),
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		}, nil
+	}
 	return &eppb.ProcessingResponse{
 		Response: &eppb.ProcessingResponse_RequestBody{
 			RequestBody: &eppb.BodyResponse{

pkg/body-based-routing/handlers/request_test.go

@@ -27,46 +27,20 @@ import (
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
-const (
-	bodyWithModel = `
-	{
-		"model": "foo",
-		"prompt": "Tell me a joke"
-	}
-	`
-	bodyWithModelNoStr = `
-	{
-		"model": 1,
-		"prompt": "Tell me a joke"
-	}
-	`
-	bodyWithoutModel = `
-	{
-		"prompt": "Tell me a joke"
-	}
-	`
-)
-
 func TestHandleRequestBody(t *testing.T) {
 	ctx := logutil.NewTestLoggerIntoContext(context.Background())
 
 	tests := []struct {
-		name    string
-		body    *extProcPb.HttpBody
-		want    *extProcPb.ProcessingResponse
-		wantErr bool
+		name      string
+		body      map[string]any
+		streaming bool
+		want      *extProcPb.ProcessingResponse
+		wantErr   bool
 	}{
-		{
-			name: "malformed body",
-			body: &extProcPb.HttpBody{
-				Body: []byte("malformed json"),
-			},
-			wantErr: true,
-		},
 		{
 			name: "model not found",
-			body: &extProcPb.HttpBody{
-				Body: []byte(bodyWithoutModel),
+			body: map[string]any{
+				"prompt": "Tell me a joke",
 			},
 			want: &extProcPb.ProcessingResponse{
 				Response: &extProcPb.ProcessingResponse_RequestBody{
@@ -76,15 +50,17 @@ func TestHandleRequestBody(t *testing.T) {
 		},
 		{
 			name: "model is not string",
-			body: &extProcPb.HttpBody{
-				Body: []byte(bodyWithModelNoStr),
+			body: map[string]any{
+				"model":  1,
+				"prompt": "Tell me a joke",
 			},
 			wantErr: true,
 		},
 		{
 			name: "success",
-			body: &extProcPb.HttpBody{
-				Body: []byte(bodyWithModel),
+			body: map[string]any{
+				"model":  "foo",
+				"prompt": "Tell me a joke",
 			},
 			want: &extProcPb.ProcessingResponse{
 				Response: &extProcPb.ProcessingResponse_RequestBody{
@@ -107,11 +83,38 @@ func TestHandleRequestBody(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "success-with-streaming",
+			body: map[string]any{
+				"model":  "foo",
+				"prompt": "Tell me a joke",
+			},
+			streaming: true,
+			want: &extProcPb.ProcessingResponse{
+				Response: &extProcPb.ProcessingResponse_RequestHeaders{
+					RequestHeaders: &extProcPb.HeadersResponse{
+						Response: &extProcPb.CommonResponse{
+							ClearRouteCache: true,
+							HeaderMutation: &extProcPb.HeaderMutation{
+								SetHeaders: []*basepb.HeaderValueOption{
+									{
+										Header: &basepb.HeaderValue{
+											Key:      "X-Gateway-Model-Name",
+											RawValue: []byte("foo"),
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			server := &Server{}
+			server := &Server{streaming: test.streaming}
 			resp, err := server.HandleRequestBody(ctx, test.body)
 			if err != nil {
 				if !test.wantErr {

pkg/body-based-routing/handlers/server.go

@@ -18,30 +18,36 @@ package handlers
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"io"
 
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+	"github.com/go-logr/logr"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
-func NewServer() *Server {
-	return &Server{}
+func NewServer(streaming bool) *Server {
+	return &Server{streaming: streaming}
 }
 
 // Server implements the Envoy external processing server.
 // https://www.envoyproxy.io/docs/envoy/latest/api-v3/service/ext_proc/v3/external_processor.proto
-type Server struct{}
+type Server struct {
+	streaming bool
+}
 
 func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 	ctx := srv.Context()
 	logger := log.FromContext(ctx)
 	loggerVerbose := logger.V(logutil.VERBOSE)
 	loggerVerbose.Info("Processing")
 
+	reader, writer := io.Pipe()
+
 	for {
 		select {
 		case <-ctx.Done():
@@ -61,12 +67,21 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 		}
 
 		var resp *extProcPb.ProcessingResponse
+		var requestBody []byte
 		var err error
 		switch v := req.Request.(type) {
 		case *extProcPb.ProcessingRequest_RequestHeaders:
-			resp, err = s.HandleRequestHeaders(req.GetRequestHeaders())
+			if !s.streaming {
+				// If streaming, then headers are handled when processing request body.
+				resp, err = s.HandleRequestHeaders(req.GetRequestHeaders())
+			} else {
+				loggerVerbose.Info("Received headers, passing off header processing until body arrives...")
+			}
 		case *extProcPb.ProcessingRequest_RequestBody:
-			resp, err = s.HandleRequestBody(ctx, req.GetRequestBody())
+			loggerVerbose.Info("Incoming body chunk", "body", string(v.RequestBody.Body), "EoS", v.RequestBody.EndOfStream)
+			resp, requestBody, err = s.processRequestBody(ctx, req.GetRequestBody(), writer, reader, logger)
+		case *extProcPb.ProcessingRequest_RequestTrailers:
+			resp, err = s.HandleRequestTrailers(req.GetRequestTrailers())
 		case *extProcPb.ProcessingRequest_ResponseHeaders:
 			resp, err = s.HandleResponseHeaders(req.GetResponseHeaders())
 		case *extProcPb.ProcessingRequest_ResponseBody:
@@ -81,10 +96,84 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 			return status.Errorf(status.Code(err), "failed to handle request: %v", err)
 		}
 
-		loggerVerbose.Info("Response generated", "response", resp)
-		if err := srv.Send(resp); err != nil {
-			logger.V(logutil.DEFAULT).Error(err, "Send failed")
-			return status.Errorf(codes.Unknown, "failed to send response back to Envoy: %v", err)
+		if resp != nil {
+			loggerVerbose.Info("Response generated", "response", resp)
+			if err := srv.Send(resp); err != nil {
+				logger.V(logutil.DEFAULT).Error(err, "Send failed")
+				return status.Errorf(codes.Unknown, "failed to send response back to Envoy: %v", err)
+			}
+
+			if s.streaming {
+				bodyResp := &extProcPb.ProcessingResponse{
+					Response: &extProcPb.ProcessingResponse_RequestBody{
+						RequestBody: &extProcPb.BodyResponse{
+							Response: &extProcPb.CommonResponse{
+								BodyMutation: &extProcPb.BodyMutation{
+									Mutation: &extProcPb.BodyMutation_StreamedResponse{
+										StreamedResponse: &extProcPb.StreamedBodyResponse{
+											Body:        requestBody,
+											EndOfStream: true,
+										},
+									},
+								},
+							},
+						},
+					},
+				}
+				loggerVerbose.Info("Response generated", "response", bodyResp)
+				if err := srv.Send(bodyResp); err != nil {
+					logger.V(logutil.DEFAULT).Error(err, "Send failed")
+					return status.Errorf(codes.Unknown, "failed to send response back to Envoy: %v", err)
+				}
+			}
 		}
 	}
 }
+
+func (s *Server) processRequestBody(ctx context.Context, body *extProcPb.HttpBody, bufferWriter *io.PipeWriter, bufferReader *io.PipeReader, logger logr.Logger) (*extProcPb.ProcessingResponse, []byte, error) {
+	loggerVerbose := logger.V(logutil.VERBOSE)
+
+	var requestBody map[string]interface{}
+	if s.streaming {
+		// In the stream case, we can receive multiple request bodies.
+		// To buffer the full message, we create a goroutine with a writer.Write()
+		// call, which will block until the corresponding reader reads from it.
+		// We do not read until we receive the EndofStream signal, and then
+		// decode the entire JSON body.
+		if !body.EndOfStream {
+			go func() {
+				loggerVerbose.Info("Writing to stream buffer")
+				_, err := bufferWriter.Write(body.Body)
+				if err != nil {
+					logger.V(logutil.DEFAULT).Error(err, "Error populating writer")
+				}
+			}()
+
+			return nil, nil, nil
+		}
+
+		if body.EndOfStream {
+			loggerVerbose.Info("Flushing stream buffer")
+			decoder := json.NewDecoder(bufferReader)
+			if err := decoder.Decode(&requestBody); err != nil {
+				logger.V(logutil.DEFAULT).Error(err, "Error unmarshaling request body")
+			}
+			bufferReader.Close()
+		}
+	} else {
+		if err := json.Unmarshal(body.GetBody(), &requestBody); err != nil {
+			return nil, nil, err
+		}
+	}
+
+	requestBodyResp, err := s.HandleRequestBody(ctx, requestBody)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	requestBodyBytes, err := json.Marshal(requestBody)
+	if err != nil {
+		return nil, nil, err
+	}
+	return requestBodyResp, requestBodyBytes, nil
+}

pkg/body-based-routing/server/runserver.go

@@ -32,17 +32,19 @@ import (
 
 // ExtProcServerRunner provides methods to manage an external process server.
 type ExtProcServerRunner struct {
-	GrpcPort int
+	GrpcPort  int
+	Streaming bool
 }
 
 // Default values for CLI flags in main
 const (
 	DefaultGrpcPort = 9002 // default for --grpcPort
 )
 
-func NewDefaultExtProcServerRunner() *ExtProcServerRunner {
+func NewDefaultExtProcServerRunner(streaming bool) *ExtProcServerRunner {
 	return &ExtProcServerRunner{
-		GrpcPort: DefaultGrpcPort,
+		GrpcPort:  DefaultGrpcPort,
+		Streaming: streaming,
 	}
 }
 
@@ -60,7 +62,7 @@ func (r *ExtProcServerRunner) AsRunnable(logger logr.Logger) manager.Runnable {
 		srv := grpc.NewServer(grpc.Creds(creds))
 		extProcPb.RegisterExternalProcessorServer(
 			srv,
-			handlers.NewServer(),
+			handlers.NewServer(r.Streaming),
 		)
 
 		// Forward to the gRPC runnable.