From 1c4df114440ed3ec7ad584d2a9e73bb43f4ca89d Mon Sep 17 00:00:00 2001
From: "mark.lin" <mrak@maicoin.com>
Date: Fri, 15 Feb 2019 13:54:00 +0800
Subject: [PATCH] health: add liveness and readiness in http server mux

---
 health/server.go | 17 +++--------------
 health/utils.go  | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/health/server.go b/health/server.go
index 74d64b7a..c93853e4 100644
--- a/health/server.go
+++ b/health/server.go
@@ -58,20 +58,9 @@ func (s *server) Liveness(ctx context.Context, req *EmptyRequest) (*EmptyRespons
 
 // Readiness is represented that whether application is ready to start accepting traffic or not.
 func (s *server) Readiness(ctx context.Context, req *EmptyRequest) (*EmptyResponse, error) {
-	if len(s.checkFns) == 0 {
-		return nil, nil
-	}
-	errCh := make(chan error, len(s.checkFns))
-	for _, checker := range s.checkFns {
-		go func(checker CheckFn) {
-			errCh <- checker(ctx)
-		}(checker)
-	}
-	for range s.checkFns {
-		if err := <-errCh; err != nil {
-			log.Error("Failed to check readiness", "err", err)
-			return nil, status.Error(codes.Unavailable, err.Error())
-		}
+	if err := CheckHealth(ctx, s.checkFns); err != nil {
+		log.Error("Failed to check readiness", "err", err)
+		return nil, status.Error(codes.Unavailable, err.Error())
 	}
 	return nil, nil
 }
diff --git a/health/utils.go b/health/utils.go
index d8db04f1..b6757892 100644
--- a/health/utils.go
+++ b/health/utils.go
@@ -16,6 +16,7 @@ package health
 
 import (
 	"context"
+	"net/http"
 	"time"
 
 	"github.com/getamis/sirius/log"
@@ -41,3 +42,37 @@ func GRPCServerHealthChecker(addr, serviceName string) CheckFn {
 		return nil
 	}
 }
+
+func CheckHealth(ctx context.Context, checkFns []CheckFn) error {
+	if len(checkFns) == 0 {
+		return nil
+	}
+	errCh := make(chan error, len(checkFns))
+	for _, checker := range checkFns {
+		go func(checker CheckFn) {
+			errCh <- checker(ctx)
+		}(checker)
+	}
+	for range checkFns {
+		if err := <-errCh; err != nil {
+			log.Error("Failed to check readiness", "err", err)
+			return err
+		}
+	}
+	return nil
+}
+
+// SetLivenessAndReadiness sets the liveness and readiness route in http server mux
+func SetLivenessAndReadiness(mux *http.ServeMux, checkFns ...CheckFn) {
+	mux.HandleFunc("/readiness", func(rw http.ResponseWriter, r *http.Request) {
+		if err := CheckHealth(r.Context(), checkFns); err != nil {
+			log.Error("Failed to check readiness", "err", err)
+			rw.WriteHeader(http.StatusServiceUnavailable)
+			return
+		}
+		rw.WriteHeader(http.StatusOK)
+	})
+	mux.HandleFunc("/liveness", func(rw http.ResponseWriter, r *http.Request) {
+		rw.WriteHeader(http.StatusOK)
+	})
+}