apply suggestions

Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>

Author: evacchi

apix/config/v1alpha1/endpointpickerconfig_types.go

@@ -302,15 +302,15 @@ type FlowControlConfig struct {
 	PriorityBands []PriorityBandConfig `json:"priorityBands,omitempty"`
 
 	// +optional
-	// UsageLimitPolicyType is the plugin type name for the UsageLimitPolicy.
-	// The named plugin must be registered in the plugin registry before the runner starts.
-	// If omitted or the plugin is not found, a default no-op policy (limit=1.0) is used.
-	UsageLimitPolicyType string `json:"usageLimitPolicyType,omitempty"`
+	// UsageLimit specifies the UsageLimitPolicy plugin to use for adaptive capacity management.
+	// The PluginRef must reference a named plugin instance defined in the top-level Plugins section.
+	// If omitted, a default no-op policy (limit=1.0) is used.
+	UsageLimit *UsageLimitConfig `json:"usageLimit,omitempty"`
 }
 
 func (fcc *FlowControlConfig) String() string {
-	return fmt.Sprintf("{MaxBytes: %v, DefaultPriorityBand: %v, PriorityBands: %v, UsageLimitPolicyType: %v}",
-		fcc.MaxBytes, fcc.DefaultPriorityBand, fcc.PriorityBands, fcc.UsageLimitPolicyType)
+	return fmt.Sprintf("{MaxBytes: %v, DefaultPriorityBand: %v, PriorityBands: %v, UsageLimit: %v}",
+		fcc.MaxBytes, fcc.DefaultPriorityBand, fcc.PriorityBands, fcc.UsageLimit)
 }
 
 // PriorityBandConfig configures a single priority band.
@@ -340,3 +340,17 @@ func (pbc PriorityBandConfig) String() string {
 	return fmt.Sprintf("{Priority: %d, MaxBytes: %v, FairnessPolicyRef: %s, OrderingPolicyRef: %s}",
 		pbc.Priority, pbc.MaxBytes, pbc.FairnessPolicyRef, pbc.OrderingPolicyRef)
 }
+
+// UsageLimitConfig contains the configuration for a UsageLimitPolicy plugin.
+type UsageLimitConfig struct {
+	// +required
+	// +kubebuilder:validation:Required
+	// PluginRef specifies a particular Plugin instance to be used as the UsageLimitPolicy.
+	// The reference is to the name of an entry of the Plugins defined in the configuration's
+	// Plugins section.
+	PluginRef string `json:"pluginRef"`
+}
+
+func (ulc UsageLimitConfig) String() string {
+	return fmt.Sprintf("{PluginRef: %s}", ulc.PluginRef)
+}

cmd/epp/runner/runner.go

@@ -59,7 +59,6 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontrol/framework/plugins/usagelimits"
 	fcregistry "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontrol/registry"
 	fwkdl "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/datalayer"
-	flowcontrolplugins "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/flowcontrol"
 	fwkplugin "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/plugin"
 	fwkrh "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/requesthandling"
 	extractormetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/plugins/datalayer/extractor/metrics"
