Package list golang-github-go-kit-kit / efd10d8
added client finalizer based on the server finazlizer TRAVIS ALLEN SALAS COX 4 years ago
2 changed file(s) with 87 addition(s) and 2 deletion(s). Raw diff Collapse all Expand all
2222 dec DecodeResponseFunc
2323 before []RequestFunc
2424 after []ClientResponseFunc
25 finalizer ClientFinalizerFunc
2526 bufferedStream bool
2627 }
2728
7172 return func(c *Client) { c.after = append(c.after, after...) }
7273 }
7374
75 // ClientFinalizer is executed at the end of every HTTP request.
76 // By default, no finalizer is registered.
77 func ClientFinalizer(f ClientFinalizerFunc) ClientOption {
78 return func(s *Client) { s.finalizer = f }
79 }
80
7481 // BufferedStream sets whether the Response.Body is left open, allowing it
7582 // to be read from later. Useful for transporting a file as a buffered stream.
7683 func BufferedStream(buffered bool) ClientOption {
8390 ctx, cancel := context.WithCancel(ctx)
8491 defer cancel()
8592
86 req, err := http.NewRequest(c.method, c.tgt.String(), nil)
93 // Vars used for client finalizer to ensure there are no nil values
94 var (
95 req *http.Request = &http.Request{}
96 resp *http.Response = &http.Response{}
97 err error
98 )
99 if c.finalizer != nil {
100 defer func() {
101 ctx = context.WithValue(ctx, ContextKeyResponseHeaders, resp.Header)
102 ctx = context.WithValue(ctx, ContextKeyResponseSize, resp.ContentLength)
103 c.finalizer(ctx, resp.StatusCode, req)
104 }()
105 }
106
107 req, err = http.NewRequest(c.method, c.tgt.String(), nil)
87108 if err != nil {
88109 return nil, err
89110 }
96117 ctx = f(ctx, req)
97118 }
98119
99 resp, err := ctxhttp.Do(ctx, c.client, req)
120 resp, err = ctxhttp.Do(ctx, c.client, req)
100121 if err != nil {
101122 return nil, err
102123 }
124
103125 if !c.bufferedStream {
104126 defer resp.Body.Close()
105127 }
116138 return response, nil
117139 }
118140 }
141
142 // ClientFinalizerFunc can be used to perform work at the end of a client HTTP
143 // request, after the response is returned. The principal
144 // intended use is for request logging. In addition to the response code
145 // provided in the function signature, additional response parameters are
146 // provided in the context under keys with the ContextKeyResponse prefix.
147 type ClientFinalizerFunc func(ctx context.Context, code int, r *http.Request)
119148
120149 // EncodeJSONRequest is an EncodeRequestFunc that serializes the request as a
121150 // JSON object to the Request body. Many JSON-over-HTTP services can use it as
136136 }
137137 if want, have := testbody, string(b); want != have {
138138 t.Errorf("want %q, have %q", want, have)
139 }
140 }
141
142 func TestClientFinalizer(t *testing.T) {
143 var (
144 headerKey = "X-Henlo-Lizer"
145 headerVal = "Helllo you stinky lizard"
146 statusCode = http.StatusTeapot
147 responseBody = "go eat a fly ugly\n"
148 done = make(chan struct{})
149 encode = func(context.Context, *http.Request, interface{}) error { return nil }
150 decode = func(_ context.Context, r *http.Response) (interface{}, error) {
151 return TestResponse{r.Body, ""}, nil
152 }
153 )
154
155 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
156 w.Header().Set(headerKey, headerVal)
157 w.WriteHeader(statusCode)
158 w.Write([]byte(responseBody))
159 }))
160 defer server.Close()
161
162 client := httptransport.NewClient(
163 "GET",
164 mustParse(server.URL),
165 encode,
166 decode,
167 httptransport.ClientFinalizer(func(ctx context.Context, code int, _ *http.Request) {
168 if want, have := statusCode, code; want != have {
169 t.Errorf("StatusCode: want %d, have %d", want, have)
170 }
171
172 responseHeader := ctx.Value(httptransport.ContextKeyResponseHeaders).(http.Header)
173 if want, have := headerVal, responseHeader.Get(headerKey); want != have {
174 t.Errorf("%s: want %q, have %q", headerKey, want, have)
175 }
176
177 responseSize := ctx.Value(httptransport.ContextKeyResponseSize).(int64)
178 if want, have := int64(len(responseBody)), responseSize; want != have {
179 t.Errorf("response size: want %d, have %d", want, have)
180 }
181
182 close(done)
183 }),
184 )
185
186 _, err := client.Endpoint()(context.Background(), struct{}{})
187 if err != nil {
188 t.Fatal(err)
189 }
190
191 select {
192 case <-done:
193 case <-time.After(time.Second):
194 t.Fatal("timeout waiting for finalizer")
139195 }
140196 }
141197