Просмотр исходного кода

client: ensure Response closed on cancel

Brian Waldon 11 лет назад
Родитель
Сommit
25cf916a80
2 измененных файлов с 42 добавлено и 2 удалено
  1. 3 2
      client/client.go
  2. 39 0
      client/client_test.go

+ 3 - 2
client/client.go

@@ -270,9 +270,10 @@ func (c *simpleHTTPClient) Do(ctx context.Context, act httpAction) (*http.Respon
 	case rtresp := <-rtchan:
 		resp, err = rtresp.resp, rtresp.err
 	case <-ctx.Done():
+		// cancel and wait for request to actually exit before continuing
 		c.transport.CancelRequest(req)
-		// wait for request to actually exit before continuing
-		<-rtchan
+		rtresp := <-rtchan
+		resp = rtresp.resp
 		err = ctx.Err()
 	}
 

+ 39 - 0
client/client_test.go

@@ -16,6 +16,7 @@ package client
 
 import (
 	"errors"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"net/url"
@@ -179,6 +180,44 @@ func TestSimpleHTTPClientDoCancelContext(t *testing.T) {
 	}
 }
 
+type checkableReadCloser struct {
+	io.ReadCloser
+	closed bool
+}
+
+func (c *checkableReadCloser) Close() error {
+	c.closed = true
+	return c.ReadCloser.Close()
+}
+
+func TestSimpleHTTPClientDoCancelContextResponseBodyClosed(t *testing.T) {
+	tr := newFakeTransport()
+	c := &simpleHTTPClient{transport: tr}
+
+	// create an already-cancelled context
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	body := &checkableReadCloser{ReadCloser: ioutil.NopCloser(strings.NewReader("foo"))}
+	go func() {
+		// wait for CancelRequest to be called, informing us that simpleHTTPClient
+		// knows the context is already timed out
+		<-tr.startCancel
+
+		tr.respchan <- &http.Response{Body: body}
+		tr.finishCancel <- struct{}{}
+	}()
+
+	_, _, 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}