refactor: update the mux server to expose both healthz and metrics, add graceful shutdown

Author: dimakis

cmd/kar-controllers/app/server.go

@@ -31,13 +31,21 @@ limitations under the License.
 package app
 
 import (
+	"context"
+	"fmt"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
 	"k8s.io/client-go/rest"
 	"k8s.io/client-go/tools/clientcmd"
-	"net/http"
 
 	"github.com/project-codeflare/multi-cluster-app-dispatcher/cmd/kar-controllers/app/options"
 	"github.com/project-codeflare/multi-cluster-app-dispatcher/pkg/controller/queuejob"
 	"github.com/project-codeflare/multi-cluster-app-dispatcher/pkg/health"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 
 	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"
 )
@@ -77,13 +85,68 @@ func Run(opt *options.ServerOption) error {
 
 // Starts the health probe listener
 func listenHealthProbe(opt *options.ServerOption) error {
-	handler := http.NewServeMux()
-	handler.Handle("/healthz", &health.Handler{})
-	err := http.ListenAndServe(opt.HealthProbeListenAddr, handler)
-	if err != nil {
-		return err
-	}
-
-	return nil
+    // Start the metrics server
+    RecordMetrics()
+
+    metricsHandler := http.NewServeMux()
+    metricsHandler.Handle("/metrics", promhttp.Handler())
+    healthHandler := http.NewServeMux()
+    healthHandler.Handle("/healthz", &health.Handler{})
+
+    metricsServer := &http.Server{
+        Addr:    opt.MetricsListenAddr,
+        Handler: metricsHandler,
+    }
+
+    healthServer := &http.Server{
+        Addr:    opt.HealthProbeListenAddr,
+        Handler: healthHandler,
+    }
+
+    // make a channel for errors for each server
+    metricsServerErrChan := make(chan error)
+    healthServerErrChan := make(chan error)
+
+    // start servers in their own goroutines
+    go func() {
+        err := metricsServer.ListenAndServe()
+        if err != nil && err != http.ErrServerClosed {
+            metricsServerErrChan <- err
+        }
+    }()
+
+    go func() {
+        err := healthServer.ListenAndServe()
+        if err != nil && err != http.ErrServerClosed {
+            healthServerErrChan <- err
+        }
+    }()
+
+    // make a channel to listen for OS shutdown signal
+    stop := make(chan os.Signal, 1)
+    signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
+
+    // use select to wait for either a shutdown signal or an error
+    select {
+    case <-stop:
+        // received an OS shutdown signal, shut down servers gracefully
+        ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+        defer cancel()
+
+  		errM := metricsServer.Shutdown(ctx)
+		if errM != nil {
+			return fmt.Errorf("metrics server shutdown error: %v", errM)
+		}
+   		errH := healthServer.Shutdown(ctx)
+			if errH != nil {
+			return fmt.Errorf("health server shutdown error: %v", errH)
+		}
+    case err := <-metricsServerErrChan:
+        return fmt.Errorf("metrics server error: %v", err)
+    case err := <-healthServerErrChan:
+        return fmt.Errorf("health server error: %v", err)
+    }
+
+    return nil
 }