From fff84e84875681706007caa7f78d110d1a97b5b8 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Thu, 13 Aug 2015 10:19:46 -0700
Subject: [PATCH] scheduler: implementing driver iterator

---
 scheduler/feasible.go      | 85 +++++++++++++++++++++-----------------
 scheduler/feasible_test.go | 35 ++++++++++++++++
 2 files changed, 83 insertions(+), 37 deletions(-)

diff --git a/scheduler/feasible.go b/scheduler/feasible.go
index 442ce366f..a8df56214 100644
--- a/scheduler/feasible.go
+++ b/scheduler/feasible.go
@@ -1,6 +1,7 @@
 package scheduler
 
 import (
+	"fmt"
 	"math/rand"
 
 	"github.com/hashicorp/nomad/nomad/structs"
@@ -59,6 +60,53 @@ func NewRandomIterator(ctx Context, nodes []*structs.Node) *StaticIterator {
 	return NewStaticIterator(ctx, nodes)
 }
 
+// DriverIterator is a FeasibleIterator which returns nodes that
+// have the drivers necessary to scheduler a task group.
+type DriverIterator struct {
+	ctx     Context
+	source  FeasibleIterator
+	drivers map[string]struct{}
+}
+
+// NewDriverIterator creates a DriverIterator from a source and set of drivers
+func NewDriverIterator(ctx Context, source FeasibleIterator, drivers map[string]struct{}) *DriverIterator {
+	iter := &DriverIterator{
+		ctx:     ctx,
+		source:  source,
+		drivers: drivers,
+	}
+	return iter
+}
+
+func (iter *DriverIterator) Next() *structs.Node {
+	for {
+		// Get the next option from the source
+		option := iter.source.Next()
+		if option == nil {
+			return nil
+		}
+
+		// Use this node if possible
+		if iter.hasDrivers(option) {
+			return option
+		}
+	}
+}
+
+// hasDrivers is used to check if the node has all the appropriate
+// drivers for this task group. Drivers are registered as node attribute
+// like "driver.docker=1" with their corresponding version.
+func (iter *DriverIterator) hasDrivers(option *structs.Node) bool {
+	for driver := range iter.drivers {
+		driverStr := fmt.Sprintf("driver.%s", driver)
+		_, ok := option.Attributes[driverStr]
+		if !ok {
+			return false
+		}
+	}
+	return true
+}
+
 // ConstraintIterator is a FeasibleIterator which returns nodes
 // that match a given set of constraints. This is used to filter
 // on job, task group, and task constraints.
@@ -97,40 +145,3 @@ func (iter *ConstraintIterator) meetsConstraints(option *structs.Node) bool {
 	// TODO:
 	return true
 }
-
-// DriverIterator is a FeasibleIterator which returns nodes that
-// have the drivers necessary to scheduler a task group.
-type DriverIterator struct {
-	ctx     Context
-	source  FeasibleIterator
-	drivers map[string]struct{}
-}
-
-// NewDriverIterator creates a DriverIterator from a source and set of drivers
-func NewDriverIterator(ctx Context, source FeasibleIterator, drivers map[string]struct{}) *DriverIterator {
-	iter := &DriverIterator{
-		ctx:     ctx,
-		source:  source,
-		drivers: drivers,
-	}
-	return iter
-}
-
-func (iter *DriverIterator) Next() *structs.Node {
-	for {
-		// Get the next option from the source
-		option := iter.source.Next()
-		if option == nil {
-			return nil
-		}
-
-		// Use this node if possible
-		if iter.hasDrivers(option) {
-			return option
-		}
-	}
-}
-
-func (iter *DriverIterator) hasDrivers(option *structs.Node) bool {
-	return true
-}
diff --git a/scheduler/feasible_test.go b/scheduler/feasible_test.go
index bcdc668cc..7b798669e 100644
--- a/scheduler/feasible_test.go
+++ b/scheduler/feasible_test.go
@@ -35,3 +35,38 @@ func TestRandomIterator(t *testing.T) {
 		t.Fatalf("same order")
 	}
 }
+
+func TestDriverIterator(t *testing.T) {
+	ctx := NewEvalContext()
+	nodes := []*structs.Node{
+		mock.Node(),
+		mock.Node(),
+		mock.Node(),
+	}
+	static := NewStaticIterator(ctx, nodes)
+
+	nodes[0].Attributes["driver.foo"] = "2"
+	nodes[2].Attributes["driver.foo"] = "2"
+
+	drivers := map[string]struct{}{
+		"docker": struct{}{},
+		"foo":    struct{}{},
+	}
+	driver := NewDriverIterator(ctx, static, drivers)
+
+	var out []*structs.Node
+	for {
+		next := driver.Next()
+		if next == nil {
+			break
+		}
+		out = append(out, next)
+	}
+
+	if len(out) != 2 {
+		t.Fatalf("missing nodes")
+	}
+	if out[0] != nodes[0] || out[1] != nodes[2] {
+		t.Fatalf("bad: %#v", out)
+	}
+}