From ff0c79e08d6688cfa7a14af121281e19eb0ab042 Mon Sep 17 00:00:00 2001
From: Nishant Totla <nishanttotla@gmail.com>
Date: Fri, 1 Apr 2016 15:27:18 -0700
Subject: [PATCH] Fix image pull bug (wait until download finishes)

Signed-off-by: Nishant Totla <nishanttotla@gmail.com>
---
 cluster/engine.go        | 47 ++++++++++++++++++++++++++--------------
 cluster/engine_test.go   |  4 ++--
 cluster/swarm/cluster.go |  4 +++-
 3 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/cluster/engine.go b/cluster/engine.go
index 927f7aa7c1..02619ebb32 100644
--- a/cluster/engine.go
+++ b/cluster/engine.go
@@ -2,6 +2,7 @@ package cluster
 
 import (
 	"crypto/tls"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -952,40 +953,54 @@ func (e *Engine) CreateVolume(request *types.VolumeCreateRequest) (*Volume, erro
 }
 
 // FIXME: This will become unnecessary after docker/engine-api#162 is merged
-func getTagFromNamedRef(dref reference.Named) string {
-	var tag string
-	switch x := dref.(type) {
-	case reference.Digested:
+func buildImagePullOptions(image string) (types.ImagePullOptions, error) {
+	distributionRef, err := reference.ParseNamed(image)
+	if err != nil {
+		return types.ImagePullOptions{}, err
+	}
+
+	name := distributionRef.Name()
+	tag := "latest"
+
+	switch x := distributionRef.(type) {
+	case reference.Canonical:
 		tag = x.Digest().String()
 	case reference.NamedTagged:
 		tag = x.Tag()
 	}
-	return tag
+
+	return types.ImagePullOptions{
+		ImageID: name,
+		Tag:     tag,
+	}, nil
 }
 
 // Pull an image on the engine
 func (e *Engine) Pull(image string, authConfig *types.AuthConfig) error {
-	distributionRef, err := reference.ParseNamed(image)
+	pullOpts, err := buildImagePullOptions(image)
 	if err != nil {
 		return err
 	}
-
-	repository := distributionRef.Name()
-	tag := getTagFromNamedRef(distributionRef)
-
-	pullOpts := types.ImagePullOptions{
-		ImageID: repository,
-		Tag:     tag,
-	}
-	_, err = e.apiClient.ImagePull(context.TODO(), pullOpts, nil)
+	pullResponse, err := e.apiClient.ImagePull(context.TODO(), pullOpts, nil)
 	e.CheckConnectionErr(err)
 	if err != nil {
 		return err
 	}
 
+	// wait until the image download is finished
+	dec := json.NewDecoder(pullResponse)
+	for {
+		m := map[string]interface{}{}
+		if err := dec.Decode(&m); err != nil {
+			if err == io.EOF {
+				break
+			}
+			return err
+		}
+	}
+
 	// force refresh images
 	e.RefreshImages()
-
 	return nil
 }
 
diff --git a/cluster/engine_test.go b/cluster/engine_test.go
index 1d6d74f7d1..d6eb803058 100644
--- a/cluster/engine_test.go
+++ b/cluster/engine_test.go
@@ -293,7 +293,7 @@ func TestCreateContainer(t *testing.T) {
 		engine     = NewEngine("test", 0, engOpts)
 		client     = mockclient.NewMockClient()
 		apiClient  = engineapimock.NewMockClient()
-		readCloser = nopCloser{bytes.NewBufferString("ok")}
+		readCloser = nopCloser{bytes.NewBufferString("")}
 	)
 
 	engine.setState(stateUnhealthy)
@@ -346,7 +346,7 @@ func TestCreateContainer(t *testing.T) {
 	id = "id3"
 	apiClient.On("ImageList", mock.Anything, mock.AnythingOfType("ImageListOptions")).Return([]types.Image{}, nil).Once()
 	mockConfig.CPUShares = int64(math.Ceil(float64(config.CPUShares*1024) / float64(mockInfo.NCPU)))
-	apiClient.On("ImagePull", mock.Anything, types.ImagePullOptions{ImageID: config.Image}, mock.Anything).Return(readCloser, nil).Once()
+	apiClient.On("ImagePull", mock.Anything, types.ImagePullOptions{ImageID: config.Image, Tag: "latest"}, mock.Anything).Return(readCloser, nil).Once()
 	// FIXMEENGINEAPI : below should return an engine-api error, or something custom
 	apiClient.On("ContainerCreate", mock.Anything, &mockConfig.Config, &mockConfig.HostConfig, &mockConfig.NetworkingConfig, name).Return(types.ContainerCreateResponse{}, dockerclient.ErrImageNotFound).Once()
 	// FIXMEENGINEAPI : below should return an engine-api error, or something custom
diff --git a/cluster/swarm/cluster.go b/cluster/swarm/cluster.go
index 6b40f01f8c..8bd1e9b61e 100644
--- a/cluster/swarm/cluster.go
+++ b/cluster/swarm/cluster.go
@@ -14,6 +14,7 @@ import (
 	log "github.com/Sirupsen/logrus"
 	"github.com/docker/docker/pkg/discovery"
 	"github.com/docker/docker/pkg/stringid"
+	"github.com/docker/engine-api/client"
 	"github.com/docker/engine-api/types"
 	containertypes "github.com/docker/engine-api/types/container"
 	networktypes "github.com/docker/engine-api/types/network"
@@ -148,8 +149,9 @@ func (c *Cluster) CreateContainer(config *cluster.ContainerConfig, name string,
 	if err != nil {
 		var retries int64
 		//  fails with image not found, then try to reschedule with image affinity
+		// ENGINEAPIFIXME: The first error can be removed once dockerclient is removed
 		bImageNotFoundError, _ := regexp.MatchString(`image \S* not found`, err.Error())
-		if bImageNotFoundError && !config.HaveNodeConstraint() {
+		if (bImageNotFoundError || client.IsErrImageNotFound(err)) && !config.HaveNodeConstraint() {
 			// Check if the image exists in the cluster
 			// If exists, retry with a image affinity
 			if c.Image(config.Image) != nil {