Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit e6d3ac2

Browse files
rramkumar1rramkumar1
committedMar 19, 2025·
Support full duplex streaming
1 parent 53cb18f commit e6d3ac2

File tree

7 files changed

+319
-122
lines changed

7 files changed

+319
-122
lines changed
 

‎cmd/body-based-routing/main.go

+4-2
Original file line numberDiff line numberDiff line change
@@ -44,14 +44,16 @@ import (
4444
var (
4545
grpcPort = flag.Int(
4646
"grpcPort",
47-
runserver.DefaultGrpcPort,
47+
9004,
4848
"The gRPC port used for communicating with Envoy proxy")
4949
grpcHealthPort = flag.Int(
5050
"grpcHealthPort",
5151
9005,
5252
"The port used for gRPC liveness and readiness probes")
5353
metricsPort = flag.Int(
5454
"metricsPort", 9090, "The metrics port")
55+
streaming = flag.Bool(
56+
"streaming", false, "Enables streaming support for Envoy full-duplex streaming mode")
5557
logVerbosity = flag.Int("v", logging.DEFAULT, "number for the log level verbosity")
5658

5759
setupLog = ctrl.Log.WithName("setup")
@@ -92,7 +94,7 @@ func run() error {
9294
ctx := ctrl.SetupSignalHandler()
9395

9496
// Setup runner.
95-
serverRunner := &runserver.ExtProcServerRunner{GrpcPort: *grpcPort}
97+
serverRunner := runserver.NewDefaultExtProcServerRunner(*grpcPort, *streaming)
9698

9799
// Register health server.
98100
if err := registerHealthServer(mgr, ctrl.Log.WithName("health"), *grpcHealthPort); err != nil {

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

+94-33
Original file line numberDiff line numberDiff line change
@@ -23,55 +23,93 @@ import (
2323

2424
basepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
2525
eppb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
26+
extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
2627
"sigs.k8s.io/controller-runtime/pkg/log"
2728
"sigs.k8s.io/gateway-api-inference-extension/pkg/body-based-routing/metrics"
2829
logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
2930
)
3031

32+
const modelHeader = "X-Gateway-Model-Name"
33+
3134
// HandleRequestBody handles request bodies.
32-
func (s *Server) HandleRequestBody(ctx context.Context, body *eppb.HttpBody) (*eppb.ProcessingResponse, error) {
35+
func (s *Server) HandleRequestBody(ctx context.Context, data map[string]any) ([]*eppb.ProcessingResponse, error) {
3336
logger := log.FromContext(ctx)
37+
var ret []*eppb.ProcessingResponse
3438

35-
var data map[string]any
36-
if err := json.Unmarshal(body.GetBody(), &data); err != nil {
39+
requestBodyBytes, err := json.Marshal(data)
40+
if err != nil {
3741
return nil, err
3842
}
3943

4044
modelVal, ok := data["model"]
4145
if !ok {
4246
metrics.RecordModelNotInBodyCounter()
4347
logger.V(logutil.DEFAULT).Info("Request body does not contain model parameter")
44-
return &eppb.ProcessingResponse{
45-
Response: &eppb.ProcessingResponse_RequestBody{
46-
RequestBody: &eppb.BodyResponse{},
47-
},
48-
}, nil
48+
if s.streaming {
49+
ret = append(ret, &eppb.ProcessingResponse{
50+
Response: &eppb.ProcessingResponse_RequestHeaders{
51+
RequestHeaders: &eppb.HeadersResponse{},
52+
},
53+
})
54+
ret = addStreamedBodyResponse(ret, requestBodyBytes)
55+
return ret, nil
56+
} else {
57+
ret = append(ret, &eppb.ProcessingResponse{
58+
Response: &eppb.ProcessingResponse_RequestBody{
59+
RequestBody: &eppb.BodyResponse{},
60+
},
61+
})
62+
}
63+
return ret, nil
4964
}
5065

5166
modelStr, ok := modelVal.(string)
5267
if !ok {
5368
metrics.RecordModelNotParsedCounter()
5469
logger.V(logutil.DEFAULT).Info("Model parameter value is not a string")
55-
return &eppb.ProcessingResponse{
56-
Response: &eppb.ProcessingResponse_RequestBody{
57-
RequestBody: &eppb.BodyResponse{},
58-
},
59-
}, fmt.Errorf("the model parameter value %v is not a string", modelVal)
70+
return nil, fmt.Errorf("the model parameter value %v is not a string", modelVal)
6071
}
6172

6273
metrics.RecordSuccessCounter()
63-
return &eppb.ProcessingResponse{
64-
Response: &eppb.ProcessingResponse_RequestBody{
65-
RequestBody: &eppb.BodyResponse{
66-
Response: &eppb.CommonResponse{
67-
// Necessary so that the new headers are used in the routing decision.
68-
ClearRouteCache: true,
69-
HeaderMutation: &eppb.HeaderMutation{
70-
SetHeaders: []*basepb.HeaderValueOption{
71-
{
72-
Header: &basepb.HeaderValue{
73-
Key: "X-Gateway-Model-Name",
74-
RawValue: []byte(modelStr),
74+
75+
if s.streaming {
76+
ret = append(ret, &eppb.ProcessingResponse{
77+
Response: &eppb.ProcessingResponse_RequestHeaders{
78+
RequestHeaders: &eppb.HeadersResponse{
79+
Response: &eppb.CommonResponse{
80+
ClearRouteCache: true,
81+
HeaderMutation: &eppb.HeaderMutation{
82+
SetHeaders: []*basepb.HeaderValueOption{
83+
{
84+
Header: &basepb.HeaderValue{
85+
Key: modelHeader,
86+
RawValue: []byte(modelStr),
87+
},
88+
},
89+
},
90+
},
91+
},
92+
},
93+
},
94+
})
95+
ret = addStreamedBodyResponse(ret, requestBodyBytes)
96+
return ret, nil
97+
}
98+
99+
return []*eppb.ProcessingResponse{
100+
{
101+
Response: &eppb.ProcessingResponse_RequestBody{
102+
RequestBody: &eppb.BodyResponse{
103+
Response: &eppb.CommonResponse{
104+
// Necessary so that the new headers are used in the routing decision.
105+
ClearRouteCache: true,
106+
HeaderMutation: &eppb.HeaderMutation{
107+
SetHeaders: []*basepb.HeaderValueOption{
108+
{
109+
Header: &basepb.HeaderValue{
110+
Key: modelHeader,
111+
RawValue: []byte(modelStr),
112+
},
75113
},
76114
},
77115
},
@@ -82,20 +120,43 @@ func (s *Server) HandleRequestBody(ctx context.Context, body *eppb.HttpBody) (*e
82120
}, nil
83121
}
84122

123+
func addStreamedBodyResponse(responses []*eppb.ProcessingResponse, requestBodyBytes []byte) []*eppb.ProcessingResponse {
124+
return append(responses, &extProcPb.ProcessingResponse{
125+
Response: &extProcPb.ProcessingResponse_RequestBody{
126+
RequestBody: &extProcPb.BodyResponse{
127+
Response: &extProcPb.CommonResponse{
128+
BodyMutation: &extProcPb.BodyMutation{
129+
Mutation: &extProcPb.BodyMutation_StreamedResponse{
130+
StreamedResponse: &extProcPb.StreamedBodyResponse{
131+
Body: requestBodyBytes,
132+
EndOfStream: true,
133+
},
134+
},
135+
},
136+
},
137+
},
138+
},
139+
})
140+
}
141+
85142
// HandleRequestHeaders handles request headers.
86-
func (s *Server) HandleRequestHeaders(headers *eppb.HttpHeaders) (*eppb.ProcessingResponse, error) {
87-
return &eppb.ProcessingResponse{
88-
Response: &eppb.ProcessingResponse_RequestHeaders{
89-
RequestHeaders: &eppb.HeadersResponse{},
143+
func (s *Server) HandleRequestHeaders(headers *eppb.HttpHeaders) ([]*eppb.ProcessingResponse, error) {
144+
return []*eppb.ProcessingResponse{
145+
{
146+
Response: &eppb.ProcessingResponse_RequestHeaders{
147+
RequestHeaders: &eppb.HeadersResponse{},
148+
},
90149
},
91150
}, nil
92151
}
93152

94153
// HandleRequestTrailers handles request trailers.
95-
func (s *Server) HandleRequestTrailers(trailers *eppb.HttpTrailers) (*eppb.ProcessingResponse, error) {
96-
return &eppb.ProcessingResponse{
97-
Response: &eppb.ProcessingResponse_RequestTrailers{
98-
RequestTrailers: &eppb.TrailersResponse{},
154+
func (s *Server) HandleRequestTrailers(trailers *eppb.HttpTrailers) ([]*eppb.ProcessingResponse, error) {
155+
return []*eppb.ProcessingResponse{
156+
{
157+
Response: &eppb.ProcessingResponse_RequestTrailers{
158+
RequestTrailers: &eppb.TrailersResponse{},
159+
},
99160
},
100161
}, nil
101162
}

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

+121-51
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ package handlers
1818

1919
import (
2020
"context"
21+
"encoding/json"
2122
"strings"
2223
"testing"
2324

@@ -31,78 +32,138 @@ import (
3132
logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
3233
)
3334

34-
const (
35-
bodyWithModel = `
36-
{
37-
"model": "foo",
38-
"prompt": "Tell me a joke"
39-
}
40-
`
41-
bodyWithModelNoStr = `
42-
{
43-
"model": 1,
44-
"prompt": "Tell me a joke"
45-
}
46-
`
47-
bodyWithoutModel = `
48-
{
49-
"prompt": "Tell me a joke"
50-
}
51-
`
52-
)
53-
5435
func TestHandleRequestBody(t *testing.T) {
5536
metrics.Register()
5637
ctx := logutil.NewTestLoggerIntoContext(context.Background())
5738

5839
tests := []struct {
59-
name string
60-
body *extProcPb.HttpBody
61-
want *extProcPb.ProcessingResponse
62-
wantErr bool
40+
name string
41+
body map[string]any
42+
streaming bool
43+
want []*extProcPb.ProcessingResponse
44+
wantErr bool
6345
}{
6446
{
65-
name: "malformed body",
66-
body: &extProcPb.HttpBody{
67-
Body: []byte("malformed json"),
47+
name: "model not found",
48+
body: map[string]any{
49+
"prompt": "Tell me a joke",
50+
},
51+
want: []*extProcPb.ProcessingResponse{
52+
{
53+
Response: &extProcPb.ProcessingResponse_RequestBody{
54+
RequestBody: &extProcPb.BodyResponse{},
55+
},
56+
},
6857
},
69-
wantErr: true,
7058
},
7159
{
72-
name: "model not found",
73-
body: &extProcPb.HttpBody{
74-
Body: []byte(bodyWithoutModel),
60+
name: "model not found with streaming",
61+
body: map[string]any{
62+
"prompt": "Tell me a joke",
7563
},
76-
want: &extProcPb.ProcessingResponse{
77-
Response: &extProcPb.ProcessingResponse_RequestBody{
78-
RequestBody: &extProcPb.BodyResponse{},
64+
streaming: true,
65+
want: []*extProcPb.ProcessingResponse{
66+
{
67+
Response: &extProcPb.ProcessingResponse_RequestHeaders{
68+
RequestHeaders: &extProcPb.HeadersResponse{},
69+
},
70+
},
71+
{
72+
Response: &extProcPb.ProcessingResponse_RequestBody{
73+
RequestBody: &extProcPb.BodyResponse{
74+
Response: &extProcPb.CommonResponse{
75+
BodyMutation: &extProcPb.BodyMutation{
76+
Mutation: &extProcPb.BodyMutation_StreamedResponse{
77+
StreamedResponse: &extProcPb.StreamedBodyResponse{
78+
Body: mapToBytes(t, map[string]any{
79+
"prompt": "Tell me a joke",
80+
}),
81+
EndOfStream: true,
82+
},
83+
},
84+
},
85+
},
86+
},
87+
},
7988
},
8089
},
8190
},
8291
{
8392
name: "model is not string",
84-
body: &extProcPb.HttpBody{
85-
Body: []byte(bodyWithModelNoStr),
93+
body: map[string]any{
94+
"model": 1,
95+
"prompt": "Tell me a joke",
8696
},
8797
wantErr: true,
8898
},
8999
{
90100
name: "success",
91-
body: &extProcPb.HttpBody{
92-
Body: []byte(bodyWithModel),
101+
body: map[string]any{
102+
"model": "foo",
103+
"prompt": "Tell me a joke",
93104
},
94-
want: &extProcPb.ProcessingResponse{
95-
Response: &extProcPb.ProcessingResponse_RequestBody{
96-
RequestBody: &extProcPb.BodyResponse{
97-
Response: &extProcPb.CommonResponse{
98-
// Necessary so that the new headers are used in the routing decision.
99-
ClearRouteCache: true,
100-
HeaderMutation: &extProcPb.HeaderMutation{
101-
SetHeaders: []*basepb.HeaderValueOption{
102-
{
103-
Header: &basepb.HeaderValue{
104-
Key: "X-Gateway-Model-Name",
105-
RawValue: []byte("foo"),
105+
want: []*extProcPb.ProcessingResponse{
106+
{
107+
Response: &extProcPb.ProcessingResponse_RequestBody{
108+
RequestBody: &extProcPb.BodyResponse{
109+
Response: &extProcPb.CommonResponse{
110+
// Necessary so that the new headers are used in the routing decision.
111+
ClearRouteCache: true,
112+
HeaderMutation: &extProcPb.HeaderMutation{
113+
SetHeaders: []*basepb.HeaderValueOption{
114+
{
115+
Header: &basepb.HeaderValue{
116+
Key: "X-Gateway-Model-Name",
117+
RawValue: []byte("foo"),
118+
},
119+
},
120+
},
121+
},
122+
},
123+
},
124+
},
125+
},
126+
},
127+
},
128+
{
129+
name: "success-with-streaming",
130+
body: map[string]any{
131+
"model": "foo",
132+
"prompt": "Tell me a joke",
133+
},
134+
streaming: true,
135+
want: []*extProcPb.ProcessingResponse{
136+
{
137+
Response: &extProcPb.ProcessingResponse_RequestHeaders{
138+
RequestHeaders: &extProcPb.HeadersResponse{
139+
Response: &extProcPb.CommonResponse{
140+
ClearRouteCache: true,
141+
HeaderMutation: &extProcPb.HeaderMutation{
142+
SetHeaders: []*basepb.HeaderValueOption{
143+
{
144+
Header: &basepb.HeaderValue{
145+
Key: "X-Gateway-Model-Name",
146+
RawValue: []byte("foo"),
147+
},
148+
},
149+
},
150+
},
151+
},
152+
},
153+
},
154+
},
155+
{
156+
Response: &extProcPb.ProcessingResponse_RequestBody{
157+
RequestBody: &extProcPb.BodyResponse{
158+
Response: &extProcPb.CommonResponse{
159+
BodyMutation: &extProcPb.BodyMutation{
160+
Mutation: &extProcPb.BodyMutation_StreamedResponse{
161+
StreamedResponse: &extProcPb.StreamedBodyResponse{
162+
Body: mapToBytes(t, map[string]any{
163+
"model": "foo",
164+
"prompt": "Tell me a joke",
165+
}),
166+
EndOfStream: true,
106167
},
107168
},
108169
},
@@ -116,7 +177,7 @@ func TestHandleRequestBody(t *testing.T) {
116177

117178
for _, test := range tests {
118179
t.Run(test.name, func(t *testing.T) {
119-
server := &Server{}
180+
server := &Server{streaming: test.streaming}
120181
resp, err := server.HandleRequestBody(ctx, test.body)
121182
if err != nil {
122183
if !test.wantErr {
@@ -147,3 +208,12 @@ func TestHandleRequestBody(t *testing.T) {
147208
t.Error(err)
148209
}
149210
}
211+
212+
func mapToBytes(t *testing.T, m map[string]any) []byte {
213+
// Convert map to JSON byte array
214+
bytes, err := json.Marshal(m)
215+
if err != nil {
216+
t.Fatalf("Marshal(): %v", err)
217+
}
218+
return bytes
219+
}

‎pkg/body-based-routing/handlers/response.go

+18-12
Original file line numberDiff line numberDiff line change
@@ -21,28 +21,34 @@ import (
2121
)
2222

2323
// HandleResponseHeaders handles response headers.
24-
func (s *Server) HandleResponseHeaders(headers *eppb.HttpHeaders) (*eppb.ProcessingResponse, error) {
25-
return &eppb.ProcessingResponse{
26-
Response: &eppb.ProcessingResponse_ResponseHeaders{
27-
ResponseHeaders: &eppb.HeadersResponse{},
24+
func (s *Server) HandleResponseHeaders(headers *eppb.HttpHeaders) ([]*eppb.ProcessingResponse, error) {
25+
return []*eppb.ProcessingResponse{
26+
&eppb.ProcessingResponse{
27+
Response: &eppb.ProcessingResponse_ResponseHeaders{
28+
ResponseHeaders: &eppb.HeadersResponse{},
29+
},
2830
},
2931
}, nil
3032
}
3133

3234
// HandleResponseBody handles response bodies.
33-
func (s *Server) HandleResponseBody(body *eppb.HttpBody) (*eppb.ProcessingResponse, error) {
34-
return &eppb.ProcessingResponse{
35-
Response: &eppb.ProcessingResponse_ResponseBody{
36-
ResponseBody: &eppb.BodyResponse{},
35+
func (s *Server) HandleResponseBody(body *eppb.HttpBody) ([]*eppb.ProcessingResponse, error) {
36+
return []*eppb.ProcessingResponse{
37+
&eppb.ProcessingResponse{
38+
Response: &eppb.ProcessingResponse_ResponseBody{
39+
ResponseBody: &eppb.BodyResponse{},
40+
},
3741
},
3842
}, nil
3943
}
4044

4145
// HandleResponseTrailers handles response trailers.
42-
func (s *Server) HandleResponseTrailers(trailers *eppb.HttpTrailers) (*eppb.ProcessingResponse, error) {
43-
return &eppb.ProcessingResponse{
44-
Response: &eppb.ProcessingResponse_ResponseTrailers{
45-
ResponseTrailers: &eppb.TrailersResponse{},
46+
func (s *Server) HandleResponseTrailers(trailers *eppb.HttpTrailers) ([]*eppb.ProcessingResponse, error) {
47+
return []*eppb.ProcessingResponse{
48+
&eppb.ProcessingResponse{
49+
Response: &eppb.ProcessingResponse_ResponseTrailers{
50+
ResponseTrailers: &eppb.TrailersResponse{},
51+
},
4652
},
4753
}, nil
4854
}

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

+71-13
Original file line numberDiff line numberDiff line change
@@ -18,30 +18,36 @@ package handlers
1818

1919
import (
2020
"context"
21+
"encoding/json"
2122
"errors"
2223
"io"
2324

2425
extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
26+
"github.com/go-logr/logr"
2527
"google.golang.org/grpc/codes"
2628
"google.golang.org/grpc/status"
2729
"sigs.k8s.io/controller-runtime/pkg/log"
2830
logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
2931
)
3032

31-
func NewServer() *Server {
32-
return &Server{}
33+
func NewServer(streaming bool) *Server {
34+
return &Server{streaming: streaming}
3335
}
3436

3537
// Server implements the Envoy external processing server.
3638
// https://www.envoyproxy.io/docs/envoy/latest/api-v3/service/ext_proc/v3/external_processor.proto
37-
type Server struct{}
39+
type Server struct {
40+
streaming bool
41+
}
3842

3943
func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
4044
ctx := srv.Context()
4145
logger := log.FromContext(ctx)
4246
loggerVerbose := logger.V(logutil.VERBOSE)
4347
loggerVerbose.Info("Processing")
4448

49+
reader, writer := io.Pipe()
50+
4551
for {
4652
select {
4753
case <-ctx.Done():
@@ -60,19 +66,25 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
6066
return status.Errorf(codes.Unknown, "cannot receive stream request: %v", recvErr)
6167
}
6268

63-
var resp *extProcPb.ProcessingResponse
69+
var responses []*extProcPb.ProcessingResponse
6470
var err error
6571
switch v := req.Request.(type) {
6672
case *extProcPb.ProcessingRequest_RequestHeaders:
67-
resp, err = s.HandleRequestHeaders(req.GetRequestHeaders())
73+
if s.streaming && !req.GetRequestHeaders().GetEndOfStream() {
74+
// If streaming and the body is not empty, then headers are handled when processing request body.
75+
loggerVerbose.Info("Received headers, passing off header processing until body arrives...")
76+
} else {
77+
responses, err = s.HandleRequestHeaders(req.GetRequestHeaders())
78+
}
6879
case *extProcPb.ProcessingRequest_RequestBody:
69-
resp, err = s.HandleRequestBody(ctx, req.GetRequestBody())
80+
loggerVerbose.Info("Incoming body chunk", "body", string(v.RequestBody.Body), "EoS", v.RequestBody.EndOfStream)
81+
responses, err = s.processRequestBody(ctx, req.GetRequestBody(), writer, reader, logger)
7082
case *extProcPb.ProcessingRequest_RequestTrailers:
71-
resp, err = s.HandleRequestTrailers(req.GetRequestTrailers())
83+
responses, err = s.HandleRequestTrailers(req.GetRequestTrailers())
7284
case *extProcPb.ProcessingRequest_ResponseHeaders:
73-
resp, err = s.HandleResponseHeaders(req.GetResponseHeaders())
85+
responses, err = s.HandleResponseHeaders(req.GetResponseHeaders())
7486
case *extProcPb.ProcessingRequest_ResponseBody:
75-
resp, err = s.HandleResponseBody(req.GetResponseBody())
87+
responses, err = s.HandleResponseBody(req.GetResponseBody())
7688
default:
7789
logger.V(logutil.DEFAULT).Error(nil, "Unknown Request type", "request", v)
7890
return status.Error(codes.Unknown, "unknown request type")
@@ -83,10 +95,56 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
8395
return status.Errorf(status.Code(err), "failed to handle request: %v", err)
8496
}
8597

86-
loggerVerbose.Info("Response generated", "response", resp)
87-
if err := srv.Send(resp); err != nil {
88-
logger.V(logutil.DEFAULT).Error(err, "Send failed")
89-
return status.Errorf(codes.Unknown, "failed to send response back to Envoy: %v", err)
98+
for _, resp := range responses {
99+
loggerVerbose.Info("Response generated", "response", resp)
100+
if err := srv.Send(resp); err != nil {
101+
logger.V(logutil.DEFAULT).Error(err, "Send failed")
102+
return status.Errorf(codes.Unknown, "failed to send response back to Envoy: %v", err)
103+
}
90104
}
91105
}
92106
}
107+
108+
func (s *Server) processRequestBody(ctx context.Context, body *extProcPb.HttpBody, bufferWriter *io.PipeWriter, bufferReader *io.PipeReader, logger logr.Logger) ([]*extProcPb.ProcessingResponse, error) {
109+
loggerVerbose := logger.V(logutil.VERBOSE)
110+
111+
var requestBody map[string]interface{}
112+
if s.streaming {
113+
// In the stream case, we can receive multiple request bodies.
114+
// To buffer the full message, we create a goroutine with a writer.Write()
115+
// call, which will block until the corresponding reader reads from it.
116+
// We do not read until we receive the EndofStream signal, and then
117+
// decode the entire JSON body.
118+
if !body.EndOfStream {
119+
go func() {
120+
loggerVerbose.Info("Writing to stream buffer")
121+
_, err := bufferWriter.Write(body.Body)
122+
if err != nil {
123+
logger.V(logutil.DEFAULT).Error(err, "Error populating writer")
124+
}
125+
}()
126+
127+
return nil, nil
128+
}
129+
130+
if body.EndOfStream {
131+
loggerVerbose.Info("Flushing stream buffer")
132+
decoder := json.NewDecoder(bufferReader)
133+
if err := decoder.Decode(&requestBody); err != nil {
134+
logger.V(logutil.DEFAULT).Error(err, "Error unmarshaling request body")
135+
}
136+
bufferReader.Close()
137+
}
138+
} else {
139+
if err := json.Unmarshal(body.GetBody(), &requestBody); err != nil {
140+
return nil, err
141+
}
142+
}
143+
144+
requestBodyResp, err := s.HandleRequestBody(ctx, requestBody)
145+
if err != nil {
146+
return nil, err
147+
}
148+
149+
return requestBodyResp, nil
150+
}

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

+8-8
Original file line numberDiff line numberDiff line change
@@ -34,17 +34,14 @@ import (
3434
type ExtProcServerRunner struct {
3535
GrpcPort int
3636
SecureServing bool
37+
Streaming bool
3738
}
3839

39-
// Default values for CLI flags in main
40-
const (
41-
DefaultGrpcPort = 9004 // default for --grpcPort
42-
)
43-
44-
func NewDefaultExtProcServerRunner() *ExtProcServerRunner {
40+
func NewDefaultExtProcServerRunner(port int, streaming bool) *ExtProcServerRunner {
4541
return &ExtProcServerRunner{
46-
GrpcPort: DefaultGrpcPort,
42+
GrpcPort: port,
4743
SecureServing: true,
44+
Streaming: streaming,
4845
}
4946
}
5047

@@ -65,7 +62,10 @@ func (r *ExtProcServerRunner) AsRunnable(logger logr.Logger) manager.Runnable {
6562
srv = grpc.NewServer()
6663
}
6764

68-
extProcPb.RegisterExternalProcessorServer(srv, handlers.NewServer())
65+
extProcPb.RegisterExternalProcessorServer(
66+
srv,
67+
handlers.NewServer(r.Streaming),
68+
)
6969

7070
// Forward to the gRPC runnable.
7171
return runnable.GRPCServer("ext-proc", srv, r.GrpcPort).Start(ctx)

‎test/integration/bbr/hermetic_test.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,6 @@ import (
3535
logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
3636
)
3737

38-
const port = runserver.DefaultGrpcPort
39-
4038
var logger = logutil.NewTestLogger().V(logutil.VERBOSE)
4139

4240
func TestBodyBasedRouting(t *testing.T) {
@@ -102,8 +100,10 @@ func TestBodyBasedRouting(t *testing.T) {
102100
}
103101

104102
func setUpHermeticServer() (client extProcPb.ExternalProcessor_ProcessClient, cleanup func()) {
103+
port := 9004
104+
105105
serverCtx, stopServer := context.WithCancel(context.Background())
106-
serverRunner := runserver.NewDefaultExtProcServerRunner()
106+
serverRunner := runserver.NewDefaultExtProcServerRunner(port, false)
107107
serverRunner.SecureServing = false
108108

109109
go func() {

0 commit comments

Comments
 (0)
Please sign in to comment.