diff --git a/prometheus/http.go b/prometheus/http.go index e391af4..eabe602 100644 --- a/prometheus/http.go +++ b/prometheus/http.go @@ -14,6 +14,9 @@ package prometheus import ( + "bufio" + "io" + "net" "net/http" "strconv" "strings" @@ -141,7 +144,18 @@ func InstrumentHandlerFuncWithOpts(opts SummaryOpts, handlerFunc func(http.Respo urlLen = len(r.URL.String()) } go computeApproximateRequestSize(r, out, urlLen) - handlerFunc(delegate, r) + + _, cn := w.(http.CloseNotifier) + _, fl := w.(http.Flusher) + _, hj := w.(http.Hijacker) + _, rf := w.(io.ReaderFrom) + var rw http.ResponseWriter + if cn && fl && hj && rf { + rw = &fancyResponseWriterDelegator{delegate} + } else { + rw = delegate + } + handlerFunc(rw, r) elapsed := float64(time.Since(now)) / float64(time.Microsecond) @@ -197,6 +211,31 @@ func (r *responseWriterDelegator) Write(b []byte) (int, error) { return n, err } +type fancyResponseWriterDelegator struct { + *responseWriterDelegator +} + +func (f *fancyResponseWriterDelegator) CloseNotify() <-chan bool { + return f.ResponseWriter.(http.CloseNotifier).CloseNotify() +} + +func (f *fancyResponseWriterDelegator) Flush() { + f.ResponseWriter.(http.Flusher).Flush() +} + +func (f *fancyResponseWriterDelegator) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return f.ResponseWriter.(http.Hijacker).Hijack() +} + +func (f *fancyResponseWriterDelegator) ReadFrom(r io.Reader) (int64, error) { + if !f.wroteHeader { + f.WriteHeader(http.StatusOK) + } + n, err := f.ResponseWriter.(io.ReaderFrom).ReadFrom(r) + f.written += n + return n, err +} + func sanitizeMethod(m string) string { switch m { case "GET", "get":