From cd55450ef60d2e86f8e0a6b51daaa35851245f5e Mon Sep 17 00:00:00 2001
From: Stefan Prodan <stefan.prodan@gmail.com>
Date: Tue, 9 Jan 2018 01:51:08 +0200
Subject: [PATCH] added http_requests_total Prometheus counter

---
 Gopkg.lock                |  2 +-
 pkg/server/handlers.go    |  2 +-
 pkg/server/interceptor.go | 30 +++++++++++++++++++++++++++++
 pkg/server/server.go      | 40 +++++++++++++++++++++++++++++----------
 4 files changed, 62 insertions(+), 12 deletions(-)
 create mode 100644 pkg/server/interceptor.go

diff --git a/Gopkg.lock b/Gopkg.lock
index 3815e1a..16ffb7c 100644
--- a/Gopkg.lock
+++ b/Gopkg.lock
@@ -74,6 +74,6 @@
 [solve-meta]
   analyzer-name = "dep"
   analyzer-version = 1
-  inputs-digest = "ee024b4cf0bcf20b65cc49aa39fd24d711cb1cea1634a739f5a5c45e06be34fa"
+  inputs-digest = "0167c7d2c6f05b3a2c3e382930866dc3c36865afba7940c620d2aef85e6f83e6"
   solver-name = "gps-cdcl"
   solver-version = 1
diff --git a/pkg/server/handlers.go b/pkg/server/handlers.go
index 476b6a3..7f7daf8 100644
--- a/pkg/server/handlers.go
+++ b/pkg/server/handlers.go
@@ -40,8 +40,8 @@ func (s *Server) echo(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 		glog.Infof("Payload received from %s: %s", r.RemoteAddr, string(body))
+		w.WriteHeader(http.StatusAccepted)
 		w.Write(body)
-		w.WriteHeader(http.StatusOK)
 	default:
 		w.WriteHeader(http.StatusNotAcceptable)
 	}
diff --git a/pkg/server/interceptor.go b/pkg/server/interceptor.go
new file mode 100644
index 0000000..5925f42
--- /dev/null
+++ b/pkg/server/interceptor.go
@@ -0,0 +1,30 @@
+package server
+
+import (
+	"net/http"
+	"net"
+	"bufio"
+	"fmt"
+)
+
+type interceptor struct {
+	http.ResponseWriter
+	statusCode int
+	recorded   bool
+}
+
+func (i *interceptor) WriteHeader(code int) {
+	if !i.recorded {
+		i.statusCode = code
+		i.recorded = true
+	}
+	i.ResponseWriter.WriteHeader(code)
+}
+
+func (i *interceptor) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hj, ok := i.ResponseWriter.(http.Hijacker)
+	if !ok {
+		return nil, nil, fmt.Errorf("interceptor: can't cast parent ResponseWriter to Hijacker")
+	}
+	return hj.Hijack()
+}
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 6f6cfba..9bba1f0 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -4,23 +4,32 @@ import (
 	"context"
 	"net/http"
 	"runtime"
+	"strconv"
 	"sync/atomic"
 	"time"
 
 	"github.com/golang/glog"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
 var (
-	healthy int32
-	ready   int32
+	healthy             int32
+	ready               int32
+	httpRequestsCounter = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Name: "http_requests_total",
+			Help: "The total number of HTTP requests.",
+		},
+		[]string{"method", "path", "status"},
+	)
 )
 
 type Server struct {
 	mux *http.ServeMux
 }
 
-func New(options ...func(*Server)) *Server {
+func NewServer(options ...func(*Server)) *Server {
 	s := &Server{mux: http.NewServeMux()}
 
 	for _, f := range options {
@@ -28,12 +37,12 @@ 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.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())
 
 	return s
@@ -45,10 +54,19 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	s.mux.ServeHTTP(w, r)
 }
 
+func instrument(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		interceptor := &interceptor{ResponseWriter: w, statusCode: http.StatusOK}
+		next.ServeHTTP(interceptor, r)
+		var status = strconv.Itoa(interceptor.statusCode)
+		httpRequestsCounter.WithLabelValues(r.Method, r.URL.Path, status).Inc()
+	})
+}
+
 func ListenAndServe(port string, timeout time.Duration, stopCh <-chan struct{}) {
 	srv := &http.Server{
 		Addr:         ":" + port,
-		Handler:      New(),
+		Handler:      instrument(NewServer()),
 		ReadTimeout:  5 * time.Second,
 		WriteTimeout: 10 * time.Second,
 		IdleTimeout:  15 * time.Second,
@@ -57,6 +75,8 @@ func ListenAndServe(port string, timeout time.Duration, stopCh <-chan struct{})
 	atomic.StoreInt32(&healthy, 1)
 	atomic.StoreInt32(&ready, 1)
 
+	prometheus.MustRegister(httpRequestsCounter)
+
 	// run server in background
 	go func() {
 		if err := srv.ListenAndServe(); err != http.ErrServerClosed {