add affinity:

Signed-off-by: Victor Vieux <vieux@docker.com>

flags.go

@@ -62,7 +62,7 @@ var (
 	}
 	flFilter = cli.StringSliceFlag{
 		Name:  "filter, f",
-		Usage: "filter to use [constraint, health, port]",
-		Value: &cli.StringSlice{"constraint", "health", "port"},
+		Usage: "filter to use [constraint, affinity, health, port]",
+		Value: &cli.StringSlice{"constraint", "affinity", "health", "port"},
 	}
 )

scheduler/filter/affinity.go

@@ -0,0 +1,50 @@
+package filter
+
+import (
+	"fmt"
+
+	log "github.com/Sirupsen/logrus"
+	"github.com/docker/swarm/cluster"
+	"github.com/samalba/dockerclient"
+)
+
+// AffinityFilter selects only nodes based on other containers on the node.
+type AffinityFilter struct {
+}
+
+func (f *AffinityFilter) Filter(config *dockerclient.ContainerConfig, nodes []*cluster.Node) ([]*cluster.Node, error) {
+	affinities := extractEnv("affinity", config.Env)
+	for k, v := range affinities {
+		log.Debugf("matching affinity: %s=%s", k, v)
+		candidates := []*cluster.Node{}
+		for _, node := range nodes {
+			switch k {
+			case "container":
+				for _, container := range node.Containers() {
+					// "node" label is a special case pinning a container to a specific node.
+					if match(v, container.Id) || match(v, container.Names[0]) {
+						candidates = append(candidates, node)
+						break
+					}
+				}
+			case "image":
+				//TODO use cache
+				images, err := node.ListImages()
+				if err != nil {
+					break
+				}
+				for _, image := range images {
+					if match(v, image) {
+						candidates = append(candidates, node)
+						break
+					}
+				}
+			}
+		}
+		if len(candidates) == 0 {
+			return nil, fmt.Errorf("unable to find a node that satisfies %s == %s", k, v)
+		}
+		nodes = candidates
+	}
+	return nodes, nil
+}

scheduler/filter/affinity_test.go

