forked from go-kivik/couchdb
/
trace.go
136 lines (120 loc) · 3.67 KB
/
trace.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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
package chttp
import (
"bytes"
"context"
"io"
"io/ioutil"
"net/http"
)
var clientTraceContextKey = &struct{ name string }{"client trace"}
// ContextClientTrace returns the ClientTrace associated with the
// provided context. If none, it returns nil.
func ContextClientTrace(ctx context.Context) *ClientTrace {
trace, _ := ctx.Value(clientTraceContextKey).(*ClientTrace)
return trace
}
// ClientTrace is a set of hooks to run at various stages of an outgoing
// HTTP request. Any particular hook may be nil. Functions may be
// called concurrently from different goroutines and some may be called
// after the request has completed or failed.
type ClientTrace struct {
// HTTPResponse returns a cloe of the *http.Response received from the
// server, with the body set to nil. If you need the body, use the more
// expensive HTTPResponseBody.
HTTPResponse func(*http.Response)
// HTTPResponseBody returns a clone of the *http.Response received from the
// server, with the body cloned. This can be expensive for responses
// with large bodies.
HTTPResponseBody func(*http.Response)
// HTTPRequest returns a clone of the *http.Request sent to the server, with
// the body set to nil. If you need the body, use the more expensive
// HTTPRequestBody.
HTTPRequest func(*http.Request)
// HTTPRequestBody returns a clone of the *http.Request sent to the server,
// with the body cloned, if it is set. This can be expensive for requests
// with large bodies.
HTTPRequestBody func(*http.Request)
}
// WithClientTrace returns a new context based on the provided parent
// ctx. HTTP client requests made with the returned context will use
// the provided trace hooks, in addition to any previous hooks
// registered with ctx. Any hooks defined in the provided trace will
// be called first.
func WithClientTrace(ctx context.Context, trace *ClientTrace) context.Context {
if trace == nil {
panic("nil trace")
}
return context.WithValue(ctx, clientTraceContextKey, trace)
}
func (t *ClientTrace) httpResponse(r *http.Response) {
if t.HTTPResponse == nil || r == nil {
return
}
clone := new(http.Response)
*clone = *r
clone.Body = nil
t.HTTPResponse(clone)
}
func (t *ClientTrace) httpResponseBody(r *http.Response) {
if t.HTTPResponseBody == nil || r == nil {
return
}
clone := new(http.Response)
*clone = *r
rBody := r.Body
body, readErr := ioutil.ReadAll(rBody)
closeErr := rBody.Close()
r.Body = newReplay(body, readErr, closeErr)
clone.Body = newReplay(body, readErr, closeErr)
t.HTTPResponseBody(clone)
}
func (t *ClientTrace) httpRequest(r *http.Request) {
if t.HTTPRequest == nil {
return
}
clone := new(http.Request)
*clone = *r
clone.Body = nil
t.HTTPRequest(clone)
}
func (t *ClientTrace) httpRequestBody(r *http.Request) {
if t.HTTPRequestBody == nil {
return
}
clone := new(http.Request)
*clone = *r
if r.Body != nil {
rBody := r.Body
body, readErr := ioutil.ReadAll(rBody)
closeErr := rBody.Close()
r.Body = newReplay(body, readErr, closeErr)
clone.Body = newReplay(body, readErr, closeErr)
}
t.HTTPRequestBody(clone)
}
func newReplay(body []byte, readErr, closeErr error) io.ReadCloser {
if readErr == nil && closeErr == nil {
return ioutil.NopCloser(bytes.NewReader(body))
}
return &replayReadCloser{
Reader: ioutil.NopCloser(bytes.NewReader(body)),
readErr: readErr,
closeErr: closeErr,
}
}
// replayReadCloser replays read and close errors
type replayReadCloser struct {
io.Reader
readErr error
closeErr error
}
func (r *replayReadCloser) Read(p []byte) (int, error) {
c, err := r.Reader.Read(p)
if err == io.EOF && r.readErr != nil {
err = r.readErr
}
return c, err
}
func (r *replayReadCloser) Close() error {
return r.closeErr
}