From 0e52b7cd0286c2370193cfd53185a835a1cebd97 Mon Sep 17 00:00:00 2001
From: Piotr Kowalczuk <p.kowalczuk.priv@gmail.com>
Date: Sun, 14 Aug 2016 21:35:19 +0200
Subject: [PATCH] health check endpoint

---
 mnemosyned/daemon.go      | 10 +++++---
 mnemosyned/health.go      | 22 +++++++++++++++++
 mnemosyned/health_test.go | 51 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 4 deletions(-)
 create mode 100644 mnemosyned/health.go
 create mode 100644 mnemosyned/health_test.go

diff --git a/mnemosyned/daemon.go b/mnemosyned/daemon.go
index d22480b..f6d1240 100644
--- a/mnemosyned/daemon.go
+++ b/mnemosyned/daemon.go
@@ -50,6 +50,7 @@ type Daemon struct {
 	opts          *DaemonOpts
 	monitor       *monitoring
 	rpcOptions    []grpc.ServerOption
+	postgres      *sql.DB
 	storage       storage
 	logger        log.Logger
 	rpcListener   net.Listener
@@ -157,6 +158,9 @@ func (d *Daemon) Run() (err error) {
 			mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
 			mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
 			mux.Handle("/metrics", prometheus.Handler())
+			mux.Handle("/health", &healthHandler{
+				postgres: d.postgres,
+			})
 			sklog.Error(d.logger, http.Serve(d.debugListener, mux))
 		}()
 	}
@@ -184,13 +188,11 @@ func (d *Daemon) Addr() net.Addr {
 }
 
 func (d *Daemon) initStorage() (err error) {
-	var db *sql.DB
-
 	switch d.opts.Storage {
 	case StorageEngineInMemory:
 		return errors.New("in memory storage is not implemented yet")
 	case StorageEnginePostgres:
-		db, err = initPostgres(
+		d.postgres, err = initPostgres(
 			d.opts.PostgresAddress,
 			d.logger,
 		)
@@ -199,7 +201,7 @@ func (d *Daemon) initStorage() (err error) {
 		}
 		if d.storage, err = initStorage(
 			d.opts.IsTest,
-			newPostgresStorage("session", db, d.monitor, d.opts.SessionTTL),
+			newPostgresStorage("session", d.postgres, d.monitor, d.opts.SessionTTL),
 			d.logger,
 		); err != nil {
 			return
diff --git a/mnemosyned/health.go b/mnemosyned/health.go
new file mode 100644
index 0000000..e4ed1a0
--- /dev/null
+++ b/mnemosyned/health.go
@@ -0,0 +1,22 @@
+package mnemosyned
+
+import (
+	"database/sql"
+	"net/http"
+)
+
+type healthHandler struct {
+	postgres *sql.DB
+}
+
+func (hh *healthHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
+	if hh.postgres != nil {
+		if err := hh.postgres.Ping(); err != nil {
+			http.Error(rw, "postgres ping failure", http.StatusServiceUnavailable)
+			return
+		}
+	}
+
+	rw.WriteHeader(http.StatusOK)
+	rw.Write([]byte("1"))
+}
diff --git a/mnemosyned/health_test.go b/mnemosyned/health_test.go
new file mode 100644
index 0000000..e3f3bd7
--- /dev/null
+++ b/mnemosyned/health_test.go
@@ -0,0 +1,51 @@
+package mnemosyned
+
+import (
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestHealthHandler_ServeHTTP(t *testing.T) {
+	s := &postgresSuite{}
+	s.setup(t)
+
+	var (
+		res *http.Response
+		pay []byte
+		err error
+	)
+	srv := httptest.NewServer(&healthHandler{
+		postgres: s.db,
+	})
+	defer srv.Close()
+
+	if res, err = http.Get(srv.URL); err != nil {
+		t.Fatalf("unexpected error: %s", err.Error())
+	}
+	if res.StatusCode != http.StatusOK {
+		t.Fatalf("wrong status code, expected %d but got %d", http.StatusOK, res.StatusCode)
+	}
+	if pay, err = ioutil.ReadAll(res.Body); err != nil {
+		t.Fatalf("unexpected error: %s", err.Error())
+	}
+	if string(pay) != "1" {
+		t.Errorf("wrong payload, expected %s but got %s", "1", string(pay))
+	}
+
+	s.teardown(t)
+
+	if res, err = http.Get(srv.URL); err != nil {
+		t.Fatalf("unexpected error: %s", err.Error())
+	}
+	if res.StatusCode != http.StatusServiceUnavailable {
+		t.Fatalf("wrong status code, expected %d but got %d", http.StatusServiceUnavailable, res.StatusCode)
+	}
+	if pay, err = ioutil.ReadAll(res.Body); err != nil {
+		t.Fatalf("unexpected error: %s", err.Error())
+	}
+	if string(pay) != "postgres ping failure\n" {
+		t.Errorf("wrong payload, expected '%s' but got '%s'", "postgres ping failure\n", string(pay))
+	}
+}