@@ -0,0 +1,88 @@
+package filter
+
+import (
+	"testing"
+
+	"github.com/docker/swarm/cluster"
+	"github.com/samalba/dockerclient"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAffinityFilter(t *testing.T) {
+	//TODO: add test for images
+
+	var (
+		f     = AffinityFilter{}
+		nodes = []*cluster.Node{
+			cluster.NewNode("node-0"),
+			cluster.NewNode("node-1"),
+			cluster.NewNode("node-2"),
+		}
+		result []*cluster.Node
+		err    error
+	)
+
+	nodes[0].ID = "node-0-id"
+	nodes[0].Name = "node-0-name"
+	nodes[0].AddContainer(&cluster.Container{
+		Container: dockerclient.Container{
+			Id:    "container-0-id",
+			Names: []string{"container-0-name"},
+		},
+	})
+
+	nodes[1].ID = "node-1-id"
+	nodes[1].Name = "node-1-name"
+	nodes[1].AddContainer(&cluster.Container{
+		Container: dockerclient.Container{
+			Id:    "container-1-id",
+			Names: []string{"container-1-name"},
+		},
+	})
+
+	nodes[2].ID = "node-2-id"
+	nodes[2].Name = "node-2-name"
+
+	// Without constraints we should get the unfiltered list of nodes back.
+	result, err = f.Filter(&dockerclient.ContainerConfig{}, nodes)
+	assert.NoError(t, err)
+	assert.Equal(t, result, nodes)
+
+	// Set a constraint that cannot be fullfilled and expect an error back.
+	result, err = f.Filter(&dockerclient.ContainerConfig{
+		Env: []string{"affinity:container=does_not_exsits"},
+	}, nodes)
+	assert.Error(t, err)
+
+	// Set a contraint that can only be filled by a single node.
+	result, err = f.Filter(&dockerclient.ContainerConfig{
+		Env: []string{"affinity:container=container-0*"},
+	}, nodes)
+	assert.NoError(t, err)
+	assert.Len(t, result, 1)
+	assert.Equal(t, result[0], nodes[0])
+
+	// This constraint can only be fullfilled by a subset of nodes.
+	result, err = f.Filter(&dockerclient.ContainerConfig{
+		Env: []string{"affinity:container=container-*"},
+	}, nodes)
+	assert.NoError(t, err)
+	assert.Len(t, result, 2)
+	assert.NotContains(t, result, nodes[2])
+
+	// Validate node pinning by id.
+	result, err = f.Filter(&dockerclient.ContainerConfig{
+		Env: []string{"affinity:container=container-0-id"},
+	}, nodes)
+	assert.NoError(t, err)
+	assert.Len(t, result, 1)
+	assert.Equal(t, result[0], nodes[0])
+
+	// Validate node pinning by name.
+	result, err = f.Filter(&dockerclient.ContainerConfig{
+		Env: []string{"affinity:container=container-1-name"},
+	}, nodes)
+	assert.NoError(t, err)
+	assert.Len(t, result, 1)
+	assert.Equal(t, result[0], nodes[1])
+}

scheduler/filter/constraint.go

@@ -2,8 +2,6 @@ package filter
 
 import (
 	"fmt"
-	"regexp"
-	"strings"
 
 	log "github.com/Sirupsen/logrus"
 	"github.com/docker/swarm/cluster"
@@ -14,31 +12,8 @@ import (
 type ConstraintFilter struct {
 }
 
-func (f *ConstraintFilter) extractConstraints(env []string) map[string]string {
-	constraints := make(map[string]string)
-	for _, e := range env {
-		if strings.HasPrefix(e, "constraint:") {
-			constraint := strings.TrimPrefix(e, "constraint:")
-			parts := strings.SplitN(constraint, "=", 2)
-			constraints[strings.ToLower(parts[0])] = strings.ToLower(parts[1])
-		}
-	}
-	return constraints
-}
-
-// Create the regex for globbing (ex: ub*t* -> ^ub.*t.*$)
-// and match.
-func (f *ConstraintFilter) match(pattern, s string) bool {
-	regex := "^" + strings.Replace(pattern, "*", ".*", -1) + "$"
-	matched, err := regexp.MatchString(regex, strings.ToLower(s))
-	if err != nil {
-		log.Error(err)
-	}
-	return matched
-}
-
 func (f *ConstraintFilter) Filter(config *dockerclient.ContainerConfig, nodes []*cluster.Node) ([]*cluster.Node, error) {
-	constraints := f.extractConstraints(config.Env)
+	constraints := extractEnv("constraint", config.Env)
 	for k, v := range constraints {
 		log.Debugf("matching constraint: %s=%s", k, v)
 		candidates := []*cluster.Node{}
@@ -46,13 +21,13 @@ func (f *ConstraintFilter) Filter(config *dockerclient.ContainerConfig, nodes []
 			switch k {
 			case "node":
 				// "node" label is a special case pinning a container to a specific node.
-				if f.match(v, node.ID) || f.match(v, node.Name) {
+				if match(v, node.ID) || match(v, node.Name) {
 					candidates = append(candidates, node)
 				}
 			default:
 				// By default match the node labels.
 				if label, ok := node.Labels[k]; ok {
-					if f.match(v, label) {
+					if match(v, label) {
 						candidates = append(candidates, node)
 					}
 				}

scheduler/filter/constraint_test.go

@@ -8,7 +8,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestConstrainteFilter(t *testing.T) {
+func TestConstraintFilter(t *testing.T) {
 	var (
 		f     = ConstraintFilter{}
 		nodes = []*cluster.Node{

scheduler/filter/filter.go

@@ -20,6 +20,7 @@ var (
 
 func init() {
 	filters = map[string]Filter{
+		"affinity":   &AffinityFilter{},
 		"health":     &HealthFilter{},
 		"constraint": &ConstraintFilter{},
 		"port":       &PortFilter{},

scheduler/filter/utils.go

@@ -0,0 +1,34 @@
+package filter
+
+import (
+	"regexp"
+	"strings"
+
+	log "github.com/Sirupsen/logrus"
+)
+
+func extractEnv(key string, env []string) map[string]string {
+	values := make(map[string]string)
+	for _, e := range env {
+		if strings.HasPrefix(e, key+":") {
+			value := strings.TrimPrefix(e, key+":")
+			parts := strings.SplitN(value, "=", 2)
+			if len(parts) == 2 {
+				values[strings.ToLower(parts[0])] = strings.ToLower(parts[1])
+			} else {
+				values[strings.ToLower(parts[0])] = ""
+			}
+		}
+	}
+	return values
+}
+
+// Create the regex for globbing (ex: ub*t* -> ^ub.*t.*$) and match.
+func match(pattern, s string) bool {
+	regex := "^" + strings.Replace(pattern, "*", ".*", -1) + "$"
+	matched, err := regexp.MatchString(regex, strings.ToLower(s))
+	if err != nil {
+		log.Error(err)
+	}
+	return matched
+}