Pluggable BBR framework (WIP) -- iteration 2

Author: davidbreitgand

cmd/bbr/runner/runner.go

@@ -34,6 +34,7 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/internal/runnable"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/metrics"
 	runserver "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/server"
+	bbrutils "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/utils"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
@@ -102,8 +103,20 @@ func (r *Runner) Run(ctx context.Context) error {
 		return err
 	}
 
+	//Initialize PluginRegistry and request/response PluginsChain instances
+	registry, requestChain, responseChain, metaDataKeys, err := bbrutils.InitPlugins()
+	if err != nil {
+		setupLog.Error(err, "Failed to initialize plugins")
+		return err
+	}
+
 	// Setup runner.
-	serverRunner := runserver.NewDefaultExtProcServerRunner(*grpcPort, *streaming)
+	serverRunner := runserver.NewDefaultExtProcServerRunner(*grpcPort,
+		*streaming,
+		registry,
+		requestChain,
+		responseChain,
+		metaDataKeys)
 
 	// Register health server.
 	if err := registerHealthServer(mgr, *grpcHealthPort); err != nil {

pkg/bbr/handlers/request.go

@@ -18,54 +18,49 @@ package handlers
 
 import (
 	"context"
-	"encoding/json"
+	"strings"
 
 	basepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	eppb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/metrics"
+	bbrplugins "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/plugins"
+	helpers "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/utils"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
-const modelHeader = "X-Gateway-Model-Name"
-
-type RequestBody struct {
-	Model string `json:"model"`
-}
-
 // HandleRequestBody handles request bodies.
 func (s *Server) HandleRequestBody(ctx context.Context, requestBodyBytes []byte) ([]*eppb.ProcessingResponse, error) {
 	logger := log.FromContext(ctx)
 	var ret []*eppb.ProcessingResponse
 
-	var requestBody RequestBody
-	if err := json.Unmarshal(requestBodyBytes, &requestBody); err != nil {
-		metrics.RecordModelNotParsedCounter()
-		return nil, err
+	allHeaders, mutatedBodyBytes, err := s.requestChain.Run(requestBodyBytes, s.metaDataKeys, s.registry)
+
+	if err != nil {
+		//TODO: add metric in metrics.go to count "other errors"
+		logger.V(logutil.DEFAULT).Info("error processing body", "error", err)
+		ret, _ := buildEmptyResponsesForMissingModel(s.streaming, requestBodyBytes)
+		return ret, nil
+	}
+
+	model, ok := allHeaders[bbrplugins.ModelHeader]
+
+	if !ok {
+		//TODO: add metric in metrics.go to count "other errors"
+		logger.V(logutil.DEFAULT).Info("manadatory header X-Gateway-Model-Name value is undetermined")
+		ret, _ := buildEmptyResponsesForMissingModel(s.streaming, requestBodyBytes)
+		return ret, nil
 	}
 
-	if requestBody.Model == "" {
+	if strings.TrimSpace(model) == "" {
 		metrics.RecordModelNotInBodyCounter()
-		logger.V(logutil.DEFAULT).Info("Request body does not contain model parameter")
-		if s.streaming {
-			ret = append(ret, &eppb.ProcessingResponse{
-				Response: &eppb.ProcessingResponse_RequestHeaders{
-					RequestHeaders: &eppb.HeadersResponse{},
-				},
-			})
-			ret = addStreamedBodyResponse(ret, requestBodyBytes)
-			return ret, nil
-		} else {
-			ret = append(ret, &eppb.ProcessingResponse{
-				Response: &eppb.ProcessingResponse_RequestBody{
-					RequestBody: &eppb.BodyResponse{},
-				},
-			})
-		}
+		ret, _ := buildEmptyResponsesForMissingModel(s.streaming, requestBodyBytes)
 		return ret, nil
 	}
 
+	logger.V(logutil.DEFAULT).Info("model extracted from request body", "model", model)
+
 	metrics.RecordSuccessCounter()
 
 	if s.streaming {
@@ -78,8 +73,8 @@ func (s *Server) HandleRequestBody(ctx context.Context, requestBodyBytes []byte)
 							SetHeaders: []*basepb.HeaderValueOption{
 								{
 									Header: &basepb.HeaderValue{
-										Key:      modelHeader,
-										RawValue: []byte(requestBody.Model),
+										Key:      bbrplugins.ModelHeader,
+										RawValue: []byte(model),
 									},
 								},
 							},
@@ -88,7 +83,10 @@ func (s *Server) HandleRequestBody(ctx context.Context, requestBodyBytes []byte)
 				},
 			},
 		})
-		ret = addStreamedBodyResponse(ret, requestBodyBytes)
+		ret = addStreamedBodyResponse(ret, mutatedBodyBytes)
+
+		logger.V(logutil.DEFAULT).Info("RESPONSE", "response", helpers.PrettyPrintResponses(ret))
+
 		return ret, nil
 	}
 
@@ -103,28 +101,33 @@ func (s *Server) HandleRequestBody(ctx context.Context, requestBodyBytes []byte)
 							SetHeaders: []*basepb.HeaderValueOption{
 								{
 									Header: &basepb.HeaderValue{
-										Key:      modelHeader,
-										RawValue: []byte(requestBody.Model),
+										Key:      bbrplugins.ModelHeader,
+										RawValue: []byte(model),
 									},
 								},
 							},
 						},
