Stop the server when the context is cancelled

Signed-off-by: Richard Wall <[email protected]>

Author: wallrj

pkg/agent/run.go

@@ -25,6 +25,7 @@ import (
 	"k8s.io/client-go/kubernetes"
 	clientgocorev1 "k8s.io/client-go/kubernetes/typed/core/v1"
 	"k8s.io/client-go/tools/record"
+	"k8s.io/klog/v2"
 	"sigs.k8s.io/controller-runtime/pkg/manager"
 
 	"github.com/jetstack/preflight/api"
@@ -50,7 +51,12 @@ const schemaVersion string = "v2.0.0"
 // Run starts the agent process
 func Run(cmd *cobra.Command, args []string) {
 	logs.Log.Printf("Preflight agent version: %s (%s)", version.PreflightVersion, version.Commit)
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(
+		klog.NewContext(
+			context.Background(),
+			klog.Background(),
+		),
+	)
 	defer cancel()
 
 	file, err := os.Open(Flags.ConfigFilePath)
@@ -83,11 +89,13 @@ func Run(cmd *cobra.Command, args []string) {
 		}
 	}()
 
-	group.Go(func() error {
+	{
 		server := http.NewServeMux()
+		const serverAddress = ":8081"
+		log := klog.FromContext(ctx).WithName("APIServer").WithValues("addr", serverAddress)
 
 		if Flags.Profiling {
-			logs.Log.Printf("pprof profiling was enabled.")
+			log.Info("Profiling endpoints enabled", "path", "/debug/pprof")
 			server.HandleFunc("/debug/pprof/", pprof.Index)
 			server.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
 			server.HandleFunc("/debug/pprof/profile", pprof.Profile)
@@ -96,7 +104,7 @@ func Run(cmd *cobra.Command, args []string) {
 		}
 
 		if Flags.Prometheus {
-			logs.Log.Printf("Prometheus was enabled.\nRunning prometheus on port :8081")
+			log.Info("Metrics endpoints enabled", "path", "/metrics")
 			prometheus.MustRegister(metricPayloadSize)
 			server.Handle("/metrics", promhttp.Handler())
 		}
@@ -105,21 +113,23 @@ func Run(cmd *cobra.Command, args []string) {
 		// what "ready" means for the agent, we just return 200 OK inconditionally.
 		// The goal is to satisfy some Kubernetes distributions, like OpenShift,
 		// that require a liveness and health probe to be present for each pod.
+		log.Info("Healthz endpoints enabled", "path", "/healthz")
 		server.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(http.StatusOK)
 		})
+		log.Info("Readyz endpoints enabled", "path", "/readyz")
 		server.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) {
 			w.WriteHeader(http.StatusOK)
 		})
 
-		err := http.ListenAndServe(":8081", server)
-		if err != nil && !errors.Is(err, http.ErrServerClosed) {
-			return fmt.Errorf("failed to run the health check server: %s", err)
-		}
-		// The agent must stop if the management server stops
-		cancel()
-		return nil
-	})
+		group.Go(func() error {
+			return listenAndServe(
+				klog.NewContext(gctx, log),
+				serverAddress,
+				server,
+			)
+		})
+	}
 
 	_, isVenConn := preflightClient.(*client.VenConnClient)
 	if isVenConn {
@@ -412,3 +422,41 @@ func postData(config CombinedConfig, preflightClient client.Client, readings []*
 
 	return nil
 }
+
+// listenAndServe starts an HTTP server at the given addr with the given handler
+// and stops it gracefully when the supplied context is cancelled.
+// It returns when the graceful server shutdown is complete or when the server
+// exits with an error.
+// If the server fails to start, it returns the server error.
+// If the server fails to shutdown gracefully, it returns the shutdown error.
+// The server is given 1s to shutdown gracefully before it is stopped
+// forcefully.
+func listenAndServe(ctx context.Context, addr string, handler http.Handler) error {
+	log := klog.FromContext(ctx).WithName("ListenAndServe")
+	server := &http.Server{Addr: addr, Handler: handler}
+	group, gCTX := errgroup.WithContext(ctx)
+	group.Go(func() error {
+		log.V(1).Info("Starting")
+		err := server.ListenAndServe()
+		log.V(1).Info("Stopped", "reason", err)
+		if errors.Is(err, http.ErrServerClosed) {
+			err = nil
+		}
+		return err
+	})
+	group.Go(func() error {
+		select {
+		case <-ctx.Done():
+			log.V(1).Info("Shutting down")
+			shutdownCTX, cancel := context.WithTimeout(context.Background(), time.Second)
+			defer cancel()
+			err := server.Shutdown(shutdownCTX)
+			log.V(1).Info("Shutdown complete", "err", err)
+			return err
+		case <-gCTX.Done():
+			log.V(1).Info("Shutdown skipped", "addr", addr, "reason", "Server already stopped")
+			return nil
+		}
+	})
+	return group.Wait()
+}