diff --git a/transport/http/client.go b/transport/http/client.go index 08f1b886e..c5dace97b 100644 --- a/transport/http/client.go +++ b/transport/http/client.go @@ -23,6 +23,7 @@ type Client struct { dec DecodeResponseFunc before []RequestFunc after []ClientResponseFunc + finalizer ClientFinalizerFunc bufferedStream bool } @@ -72,6 +73,12 @@ func ClientAfter(after ...ClientResponseFunc) ClientOption { return func(c *Client) { c.after = append(c.after, after...) } } +// ClientFinalizer is executed at the end of every HTTP request. +// By default, no finalizer is registered. +func ClientFinalizer(f ClientFinalizerFunc) ClientOption { + return func(s *Client) { s.finalizer = f } +} + // BufferedStream sets whether the Response.Body is left open, allowing it // to be read from later. Useful for transporting a file as a buffered stream. func BufferedStream(buffered bool) ClientOption { @@ -84,6 +91,20 @@ func (c Client) Endpoint() endpoint.Endpoint { ctx, cancel := context.WithCancel(ctx) defer cancel() + var ( + resp *http.Response + err error + ) + if c.finalizer != nil { + defer func() { + if resp != nil { + ctx = context.WithValue(ctx, ContextKeyResponseHeaders, resp.Header) + ctx = context.WithValue(ctx, ContextKeyResponseSize, resp.ContentLength) + } + c.finalizer(ctx, err) + }() + } + req, err := http.NewRequest(c.method, c.tgt.String(), nil) if err != nil { return nil, err @@ -97,10 +118,11 @@ func (c Client) Endpoint() endpoint.Endpoint { ctx = f(ctx, req) } - resp, err := ctxhttp.Do(ctx, c.client, req) + resp, err = ctxhttp.Do(ctx, c.client, req) if err != nil { return nil, err } + if !c.bufferedStream { defer resp.Body.Close() } @@ -118,6 +140,14 @@ func (c Client) Endpoint() endpoint.Endpoint { } } +// ClientFinalizerFunc can be used to perform work at the end of a client HTTP +// request, after the response is returned. The principal +// intended use is for error logging. Additional response parameters are +// provided in the context under keys with the ContextKeyResponse prefix. +// Note: err may be nil. There maybe also no additional response parameters depending on +// when an error occurs. +type ClientFinalizerFunc func(ctx context.Context, err error) + // EncodeJSONRequest is an EncodeRequestFunc that serializes the request as a // JSON object to the Request body. Many JSON-over-HTTP services can use it as // a sensible default. If the request implements Headerer, the provided headers diff --git a/transport/http/client_test.go b/transport/http/client_test.go index ad366d1ac..d66381000 100644 --- a/transport/http/client_test.go +++ b/transport/http/client_test.go @@ -140,6 +140,56 @@ func TestHTTPClientBufferedStream(t *testing.T) { } } +func TestClientFinalizer(t *testing.T) { + var ( + headerKey = "X-Henlo-Lizer" + headerVal = "Helllo you stinky lizard" + responseBody = "go eat a fly ugly\n" + done = make(chan struct{}) + encode = func(context.Context, *http.Request, interface{}) error { return nil } + decode = func(_ context.Context, r *http.Response) (interface{}, error) { + return TestResponse{r.Body, ""}, nil + } + ) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set(headerKey, headerVal) + w.Write([]byte(responseBody)) + })) + defer server.Close() + + client := httptransport.NewClient( + "GET", + mustParse(server.URL), + encode, + decode, + httptransport.ClientFinalizer(func(ctx context.Context, err error) { + responseHeader := ctx.Value(httptransport.ContextKeyResponseHeaders).(http.Header) + if want, have := headerVal, responseHeader.Get(headerKey); want != have { + t.Errorf("%s: want %q, have %q", headerKey, want, have) + } + + responseSize := ctx.Value(httptransport.ContextKeyResponseSize).(int64) + if want, have := int64(len(responseBody)), responseSize; want != have { + t.Errorf("response size: want %d, have %d", want, have) + } + + close(done) + }), + ) + + _, err := client.Endpoint()(context.Background(), struct{}{}) + if err != nil { + t.Fatal(err) + } + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("timeout waiting for finalizer") + } +} + func TestEncodeJSONRequest(t *testing.T) { var header http.Header var body string