feat: add response time middleware

examples/middlewares/main.go

@@ -438,6 +438,36 @@ func ExampleSessionID() {
 	svr.Run()
 }
 
+func ExmampleResponseTime() {
+	svr := keel.NewServer()
+
+	// get logger
+	l := svr.Logger()
+
+	// create demo service
+	svs := http.NewServeMux()
+	svs.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("OK"))
+	})
+	svs.HandleFunc("/sleep", func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(time.Second)
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("OK"))
+	})
+
+	svr.AddService(
+		keel.NewServiceHTTP(l, "demo", ":8080", svs,
+			middleware.ResponseTime(
+				// automatically set cookie if not exists
+				middleware.ResponseTimeWithMaxDuration(time.Millisecond * 500),
+			),
+		),
+	)
+
+	svr.Run()
+}
+
 func ExampleSkip() {
 	svr := keel.NewServer()
 

net/http/header.go

@@ -24,6 +24,7 @@ const (
 	HeaderXUrlScheme          = "X-Url-Scheme"
 	HeaderXHTTPMethodOverride = "X-HTTP-Method-Override"
 	HeaderXRealIP             = "X-Real-IP"
+	HeaderXResponseTime       = "X-Response-Time"
 	HeaderXRequestID          = "X-Request-ID"
 	HeaderXRequestedWith      = "X-Requested-With"
 	HeaderXSessionID          = "X-Session-ID"

net/http/middleware/logger.go

@@ -37,7 +37,7 @@ func LoggerWithOptions(opts LoggerOptions) Middleware {
 			start := time.Now()
 
 			// wrap response write to get access to status & size
-			wr := wrapResponseWriter(w)
+			wr := WrapResponseWriter(w)
 
 			next.ServeHTTP(wr, r)
 

net/http/middleware/responsetime.go

@@ -0,0 +1,75 @@
+package middleware
+
+import (
+	"net/http"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/foomo/keel/log"
+)
+
+type (
+	ResponseTimeOptions struct {
+		SetHeader          bool
+		MaxDuration        time.Duration
+		MaxDurationMessage string
+	}
+	ResponseTimeOption func(*ResponseTimeOptions)
+)
+
+// GetDefaultResponseTimeOptions returns the default options
+func GetDefaultResponseTimeOptions() ResponseTimeOptions {
+	return ResponseTimeOptions{
+		SetHeader:          true,
+		MaxDurationMessage: "max response time exceeded",
+	}
+}
+
+// ResponseTimeWithMaxDurationMessage middleware option
+func ResponseTimeWithMaxDurationMessage(v string) ResponseTimeOption {
+	return func(o *ResponseTimeOptions) {
+		o.MaxDurationMessage = v
+	}
+}
+
+// ResponseTimeWithMaxDuration middleware option
+func ResponseTimeWithMaxDuration(v time.Duration) ResponseTimeOption {
+	return func(o *ResponseTimeOptions) {
+		o.MaxDuration = v
+	}
+}
+
+// ResponseTimeWithSetHeader middleware option
+func ResponseTimeWithSetHeader(v bool) ResponseTimeOption {
+	return func(o *ResponseTimeOptions) {
+		o.SetHeader = v
+	}
+}
+
+// ResponseTime middleware
+func ResponseTime(opts ...ResponseTimeOption) Middleware {
+	options := GetDefaultResponseTimeOptions()
+	for _, opt := range opts {
+		if opt != nil {
+			opt(&options)
+		}
+	}
+	return ResponseTimeWithOptions(options)
+}
+
+// ResponseTimeWithOptions middleware
+func ResponseTimeWithOptions(opts ResponseTimeOptions) Middleware {
+	return func(l *zap.Logger, next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			start := time.Now()
+			rw := WrapResponseWriter(w)
+			rw.SetWriteResponseTimeHeader(opts.SetHeader)
+			next.ServeHTTP(rw, r)
+			duration := time.Since(start)
+			if opts.MaxDuration > 0 && duration > opts.MaxDuration {
+				l.Warn(opts.MaxDurationMessage, log.FDuration(opts.MaxDuration), log.FValue(duration.Microseconds()))
+			}
+		})
+	}
+}

net/http/middleware/responsewriter.go

@@ -3,25 +3,36 @@ package middleware
 import (
 	"fmt"
 	"net/http"
+	"strconv"
+	"time"
+
+	http2 "github.com/foomo/keel/net/http"
 )
 
 // responseWriter is a wrapper that includes that http statusCode and size for logging
 type responseWriter struct {
 	http.ResponseWriter
-	wroteHeader bool
-	statusCode  int
-	size        int
+	writeResponseTimeHeader bool
+	wroteHeader             bool
+	statusCode              int
+	start                   time.Time
+	size                    int
 }
 
-func wrapResponseWriter(w http.ResponseWriter) *responseWriter {
+func WrapResponseWriter(w http.ResponseWriter) *responseWriter {
 	if wr, ok := w.(*responseWriter); ok {
 		return wr
 	}
 	return &responseWriter{
 		ResponseWriter: w,
+		start:          time.Now(),
 	}
 }
 
+func (w *responseWriter) SetWriteResponseTimeHeader(write bool) {
+	w.writeResponseTimeHeader = write
+}
+
 func (w *responseWriter) Size() int {
 	return w.size
 }
@@ -38,9 +49,12 @@ func (w *responseWriter) Status() string {
 }
 
 func (w *responseWriter) WriteHeader(statusCode int) {
-	w.wroteHeader = true
-	w.statusCode = statusCode
+	if w.writeResponseTimeHeader {
+		w.Header().Set(http2.HeaderXResponseTime, strconv.FormatInt(time.Since(w.start).Microseconds(), 10))
+	}
 	w.ResponseWriter.WriteHeader(statusCode)
+	w.statusCode = statusCode
+	w.wroteHeader = true
 }
 
 func (w *responseWriter) Write(b []byte) (int, error) {

net/http/middleware/telemetry.go

@@ -36,7 +36,7 @@ func TelemetryWithOptions(name string, opts TelemetryOptions) Middleware {
 	return func(l *zap.Logger, next http.Handler) http.Handler {
 		return otelhttp.NewHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			// wrap response write to get access to status & size
-			wr := wrapResponseWriter(w)
+			wr := WrapResponseWriter(w)
 
 			start := time.Now()
 			next.ServeHTTP(wr, r)