Manage all runnable components using manager

Instead of managing all runnable components using errgroup.Group,
they are all wrapped in manager.Runnable interface as the manager
is being used to govern the init and termination process.

Author: tchap

internal/runnable/grpc.go

@@ -0,0 +1,50 @@
+package runnable
+
+import (
+	"context"
+	"fmt"
+	"net"
+
+	"google.golang.org/grpc"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/manager"
+)
+
+func GRPCServer(name string, srv *grpc.Server, port int) manager.Runnable {
+	return manager.RunnableFunc(func(ctx context.Context) error {
+		// Use "name" key as that is what manager.Server does as well.
+		log := ctrl.Log.WithValues("name", name)
+		log.Info("gRPC server starting")
+
+		// Start listening.
+		lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
+		if err != nil {
+			log.Error(err, "gRPC server failed to listen")
+			return err
+		}
+
+		log.Info("gRPC server listening", "port", port)
+
+		// Shutdown on context closed.
+		// Terminate the server on context closed.
+		// Make sure the goroutine does not leak.
+		doneCh := make(chan struct{})
+		defer close(doneCh)
+		go func() {
+			select {
+			case <-ctx.Done():
+				log.Info("gRPC server shutting down")
+				srv.GracefulStop()
+			case <-doneCh:
+			}
+		}()
+
+		// Keep serving until terminated.
+		if err := srv.Serve(lis); err != nil && err != grpc.ErrServerStopped {
+			log.Error(err, "gRPC server failed")
+			return err
+		}
+		log.Info("gRPC server terminated")
+		return nil
+	})
+}

internal/runnable/leader_election.go

@@ -0,0 +1,27 @@
+package runnable
+
+import "sigs.k8s.io/controller-runtime/pkg/manager"
+
+type leaderElection struct {
+	manager.Runnable
+	needsLeaderElection bool
+}
+
+func LeaderElection(runnable manager.Runnable, needsLeaderElection bool) manager.Runnable {
+	return &leaderElection{
+		Runnable:            runnable,
+		needsLeaderElection: needsLeaderElection,
+	}
+}
+
+func RequireLeaderElection(runnable manager.Runnable) manager.Runnable {
+	return LeaderElection(runnable, true)
+}
+
+func NoLeaderElection(runnable manager.Runnable) manager.Runnable {
+	return LeaderElection(runnable, false)
+}
+
+func (r *leaderElection) NeedLeaderElection() bool {
+	return r.needsLeaderElection
+}

pkg/ext-proc/main.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"context"
 	"flag"
 	"fmt"
 	"net"
