Merge pull request #533 from travissalascox/master
added client finalizer
Peter Bourgon authored 6 years ago
GitHub committed 6 years ago
| 22 | 22 | dec DecodeResponseFunc |
| 23 | 23 | before []RequestFunc |
| 24 | 24 | after []ClientResponseFunc |
| 25 | finalizer ClientFinalizerFunc | |
| 25 | 26 | bufferedStream bool |
| 26 | 27 | } |
| 27 | 28 | |
| 71 | 72 | return func(c *Client) { c.after = append(c.after, after...) } |
| 72 | 73 | } |
| 73 | 74 | |
| 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 | ||
| 74 | 81 | // BufferedStream sets whether the Response.Body is left open, allowing it |
| 75 | 82 | // to be read from later. Useful for transporting a file as a buffered stream. |
| 76 | 83 | func BufferedStream(buffered bool) ClientOption { |
| 82 | 89 | return func(ctx context.Context, request interface{}) (interface{}, error) { |
| 83 | 90 | ctx, cancel := context.WithCancel(ctx) |
| 84 | 91 | defer cancel() |
| 92 | ||
| 93 | var ( | |
| 94 | resp *http.Response | |
| 95 | err error | |
| 96 | ) | |
| 97 | if c.finalizer != nil { | |
| 98 | defer func() { | |
| 99 | if resp != nil { | |
| 100 | ctx = context.WithValue(ctx, ContextKeyResponseHeaders, resp.Header) | |
| 101 | ctx = context.WithValue(ctx, ContextKeyResponseSize, resp.ContentLength) | |
| 102 | } | |
| 103 | c.finalizer(ctx, err) | |
| 104 | }() | |
| 105 | } | |
| 85 | 106 | |
| 86 | 107 | req, err := http.NewRequest(c.method, c.tgt.String(), nil) |
| 87 | 108 | if err != nil { |
| 96 | 117 | ctx = f(ctx, req) |
| 97 | 118 | } |
| 98 | 119 | |
| 99 | resp, err := ctxhttp.Do(ctx, c.client, req) | |
| 120 | resp, err = ctxhttp.Do(ctx, c.client, req) | |
| 100 | 121 | if err != nil { |
| 101 | 122 | return nil, err |
| 102 | 123 | } |
| 124 | ||
| 103 | 125 | if !c.bufferedStream { |
| 104 | 126 | defer resp.Body.Close() |
| 105 | 127 | } |
| 116 | 138 | return response, nil |
| 117 | 139 | } |
| 118 | 140 | } |
| 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 error logging. Additional response parameters are | |
| 145 | // provided in the context under keys with the ContextKeyResponse prefix. | |
| 146 | // Note: err may be nil. There maybe also no additional response parameters depending on | |
| 147 | // when an error occurs. | |
| 148 | type ClientFinalizerFunc func(ctx context.Context, err error) | |
| 119 | 149 | |
| 120 | 150 | // EncodeJSONRequest is an EncodeRequestFunc that serializes the request as a |
| 121 | 151 | // JSON object to the Request body. Many JSON-over-HTTP services can use it as |
| 136 | 136 | } |
| 137 | 137 | if want, have := testbody, string(b); want != have { |
| 138 | 138 | 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 | responseBody = "go eat a fly ugly\n" | |
| 147 | done = make(chan struct{}) | |
| 148 | encode = func(context.Context, *http.Request, interface{}) error { return nil } | |
| 149 | decode = func(_ context.Context, r *http.Response) (interface{}, error) { | |
| 150 | return TestResponse{r.Body, ""}, nil | |
| 151 | } | |
| 152 | ) | |
| 153 | ||
| 154 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { | |
| 155 | w.Header().Set(headerKey, headerVal) | |
| 156 | w.Write([]byte(responseBody)) | |
| 157 | })) | |
| 158 | defer server.Close() | |
| 159 | ||
| 160 | client := httptransport.NewClient( | |
| 161 | "GET", | |
| 162 | mustParse(server.URL), | |
| 163 | encode, | |
| 164 | decode, | |
| 165 | httptransport.ClientFinalizer(func(ctx context.Context, err error) { | |
| 166 | responseHeader := ctx.Value(httptransport.ContextKeyResponseHeaders).(http.Header) | |
| 167 | if want, have := headerVal, responseHeader.Get(headerKey); want != have { | |
| 168 | t.Errorf("%s: want %q, have %q", headerKey, want, have) | |
| 169 | } | |
| 170 | ||
| 171 | responseSize := ctx.Value(httptransport.ContextKeyResponseSize).(int64) | |
| 172 | if want, have := int64(len(responseBody)), responseSize; want != have { | |
| 173 | t.Errorf("response size: want %d, have %d", want, have) | |
| 174 | } | |
| 175 | ||
| 176 | close(done) | |
| 177 | }), | |
| 178 | ) | |
| 179 | ||
| 180 | _, err := client.Endpoint()(context.Background(), struct{}{}) | |
| 181 | if err != nil { | |
| 182 | t.Fatal(err) | |
| 183 | } | |
| 184 | ||
| 185 | select { | |
| 186 | case <-done: | |
| 187 | case <-time.After(time.Second): | |
| 188 | t.Fatal("timeout waiting for finalizer") | |
| 139 | 189 | } |
| 140 | 190 | } |
| 141 | 191 |