From 8c7c0b248b9639b2a9eefaa098baaf94e35ef466 Mon Sep 17 00:00:00 2001
From: Xian Chaobo <xianchaobo@huawei.com>
Date: Tue, 16 Jun 2015 07:52:00 +0800
Subject: [PATCH] reschedule with soft image affinity

Signed-off-by: Xian Chaobo <xianchaobo@huawei.com>
---
 cluster/config.go        | 15 +++++++++++++++
 cluster/swarm/cluster.go | 23 +++++++++++++++++++++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/cluster/config.go b/cluster/config.go
index 69c547a907..2041cf440f 100644
--- a/cluster/config.go
+++ b/cluster/config.go
@@ -146,3 +146,18 @@ func (c *ContainerConfig) Affinities() []string {
 func (c *ContainerConfig) Constraints() []string {
 	return c.extractExprs("constraints")
 }
+
+// AddSoftImageAffinity
+func (c *ContainerConfig) AddSoftImageAffinity(image string) error {
+	affnities_json := c.Labels[SwarmLabelNamespace+".affinities"]
+	var affinities []string
+	json.Unmarshal([]byte(affnities_json), &affinities)
+	affinities = append(affinities, "image==~"+image)
+
+	labels, err := json.Marshal(affinities)
+	if err != nil {
+		return err
+	}
+	c.Labels[SwarmLabelNamespace+".affinities"] = string(labels)
+	return nil
+}
diff --git a/cluster/swarm/cluster.go b/cluster/swarm/cluster.go
index 8f2c7d1f65..1acf60d445 100644
--- a/cluster/swarm/cluster.go
+++ b/cluster/swarm/cluster.go
@@ -89,6 +89,20 @@ func (c *Cluster) generateUniqueID() string {
 
 // CreateContainer aka schedule a brand new container into the cluster.
 func (c *Cluster) CreateContainer(config *cluster.ContainerConfig, name string) (*cluster.Container, error) {
+	container, err := c.createContainer(config, name, false)
+
+	//  fails with image not found, then try to reschedule with soft-image-affinity
+	if err != nil && strings.HasSuffix(err.Error(), "not found") {
+		// Check if the image exists in the cluster
+		// If exists, retry with a soft-image-affinity
+		if image := c.Image(config.Image); image != nil {
+			container, err = c.createContainer(config, name, true)
+		}
+	}
+	return container, err
+}
+
+func (c *Cluster) createContainer(config *cluster.ContainerConfig, name string, with_soft_image_affinity bool) (*cluster.Container, error) {
 	c.scheduler.Lock()
 	defer c.scheduler.Unlock()
 
@@ -100,13 +114,18 @@ func (c *Cluster) CreateContainer(config *cluster.ContainerConfig, name string)
 	// Associate a Swarm ID to the container we are creating.
 	config.SetSwarmID(c.generateUniqueID())
 
-	n, err := c.scheduler.SelectNodeForContainer(c.listNodes(), config)
+	config_temp := config
+	if with_soft_image_affinity {
+		config_temp.AddSoftImageAffinity(config.Image)
+	}
+
+	n, err := c.scheduler.SelectNodeForContainer(c.listNodes(), config_temp)
 	if err != nil {
 		return nil, err
 	}
 
 	if nn, ok := c.engines[n.ID]; ok {
-		container, err := nn.Create(config, name, true)
+		container, err := nn.Create(config_temp, name, true)
 		if err != nil {
 			return nil, err
 		}