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	m := Metrics{Code: http.StatusOK}
39	m.CaptureMetrics(w, fn)
40	return m
41}
42
43// CaptureMetrics wraps w and calls fn with the wrapped w and updates
44// Metrics m with the resulting metrics. This is similar to CaptureMetricsFn,
45// but allows one to customize starting Metrics object.
46func (m *Metrics) CaptureMetrics(w http.ResponseWriter, fn func(http.ResponseWriter)) {
47	var (
48		start         = time.Now()
49		headerWritten bool
50		hooks         = Hooks{
51			WriteHeader: func(next WriteHeaderFunc) WriteHeaderFunc {
52				return func(code int) {
53					next(code)
54
55					if !(code >= 100 && code <= 199) && !headerWritten {
56						m.Code = code
57						headerWritten = true
58					}
59				}
60			},
61
62			Write: func(next WriteFunc) WriteFunc {
63				return func(p []byte) (int, error) {
64					n, err := next(p)
65
66					m.Written += int64(n)
67					headerWritten = true
68					return n, err
69				}
70			},
71
72			ReadFrom: func(next ReadFromFunc) ReadFromFunc {
73				return func(src io.Reader) (int64, error) {
74					n, err := next(src)
75
76					headerWritten = true
77					m.Written += n
78					return n, err
79				}
80			},
81		}
82	)
83
84	fn(Wrap(w, hooks))
85	m.Duration += time.Since(start)
86}