From 2a3b5ce448caa0095970b556717a14a795cb50aa Mon Sep 17 00:00:00 2001
From: Ludvig Lundgren <hello@ludviglundgren.se>
Date: Tue, 12 Apr 2022 18:19:07 +0200
Subject: [PATCH] feat(http): add healthcheck endpoints liveness and readiness
 (#240)

* feat(http): add liveness and readiness endpoints

* feat(http): improve unhealthy msg
---
 cmd/autobrr/main.go           |  2 +-
 internal/database/database.go |  4 +++
 internal/http/health.go       | 51 +++++++++++++++++++++++++++++++++++
 internal/http/server.go       |  6 ++++-
 4 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 internal/http/health.go

diff --git a/cmd/autobrr/main.go b/cmd/autobrr/main.go
index fe0a917..701c379 100644
--- a/cmd/autobrr/main.go
+++ b/cmd/autobrr/main.go
@@ -100,7 +100,7 @@ func main() {
 	errorChannel := make(chan error)
 
 	go func() {
-		httpServer := http.NewServer(cfg, serverEvents, version, commit, date, actionService, authService, downloadClientService, filterService, indexerService, ircService, notificationService, releaseService)
+		httpServer := http.NewServer(cfg, serverEvents, db, version, commit, date, actionService, authService, downloadClientService, filterService, indexerService, ircService, notificationService, releaseService)
 		errorChannel <- httpServer.Open()
 	}()
 
diff --git a/internal/database/database.go b/internal/database/database.go
index 027a274..7fbbea0 100644
--- a/internal/database/database.go
+++ b/internal/database/database.go
@@ -82,6 +82,10 @@ func (db *DB) Close() error {
 	return nil
 }
 
+func (db *DB) Ping() error {
+	return db.handler.Ping()
+}
+
 func (db *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) (*Tx, error) {
 	tx, err := db.handler.BeginTx(ctx, opts)
 	if err != nil {
diff --git a/internal/http/health.go b/internal/http/health.go
new file mode 100644
index 0000000..1828d6c
--- /dev/null
+++ b/internal/http/health.go
@@ -0,0 +1,51 @@
+package http
+
+import (
+	"net/http"
+
+	"github.com/autobrr/autobrr/internal/database"
+
+	"github.com/go-chi/chi"
+)
+
+type healthHandler struct {
+	encoder encoder
+	db      *database.DB
+}
+
+func newHealthHandler(encoder encoder, db *database.DB) *healthHandler {
+	return &healthHandler{
+		encoder: encoder,
+		db:      db,
+	}
+}
+
+func (h healthHandler) Routes(r chi.Router) {
+	r.Get("/liveness", h.handleLiveness)
+	r.Get("/readiness", h.handleReadiness)
+}
+
+func (h healthHandler) handleLiveness(w http.ResponseWriter, _ *http.Request) {
+	writeHealthy(w)
+}
+
+func (h healthHandler) handleReadiness(w http.ResponseWriter, _ *http.Request) {
+	if err := h.db.Ping(); err != nil {
+		writeUnhealthy(w)
+		return
+	}
+
+	writeHealthy(w)
+}
+
+func writeHealthy(w http.ResponseWriter) {
+	w.Header().Set("Content-Type", "text/plain")
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("OK"))
+}
+
+func writeUnhealthy(w http.ResponseWriter) {
+	w.Header().Set("Content-Type", "text/plain")
+	w.WriteHeader(http.StatusInternalServerError)
+	w.Write([]byte("Unhealthy. Database unreachable"))
+}
diff --git a/internal/http/server.go b/internal/http/server.go
index aabeff4..398abff 100644
--- a/internal/http/server.go
+++ b/internal/http/server.go
@@ -6,6 +6,7 @@ import (
 	"net"
 	"net/http"
 
+	"github.com/autobrr/autobrr/internal/database"
 	"github.com/autobrr/autobrr/internal/domain"
 	"github.com/autobrr/autobrr/web"
 
@@ -17,6 +18,7 @@ import (
 
 type Server struct {
 	sse *sse.Server
+	db  *database.DB
 
 	config      domain.Config
 	cookieStore *sessions.CookieStore
@@ -35,10 +37,11 @@ type Server struct {
 	releaseService        releaseService
 }
 
-func NewServer(config domain.Config, sse *sse.Server, version string, commit string, date string, actionService actionService, authService authService, downloadClientSvc downloadClientService, filterSvc filterService, indexerSvc indexerService, ircSvc ircService, notificationSvc notificationService, releaseSvc releaseService) Server {
+func NewServer(config domain.Config, sse *sse.Server, db *database.DB, version string, commit string, date string, actionService actionService, authService authService, downloadClientSvc downloadClientService, filterSvc filterService, indexerSvc indexerService, ircSvc ircService, notificationSvc notificationService, releaseSvc releaseService) Server {
 	return Server{
 		config:  config,
 		sse:     sse,
+		db:      db,
 		version: version,
 		commit:  commit,
 		date:    date,
@@ -98,6 +101,7 @@ func (s Server) Handler() http.Handler {
 	})
 
 	r.Route("/api/auth", newAuthHandler(encoder, s.config, s.cookieStore, s.authService).Routes)
+	r.Route("/api/healthz", newHealthHandler(encoder, s.db).Routes)
 
 	r.Group(func(r chi.Router) {
 		r.Use(s.IsAuthenticated)