Adding util func for splitting large bodies into chunks (#859)

Author: kfswain

pkg/epp/handlers/server.go

@@ -37,6 +37,11 @@ import (
 	requtil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/request"
 )
 
+const (
+	// Certain envoy implementations set a max limit of 64Kb per streamed chunk, intentionally setting this lower for a safe margin.
+	bodyByteLimit = 62000
+)
+
 func NewStreamingServer(destinationEndpointHintMetadataNamespace, destinationEndpointHintKey string, datastore Datastore, director Director) *StreamingServer {
 	return &StreamingServer{
 		destinationEndpointHintMetadataNamespace: destinationEndpointHintMetadataNamespace,
@@ -460,3 +465,32 @@ func BuildErrResponse(err error) (*extProcPb.ProcessingResponse, error) {
 
 	return resp, nil
 }
+
+func buildCommonResponses(bodyBytes []byte, byteLimit int) []*extProcPb.CommonResponse {
+	responses := []*extProcPb.CommonResponse{}
+	startingIndex := 0
+	bodyLen := len(bodyBytes)
+
+	for startingIndex < bodyLen {
+		eos := false
+		len := min(bodyLen-startingIndex, byteLimit)
+		chunk := bodyBytes[startingIndex : len+startingIndex]
+		if len+startingIndex == bodyLen {
+			eos = true
+		}
+
+		commonResp := &extProcPb.CommonResponse{
+			BodyMutation: &extProcPb.BodyMutation{
+				Mutation: &extProcPb.BodyMutation_StreamedResponse{
+					StreamedResponse: &extProcPb.StreamedBodyResponse{
+						Body:        chunk,
+						EndOfStream: eos,
+					},
+				},
+			},
+		}
+		responses = append(responses, commonResp)
+		startingIndex += len
+	}
+	return responses
+}

pkg/epp/handlers/server_test.go

@@ -0,0 +1,64 @@
+package handlers
+
+import (
+	"crypto/rand"
+	"testing"
+)
+
+func TestBuildCommonResponses(t *testing.T) {
+	tests := []struct {
+		name                 string
+		count                int
+		expectedMessageCount int
+	}{
+		{
+			name:                 "below limit",
+			count:                bodyByteLimit - 1000,
+			expectedMessageCount: 1,
+		},
+		{
+			name:                 "at limit",
+			count:                bodyByteLimit,
+			expectedMessageCount: 1,
+		},
+		{
+			name:                 "off by one error?",
+			count:                bodyByteLimit + 1,
+			expectedMessageCount: 2,
+		},
+		{
+			name:                 "above limit",
+			count:                bodyByteLimit + 1000,
+			expectedMessageCount: 2,
+		},
+		{
+			name:                 "above limit",
+			count:                (bodyByteLimit * 2) + 1000,
+			expectedMessageCount: 3,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			arr := generateBytes(test.count)
+			responses := buildCommonResponses(arr, bodyByteLimit)
+			for i, response := range responses {
+				eos := response.BodyMutation.GetStreamedResponse().GetEndOfStream()
+				if eos == true && i+1 != len(responses) {
+					t.Fatalf("EoS should not be set")
+				}
+				if eos == false && i+1 == len(responses) {
+					t.Fatalf("EoS should be set")
+				}
+			}
+			if len(responses) != test.expectedMessageCount {
+				t.Fatalf("Expected: %v, Got %v", test.expectedMessageCount, len(responses))
+			}
+		})
+	}
+}
+
+func generateBytes(count int) []byte {
+	arr := make([]byte, count)
+	_, _ = rand.Read(arr)
+	return arr
+}