From 5481cfe7efa9c14de44f33cb73ef7eefe2759bc2 Mon Sep 17 00:00:00 2001
From: Stefan Prodan <stefan.prodan@gmail.com>
Date: Sun, 7 Jan 2018 17:58:41 +0200
Subject: [PATCH] use an atomic int to determine health status

---
 pkg/server/response.go | 10 +++++-----
 pkg/server/server.go   | 25 ++++++++++++++++++++-----
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/pkg/server/response.go b/pkg/server/response.go
index bd4e605..cfea3ad 100644
--- a/pkg/server/response.go
+++ b/pkg/server/response.go
@@ -17,11 +17,11 @@ import (
 )
 
 type Response struct {
-	Runtime     map[string]string `json:"runtime"`
-	Labels      map[string]string `json:"labels"`
-	Annotations map[string]string `json:"annotations"`
-	Environment map[string]string `json:"environment"`
-	ExternalIP  map[string]string `json:"externalips"`
+	Runtime     map[string]string `json:"runtime" yaml:"runtime"`
+	Labels      map[string]string `json:"labels,omitempty" yaml:"labels,omitempty"`
+	Annotations map[string]string `json:"annotations,omitempty" yaml:"annotations,omitempty"`
+	Environment map[string]string `json:"environment" yaml:"environment"`
+	ExternalIP  map[string]string `json:"externalIP" yaml:"externalIP"`
 }
 
 func makeResponse() (*Response, error) {
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 6663591..281ca80 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -9,8 +9,11 @@ import (
 	"github.com/golang/glog"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"gopkg.in/yaml.v2"
+	"sync/atomic"
 )
 
+var status int32
+
 type Server struct {
 	mux *http.ServeMux
 }
@@ -42,13 +45,19 @@ func (s *Server) index(w http.ResponseWriter, r *http.Request) {
 		w.Write([]byte(err.Error()))
 	}
 
-	w.Header().Set("Content-Type", "text/x-yaml")
+	w.Header().Set("Content-Type", "text/x-yaml; charset=utf-8")
+	w.Header().Set("X-Content-Type-Options", "nosniff")
+	w.WriteHeader(http.StatusOK)
 	w.Write(d)
 }
 
 func (s *Server) healthz(w http.ResponseWriter, r *http.Request) {
-	w.WriteHeader(http.StatusOK)
-	w.Write([]byte("OK"))
+	if atomic.LoadInt32(&status) == 1 {
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte("OK"))
+		return
+	}
+	w.WriteHeader(http.StatusServiceUnavailable)
 }
 
 func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -61,8 +70,13 @@ func ListenAndServe(port string, timeout time.Duration, stopCh <-chan struct{})
 	srv := &http.Server{
 		Addr:    ":" + port,
 		Handler: New(),
+		ReadTimeout:  5 * time.Second,
+		WriteTimeout: 10 * time.Second,
+		IdleTimeout:  15 * time.Second,
 	}
 
+	atomic.StoreInt32(&status, 1)
+
 	// run server in background
 	go func() {
 		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
@@ -72,13 +86,14 @@ func ListenAndServe(port string, timeout time.Duration, stopCh <-chan struct{})
 
 	// wait for SIGTERM or SIGINT
 	<-stopCh
+	atomic.StoreInt32(&status, 0)
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 
-	glog.Infof("Shutdown HTTP server with timeout: %s", timeout)
+	glog.Infof("Shutting down HTTP server with timeout: %v", timeout)
 
 	if err := srv.Shutdown(ctx); err != nil {
-		glog.Error(err)
+		glog.Errorf("HTTP server graceful shutdown failed with error: %v",err)
 	} else {
 		glog.Info("HTTP server stopped")
 	}