1package httpsnoop
2
3import (
4	"io"
5	"net/http"
6	"time"
7)
8
9// Metrics holds metrics captured from CaptureMetrics.
10type Metrics struct {
11	// Code is the first http response code passed to the WriteHeader func of
12	// the ResponseWriter. If no such call is made, a default code of 200 is
13	// assumed instead.
14	Code int
15	// Duration is the time it took to execute the handler.
16	Duration time.Duration
17	// Written is the number of bytes successfully written by the Write or
18	// ReadFrom function of the ResponseWriter. ResponseWriters may also write
19	// data to their underlaying connection directly (e.g. headers), but those
20	// are not tracked. Therefor the number of Written bytes will usually match
21	// the size of the response body.
22	Written int64
23}
24
25// CaptureMetrics wraps the given hnd, executes it with the given w and r, and
26// returns the metrics it captured from it.
27func CaptureMetrics(hnd http.Handler, w http.ResponseWriter, r *http.Request) Metrics {
28	return CaptureMetricsFn(w, func(ww http.ResponseWriter) {
29		hnd.ServeHTTP(ww, r)
30	})
31}
32
33// CaptureMetricsFn wraps w and calls fn with the wrapped w and returns the
34// resulting metrics. This is very similar to CaptureMetrics (which is just
35// sugar on top of this func), but is a more usable interface if your
36// application doesn't use the Go http.Handler interface.
37func CaptureMetricsFn(w http.ResponseWriter, fn func(http.ResponseWriter)) Metrics {
38	var (
39		start         = time.Now()
40		m             = Metrics{Code: http.StatusOK}
41		headerWritten bool
42		hooks         = Hooks{
43			WriteHeader: func(next WriteHeaderFunc) WriteHeaderFunc {
44				return func(code int) {
45					next(code)
46
47					if !headerWritten {
48						m.Code = code
49						headerWritten = true
50					}
51				}
52			},
53
54			Write: func(next WriteFunc) WriteFunc {
55				return func(p []byte) (int, error) {
56					n, err := next(p)
57
58					m.Written += int64(n)
59					headerWritten = true
60					return n, err
61				}
62			},
63
64			ReadFrom: func(next ReadFromFunc) ReadFromFunc {
65				return func(src io.Reader) (int64, error) {
66					n, err := next(src)
67
68					headerWritten = true
69					m.Written += n
70					return n, err
71				}
72			},
73		}
74	)
75
76	fn(Wrap(w, hooks))
77	m.Duration = time.Since(start)
78	return m
79}
80