/
httpbrake.go
64 lines (52 loc) · 1.72 KB
/
httpbrake.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
package http
import (
"net/http"
"github.com/airbrake/gobrake/v5"
)
// A Handler is an HTTP middleware that provides integration with
// Airbrake.
type Handler struct {
Notifier *gobrake.Notifier
}
type airbrakeResponseWriter struct {
http.ResponseWriter
statusCode int
}
// New returns a new Handler. Use the Handler and HandleFunc methods to wrap
// existing HTTP handlers.
func New(notifier *gobrake.Notifier) *Handler {
h := Handler{notifier}
return &h
}
// Handle works as a middleware that wraps an existing http.Handler and sends route performance stats
func (h *Handler) Handle(handler http.Handler) http.Handler {
return h.handle(handler)
}
// HandleFunc is like Handler, but with a handler function parameter for cases
// where that is convenient. In particular, use it to wrap a handler function
// literal.
//
// http.Handle(pattern, h.HandleFunc(func (w http.ResponseWriter, r *http.Request) {
// // handler code here
// }))
func (h *Handler) HandleFunc(handler http.HandlerFunc) http.HandlerFunc {
return h.handle(handler)
}
func (h *Handler) handle(handler http.Handler) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
ctx, routeMetric := gobrake.NewRouteMetric(ctx, r.Method, r.URL.Path) // Starts the timing
arw := newAirbrakeResponseWriter(w)
handler.ServeHTTP(arw, r)
routeMetric.StatusCode = arw.statusCode
_ = h.Notifier.Routes.Notify(ctx, routeMetric)
}
}
func newAirbrakeResponseWriter(w http.ResponseWriter) *airbrakeResponseWriter {
// Returns 200 OK if WriteHeader isn't called
return &airbrakeResponseWriter{w, http.StatusOK}
}
func (arw *airbrakeResponseWriter) WriteHeader(code int) {
arw.statusCode = code
arw.ResponseWriter.WriteHeader(code)
}