@@ -10,10 +9,10 @@ import (
 	"strconv"
 
 	"github.com/prometheus/client_golang/prometheus/promhttp"
-	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	healthPb "google.golang.org/grpc/health/grpc_health_v1"
 	"inference.networking.x-k8s.io/gateway-api-inference-extension/api/v1alpha1"
+	"inference.networking.x-k8s.io/gateway-api-inference-extension/internal/runnable"
 	"inference.networking.x-k8s.io/gateway-api-inference-extension/pkg/ext-proc/backend"
 	"inference.networking.x-k8s.io/gateway-api-inference-extension/pkg/ext-proc/backend/vllm"
 	"inference.networking.x-k8s.io/gateway-api-inference-extension/pkg/ext-proc/metrics"
@@ -26,6 +25,7 @@ import (
 	"k8s.io/component-base/metrics/legacyregistry"
 	klog "k8s.io/klog/v2"
 	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/controller-runtime/pkg/metrics/filters"
 )
 
@@ -119,123 +119,73 @@ func run() error {
 		klog.ErrorS(err, "Failed to setup ext-proc server")
 		return err
 	}
+	mgr := serverRunner.Manager
 
-	k8sClient, err := kubernetes.NewForConfigAndClient(cfg, serverRunner.Manager.GetHTTPClient())
+	k8sClient, err := kubernetes.NewForConfigAndClient(cfg, mgr.GetHTTPClient())
 	if err != nil {
 		klog.ErrorS(err, "Failed to create client")
 		return err
 	}
 	datastore.SetClient(k8sClient)
 
-	if err := serverRunner.Setup(); err != nil {
-		klog.ErrorS(err, "Failed to setup server runner")
+	// Register health server.
+	if err := registerHealthServer(mgr, datastore, *grpcHealthPort); err != nil {
 		return err
 	}
 
-	// Start processing signals and init the group to manage goroutines.
-	g, ctx := errgroup.WithContext(ctrl.SetupSignalHandler())
-
-	// Start health server.
-	startHealthServer(ctx, g, datastore, *grpcHealthPort)
-
-	// Start ext-proc server.
-	g.Go(func() error {
-		return serverRunner.Start(ctx, &vllm.PodMetricsClientImpl{})
-	})
-
-	// Start metrics handler.
-	startMetricsHandler(ctx, g, *metricsPort, cfg)
+	// Register ext-proc server.
+	if err := mgr.Add(serverRunner.AsRunnable(&vllm.PodMetricsClientImpl{})); err != nil {
+		klog.ErrorS(err, "Failed to register ext-proc server")
+		return err
+	}
 
-	// Start manager.
-	g.Go(func() error {
-		return serverRunner.StartManager(ctx)
-	})
+	// Register metrics handler.
+	if err := registerMetricsHandler(mgr, *metricsPort, cfg); err != nil {
+		return err
+	}
 
-	err = g.Wait()
-	klog.InfoS("All components terminated")
-	return err
+	// Start the manager.
+	return serverRunner.StartManager(ctrl.SetupSignalHandler())
 }
 
-// startHealthServer starts the gRPC health probe server using the given errgroup.
-func startHealthServer(ctx context.Context, g *errgroup.Group, ds *backend.K8sDatastore, port int) {
-	g.Go(func() error {
-		klog.InfoS("Health server starting...")
-
-		// Start listening.
-		lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
-		if err != nil {
-			klog.ErrorS(err, "Health server failed to listen")
-			return err
-		}
-
-		klog.InfoS("Health server listening", "port", port)
-
-		svr := grpc.NewServer()
-		healthPb.RegisterHealthServer(svr, &healthServer{datastore: ds})
-
-		// Shutdown on context closed.
-		g.Go(func() error {
-			<-ctx.Done()
-			klog.InfoS("Health server shutting down...")
-			svr.GracefulStop()
-			return nil
-		})
-
-		// Keep serving until terminated.
-		if err := svr.Serve(lis); err != nil && err != grpc.ErrServerStopped {
-			klog.ErrorS(err, "Health server failed")
-			return err
-		}
-		klog.InfoS("Health server terminated")
-		return nil
-	})
+// registerHealthServer adds the Health gRPC server as a Runnable to the given manager.
+func registerHealthServer(mgr manager.Manager, ds *backend.K8sDatastore, port int) error {
+	srv := grpc.NewServer()
+	healthPb.RegisterHealthServer(srv, &healthServer{datastore: ds})
+	if err := mgr.Add(
+		runnable.NoLeaderElection(runnable.GRPCServer("health", srv, port))); err != nil {
+		klog.ErrorS(err, "Failed to register health server")
+		return err
+	}
+	return nil
 }
 
-// startMetricsHandler starts the metrics HTTP handler using the given errgroup.
-func startMetricsHandler(ctx context.Context, g *errgroup.Group, port int, cfg *rest.Config) {
-	g.Go(func() error {
-		metrics.Register()
-		klog.InfoS("Metrics HTTP handler starting...")
+// registerMetricsHandler adds the metrics HTTP handler as a Runnable to the given manager.
+func registerMetricsHandler(mgr manager.Manager, port int, cfg *rest.Config) error {
+	metrics.Register()
 
-		// Start listening.
-		lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
-		if err != nil {
-			klog.ErrorS(err, "Metrics HTTP handler failed to listen")
-			return err
-		}
-
-		klog.InfoS("Metrics HTTP handler listening", "port", port)
-
-		// Init HTTP server.
-		h, err := metricsHandlerWithAuthenticationAndAuthorization(cfg)
-		if err != nil {
-			return err
-		}
-
-		mux := http.NewServeMux()
-		mux.Handle(defaultMetricsEndpoint, h)
+	// Init HTTP server.
+	h, err := metricsHandlerWithAuthenticationAndAuthorization(cfg)
+	if err != nil {
+		return err
+	}
 
-		svr := &http.Server{
-			Addr:    net.JoinHostPort("", strconv.Itoa(port)),
-			Handler: mux,
-		}
+	mux := http.NewServeMux()
+	mux.Handle(defaultMetricsEndpoint, h)
 
-		// Shutdown on interrupt.
-		g.Go(func() error {
-			<-ctx.Done()
-			klog.InfoS("Metrics HTTP handler shutting down...")
-			_ = svr.Shutdown(context.Background())
-			return nil
-		})
+	srv := &http.Server{
+		Addr:    net.JoinHostPort("", strconv.Itoa(port)),
+		Handler: mux,
+	}
 
-		// Keep serving until terminated.
-		if err := svr.Serve(lis); err != http.ErrServerClosed {
-			klog.ErrorS(err, "Metrics HTTP handler failed")
-			return err
-		}
-		klog.InfoS("Metrics HTTP handler terminated")
-		return nil
-	})
+	if err := mgr.Add(&manager.Server{
+		Name:   "metrics",
+		Server: srv,
+	}); err != nil {
+		klog.ErrorS(err, "Failed to register metrics HTTP handler")
+		return err
+	}
+	return nil
 }
 
 func metricsHandlerWithAuthenticationAndAuthorization(cfg *rest.Config) (http.Handler, error) {

pkg/ext-proc/server/runserver.go

@@ -4,11 +4,11 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"net"
 	"time"
 
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"google.golang.org/grpc"
+	"inference.networking.x-k8s.io/gateway-api-inference-extension/internal/runnable"
 	"inference.networking.x-k8s.io/gateway-api-inference-extension/pkg/ext-proc/backend"
 	"inference.networking.x-k8s.io/gateway-api-inference-extension/pkg/ext-proc/handlers"
 	"inference.networking.x-k8s.io/gateway-api-inference-extension/pkg/ext-proc/scheduling"
@@ -17,6 +17,7 @@ import (
 	"k8s.io/client-go/rest"
 	klog "k8s.io/klog/v2"
 	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/manager"
 )
 
 // ExtProcServerRunner provides methods to manage an external process server.
@@ -93,60 +94,27 @@ func (r *ExtProcServerRunner) Setup() error {
 	return nil
 }
 
-// Start starts the Envoy external processor server and blocks
-// until the context is canceled or an error encountered.
-func (r *ExtProcServerRunner) Start(
-	ctx context.Context,
-	podMetricsClient backend.PodMetricsClient,
-) error {
-	klog.InfoS("Ext-proc server starting...")
-
-	// Start listening.
-	svr := grpc.NewServer()
-	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", r.GrpcPort))
-	if err != nil {
-		klog.ErrorS(err, "Ext-proc server failed to listen", "port", r.GrpcPort)
-		return err
-	}
-	// The listener will be closed by the server,
-	// but the function may also return earlier on error.
-	defer lis.Close()
-
-	klog.InfoS("Ext-proc server listening", "port", r.GrpcPort)
-
-	// Initialize backend provider
-	pp := backend.NewProvider(podMetricsClient, r.Datastore)
-	if err := pp.Init(r.RefreshPodsInterval, r.RefreshMetricsInterval); err != nil {
-		klog.ErrorS(err, "Failed to initialize backend provider")
-		return err
-	}
-
-	// Register ext_proc handlers
-	extProcPb.RegisterExternalProcessorServer(
-		svr,
-		handlers.NewServer(pp, scheduling.NewScheduler(pp), r.TargetEndpointKey, r.Datastore),
-	)
-
-	// Terminate the server on context closed.
-	// Make sure the goroutine does not leak.
-	doneCh := make(chan struct{})
-	defer close(doneCh)
-	go func() {
-		select {
-		case <-ctx.Done():
-			klog.InfoS("Ext-proc server shutting down...")
-			svr.GracefulStop()
-		case <-doneCh:
+// AsRunnable returns a Runnable that can be used to start the ext-proc gRPC server.
+// The runnable implements LeaderElectionRunnable with leader election disabled.
+func (r *ExtProcServerRunner) AsRunnable(podMetricsClient backend.PodMetricsClient) manager.Runnable {
+	return runnable.NoLeaderElection(manager.RunnableFunc(func(ctx context.Context) error {
+		// Initialize backend provider
+		pp := backend.NewProvider(podMetricsClient, r.Datastore)
+		if err := pp.Init(r.RefreshPodsInterval, r.RefreshMetricsInterval); err != nil {
+			klog.ErrorS(err, "Failed to initialize backend provider")
+			return err
 		}
-	}()
 
-	// Block until terminated.
-	if err := svr.Serve(lis); err != nil && err != grpc.ErrServerStopped {
-		klog.ErrorS(err, "Ext-proc server failed")
-		return err
-	}
-	klog.InfoS("Ext-proc server terminated")
-	return nil
+		// Init the server.
+		srv := grpc.NewServer()
+		extProcPb.RegisterExternalProcessorServer(
+			srv,
+			handlers.NewServer(pp, scheduling.NewScheduler(pp), r.TargetEndpointKey, r.Datastore),
+		)
+
+		// Forward to the gRPC runnable.
+		return runnable.GRPCServer("ext-proc", srv, r.GrpcPort).Start(ctx)
+	}))
 }
 
 func (r *ExtProcServerRunner) StartManager(ctx context.Context) error {
@@ -157,7 +125,7 @@ func (r *ExtProcServerRunner) StartManager(ctx context.Context) error {
 	}
 
 	// Start the controller manager. Blocking and will return when shutdown is complete.
-	klog.InfoS("Controller manager starting...")
+	klog.InfoS("Controller manager starting")
 	if err := r.Manager.Start(ctx); err != nil {
 		klog.ErrorS(err, "Error starting controller manager")
 		return err

test/integration/hermetic_test.go

@@ -343,7 +343,7 @@ func setUpHermeticServer(t *testing.T, pods []*backend.PodMetrics) (client extPr
 
 	serverCtx, stopServer := context.WithCancel(context.Background())
 	go func() {
-		if err := serverRunner.Start(serverCtx, pmc); err != nil {
+		if err := serverRunner.AsRunnable(pmc).Start(serverCtx); err != nil {
 			log.Fatalf("Failed to start ext-proc server: %v", err)
 		}
 	}()