Golang net/http X-Response-Time header middleware

main.go

package main

import (
	"fmt"
	"log"
	"net/http"
	"time"

	"httptimer"
)

func main() {
	h := http.NewServeMux()
	h.HandleFunc("/", index)

	log.Fatal(http.ListenAndServe("127.0.0.1:18080", httptimer.Timed(h)))
}

func index(w http.ResponseWriter, r *http.Request) {
	wait, err := time.ParseDuration(r.URL.Query().Get("wait"))
	if err != nil {
		wait = 2333 * time.Nanosecond
	}
	time.Sleep(wait)

	w.Header().Set("Content-Type", "text/plain")
	fmt.Fprintf(w, "%v: %v\n", r.Method, r.URL.Path)
	fmt.Fprintln(w, "Check X-Response-Time for the processing time of this request.")
}

timer.go

package httptimer

import (
	"net/http"
	"strconv"
	"time"
)

type responseWriterWithTimer struct {
	http.ResponseWriter

	isHeaderWritten bool
	start           time.Time
}

func Timed(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		h.ServeHTTP(&responseWriterWithTimer{w, false, time.Now()}, r)
	})
}

func (w *responseWriterWithTimer) WriteHeader(statusCode int) {
	duration := time.Now().Sub(w.start)
	us := int(duration.Truncate(1000*time.Nanosecond).Nanoseconds() / 1000)
	w.Header().Set("X-Response-Time", strconv.Itoa(us)+" us")

	w.ResponseWriter.WriteHeader(statusCode)
	w.isHeaderWritten = true
}

func (w *responseWriterWithTimer) Write(b []byte) (int, error) {
	if !w.isHeaderWritten {
		w.WriteHeader(200)
	}

	return w.ResponseWriter.Write(b)
}