etcdserver/api/metrics: hijack Prometheus handler, add 'health' key

Hijacks Prometheus handler, and add health information

Signed-off-by: Gyu-Ho Lee <gyuhox@gmail.com>

etcdserver/api/metrics/metrics.go

@@ -16,7 +16,9 @@
 package metrics
 
 import (
+	"bufio"
 	"context"
+	"fmt"
 	"net/http"
 	"time"
 
@@ -72,3 +74,46 @@ func RegisterPrometheus(mux *http.ServeMux) {
 func RegisterHealth(mux *http.ServeMux, srv *etcdserver.EtcdServer) {
 	mux.Handle(PathHealth, HealthHandler(srv))
 }
+
+// Handler serves health and metrics information.
+func Handler(srv *etcdserver.EtcdServer) http.HandlerFunc {
+	h := prometheus.Handler()
+	return func(w http.ResponseWriter, r *http.Request) {
+		hj, ok := w.(http.Hijacker)
+		if !ok { // webserver doesn't support hijacking
+			h.ServeHTTP(w, r)
+			return
+		}
+		conn, buf, err := hj.Hijack()
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		defer conn.Close()
+		hv := 0
+		if Health(srv) {
+			hv = 1
+		}
+		buf.WriteString(fmt.Sprintf("health %d\n", hv))
+		wr := &responseWriter{w: w, buf: buf}
+		h.ServeHTTP(wr, r)
+		buf.Flush()
+	}
+}
+
+type responseWriter struct {
+	w   http.ResponseWriter
+	buf *bufio.ReadWriter
+}
+
+func (w *responseWriter) Header() http.Header {
+	return w.w.Header()
+}
+
+func (w *responseWriter) Write(b []byte) (int, error) {
+	return w.buf.Write(b)
+}
+
+func (w *responseWriter) WriteHeader(h int) {
+	w.w.WriteHeader(h)
+}