feature: Add external-filter in Header for advanced routing

Signed-off-by: rayne-Li <[email protected]>

Author: rayne-Li

pkg/plugins/gateway/gateway.go

@@ -28,6 +28,8 @@ import (
 	"github.com/redis/go-redis/v9"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/klog/v2"
 
@@ -160,7 +162,7 @@ func (s *Server) Process(srv extProcPb.ExternalProcessor_ProcessServer) error {
 	}
 }
 
-func (s *Server) selectTargetPod(ctx *types.RoutingContext, pods types.PodList) (string, error) {
+func (s *Server) selectTargetPod(ctx *types.RoutingContext, pods types.PodList, externalFilterExpr string) (string, error) {
 	router, err := routing.Select(ctx)
 	if err != nil {
 		return "", err
@@ -170,6 +172,25 @@ func (s *Server) selectTargetPod(ctx *types.RoutingContext, pods types.PodList)
 		return "", fmt.Errorf("no pods for routing")
 	}
 	readyPods := utils.FilterRoutablePods(pods.All())
+
+	// filter pod by header ext-filter
+	// k8s labelSelector format, eg: "k=v"、"env in (prod,stg)"
+	if externalFilterExpr != "" {
+		sel, err := labels.Parse(externalFilterExpr)
+		if err != nil {
+			return "", err
+		}
+		out := make([]*corev1.Pod, 0, len(readyPods))
+		for _, p := range readyPods {
+			klog.V(3).InfoS("filtering pod", "pod", p.Name)
+			if sel.Matches(labels.Set(p.Labels)) {
+				out = append(out, p)
+				klog.V(3).InfoS("filter passed", "pod", p.Name)
+			}
+		}
+		readyPods = out
+	}
+
 	if len(readyPods) == 0 {
 		return "", fmt.Errorf("no ready pods for routing")
 	}

pkg/plugins/gateway/gateway_req_body.go

@@ -74,7 +74,10 @@ func (s *Server) HandleRequestBody(ctx context.Context, requestID string, req *e
 		headers = buildEnvoyProxyHeaders(headers, HeaderModel, model)
 		klog.InfoS("request start", "requestID", requestID, "requestPath", requestPath, "model", model, "stream", stream)
 	} else {
-		targetPodIP, err := s.selectTargetPod(routingCtx, podsArr)
+		// external filter in header
+		externalFilter := routingCtx.ReqHeaders[HeaderExternalFilter]
+		klog.InfoS("found external filter", "filter", externalFilter)
+		targetPodIP, err := s.selectTargetPod(routingCtx, podsArr, externalFilter)
 		if targetPodIP == "" || err != nil {
 			klog.ErrorS(err, "failed to select target pod", "requestID", requestID, "routingStrategy", routingAlgorithm, "model", model, "routingDuration", routingCtx.GetRoutingDelay())
 			return generateErrorResponse(

pkg/plugins/gateway/gateway_req_headers.go

@@ -56,6 +56,8 @@ func (s *Server) HandleRequestHeaders(ctx context.Context, requestID string, req
 			requestPath = string(n.RawValue)
 		case authorizationKey:
 			reqHeaders[n.Key] = string(n.RawValue)
+		case HeaderExternalFilter:
+			reqHeaders[n.Key] = string(n.RawValue)
 		}
 	}
 

pkg/plugins/gateway/gateway_test.go

@@ -159,11 +159,12 @@ func Test_selectTargetPod(t *testing.T) {
 
 	// Define test cases for different pod selection and error scenarios
 	tests := []struct {
-		name          string
-		pods          types.PodList
-		mockSetup     func(*mockRouter, types.RoutingAlgorithm)
-		expectedError bool
-		expectedPodIP string
+		name           string
+		pods           types.PodList
+		mockSetup      func(*mockRouter, types.RoutingAlgorithm)
+		expectedError  bool
+		expectedPodIP  string
+		externalFilter string
 	}{
 		{
 			name: "routing.Route returns error",
@@ -288,6 +289,42 @@ func Test_selectTargetPod(t *testing.T) {
 			expectedError: false,
 			expectedPodIP: "1.2.3.4:8000",
 		},
+		{
+			name: "external filter",
+			pods: &utils.PodArray{Pods: []*v1.Pod{
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Labels: map[string]string{
+							"foo": "bar",
+						},
+					},
+					Status: v1.PodStatus{
+						PodIP:      "1.2.3.4",
+						Conditions: []v1.PodCondition{{Type: v1.PodReady, Status: v1.ConditionTrue}},
+					},
+				},
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Labels: map[string]string{
+							"foo": "sad",
+						},
+					},
+					Status: v1.PodStatus{
+						PodIP:      "5.6.7.8",
+						Conditions: []v1.PodCondition{{Type: v1.PodReady, Status: v1.ConditionTrue}},
+					},
+				},
+			}},
+			mockSetup: func(mockRouter *mockRouter, algo types.RoutingAlgorithm) {
+				// Register a mock router that selects a pod from multiple ready pods
+				routing.Register(algo, func() (types.Router, error) {
+					return mockRouter, nil
+				})
+				mockRouter.On("Route", mock.Anything, mock.Anything).Return("1.2.3.4:8000", nil).Once()
+			},
+			expectedError: false,
+			expectedPodIP: "1.2.3.4:8000",
+		},
 	}
 
 	for _, tt := range tests {
@@ -305,7 +342,7 @@ func Test_selectTargetPod(t *testing.T) {
 			ctx := types.NewRoutingContext(context.Background(), routingAlgo, "test-model", "test-message", "test-request", "test-user")
 
 			// Call selectTargetPod and check the result
-			podIP, err := server.selectTargetPod(ctx, tt.pods)
+			podIP, err := server.selectTargetPod(ctx, tt.pods, tt.externalFilter)
 
 			if tt.expectedError {
 				assert.Error(subtest, err)

pkg/plugins/gateway/types.go

@@ -46,6 +46,7 @@ const (
 	HeaderRoutingStrategy    = "routing-strategy"
 	HeaderRequestID          = "request-id"
 	HeaderModel              = "model"
+	HeaderExternalFilter     = "external-filter"
 
 	// RPM & TPM Update Errors
 	HeaderUpdateTPM        = "x-update-tpm"