extproc: refactors router code

internal/extproc/mocks_test.go

@@ -104,9 +104,10 @@ type mockRouter struct {
 }
 
 // Calculate implements [router.Router.Calculate].
-func (m mockRouter) Calculate(headers map[string]string) (string, filterconfig.VersionedAPISchema, error) {
+func (m mockRouter) Calculate(headers map[string]string) (*filterconfig.Backend, error) {
 	require.Equal(m.t, m.expHeaders, headers)
-	return m.retBackendName, m.retVersionedAPISchema, m.retErr
+	b := &filterconfig.Backend{Name: m.retBackendName, OutputSchema: m.retVersionedAPISchema}
+	return b, m.retErr
 }
 
 // mockRequestBodyParser implements [router.RequestBodyParser] for testing.

internal/extproc/processor.go

@@ -73,14 +73,14 @@ func (p *Processor) ProcessRequestBody(_ context.Context, rawBody *extprocv3.Htt
 	}
 
 	p.requestHeaders[p.config.ModelNameHeaderKey] = model
-	backendName, outputSchema, err := p.config.router.Calculate(p.requestHeaders)
+	b, err := p.config.router.Calculate(p.requestHeaders)
 	if err != nil {
 		return nil, fmt.Errorf("failed to calculate route: %w", err)
 	}
 
-	factory, ok := p.config.factories[outputSchema]
+	factory, ok := p.config.factories[b.OutputSchema]
 	if !ok {
-		return nil, fmt.Errorf("failed to find factory for output schema %q", outputSchema)
+		return nil, fmt.Errorf("failed to find factory for output schema %q", b.OutputSchema)
 	}
 
 	t, err := factory(path)
@@ -101,10 +101,10 @@ func (p *Processor) ProcessRequestBody(_ context.Context, rawBody *extprocv3.Htt
 	headerMutation.SetHeaders = append(headerMutation.SetHeaders, &corev3.HeaderValueOption{
 		Header: &corev3.HeaderValue{Key: p.config.ModelNameHeaderKey, RawValue: []byte(model)},
 	}, &corev3.HeaderValueOption{
-		Header: &corev3.HeaderValue{Key: p.config.selectedBackendHeaderKey, RawValue: []byte(backendName)},
+		Header: &corev3.HeaderValue{Key: p.config.selectedBackendHeaderKey, RawValue: []byte(b.Name)},
 	})
 
-	if authHandler, ok := p.config.backendAuthHandlers[backendName]; ok {
+	if authHandler, ok := p.config.backendAuthHandlers[b.Name]; ok {
 		if err := authHandler.Do(p.requestHeaders, headerMutation, bodyMutation); err != nil {
 			return nil, fmt.Errorf("failed to do auth request: %w", err)
 		}

internal/extproc/router/router.go

@@ -13,7 +13,7 @@ import (
 type Router interface {
 	// Calculate determines the backend to route to based on the headers.
 	// Returns the backend name and the output schema.
-	Calculate(headers map[string]string) (backendName string, outputSchema filterconfig.VersionedAPISchema, err error)
+	Calculate(headers map[string]string) (backend *filterconfig.Backend, err error)
 }
 
 // router implements [Router].
@@ -28,7 +28,7 @@ func NewRouter(config *filterconfig.Config) (Router, error) {
 }
 
 // Calculate implements [Router.Calculate].
-func (r *router) Calculate(headers map[string]string) (backendName string, outputSchema filterconfig.VersionedAPISchema, err error) {
+func (r *router) Calculate(headers map[string]string) (backend *filterconfig.Backend, err error) {
 	var rule *filterconfig.RouteRule
 	for i := range r.rules {
 		_rule := &r.rules[i]
@@ -42,28 +42,28 @@ func (r *router) Calculate(headers map[string]string) (backendName string, outpu
 		}
 	}
 	if rule == nil {
-		return "", filterconfig.VersionedAPISchema{}, errors.New("no matching rule found")
+		return nil, errors.New("no matching rule found")
 	}
-	backendName, outputSchema = r.selectBackendFromRule(rule)
-	return
+	return r.selectBackendFromRule(rule), nil
 }
 
-func (r *router) selectBackendFromRule(rule *filterconfig.RouteRule) (backendName string, outputSchema filterconfig.VersionedAPISchema) {
+func (r *router) selectBackendFromRule(rule *filterconfig.RouteRule) (backend *filterconfig.Backend) {
 	// Each backend has a weight, so we randomly select depending on the weight.
 	// This is a pretty naive implementation and can be buggy, so fix it later.
 	totalWeight := 0
 	for _, b := range rule.Backends {
 		totalWeight += b.Weight
 	}
 	if totalWeight == 0 {
-		return rule.Backends[0].Name, rule.Backends[0].OutputSchema
+		return &rule.Backends[0]
 	}
 	selected := r.rng.Intn(totalWeight)
-	for _, b := range rule.Backends {
+	for i := range rule.Backends {
+		b := &rule.Backends[i]
 		if selected < b.Weight {
-			return b.Name, b.OutputSchema
+			return b
 		}
 		selected -= b.Weight
 	}
-	return rule.Backends[0].Name, rule.Backends[0].OutputSchema
+	return &rule.Backends[0]
 }

internal/extproc/router/router_test.go

@@ -36,25 +36,24 @@ func TestRouter_Calculate(t *testing.T) {
 	require.True(t, ok)
 
 	t.Run("no matching rule", func(t *testing.T) {
-		backendName, outputSchema, err := r.Calculate(map[string]string{"x-model-name": "something-quirky"})
+		b, err := r.Calculate(map[string]string{"x-model-name": "something-quirky"})
 		require.Error(t, err)
-		require.Empty(t, backendName)
-		require.Empty(t, outputSchema)
+		require.Nil(t, b)
 	})
 	t.Run("matching rule - single backend choice", func(t *testing.T) {
-		backendName, outputSchema, err := r.Calculate(map[string]string{"x-model-name": "gpt4.4444"})
+		b, err := r.Calculate(map[string]string{"x-model-name": "gpt4.4444"})
 		require.NoError(t, err)
-		require.Equal(t, "openai", backendName)
-		require.Equal(t, outSchema, outputSchema)
+		require.Equal(t, "openai", b.Name)
+		require.Equal(t, outSchema, b.OutputSchema)
 	})
 	t.Run("matching rule - multiple backend choices", func(t *testing.T) {
 		chosenNames := make(map[string]int)
 		for i := 0; i < 1000; i++ {
-			backendName, outputSchema, err := r.Calculate(map[string]string{"x-model-name": "llama3.3333"})
+			b, err := r.Calculate(map[string]string{"x-model-name": "llama3.3333"})
 			require.NoError(t, err)
-			chosenNames[backendName]++
-			require.Contains(t, []string{"foo", "bar"}, backendName)
-			require.Equal(t, outSchema, outputSchema)
+			chosenNames[b.Name]++
+			require.Contains(t, []string{"foo", "bar"}, b.Name)
+			require.Equal(t, outSchema, b.OutputSchema)
 		}
 		require.Greater(t, chosenNames["bar"], chosenNames["foo"])
 		require.Greater(t, chosenNames["bar"], 700)
@@ -79,8 +78,8 @@ func TestRouter_selectBackendFromRule(t *testing.T) {
 
 	chosenNames := make(map[string]int)
 	for i := 0; i < 1000; i++ {
-		backendName, _ := r.selectBackendFromRule(rule)
-		chosenNames[backendName]++
+		b := r.selectBackendFromRule(rule)
+		chosenNames[b.Name]++
 	}
 
 	require.Greater(t, chosenNames["bar"], chosenNames["foo"])