From c4317ac0c8afdcfde8643233ecfac2af9be087aa Mon Sep 17 00:00:00 2001
From: Stefan Prodan <stefan.prodan@gmail.com>
Date: Mon, 8 Jan 2018 23:27:48 +0200
Subject: [PATCH] readiness handler

---
 deploy/podinfo-dep.yaml |  2 +-
 pkg/server/handlers.go  | 19 ++++++++++++++++++-
 pkg/server/server.go    | 16 ++++++++++++----
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/deploy/podinfo-dep.yaml b/deploy/podinfo-dep.yaml
index b0dccdb..8b416a5 100644
--- a/deploy/podinfo-dep.yaml
+++ b/deploy/podinfo-dep.yaml
@@ -33,7 +33,7 @@ spec:
           protocol: TCP
         readinessProbe:
           httpGet:
-            path: /healthz
+            path: /readyz
             port: 9898
           initialDelaySeconds: 1
           periodSeconds: 2
diff --git a/pkg/server/handlers.go b/pkg/server/handlers.go
index c93a190..476b6a3 100644
--- a/pkg/server/handlers.go
+++ b/pkg/server/handlers.go
@@ -48,7 +48,7 @@ func (s *Server) echo(w http.ResponseWriter, r *http.Request) {
 }
 
 func (s *Server) healthz(w http.ResponseWriter, r *http.Request) {
-	if atomic.LoadInt32(&status) == 1 {
+	if atomic.LoadInt32(&healthy) == 1 {
 		w.WriteHeader(http.StatusOK)
 		w.Write([]byte("OK"))
 		return
@@ -56,6 +56,23 @@ func (s *Server) healthz(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusServiceUnavailable)
 }
 
+func (s *Server) readyz(w http.ResponseWriter, r *http.Request) {
+	if atomic.LoadInt32(&ready) == 1 {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte("OK"))
+		return
+	}
+	w.WriteHeader(http.StatusServiceUnavailable)
+}
+
+func (s *Server) enable(w http.ResponseWriter, r *http.Request) {
+	atomic.StoreInt32(&ready, 1)
+}
+
+func (s *Server) disable(w http.ResponseWriter, r *http.Request) {
+	atomic.StoreInt32(&ready, 0)
+}
+
 func (s *Server) panic(w http.ResponseWriter, r *http.Request) {
 	glog.Fatal("Kill switch triggered")
 }
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 3ecf287..6f6cfba 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -11,7 +11,10 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
-var status int32
+var (
+	healthy int32
+	ready   int32
+)
 
 type Server struct {
 	mux *http.ServeMux
@@ -26,6 +29,9 @@ func New(options ...func(*Server)) *Server {
 
 	s.mux.HandleFunc("/", s.index)
 	s.mux.HandleFunc("/healthz/", s.healthz)
+	s.mux.HandleFunc("/readyz/", s.readyz)
+	s.mux.HandleFunc("/readyz/enable/", s.enable)
+	s.mux.HandleFunc("/readyz/disable/", s.disable)
 	s.mux.HandleFunc("/echo/", s.echo)
 	s.mux.HandleFunc("/panic/", s.panic)
 	s.mux.Handle("/metrics", promhttp.Handler())
@@ -48,7 +54,8 @@ func ListenAndServe(port string, timeout time.Duration, stopCh <-chan struct{})
 		IdleTimeout:  15 * time.Second,
 	}
 
-	atomic.StoreInt32(&status, 1)
+	atomic.StoreInt32(&healthy, 1)
+	atomic.StoreInt32(&ready, 1)
 
 	// run server in background
 	go func() {
@@ -62,8 +69,9 @@ func ListenAndServe(port string, timeout time.Duration, stopCh <-chan struct{})
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 
-	// all calls to /healthz will fail from now on
-	atomic.StoreInt32(&status, 0)
+	// all calls to /healthz and /readyz will fail from now on
+	atomic.StoreInt32(&healthy, 0)
+	atomic.StoreInt32(&ready, 0)
 
 	glog.Infof("Shutting down HTTP server with timeout: %v", timeout)