added client finalizer

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

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