+						BodyMutation: &eppb.BodyMutation{
+							Mutation: &eppb.BodyMutation_Body{
+								Body: mutatedBodyBytes,
+							},
+						},
 					},
 				},
 			},
 		},
 	}, nil
 }
 
-func addStreamedBodyResponse(responses []*eppb.ProcessingResponse, requestBodyBytes []byte) []*eppb.ProcessingResponse {
+func addStreamedBodyResponse(responses []*eppb.ProcessingResponse, mutatedBodyBytes []byte) []*eppb.ProcessingResponse {
 	return append(responses, &eppb.ProcessingResponse{
 		Response: &eppb.ProcessingResponse_RequestBody{
 			RequestBody: &eppb.BodyResponse{
 				Response: &eppb.CommonResponse{
 					BodyMutation: &eppb.BodyMutation{
 						Mutation: &eppb.BodyMutation_StreamedResponse{
 							StreamedResponse: &eppb.StreamedBodyResponse{
-								Body:        requestBodyBytes,
+								Body:        mutatedBodyBytes,
 								EndOfStream: true,
 							},
 						},
@@ -156,3 +159,31 @@ func (s *Server) HandleRequestTrailers(trailers *eppb.HttpTrailers) ([]*eppb.Pro
 		},
 	}, nil
 }
+
+// buildEmptyResponsesForMissingModel is a local helper that returns the appropriate empty responses
+// for the "model not found" branch depending on streaming mode.
+// It is also used to create empty responses in case of other errors related to running plugins on the body
+// This is not very clean and MUST be segregated in the future.
+// Corresponding metrics should be defined to make different errors observable
+func buildEmptyResponsesForMissingModel(streaming bool, requestBodyBytes []byte) ([]*eppb.ProcessingResponse, error) {
+	var ret []*eppb.ProcessingResponse
+
+	if streaming {
+		// Emit empty headers response, then stream body unchanged.
+		ret = append(ret, &eppb.ProcessingResponse{
+			Response: &eppb.ProcessingResponse_RequestHeaders{
+				RequestHeaders: &eppb.HeadersResponse{},
+			},
+		})
+		ret = addStreamedBodyResponse(ret, requestBodyBytes)
+		return ret, nil
+	}
+
+	// Non-streaming: emit empty body response.
+	ret = append(ret, &eppb.ProcessingResponse{
+		Response: &eppb.ProcessingResponse_RequestBody{
+			RequestBody: &eppb.BodyResponse{},
+		},
+	})
+	return ret, nil
+}

pkg/bbr/handlers/server.go

@@ -27,18 +27,34 @@ import (
 	"google.golang.org/grpc/status"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/framework"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 	requtil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/request"
 )
 
-func NewServer(streaming bool) *Server {
-	return &Server{streaming: streaming}
+func NewServer(streaming bool,
+	reg framework.PluginRegistry,
+	reqChain framework.PluginsChain,
+	respChain framework.PluginsChain,
+	metaDataKeys []string) *Server {
+	return &Server{streaming: streaming,
+		registry:      reg,
+		requestChain:  reqChain,
+		responseChain: respChain,
+		endpoint:      "", //will be known upon receiving a request
+		metaDataKeys:  metaDataKeys,
+	}
 }
 
 // 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 {
-	streaming bool
+	streaming     bool
+	registry      framework.PluginRegistry
+	requestChain  framework.PluginsChain
+	responseChain framework.PluginsChain
+	endpoint      string //holds endpoint path of an incoming openai request (extracted from an envoy :path pseudo-header)
+	metaDataKeys  []string
 }
 
 func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {

pkg/bbr/handlers/server_test.go

@@ -26,11 +26,8 @@ import (
 	"google.golang.org/protobuf/testing/protocmp"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
-<<<<<<< HEAD
-=======
 	bbrplugins "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/plugins"
 	utils "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/utils"
->>>>>>> 92cdb00 (Pluggable BBR framework (WiP))
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
@@ -68,7 +65,7 @@ func TestProcessRequestBody(t *testing.T) {
 									SetHeaders: []*basepb.HeaderValueOption{
 										{
 											Header: &basepb.HeaderValue{
-												Key:      modelHeader,
+												Key:      bbrplugins.ModelHeader,
 												RawValue: []byte("foo"),
 											},
 										},
@@ -103,7 +100,7 @@ func TestProcessRequestBody(t *testing.T) {
 									SetHeaders: []*basepb.HeaderValueOption{
 										{
 											Header: &basepb.HeaderValue{
-												Key:      modelHeader,
+												Key:      bbrplugins.ModelHeader,
 												RawValue: []byte("foo"),
 											},
 										},
@@ -135,11 +132,6 @@ func TestProcessRequestBody(t *testing.T) {
 		},
 	}
 
-<<<<<<< HEAD
-	for _, tc := range cases {
-		t.Run(tc.desc, func(t *testing.T) {
-			srv := NewServer(tc.streaming)
-=======
 	//Initialize PluginRegistry and request/response PluginsChain instances based on the minimal configuration setting vi env vars
 	registry, requestChain, responseChain, metaDataKeys, err := utils.InitPlugins()
 	if err != nil {
@@ -149,7 +141,6 @@ func TestProcessRequestBody(t *testing.T) {
 	for _, tc := range cases {
 		t.Run(tc.desc, func(t *testing.T) {
 			srv := NewServer(tc.streaming, registry, requestChain, responseChain, metaDataKeys)
->>>>>>> 92cdb00 (Pluggable BBR framework (WiP))
 			streamedBody := &streamedBody{}
 			for i, body := range tc.bodys {
 				got, err := srv.processRequestBody(context.Background(), body, streamedBody, log.FromContext(ctx))

pkg/bbr/server/runserver.go

@@ -29,21 +29,36 @@ import (
 
 	"sigs.k8s.io/gateway-api-inference-extension/internal/runnable"
 	tlsutil "sigs.k8s.io/gateway-api-inference-extension/internal/tls"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/framework"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/handlers"
 )
 
 // ExtProcServerRunner provides methods to manage an external process server.
 type ExtProcServerRunner struct {
-	GrpcPort      int
-	SecureServing bool
-	Streaming     bool
+	GrpcPort             int
+	SecureServing        bool
+	Streaming            bool
+	registry             framework.PluginRegistry
+	requestPluginsChain  framework.PluginsChain
+	responsePluginsChain framework.PluginsChain
+	metaDataKeys         []string
 }
 
-func NewDefaultExtProcServerRunner(port int, streaming bool) *ExtProcServerRunner {
+func NewDefaultExtProcServerRunner(
+	port int,
+	streaming bool,
+	r framework.PluginRegistry,
+	reqChain framework.PluginsChain,
+	respChain framework.PluginsChain,
+	metaDataKeys []string) *ExtProcServerRunner {
 	return &ExtProcServerRunner{
-		GrpcPort:      port,
-		SecureServing: true,
-		Streaming:     streaming,
+		GrpcPort:             port,
+		SecureServing:        true,
+		Streaming:            streaming,
+		registry:             r,
+		requestPluginsChain:  reqChain,
+		responsePluginsChain: respChain,
+		metaDataKeys:         metaDataKeys,
 	}
 }
 
@@ -65,7 +80,7 @@ func (r *ExtProcServerRunner) AsRunnable(logger logr.Logger) manager.Runnable {
 
 		extProcPb.RegisterExternalProcessorServer(
 			srv,
-			handlers.NewServer(r.Streaming),
+			handlers.NewServer(r.Streaming, r.registry, r.requestPluginsChain, r.responsePluginsChain, r.metaDataKeys),
 		)
 
 		// Forward to the gRPC runnable.