From 743849a1313425f7cfbcd1626e4f3604520b4db5 Mon Sep 17 00:00:00 2001
From: matst80 <mats.tornberg@gmail.com>
Date: Sun, 10 Nov 2024 23:20:52 +0100
Subject: [PATCH] health check

---
 main.go        | 25 ++++++++++++++++++++++++-
 rpc-server.go  |  4 ++++
 synced-pool.go | 13 +++++++++++++
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/main.go b/main.go
index eebce40..1ad99cf 100644
--- a/main.go
+++ b/main.go
@@ -125,7 +125,7 @@ func main() {
 	// if local
 	//syncedPool.AddRemote("localhost")
 
-	_, err = NewGrainHandler(app.pool, ":1337")
+	hg, err := NewGrainHandler(app.pool, ":1337")
 	if err != nil {
 		log.Fatalf("Error creating handler: %v\n", err)
 	}
@@ -155,6 +155,29 @@ func main() {
 	mux.HandleFunc("/pprof/symbol", pprof.Symbol)
 	mux.HandleFunc("/pprof/trace", pprof.Trace)
 	mux.Handle("/metrics", promhttp.Handler())
+	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
+		if !hg.IsHealthy() {
+			w.WriteHeader(http.StatusInternalServerError)
+			w.Write([]byte("handler not healthy"))
+			return
+		}
+		if !syncedPool.IsHealthy() {
+			w.WriteHeader(http.StatusInternalServerError)
+			w.Write([]byte("pool not healthy"))
+			return
+		}
+
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte("ok"))
+	})
+	mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte("ok"))
+	})
+	mux.HandleFunc("/livez", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte("ok"))
+	})
 
 	sigs := make(chan os.Signal, 1)
 	done := make(chan bool, 1)
diff --git a/rpc-server.go b/rpc-server.go
index 4e87f43..bf52f4c 100644
--- a/rpc-server.go
+++ b/rpc-server.go
@@ -30,6 +30,10 @@ func NewGrainHandler(pool *GrainLocalPool, listen string) (*GrainHandler, error)
 	return handler, err
 }
 
+func (h *GrainHandler) IsHealthy() bool {
+	return len(h.pool.grains) < h.pool.PoolSize
+}
+
 func (h *GrainHandler) RemoteHandleMessageHandler(id CartId, data []byte) (uint32, []byte, error) {
 	var msg Message
 	err := ReadMessage(bytes.NewReader(data), &msg)
diff --git a/synced-pool.go b/synced-pool.go
index 940a8ab..23d9e57 100644
--- a/synced-pool.go
+++ b/synced-pool.go
@@ -17,6 +17,10 @@ type Quorum interface {
 	OwnerChanged(CartId, host string) error
 }
 
+type HealthHandler interface {
+	IsHealthy() bool
+}
+
 type RemoteHost struct {
 	*Client
 	Host        string
@@ -200,6 +204,15 @@ func NewSyncedPool(local *GrainLocalPool, hostname string, discovery Discovery)
 	return pool, nil
 }
 
+func (p *SyncedPool) IsHealthy() bool {
+	for _, r := range p.remotes {
+		if r.MissedPings > 3 {
+			return false
+		}
+	}
+	return true
+}
+
 func (p *SyncedPool) IsKnown(host string) bool {
 	for _, r := range p.remotes {
 		if r.Host == host {