From c1694f1a227f126091ba5c11cccaaf20e4fd0d0a Mon Sep 17 00:00:00 2001
From: Matthias Fechner <idefix@fechner.net>
Date: Mon, 30 Jan 2023 10:56:58 +0200
Subject: [PATCH] Feature: Add Kubernetes API health (livez/readyz) endpoints

Kubernetes checks if a pod is ok and if it can retrieve traffic using probes.
This commit add two routes to make a liveness probe and a readiness probe.
---
 server/handlers/k8healthz.go |  8 ++++++++
 server/handlers/k8sready.go  | 17 +++++++++++++++++
 server/server.go             | 24 ++++++++++++++++++------
 3 files changed, 43 insertions(+), 6 deletions(-)
 create mode 100644 server/handlers/k8healthz.go
 create mode 100644 server/handlers/k8sready.go

diff --git a/server/handlers/k8healthz.go b/server/handlers/k8healthz.go
new file mode 100644
index 0000000..818cacc
--- /dev/null
+++ b/server/handlers/k8healthz.go
@@ -0,0 +1,8 @@
+package handlers
+
+import "net/http"
+
+// Healthz is a liveness probe
+func HealthzHandler(w http.ResponseWriter, _ *http.Request) {
+	w.WriteHeader(http.StatusOK)
+}
diff --git a/server/handlers/k8sready.go b/server/handlers/k8sready.go
new file mode 100644
index 0000000..875b2b4
--- /dev/null
+++ b/server/handlers/k8sready.go
@@ -0,0 +1,17 @@
+package handlers
+
+import (
+	"net/http"
+	"sync/atomic"
+)
+
+// ReadyzHandler is a ready probe that signals k8s to be able to retrieve traffic
+func ReadyzHandler(isReady *atomic.Value) http.HandlerFunc {
+	return func(w http.ResponseWriter, _ *http.Request) {
+		if isReady == nil || !isReady.Load().(bool) {
+			http.Error(w, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable)
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+	}
+}
diff --git a/server/server.go b/server/server.go
index 021c468..129242e 100644
--- a/server/server.go
+++ b/server/server.go
@@ -3,17 +3,18 @@ package server
 import (
 	"compress/gzip"
 	"embed"
+	"github.com/axllent/mailpit/config"
+	"github.com/axllent/mailpit/server/apiv1"
+	"github.com/axllent/mailpit/server/handlers"
+	"github.com/axllent/mailpit/server/websockets"
+	"github.com/axllent/mailpit/utils/logger"
+	"github.com/gorilla/mux"
 	"io"
 	"io/fs"
 	"net/http"
 	"os"
 	"strings"
-
-	"github.com/axllent/mailpit/config"
-	"github.com/axllent/mailpit/server/apiv1"
-	"github.com/axllent/mailpit/server/websockets"
-	"github.com/axllent/mailpit/utils/logger"
-	"github.com/gorilla/mux"
+	"sync/atomic"
 )
 
 //go:embed ui
@@ -21,6 +22,9 @@ var embeddedFS embed.FS
 
 // Listen will start the httpd
 func Listen() {
+	isReady := &atomic.Value{}
+	isReady.Store(false)
+
 	serverRoot, err := fs.Sub(embeddedFS, "ui")
 	if err != nil {
 		logger.Log().Errorf("[http] %s", err)
@@ -33,6 +37,10 @@ func Listen() {
 
 	r := defaultRoutes()
 
+	// kubernetes probes
+	r.HandleFunc("/livez", handlers.HealthzHandler)
+	r.HandleFunc("/readyz", handlers.ReadyzHandler(isReady))
+
 	// web UI websocket
 	r.HandleFunc(config.Webroot+"api/events", apiWebsocket).Methods("GET")
 
@@ -51,6 +59,9 @@ func Listen() {
 		logger.Log().Info("[http] enabling web UI basic authentication")
 	}
 
+	// Mark the application here as ready
+	isReady.Store(true)
+
 	if config.UISSLCert != "" && config.UISSLKey != "" {
 		logger.Log().Infof("[http] starting secure server on https://%s%s", config.HTTPListen, config.Webroot)
 		logger.Log().Fatal(http.ListenAndServeTLS(config.HTTPListen, config.UISSLCert, config.UISSLKey, nil))
@@ -58,6 +69,7 @@ func Listen() {
 		logger.Log().Infof("[http] starting server on http://%s%s", config.HTTPListen, config.Webroot)
 		logger.Log().Fatal(http.ListenAndServe(config.HTTPListen, nil))
 	}
+
 }
 
 func defaultRoutes() *mux.Router {