container: make Stats more reliable

Specially on error cases. Should fix the intermittent tests on Travis.
yyuuttaaoo · Jul 26, 2017 · 85a022e · 85a022e
1 parent 7228e3d
commit 85a022e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/client.go b/client.go
@@ -498,6 +498,7 @@ type streamOptions struct {
 	in             io.Reader
 	stdout         io.Writer
 	stderr         io.Writer
+	reqSent        chan struct{}
 	// timeout is the initial connection timeout
 	timeout time.Duration
 	// Timeout with no data is received, it's reset every time new data
@@ -576,6 +577,9 @@ func (c *Client) stream(method, path string, streamOptions streamOptions) error
 			dial.SetDeadline(time.Now().Add(streamOptions.timeout))
 		}
 
+		if streamOptions.reqSent != nil {
+			close(streamOptions.reqSent)
+		}
 		if resp, err = http.ReadResponse(breader, req); err != nil {
 			// Cancel timeout for future I/O operations
 			if streamOptions.timeout > 0 {
@@ -594,6 +598,9 @@ func (c *Client) stream(method, path string, streamOptions streamOptions) error
 			}
 			return chooseError(subCtx, err)
 		}
+		if streamOptions.reqSent != nil {
+			close(streamOptions.reqSent)
+		}
 	}
 	defer resp.Body.Close()
 	if resp.StatusCode < 200 || resp.StatusCode >= 400 {

diff --git a/container.go b/container.go
@@ -1056,6 +1056,7 @@ func (c *Client) Stats(opts StatsOptions) (retErr error) {
 		}
 	}()
 
+	reqSent := make(chan struct{})
 	go func() {
 		err := c.stream("GET", fmt.Sprintf("/containers/%s/stats?stream=%v", opts.ID, opts.Stream), streamOptions{
 			rawJSONStream:     true,
@@ -1064,6 +1065,7 @@ func (c *Client) Stats(opts StatsOptions) (retErr error) {
 			timeout:           opts.Timeout,
 			inactivityTimeout: opts.InactivityTimeout,
 			context:           opts.Context,
+			reqSent:           reqSent,
 		})
 		if err != nil {
 			dockerError, ok := err.(*Error)
@@ -1094,6 +1096,7 @@ func (c *Client) Stats(opts StatsOptions) (retErr error) {
 
 	decoder := json.NewDecoder(readCloser)
 	stats := new(Stats)
+	<-reqSent
 	for err := decoder.Decode(stats); err != io.EOF; err = decoder.Decode(stats) {
 		if err != nil {
 			return err