From ae86223ace6e7890a35d27fa523c5e0649f417bd Mon Sep 17 00:00:00 2001
From: Jannis Mattheis <contact@jmattheis.de>
Date: Fri, 11 Oct 2024 14:56:44 +0200
Subject: [PATCH] fix: add health endpoint

---
 router/router.go   | 19 +++++++++++++++++++
 ws/event_health.go | 10 ++++++++++
 ws/rooms.go        | 16 ++++++++++++++++
 3 files changed, 45 insertions(+)
 create mode 100644 ws/event_health.go

diff --git a/router/router.go b/router/router.go
index 2cbd54d..99eaabb 100644
--- a/router/router.go
+++ b/router/router.go
@@ -16,6 +16,12 @@ import (
 	"github.com/screego/server/ws"
 )
 
+type Health struct {
+	Status  string `json:"status"`
+	Clients int    `json:"clients"`
+	Reason  string `json:"reason,omitempty"`
+}
+
 type UIConfig struct {
 	AuthMode                 string `json:"authMode"`
 	User                     string `json:"user"`
@@ -47,6 +53,19 @@ func Router(conf config.Config, rooms *ws.Rooms, users *auth.Users, version stri
 			CloseRoomWhenOwnerLeaves: conf.CloseRoomWhenOwnerLeaves,
 		})
 	})
+	router.Methods("GET").Path("/health").HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		i, err := rooms.Count()
+		status := "up"
+		if err != "" {
+			status = "down"
+			w.WriteHeader(500)
+		}
+		_ = json.NewEncoder(w).Encode(Health{
+			Status:  status,
+			Clients: i,
+			Reason:  err,
+		})
+	})
 	if conf.Prometheus {
 		log.Info().Msg("Prometheus enabled")
 		router.Methods("GET").Path("/metrics").Handler(basicAuth(promhttp.Handler(), users))
diff --git a/ws/event_health.go b/ws/event_health.go
new file mode 100644
index 0000000..1264fd4
--- /dev/null
+++ b/ws/event_health.go
@@ -0,0 +1,10 @@
+package ws
+
+type Health struct {
+	Response chan int
+}
+
+func (e *Health) Execute(rooms *Rooms, current ClientInfo) error {
+	e.Response <- len(rooms.connected)
+	return nil
+}
diff --git a/ws/rooms.go b/ws/rooms.go
index b960653..a75dba7 100644
--- a/ws/rooms.go
+++ b/ws/rooms.go
@@ -110,6 +110,22 @@ func (r *Rooms) Start() {
 	}
 }
 
+func (r *Rooms) Count() (int, string) {
+	h := Health{Response: make(chan int, 1)}
+	select {
+	case r.Incoming <- ClientMessage{SkipConnectedCheck: true, Incoming: &h}:
+	case <-time.After(5 * time.Second):
+		return -1, "main loop didn't accept a message within 5 second"
+	}
+	r.Incoming <- ClientMessage{SkipConnectedCheck: true, Incoming: &h}
+	select {
+	case count := <-h.Response:
+		return count, ""
+	case <-time.After(5 * time.Second):
+		return -1, "main loop didn't respond to a message within 5 second"
+	}
+}
+
 func (r *Rooms) closeRoom(roomID string) {
 	room, ok := r.Rooms[roomID]
 	if !ok {