refactor: do not block reading messages from docker daemon

If body responses from the docker daemon were not read fast enough,
docker would block image pulls.

This now reads the messages from the docker response body and no longer blocks.

As a bonus, context is now respected so ctrl+c will stop an image pull.

part of DEP-1801

Signed-off-by: Chris Goller <[email protected]>

pkg/load/pull.go

@@ -58,7 +58,13 @@ func ImagePullPrivileged(ctx context.Context, dockerapi docker.APIClient, imageN
 			return err
 		}
 	} else {
-		if err := printPull(ctx, responseBody, logger); err != nil {
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel() // Used to stop the go routine if printPull returns error early.
+
+		msgCh := make(chan Message, 4096)
+		go decode(ctx, responseBody, msgCh)
+		err := printPull(ctx, msgCh, logger)
+		if err != nil {
 			return err
 		}
 	}
@@ -156,7 +162,52 @@ type PullProgress struct {
 	Vtx    *client.VertexStatus
 }
 
-func printPull(_ context.Context, rc io.Reader, l progress.SubLogger) error {
+type Message struct {
+	msg *jsonmessage.JSONMessage
+	err error
+}
+
+// decode reads the body of the response from Docker and decodes it into JSON messages as fast
+// as it can.  It does not block on the channel and prefers to drop messages if the channel is full
+// to prevent Docker from blocking on the pull.
+func decode(ctx context.Context, r io.Reader, msgCh chan<- Message) {
+	defer close(msgCh)
+
+	dec := json.NewDecoder(r)
+	for {
+		select {
+		case <-ctx.Done():
+			select {
+			case msgCh <- Message{err: ctx.Err()}:
+			default:
+			}
+			return
+		default:
+		}
+
+		var msg jsonmessage.JSONMessage
+		if err := dec.Decode(&msg); err != nil {
+			if err == io.EOF {
+				return
+			}
+
+			select {
+			case msgCh <- Message{err: err}:
+			default:
+			}
+		}
+
+		// If we block here it is possible for Docker to block on the pull.
+		select {
+		case msgCh <- Message{msg: &msg}:
+		default:
+		}
+	}
+}
+
+// printPull will convert the messages to useful on screen content.
+// we want to read as fast as possible as docker will block if the body buffer becomes too full.
+func printPull(ctx context.Context, msgCh <-chan Message, l progress.SubLogger) error {
 	started := map[string]PullProgress{}
 
 	defer func() {
@@ -170,26 +221,29 @@ func printPull(_ context.Context, rc io.Reader, l progress.SubLogger) error {
 		}
 	}()
 
-	dec := json.NewDecoder(rc)
-
 	var (
-		parsedError error
-		jm          jsonmessage.JSONMessage
+		msg Message
+		ok  bool
 	)
 
 	for {
-		if err := dec.Decode(&jm); err != nil {
-			if parsedError != nil {
-				return parsedError
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case msg, ok = <-msgCh:
+			if !ok {
+				return nil
 			}
-			if err == io.EOF {
-				break
-			}
-			return err
 		}
 
+		if msg.err != nil {
+			return msg.err
+		}
+
+		jm := msg.msg
+
 		if jm.Error != nil {
-			parsedError = jm.Error
+			return jm.Error
 		}
 
 		if jm.ID == "" {
@@ -270,5 +324,4 @@ func printPull(_ context.Context, rc io.Reader, l progress.SubLogger) error {
 
 		l.SetStatus(st.Vtx)
 	}
-	return nil
 }