From 346c6dfcf89af9e28848d4d781c1b7a09e2f0e4f Mon Sep 17 00:00:00 2001
From: Michal Wasilewski <michal@mwasilewski.net>
Date: Mon, 2 Dec 2024 17:28:04 +0100
Subject: [PATCH] Expose healthcheck server (#643)

This exposes a basic healtcheck server. Internally, the healtcheck handler doesn't do anything complex.
The idea is to just check if the process is responding.

Signed-off-by: Michal Wasilewski <michal@mwasilewski.net>
---
 main.go             | 11 ++++++++++
 util/healthcheck.go | 49 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)
 create mode 100644 util/healthcheck.go

diff --git a/main.go b/main.go
index 7d9137f77..9f8df34d0 100644
--- a/main.go
+++ b/main.go
@@ -64,6 +64,8 @@ func main() {
 	var reloadRun bool
 	var noReload bool
 	var benchmark bool
+	var collectorHealthCheckEnabled bool
+	var collectorHealthCheckAddress string
 
 	logFlags := log.LstdFlags
 	logger := &util.Logger{}
@@ -107,6 +109,8 @@ func main() {
 	flag.StringVar(&stateFilename, "statefile", defaultStateFile, "Specify alternative path for state file")
 	flag.StringVar(&pidFilename, "pidfile", "", "Specifies a path that a pidfile should be written to (default is no pidfile being written)")
 	flag.BoolVar(&benchmark, "benchmark", false, "Runs collector in benchmark mode (skip submitting the statistics to the server)")
+	flag.BoolVar(&collectorHealthCheckEnabled, "collector-health-check-enabled", true, "Enable the health check endpoint on the collector")
+	flag.StringVar(&collectorHealthCheckAddress, "collector-health-check-address", ":8080", "Address the health check webserver should listen on")
 	flag.Parse()
 
 	// Automatically reload the configuration after a successful test run.
@@ -281,6 +285,13 @@ ReadConfigAndRun:
 	exitCode := 0
 	keepRunning, testRunSuccess, writeStateFile, shutdown := runner.Run(ctx, &wg, globalCollectionOpts, logger, configFilename)
 
+	if collectorHealthCheckEnabled {
+		err := util.SetupHealthCheck(ctx, logger, &wg, collectorHealthCheckAddress)
+		if err != nil {
+			logger.PrintError("Failed to setup health check server: %s", err)
+		}
+	}
+
 	if keepRunning {
 		// Block here until we get any of the registered signals
 		s := <-sigs
diff --git a/util/healthcheck.go b/util/healthcheck.go
new file mode 100644
index 000000000..d3a015b36
--- /dev/null
+++ b/util/healthcheck.go
@@ -0,0 +1,49 @@
+package util
+
+import (
+	"context"
+	"net/http"
+	"sync"
+	"time"
+)
+
+var (
+	healthCheckServerShutdownTimeout = 1 * time.Second
+)
+
+func SetupHealthCheck(ctx context.Context, logger *Logger, wg *sync.WaitGroup, address string) error {
+	var srv http.Server
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		srv = http.Server{
+			Addr: address,
+		}
+		http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+		})
+
+		err := srv.ListenAndServe()
+		if err != nil && err != http.ErrServerClosed {
+			logger.PrintError("error when running the healthcheck server: %s", err)
+		}
+
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		<-ctx.Done()
+		ctxWithTimeout, _ := context.WithTimeout(context.Background(), healthCheckServerShutdownTimeout)
+		err := srv.Shutdown(ctxWithTimeout)
+		if err != nil {
+			logger.PrintError("failed to shutdown the health check server: %s", err)
+		}
+
+	}()
+
+	return nil
+}