Add flow controller.

Author: LukeAVanDrie

cmd/epp/main.go

@@ -17,6 +17,7 @@ limitations under the License.
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"net"
@@ -41,8 +42,10 @@ import (
 	"sigs.k8s.io/gateway-api-inference-extension/internal/runnable"
 	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/datastore"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/flowcontroller"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics/collectors"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/saturationdetector"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/plugins/filter"
@@ -157,13 +160,20 @@ func run() error {
 	})
 	setupLog.Info("Flags processed", "flags", flags)
 
-	// Init runtime.
+	// --- Load Configurations from Environment Variables ---
+	// Note: Scheduler config is loaded via its package init currently. We may
+	// want to load it here explicitly:
+	fcConfig, flowControllerEnabled := flowcontroller.LoadConfigFromEnv()
+	sdConfig := saturationdetector.LoadConfigFromEnv()
+
+	// --- Get Kubernetes Config ---
 	cfg, err := ctrl.GetConfig()
 	if err != nil {
-		setupLog.Error(err, "Failed to get rest config")
+		setupLog.Error(err, "Failed to get Kubernetes rest config")
 		return err
 	}
 
+	// --- Setup Manager ---
 	poolNamespacedName := types.NamespacedName{
 		Name:      *poolName,
 		Namespace: *poolNamespace,
@@ -174,7 +184,7 @@ func run() error {
 		return err
 	}
 
-	// Set up mapper for metric scraping.
+	// --- Setup Datastore ---
 	mapping, err := backendmetrics.NewMetricMapping(
 		*totalQueuedRequestsMetric,
 		*kvCacheUsagePercentageMetric,
@@ -185,14 +195,12 @@ func run() error {
 		return err
 	}
 	verifyMetricMapping(*mapping, setupLog)
-
 	pmf := backendmetrics.NewPodMetricsFactory(&backendmetrics.PodMetricsClientImpl{MetricMapping: mapping}, *refreshMetricsInterval)
-	// Setup runner.
 	ctx := ctrl.SetupSignalHandler()
+	appDatastore := datastore.NewDatastore(ctx, pmf)
 
-	datastore := datastore.NewDatastore(ctx, pmf)
-
-	scheduler := scheduling.NewScheduler(datastore)
+	// --- Initialize EPP Components ---
+	appScheduler := scheduling.NewScheduler(appDatastore)
 	if schedulerV2 == "true" {
 		queueScorerWeight := envutil.GetEnvInt("QUEUE_SCORE_WEIGHT", scorer.DefaultQueueScorerWeight, setupLog)
 		kvCacheScorerWeight := envutil.GetEnvInt("KV_CACHE_SCORE_WEIGHT", scorer.DefaultKVCacheScorerWeight, setupLog)
@@ -213,41 +221,78 @@ func run() error {
 			[]plugins.PostSchedule{},
 			[]plugins.PostResponse{},
 			schedConfigOpts...)
-		scheduler = scheduling.NewSchedulerWithConfig(datastore, schedulerConfig)
+			appScheduler = scheduling.NewSchedulerWithConfig(appDatastore, schedulerConfig)
+	}
+
+	appSaturationDetector, err := saturationdetector.NewDetector(
+		*sdConfig,
+		appDatastore,
+		ctrl.Log.WithName("saturation-detector"),
+	)
+	if err != nil {
+		setupLog.Error(err, "Failed to create SaturationDetector")
+		return err
+	}
+
+	var appFlowController *flowcontroller.FlowController
+	if flowControllerEnabled {
+		appFlowController, err = flowcontroller.NewFlowController(
+			appSaturationDetector,
+			fcConfig,
+		)
+		if err != nil {
+			setupLog.Error(err, "Failed to create FlowController")
+			return err
+		}
+		setupLog.Info("FlowController enabled and initialized.")
+	} else {
+		setupLog.Info("FlowController is disabled via configuration.")
 	}
+
+	// --- Setup ExtProc Server Runner ---
 	serverRunner := &runserver.ExtProcServerRunner{
 		GrpcPort:                                 *grpcPort,
 		DestinationEndpointHintMetadataNamespace: *destinationEndpointHintMetadataNamespace,
 		DestinationEndpointHintKey:               *destinationEndpointHintKey,
 		PoolNamespacedName:                       poolNamespacedName,
-		Datastore:                                datastore,
+		Datastore:                                appDatastore,
 		SecureServing:                            *secureServing,
 		CertPath:                                 *certPath,
 		RefreshPrometheusMetricsInterval:         *refreshPrometheusMetricsInterval,
-		Scheduler:                                scheduler,
+		Scheduler:                                appScheduler,
+		FlowController:                           appFlowController, // Pass instance (can be nil)
+		SaturationDetector:                       appSaturationDetector,
+		FlowControllerEnabled:                    flowControllerEnabled,
 	}
 	if err := serverRunner.SetupWithManager(ctx, mgr); err != nil {
-		setupLog.Error(err, "Failed to setup ext-proc controllers")
+		setupLog.Error(err, "Failed to setup EPP controllers")
+		return err
+	}
+
+	// --- Add Runnables to Manager ---
+
+	// Register FlowController Run loop.
+	if err := registerFlowController(mgr, appFlowController); err != nil {
 		return err
 	}
 
 	// Register health server.
-	if err := registerHealthServer(mgr, ctrl.Log.WithName("health"), datastore, *grpcHealthPort); err != nil {
+	if err := registerHealthServer(mgr, ctrl.Log.WithName("health"), appDatastore, *grpcHealthPort); err != nil {
 		return err
 	}
 
 	// Register ext-proc server.
-	if err := mgr.Add(serverRunner.AsRunnable(ctrl.Log.WithName("ext-proc"))); err != nil {
-		setupLog.Error(err, "Failed to register ext-proc gRPC server")
+	if err := registerExtProcServer(mgr, serverRunner, ctrl.Log.WithName("ext-proc")); err != nil {
 		return err
 	}
 
 	// Register metrics handler.
-	if err := registerMetricsHandler(mgr, *metricsPort, cfg, datastore); err != nil {
+	if err := registerMetricsHandler(mgr, *metricsPort, cfg, appDatastore); err != nil {
 		return err
 	}
 
-	// Start the manager. This blocks until a signal is received.
+	// --- Start Manager ---
+	// This blocks until a signal is received.
 	setupLog.Info("Controller manager starting")
 	if err := mgr.Start(ctx); err != nil {
 		setupLog.Error(err, "Error starting controller manager")
@@ -275,6 +320,39 @@ func initLogging(opts *zap.Options) {
 	ctrl.SetLogger(logger)
 }
 
+// registerFlowController adds the FlowController Run loop as a Runnable to the
+// manager.
+func registerFlowController(mgr manager.Manager, fc *flowcontroller.FlowController) error {
+	if fc == nil {
+		setupLog.Info("FlowController is nil (disabled), skipping registration.")
+		return nil // Not an error if it's intentionally disabled
+	}
+	if err := mgr.Add(manager.RunnableFunc(func(runCtx context.Context) error {
+		fcLog := ctrl.Log.WithName("flowcontroller-runnable")
+		fcLog.Info("Starting FlowController Run loop")
+		// Run the FlowController; it handles context cancellation internally.
+		fc.Run(runCtx)
+		fcLog.Info("FlowController Run loop stopped")
+		return nil
+	})); err != nil {
+		setupLog.Error(err, "Failed to add FlowController runnable to manager")
+		return err
+	}
+	setupLog.Info("FlowController Run loop added to manager.")
+	return nil
+}
+
+// registerExtProcServer adds the ExtProcServerRunner as a Runnable to the
+// manager.
+func registerExtProcServer(mgr manager.Manager, runner *runserver.ExtProcServerRunner, logger logr.Logger) error {
+	if err := mgr.Add(runner.AsRunnable(logger)); err != nil {
+		setupLog.Error(err, "Failed to register ext-proc gRPC server runnable")
+		return err
+	}
+	setupLog.Info("ExtProc server runner added to manager.")
+	return nil
+}
+
 // registerHealthServer adds the Health gRPC server as a Runnable to the given manager.
 func registerHealthServer(mgr manager.Manager, logger logr.Logger, ds datastore.Datastore, port int) error {
 	srv := grpc.NewServer()
@@ -364,5 +442,4 @@ func verifyMetricMapping(mapping backendmetrics.MetricMapping, logger logr.Logge
 	if mapping.LoraRequestInfo == nil {
 		logger.Info("Not scraping metric: LoraRequestInfo")
 	}
-
 }