@@ -329,8 +328,6 @@ func (r *Runner) setup(ctx context.Context, cfg *rest.Config, opts *runserver.Op
 
 	saturationDetector := utilizationdetector.NewDetector(eppConfig.SaturationDetectorConfig, setupLog)
 
-	usageLimitPolicy := r.setupFlowControlPlugins(eppConfig.FlowControlConfig.UsageLimitPolicyType, setupLog)
-
 	// --- Admission Control Initialization ---
 	var admissionController requestcontrol.AdmissionController
 	var locator contracts.PodLocator
@@ -347,7 +344,7 @@ func (r *Runner) setup(ctx context.Context, cfg *rest.Config, opts *runserver.Op
 			opts.PoolName,
 			eppConfig.FlowControlConfig.Controller,
 			registry, saturationDetector,
-			locator, usageLimitPolicy,
+			locator, eppConfig.FlowControlConfig.UsageLimitPolicy,
 		)
 		if err != nil {
 			return nil, nil, fmt.Errorf("failed to initialize Flow Controller: %w", err)
@@ -463,6 +460,7 @@ func (r *Runner) registerInTreePlugins() {
 	fwkplugin.Register(fairness.RoundRobinFairnessPolicyType, fairness.RoundRobinFairnessPolicyFactory)
 	fwkplugin.Register(ordering.FCFSOrderingPolicyType, ordering.FCFSOrderingPolicyFactory)
 	fwkplugin.Register(ordering.EDFOrderingPolicyType, ordering.EDFOrderingPolicyFactory)
+	fwkplugin.Register(usagelimits.NoopUsageLimitPolicyType, usagelimits.NoopPolicyFactory)
 	// Latency predictor plugins
 	fwkplugin.Register(predictedlatency.PredictedLatencyPluginType, predictedlatency.PredictedLatencyFactory)
 	// register filter for test purpose only (used in conformance tests)
@@ -639,26 +637,6 @@ func (r *Runner) setupDataLayer(enableNewMetrics bool, cfg *datalayer.Config,
 	return nil
 }
 
-func (r *Runner) setupFlowControlPlugins(policyType string, logger logr.Logger) flowcontrolplugins.UsageLimitPolicy {
-	if factory, ok := fwkplugin.Registry[policyType]; ok {
-		p, err := factory(policyType, nil, nil)
-		if err != nil {
-			logger.Info("Failed to instantiate usage limit policy plugin, using default", "type", policyType, "error", err)
-			return usagelimits.DefaultPolicy()
-		}
-		policy, ok := p.(flowcontrolplugins.UsageLimitPolicy)
-		if !ok {
-			logger.Info("Registered plugin does not implement UsageLimitPolicy, using default", "type", policyType)
-			return usagelimits.DefaultPolicy()
-		}
-		return policy
-	}
-	if policyType != "" {
-		logger.Info("No plugin registered for usage limit policy type, using default", "type", policyType)
-	}
-	return usagelimits.DefaultPolicy()
-}
-
 func (r *Runner) setupMetricsCollection(enableNewMetrics bool, opts *runserver.Options, pmc backendmetrics.PodMetricsClient) datalayer.EndpointFactory {
 	if enableNewMetrics {
 		return datalayer.NewEndpointFactory(nil, opts.RefreshMetricsInterval)

cmd/epp/runner/runner_test.go

@@ -18,109 +18,13 @@ package runner
 
 import (
 	"context"
-	"encoding/json"
 	"testing"
 
-	"github.com/go-logr/logr"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 
-	fc "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontrol"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontrol/framework/plugins/usagelimits"
-	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/flowcontrol"
-	fwkplugin "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/plugin"
 )
 
-// TestSetupFlowControlPlugins verifies that setupFlowControlPlugins resolves a registered
-// UsageLimitPolicy from the plugin registry and that the resolved policy returns the expected limit.
-func TestSetupFlowControlPlugins(t *testing.T) {
-	// Register the factory before any parallel sub-test starts so the write
-	// to the global registry map does not race with reads in sub-tests.
-	const pluginType = "func-plugin"
-	fwkplugin.Register(pluginType, func(name string, _ json.RawMessage, _ fwkplugin.Handle) (fwkplugin.Plugin, error) {
-		return usagelimits.NewPolicyFunc(name, func(_ context.Context, _ float64, priorities []int) []float64 {
-			result := make([]float64, len(priorities))
-			for i := range result {
-				result[i] = 0.8
-			}
-			return result
-		}), nil
-	})
-	t.Cleanup(func() {
-		delete(fwkplugin.Registry, pluginType)
-	})
-
-	r := &Runner{}
-
-	t.Run("resolves registered plugin and returns 0.8", func(t *testing.T) {
-		t.Parallel()
-
-		policy := r.setupFlowControlPlugins(pluginType, logr.Discard())
-		require.NotNil(t, policy)
-
-		ctx := context.Background()
-		for _, tc := range []struct {
-			name       string
-			priority   int
-			saturation float64
-		}{
-			{"zero saturation", 0, 0.0},
-			{"half saturation", 1, 0.5},
-			{"full saturation", 5, 1.0},
-		} {
-			t.Run(tc.name, func(t *testing.T) {
-				t.Parallel()
-				assert.Equal(t, []float64{0.8}, policy.ComputeLimit(ctx, tc.saturation, []int{tc.priority}))
-			})
-		}
-	})
-
-	t.Run("falls back to default policy when type is empty", func(t *testing.T) {
-		t.Parallel()
-
-		policy := r.setupFlowControlPlugins("", logr.Discard())
-		require.NotNil(t, policy)
-
-		ctx := context.Background()
-		assert.Equal(t, []float64{1.0}, policy.ComputeLimit(ctx, 0.5, []int{0}))
-	})
-
-	t.Run("falls back to default policy when type is unknown", func(t *testing.T) {
-		t.Parallel()
-
-		policy := r.setupFlowControlPlugins("no-such-policy-type", logr.Discard())
-		require.NotNil(t, policy)
-
-		ctx := context.Background()
-		assert.Equal(t, []float64{1.0}, policy.ComputeLimit(ctx, 0.5, []int{0}))
-	})
-}
-
-const constantPointEightStructPolicyType = "test-constant-point-eight-struct-usage-limit-policy"
-
-// constantPointEightPolicy is a hand-rolled UsageLimitPolicy implementation that always returns 0.8.
-// It exists to show that any struct satisfying the interface can be registered and resolved,
-// without relying on the usagelimits.NewPolicyFunc helper.
-type constantPointEightPolicy struct{}
-
-func (p *constantPointEightPolicy) TypedName() fwkplugin.TypedName {
-	return fwkplugin.TypedName{
-		Type: constantPointEightStructPolicyType + "-type",
-		Name: constantPointEightStructPolicyType,
-	}
-}
-
-func (p *constantPointEightPolicy) ComputeLimit(_ context.Context, _ float64, priorities []int) []float64 {
-	result := make([]float64, len(priorities))
-	for i := range result {
-		result[i] = 0.8
-	}
-	return result
-}
-
-// compile-time check that constantPointEightPolicy satisfies the interface.
-var _ flowcontrol.UsageLimitPolicy = (*constantPointEightPolicy)(nil)
-
 // TestLinearSpacingPolicy demonstrates a stateless UsageLimitPolicy that dynamically spaces
 // ceilings based on the active priority domain. The highest-active priority always gets ceiling
 // 1.0, and each subsequent tier drops by a fixed step (0.2). When priorities go idle and the
@@ -154,58 +58,3 @@ func TestLinearSpacingPolicy(t *testing.T) {
 	got = linearSpacing.ComputeLimit(ctx, 0.5, []int{0})
 	assert.Equal(t, []float64{1.0}, got, "single active priority should get full ceiling")
 }
-
-// TestSetupFlowControlPlugins_WithUsageLimitPolicyType verifies that UsageLimitPolicyType on
-// flowcontrol.Config (the field a downstream project would set from a config file) is read by
-// setupFlowControlPlugins to resolve and return the correct registered plugin.
-func TestSetupFlowControlPlugins_WithUsageLimitPolicyType(t *testing.T) {
-	const pluginType = "struct-plugin-via-config-option"
-
-	fwkplugin.Register(pluginType, func(name string, _ json.RawMessage, _ fwkplugin.Handle) (fwkplugin.Plugin, error) {
-		return &constantPointEightPolicy{}, nil
-	})
-	t.Cleanup(func() {
-		delete(fwkplugin.Registry, pluginType)
-	})
-
-	cfg := &fc.Config{UsageLimitPolicyType: pluginType}
-
-	policy := (&Runner{}).setupFlowControlPlugins(cfg.UsageLimitPolicyType, logr.Discard())
-	require.NotNil(t, policy)
-
-	ctx := context.Background()
-	assert.Equal(t, []float64{0.8}, policy.ComputeLimit(ctx, 0.5, []int{0}))
-}
-
-// TestSetupFlowControlPlugins_StructPlugin verifies that a hand-rolled struct implementing
-// UsageLimitPolicy (without using usagelimits.NewPolicyFunc) can be registered and resolved.
-func TestSetupFlowControlPlugins_StructPlugin(t *testing.T) {
-
-	fwkplugin.Register(constantPointEightStructPolicyType, func(name string, _ json.RawMessage, _ fwkplugin.Handle) (fwkplugin.Plugin, error) {
-		return &constantPointEightPolicy{}, nil
-	})
-	t.Cleanup(func() {
-		delete(fwkplugin.Registry, constantPointEightStructPolicyType)
-	})
-
-	r := &Runner{}
-
-	policy := r.setupFlowControlPlugins(constantPointEightStructPolicyType, logr.Discard())
-	require.NotNil(t, policy)
-
-	ctx := context.Background()
-	for _, tc := range []struct {
-		name       string
-		priority   int
-		saturation float64
-	}{
-		{"zero saturation", 0, 0.0},
-		{"half saturation", 1, 0.5},
-		{"full saturation", 5, 1.0},
-	} {
-		t.Run(tc.name, func(t *testing.T) {
-			t.Parallel()
-			assert.Equal(t, []float64{0.8}, policy.ComputeLimit(ctx, tc.saturation, []int{tc.priority}))
-		})
-	}
-}

pkg/epp/config/loader/defaults.go

@@ -202,6 +202,9 @@ func ensureFlowControlLayer(
 	if _, ok := allPlugins[registry.DefaultFairnessPolicyRef]; !ok {
 		return registerDefaultPlugin(cfg, handle, registry.DefaultFairnessPolicyRef)
 	}
+	if _, ok := allPlugins[registry.DefaultUsageLimitPolicyRef]; !ok {
+		return registerDefaultPlugin(cfg, handle, registry.DefaultUsageLimitPolicyRef)
+	}
 	return nil
 }
 

pkg/epp/flowcontrol/config.go

@@ -22,6 +22,7 @@ import (
 	configapi "sigs.k8s.io/gateway-api-inference-extension/apix/config/v1alpha1"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontrol/controller"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontrol/registry"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/flowcontrol"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/framework/interface/plugin"
 )
 
@@ -31,12 +32,9 @@ const FeatureGate = "flowControl"
 // It embeds the configurations for the controller and the registry, providing a single point of entry for validation
 // and initialization.
 type Config struct {
-	Controller *controller.Config
-	Registry   *registry.Config
-
-	// UsageLimitPolicyType is the plugin type name used to resolve a UsageLimitPolicy from the plugin registry.
-	// Optional: Defaults to empty string (uses the default no-op policy).
-	UsageLimitPolicyType string
+	Controller       *controller.Config
+	Registry         *registry.Config
+	UsageLimitPolicy flowcontrol.UsageLimitPolicy
 }
 
 // NewConfigFromAPI creates a new Config by translating the top-level API configuration.
@@ -49,12 +47,30 @@ func NewConfigFromAPI(apiConfig *configapi.FlowControlConfig, handle plugin.Hand
 	if err != nil {
 		return nil, fmt.Errorf("failed to create controller config: %w", err)
 	}
-	cfg := &Config{
-		Controller: ctrlCfg,
-		Registry:   registryConfig,
+	usageLimitPolicy, err := ensureUsageLimitPolicy(apiConfig, handle)
+	if err != nil {
+		return nil, err
 	}
-	if apiConfig != nil {
-		cfg.UsageLimitPolicyType = apiConfig.UsageLimitPolicyType
+	cfg := &Config{
+		Controller:       ctrlCfg,
+		Registry:         registryConfig,
+		UsageLimitPolicy: usageLimitPolicy,
 	}
 	return cfg, nil
 }
+
+func ensureUsageLimitPolicy(apiConfig *configapi.FlowControlConfig, handle plugin.Handle) (flowcontrol.UsageLimitPolicy, error) {
+	ref := registry.DefaultUsageLimitPolicyRef
+	if apiConfig != nil && apiConfig.UsageLimit != nil {
+		ref = apiConfig.UsageLimit.PluginRef
+	}
+	p := handle.Plugin(ref)
+	if p == nil {
+		return nil, fmt.Errorf("usage limit policy plugin '%s' not found", ref)
+	}
+	usageLimitPolicy, ok := p.(flowcontrol.UsageLimitPolicy)
+	if !ok {
+		return nil, fmt.Errorf("plugin '%s' does not implement UsageLimitPolicy", ref)
+	}
+	return usageLimitPolicy, nil
+}