Skip to content

Commit

Permalink
Expose healthcheck server (pganalyze#643)
Browse files Browse the repository at this point in the history
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 <[email protected]>
  • Loading branch information
mwasilew2 committed Dec 2, 2024
1 parent c477eaf commit 346c6df
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 0 deletions.
11 changes: 11 additions & 0 deletions main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{}
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down
49 changes: 49 additions & 0 deletions util/healthcheck.go
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit 346c6df

Please sign in to comment.