Refine controller template and probe listeners

clitools/pkg/cmd/controller/controller.go

@@ -3,6 +3,7 @@ package controller
 import (
 	"context"
 	"errors"
+	"fmt"
 	"net"
 	"net/http"
 	"os"
@@ -57,7 +58,7 @@ func NewCmdController(flags *genericclioptions.ConfigFlags) *cobra.Command {
 			}()
 
 			go func() {
-				httpErrCh <- httpListen(ctx, clients, conf)
+				httpErrCh <- listenAndServe(ctx, clients, conf)
 			}()
 
 			select {
@@ -92,63 +93,125 @@ func NewCmdController(flags *genericclioptions.ConfigFlags) *cobra.Command {
 	return cmd
 }
 
-func httpListen(ctx context.Context, clients *kube.Clients, conf ServerConfig) error {
-	address, port := "", "8443"
-	addr := net.JoinHostPort(address, port)
-
+func listenAndServe(ctx context.Context, clients *kube.Clients, conf ServerConfig) error {
 	nodeName := os.Getenv("NODE_NAME")
 
-	server := mkscontroller.NewServer(ctx, clients, conf.Namespace, nodeName)
+	controllerServer := mkscontroller.NewServer(ctx, clients, conf.Namespace, nodeName)
 
-	s := &http.Server{
-		Addr:           addr,
-		Handler:        server,
+	healthMux := http.NewServeMux()
+	healthMux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("ok\n"))
+	})
+	healthMux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("ok\n"))
+	})
+
+	healthAddr := net.JoinHostPort("", "8080")
+	controllerAddr := net.JoinHostPort("", "8443")
+
+	healthHTTPServer := &http.Server{
+		Addr:           healthAddr,
+		Handler:        healthMux,
+		IdleTimeout:    90 * time.Second,
+		ReadTimeout:    10 * time.Second,
+		WriteTimeout:   10 * time.Second,
+		MaxHeaderBytes: 1 << 20,
+	}
+
+	controllerHTTPServer := &http.Server{
+		Addr:           controllerAddr,
+		Handler:        controllerServer,
 		IdleTimeout:    90 * time.Second,
 		ReadTimeout:    4 * time.Minute,
 		WriteTimeout:   4 * time.Minute,
 		MaxHeaderBytes: 1 << 20,
 	}
 
-	serverErrCh := make(chan error, 1)
+	serverErrCh := make(chan error, 2)
 
 	go func() {
-		if conf.TLSCertFile != "" {
-			klog.InfoS("starting HTTPS server",
-				"addr", addr,
-				"certFile", conf.TLSCertFile,
-				"keyFile", conf.TLSPrivateKeyFile,
-			)
-			serverErrCh <- s.ListenAndServeTLS(conf.TLSCertFile, conf.TLSPrivateKeyFile)
+		klog.InfoS("starting health HTTP server", "addr", healthAddr)
+
+		err := healthHTTPServer.ListenAndServe()
+		if err != nil && !errors.Is(err, http.ErrServerClosed) {
+			serverErrCh <- fmt.Errorf("health HTTP server: %w", err)
 			return
 		}
 
-		klog.InfoS("starting HTTP server", "addr", addr)
-		serverErrCh <- s.ListenAndServe()
+		serverErrCh <- nil
+	}()
+
+	go func() {
+		if conf.TLSCertFile != "" {
+			klog.InfoS("starting controller HTTPS server",
+				"addr", controllerAddr,
+				"certFile", conf.TLSCertFile,
+				"keyFile", conf.TLSPrivateKeyFile,
+			)
+
+			err := controllerHTTPServer.ListenAndServeTLS(conf.TLSCertFile, conf.TLSPrivateKeyFile)
+			if err != nil && !errors.Is(err, http.ErrServerClosed) {
+				serverErrCh <- fmt.Errorf("controller HTTPS server: %w", err)
+				return
+			}
+
+			serverErrCh <- nil
+			return
+		}
+
+		klog.InfoS("starting controller HTTP server", "addr", controllerAddr)
+
+		err := controllerHTTPServer.ListenAndServe()
+		if err != nil && !errors.Is(err, http.ErrServerClosed) {
+			serverErrCh <- fmt.Errorf("controller HTTP server: %w", err)
+			return
+		}
+
+		serverErrCh <- nil
 	}()
 
 	select {
 	case <-ctx.Done():
-		klog.InfoS("shutting down HTTP server", "addr", addr)
+		klog.InfoS("shutting down HTTP servers",
+			"healthAddr", healthAddr,
+			"controllerAddr", controllerAddr,
+		)
 
 		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 		defer cancel()
 
-		err := s.Shutdown(shutdownCtx)
-		if err != nil {
-			return err
+		var errs []error
+
+		if err := healthHTTPServer.Shutdown(shutdownCtx); err != nil {
+			errs = append(errs, fmt.Errorf("shutdown health HTTP server: %w", err))
 		}
 
-		err = <-serverErrCh
-		if err != nil && !errors.Is(err, http.ErrServerClosed) {
-			return err
+		if err := controllerHTTPServer.Shutdown(shutdownCtx); err != nil {
+			errs = append(errs, fmt.Errorf("shutdown controller HTTP server: %w", err))
 		}
+
+		for i := 0; i < 2; i++ {
+			if err := <-serverErrCh; err != nil {
+				errs = append(errs, err)
+			}
+		}
+
+		if len(errs) > 0 {
+			return errors.Join(errs...)
+		}
+
 		return context.Canceled
 
 	case err := <-serverErrCh:
-		if err != nil && !errors.Is(err, http.ErrServerClosed) {
+		if err != nil {
 			klog.ErrorS(err, "HTTP server failed")
 			return err
 		}
-		return nil
+
+		// One server exited cleanly unexpectedly. Treat that as failure because
+		// the process should keep both servers alive until ctx is canceled.
+		return fmt.Errorf("HTTP server exited unexpectedly")
 	}
 }