health check endpoint for checking both the SQL store and also influx (…

…#372)
openziti · Aug 8, 2023 · 7cc83e3 · 7cc83e3
1 parent e4a1261
commit 7cc83e3
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 3 deletions.
diff --git a/controller/controller.go b/controller/controller.go
@@ -128,10 +128,15 @@ func Run(inCfg *config.Config) error {
 	defer func() { _ = server.Shutdown() }()
 	server.Host = cfg.Endpoint.Host
 	server.Port = cfg.Endpoint.Port
+	rest_server_zrok.HealthCheck = HealthCheckHTTP
 	server.ConfigureAPI()
 	if err := server.Serve(); err != nil {
 		return errors.Wrap(err, "api server error")
 	}
 
 	return nil
 }
+
+func Store() *store.Store {
+	return str
+}
diff --git a/controller/health.go b/controller/health.go
@@ -0,0 +1,62 @@
+package controller
+
+import (
+	"context"
+	"fmt"
+	"github.com/sirupsen/logrus"
+	"net/http"
+)
+
+func HealthCheckHTTP(w http.ResponseWriter, _ *http.Request) {
+	if err := healthCheckStore(w); err != nil {
+		logrus.Error(err)
+		return
+	}
+	if err := healthCheckMetrics(w); err != nil {
+		logrus.Error(err)
+		return
+	}
+	w.Header().Set("Content-Type", "text/html; charset=utf-8")
+	w.Write([]byte("<html><body><h1>Healthy</h1></body></html>"))
+}
+
+func healthCheckStore(w http.ResponseWriter) error {
+	trx, err := str.Begin()
+	if err != nil {
+		http.Error(w, "error starting transaction", http.StatusInternalServerError)
+		return err
+	}
+	defer func() {
+		_ = trx.Rollback()
+	}()
+	count := -1
+	if err := trx.QueryRowx("select count(0) from migrations").Scan(&count); err != nil {
+		http.Error(w, "error selecting migration count", http.StatusInternalServerError)
+		return err
+	}
+	logrus.Debugf("%d migrations", count)
+	return nil
+}
+
+func healthCheckMetrics(w http.ResponseWriter) error {
+	if cfg.Metrics != nil && cfg.Metrics.Influx != nil {
+		queryApi := idb.QueryAPI(cfg.Metrics.Influx.Org)
+		query := fmt.Sprintf("from(bucket: \"%v\")\n", cfg.Metrics.Influx.Bucket) +
+			fmt.Sprintf("|> range(start: -5s)\n") +
+			"|> filter(fn: (r) => r[\"_measurement\"] == \"xfer\")\n" +
+			"|> filter(fn: (r) => r[\"_field\"] == \"rx\" or r[\"_field\"] == \"tx\")\n" +
+			"|> filter(fn: (r) => r[\"namespace\"] == \"backend\")\n" +
+			"|> sum()"
+		result, err := queryApi.Query(context.Background(), query)
+		if err != nil {
+			http.Error(w, "error querying influx", http.StatusInternalServerError)
+			return err
+		}
+		results := 0
+		for result.Next() {
+			results++
+		}
+		logrus.Debugf("%d results", results)
+	}
+	return nil
+}
diff --git a/rest_server_zrok/configure_zrok.go b/rest_server_zrok/configure_zrok.go
@@ -13,6 +13,8 @@ import (
 	"github.com/openziti/zrok/rest_server_zrok/operations"
 )
 
+var HealthCheck func(w http.ResponseWriter, r *http.Request)
+
 //go:generate swagger generate server --target ../../zrok --name Zrok --spec ../specs/zrok.yml --model-package rest_model_zrok --server-package rest_server_zrok --principal interface{} --exclude-main
 
 func configureFlags(api *operations.ZrokAPI) {
@@ -53,6 +55,5 @@ func setupMiddlewares(handler http.Handler) http.Handler {
 // The middleware configuration happens before anything, this middleware also applies to serving the swagger.json document.
 // So this is a good place to plug in a panic handling middleware, logging and metrics.
 func setupGlobalMiddleware(handler http.Handler) http.Handler {
-	logrus.Infof("configuring")
-	return ui.StaticBuilder(handler)
+	return ui.Middleware(handler, HealthCheck)
 }
diff --git a/ui/static.go → ui/middleware.go b/ui/static.go → ui/middleware.go
@@ -9,14 +9,18 @@ import (
 	"strings"
 )
 
-func StaticBuilder(handler http.Handler) http.Handler {
+func Middleware(handler http.Handler, healthCheck func(w http.ResponseWriter, r *http.Request)) http.Handler {
 	logrus.Infof("building")
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if strings.HasPrefix(r.URL.Path, "/api/v1") {
 			logrus.Debugf("directing '%v' to api handler", r.URL.Path)
 			handler.ServeHTTP(w, r)
 			return
 		}
+		if strings.HasPrefix(r.URL.Path, "/health") {
+			healthCheck(w, r)
+			return
+		}
 
 		logrus.Debugf("directing '%v' to static handler", r.URL.Path)