-
Notifications
You must be signed in to change notification settings - Fork 32
/
handler.go
122 lines (103 loc) · 2.33 KB
/
handler.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
package httpstats
import (
"bufio"
"net"
"net/http"
"time"
"github.com/segmentio/stats"
)
// NewHandler wraps h to produce metrics on the default engine for every request
// received and every response sent.
func NewHandler(h http.Handler) http.Handler {
return NewHandlerWith(stats.DefaultEngine, h)
}
// NewHandlerWith wraps h to produce metrics on eng for every request received
// and every response sent.
func NewHandlerWith(eng *stats.Engine, h http.Handler) http.Handler {
return &handler{
handler: h,
eng: eng,
}
}
type handler struct {
handler http.Handler
eng *stats.Engine
}
func (h *handler) ServeHTTP(res http.ResponseWriter, req *http.Request) {
m := &metrics{}
w := &responseWriter{
ResponseWriter: res,
eng: h.eng,
req: req,
metrics: m,
start: time.Now(),
}
defer w.complete()
b := &requestBody{
body: req.Body,
eng: h.eng,
req: req,
metrics: m,
op: "read",
}
defer b.close()
req.Body = b
h.handler.ServeHTTP(w, req)
}
type responseWriter struct {
http.ResponseWriter
start time.Time
eng *stats.Engine
req *http.Request
metrics *metrics
status int
bytes int
wroteHeader bool
wroteStats bool
}
func (w *responseWriter) WriteHeader(status int) {
if !w.wroteHeader {
w.wroteHeader = true
w.status = status
w.ResponseWriter.WriteHeader(status)
}
}
func (w *responseWriter) Write(b []byte) (n int, err error) {
if !w.wroteHeader {
w.wroteHeader = true
w.status = http.StatusOK
}
if n, err = w.ResponseWriter.Write(b); n > 0 {
w.bytes += n
}
return
}
func (w *responseWriter) Hijack() (conn net.Conn, buf *bufio.ReadWriter, err error) {
if conn, buf, err = w.ResponseWriter.(http.Hijacker).Hijack(); err == nil {
w.wroteHeader = true
w.complete()
}
return
}
func (w *responseWriter) complete() {
if w.wroteStats {
return
}
w.wroteStats = true
if !w.wroteHeader {
w.wroteHeader = true
w.status = http.StatusOK
}
now := time.Now()
res := &http.Response{
ProtoMajor: w.req.ProtoMajor,
ProtoMinor: w.req.ProtoMinor,
Proto: w.req.Proto,
StatusCode: w.status,
Header: w.Header(),
Request: w.req,
ContentLength: -1,
}
w.metrics.observeResponse(res, "write", w.bytes, now.Sub(w.start))
w.eng.ReportAt(w.start, w.metrics)
}