From 15ac4f08f83b2a9e729b3dec49b6e6267773c753 Mon Sep 17 00:00:00 2001
From: Xiang Li <xiangli.cs@gmail.com>
Date: Fri, 5 Jun 2015 14:11:54 -0700
Subject: [PATCH] client: fix cancel watch

ioutil.ReadAll is a blocking call, we need to wait cancelation
during the call.
---
 client/client.go      | 18 +++++++++++++++++-
 client/client_test.go | 44 +++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/client/client.go b/client/client.go
index c2a4010c9..bc314dde2 100644
--- a/client/client.go
+++ b/client/client.go
@@ -293,7 +293,23 @@ func (c *simpleHTTPClient) Do(ctx context.Context, act httpAction) (*http.Respon
 		return nil, nil, err
 	}
 
-	body, err := ioutil.ReadAll(resp.Body)
+	var body []byte
+	done := make(chan struct{})
+	go func() {
+		body, err = ioutil.ReadAll(resp.Body)
+		done <- struct{}{}
+	}()
+
+	select {
+	case <-ctx.Done():
+		err = resp.Body.Close()
+		<-done
+		if err == nil {
+			err = ctx.Err()
+		}
+	case <-done:
+	}
+
 	return resp, body, err
 }
 
diff --git a/client/client_test.go b/client/client_test.go
index bf2e9b1bb..97383b60f 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -186,8 +186,11 @@ type checkableReadCloser struct {
 }
 
 func (c *checkableReadCloser) Close() error {
-	c.closed = true
-	return c.ReadCloser.Close()
+	if !c.closed {
+		c.closed = true
+		return c.ReadCloser.Close()
+	}
+	return nil
 }
 
 func TestSimpleHTTPClientDoCancelContextResponseBodyClosed(t *testing.T) {
@@ -218,6 +221,43 @@ func TestSimpleHTTPClientDoCancelContextResponseBodyClosed(t *testing.T) {
 	}
 }
 
+type blockingBody struct {
+	c chan struct{}
+}
+
+func (bb *blockingBody) Read(p []byte) (n int, err error) {
+	<-bb.c
+	return 0, errors.New("closed")
+}
+
+func (bb *blockingBody) Close() error {
+	close(bb.c)
+	return nil
+}
+
+func TestSimpleHTTPClientDoCancelContextResponseBodyClosedWithBlockingBody(t *testing.T) {
+	tr := newFakeTransport()
+	c := &simpleHTTPClient{transport: tr}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	body := &checkableReadCloser{ReadCloser: &blockingBody{c: make(chan struct{})}}
+	go func() {
+		tr.respchan <- &http.Response{Body: body}
+		time.Sleep(2 * time.Millisecond)
+		// cancel after the body is received
+		cancel()
+	}()
+
+	_, _, err := c.Do(ctx, &fakeAction{})
+	if err == nil {
+		t.Fatalf("expected non-nil error, got nil")
+	}
+
+	if !body.closed {
+		t.Fatalf("expected closed body")
+	}
+}
+
 func TestSimpleHTTPClientDoCancelContextWaitForRoundTrip(t *testing.T) {
 	tr := newFakeTransport()
 	c := &simpleHTTPClient{transport: tr}