From fa4736e509f128ba716d401e5cbd1ab4629b23ff Mon Sep 17 00:00:00 2001
From: Tobias Lindberg <tobias.ehlert@gmail.com>
Date: Wed, 24 Jul 2024 22:31:44 +0200
Subject: [PATCH] Feature: adding support for k8s health endpoints (#191)

* adding support for k8s health endpoints
* fix: update readyz if mqtt is lost
---
 README.md                        |  2 ++
 src/v1_TeslaMateAPICarsStatus.go |  8 ++++++++
 src/webserver.go                 | 30 ++++++++++++++++++++++++++++++
 3 files changed, 40 insertions(+)

diff --git a/README.md b/README.md
index f619fe9..ea3b6c6 100644
--- a/README.md
+++ b/README.md
@@ -171,7 +171,9 @@ More detailed documentation of every endpoint will come..
 - GET `/api/v1/cars/:CarID/updates`
 - POST `/api/v1/cars/:CarID/wake_up`
 - GET `/api/v1/globalsettings`
+- GET `/api/healthz`
 - GET `/api/ping`
+- GET `/api/readyz`
 
 ### Authentication
 
diff --git a/src/v1_TeslaMateAPICarsStatus.go b/src/v1_TeslaMateAPICarsStatus.go
index 73aa3c4..2434d40 100644
--- a/src/v1_TeslaMateAPICarsStatus.go
+++ b/src/v1_TeslaMateAPICarsStatus.go
@@ -190,6 +190,9 @@ func startMQTT() (*statusCache, error) {
 
 	s.topicScan = fmt.Sprintf("teslamate%s/cars/%%d/%%s", getMQTTNameSpace())
 
+	// setting readyz endpoint to true (when using MQTT)
+	isReady.Store(true)
+
 	// Thats all - newMessage will be called when something new arrives
 	return &s, nil
 }
@@ -215,12 +218,17 @@ func (s *statusCache) connectedHandler(c mqtt.Client) {
 	}
 	log.Println("[info] subscribed to: " + topic)
 
+	// setting readyz endpoint to true (when using MQTT)
+	isReady.Store(true)
 }
 
 // connectionLost - called by mqtt package when the connection get lost
 func (s *statusCache) connectionLost(c mqtt.Client, err error) {
 	log.Println("[error] MQTT connection lost: " + err.Error())
 	s.mqttConnected = false
+
+	// setting readyz endpoint to false (when using MQTT)
+	isReady.Store(false)
 }
 
 // newMessage - called by mqtt package when new message received
diff --git a/src/webserver.go b/src/webserver.go
index 142806e..22e0f3f 100644
--- a/src/webserver.go
+++ b/src/webserver.go
@@ -9,6 +9,7 @@ import (
 	"os"
 	"os/signal"
 	"strconv"
+	"sync/atomic"
 	"time"
 
 	"github.com/gin-contrib/gzip"
@@ -17,6 +18,9 @@ import (
 )
 
 var (
+	// application readyz endpoint value for k8s
+	isReady *atomic.Value
+
 	// setting TeslaMateApi version number
 	apiVersion = "unspecified"
 
@@ -32,6 +36,9 @@ var (
 
 // main function
 func main() {
+	// setup of readyness endpoint code
+	isReady := &atomic.Value{}
+	isReady.Store(false)
 
 	// setting log parameters
 	log.SetFlags(log.Ldate | log.Lmicroseconds)
@@ -141,6 +148,10 @@ func main() {
 
 		// /api/ping endpoint
 		api.GET("/ping", func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"message": "pong"}) })
+
+		// health endpoints for kubernetes
+		api.GET("/healthz", healthz)
+		api.GET("/readyz", readyz)
 	}
 
 	// TeslaMateApi endpoints (before versioning)
@@ -161,6 +172,11 @@ func main() {
 		Handler: r,
 	}
 
+	// setting readyz endpoint to true (if not using MQTT)
+	if getEnvAsBool("DISABLE_MQTT", false) {
+		isReady.Store(true)
+	}
+
 	// graceful shutdown
 	quit := make(chan os.Signal, 1)
 	signal.Notify(quit, os.Interrupt)
@@ -428,3 +444,17 @@ func checkArrayContainsString(s []string, e string) bool {
 	}
 	return false
 }
+
+// healthz is a liveness probe.
+func healthz(c *gin.Context) {
+	c.JSON(http.StatusOK, gin.H{"status": http.StatusText(http.StatusOK)})
+}
+
+// readyz is a readiness probe.
+func readyz(c *gin.Context) {
+	if isReady == nil || !isReady.Load().(bool) {
+		c.JSON(http.StatusServiceUnavailable, gin.H{"error": http.StatusText(http.StatusServiceUnavailable)})
+		return
+	}
+	TeslaMateAPIHandleSuccessResponse(c, "webserver", gin.H{"status": http.StatusText(http.StatusOK)})
+}