Add "distinctHost" constraint

jobspec/parse.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"strconv"
 	"strings"
 	"time"
 
@@ -256,6 +257,20 @@ func parseConstraints(result *[]*structs.Constraint, obj *hclobj.Object) error {
 			m["RTarget"] = constraint
 		}
 
+		if value, ok := m["distinctHosts"]; ok {
+			enabled, err := strconv.ParseBool(value.(string))
+			if err != nil {
+				return err
+			}
+
+			// If it is not enabled, skip the constraint.
+			if !enabled {
+				continue
+			}
+
+			m["Operand"] = "distinctHosts"
+		}
+
 		// Build the constraint
 		var c structs.Constraint
 		if err := mapstructure.WeakDecode(m, &c); err != nil {

jobspec/parse_test.go

@@ -192,6 +192,24 @@ func TestParse(t *testing.T) {
 			false,
 		},
 
+		{
+			"distinctHosts-constraint.hcl",
+			&structs.Job{
+				ID:       "foo",
+				Name:     "foo",
+				Priority: 50,
+				Region:   "global",
+				Type:     "service",
+				Constraints: []*structs.Constraint{
+					&structs.Constraint{
+						Hard:    true,
+						Operand: "distinctHosts",
+					},
+				},
+			},
+			false,
+		},
+
 		{
 			"specify-job.hcl",
 			&structs.Job{

jobspec/test-fixtures/distinctHosts-constraint.hcl

@@ -0,0 +1,5 @@
+job "foo" {
+    constraint {
+        distinctHosts = "true"
+    }
+}

scheduler/feasible.go

@@ -150,6 +150,117 @@ func (iter *DriverIterator) hasDrivers(option *structs.Node) bool {
 	return true
 }
 
+// DynamicConstraintIterator is a FeasibleIterator which returns nodes that
+// match constraints that are not static such as Node attributes but are
+// effected by alloc placements. Examples are distinctHosts and tenancy constraints.
+// This is used to filter on job and task group constraints.
+type DynamicConstraintIterator struct {
+	ctx    Context
+	source FeasibleIterator
+	tg     *structs.TaskGroup
+	job    *structs.Job
+
+	// Store whether the Job or TaskGroup has a distinctHosts constraints so
+	// they don't have to be calculated every time Next() is called.
+	tgDistinctHosts  bool
+	jobDistinctHosts bool
+}
+
+// NewDynamicConstraintIterator creates a DynamicConstraintIterator from a
+// source.
+func NewDynamicConstraintIterator(ctx Context, source FeasibleIterator) *DynamicConstraintIterator {
+	iter := &DynamicConstraintIterator{
+		ctx:    ctx,
+		source: source,
+	}
+	return iter
+}
+
+func (iter *DynamicConstraintIterator) SetTaskGroup(tg *structs.TaskGroup) {
+	iter.tg = tg
+	iter.tgDistinctHosts = iter.hasDistinctHostsConstraint(tg.Constraints)
+}
+
+func (iter *DynamicConstraintIterator) SetJob(job *structs.Job) {
+	iter.job = job
+	iter.jobDistinctHosts = iter.hasDistinctHostsConstraint(job.Constraints)
+}
+
+func (iter *DynamicConstraintIterator) hasDistinctHostsConstraint(constraints []*structs.Constraint) bool {
+	if constraints == nil {
+		return false
+	}
+
+	for _, con := range constraints {
+		if con.Operand == "distinctHosts" {
+			return true
+		}
+	}
+	return false
+}
+
+func (iter *DynamicConstraintIterator) Next() *structs.Node {
+	if iter.job == nil {
+		iter.ctx.Logger().Printf("[ERR] scheduler.dynamic-constraint: job not set")
+		return nil
+	}
+
+	if iter.tg == nil {
+		iter.ctx.Logger().Printf("[ERR] scheduler.dynamic-constraint: task group not set")
+		return nil
+	}
+
+	for {
+		// Get the next option from the source
+		option := iter.source.Next()
+
+		// Hot-path if the option is nil or there are no distinctHosts constraints.
+		if option == nil || (!iter.jobDistinctHosts && !iter.tgDistinctHosts) {
+			return option
+		}
+
+		if !iter.satisfiesDistinctHosts(option, iter.jobDistinctHosts) {
+			iter.ctx.Metrics().FilterNode(option, "distinctHosts")
+			continue
+		}
+
+		return option
+	}
+}
+
+// satisfiesDistinctHosts checks if the node satisfies a distinctHosts
+// constraint either specified at the job level or the TaskGroup level.
+func (iter *DynamicConstraintIterator) satisfiesDistinctHosts(option *structs.Node, job bool) bool {
+	// Get the proposed allocations
+	proposed, err := iter.ctx.ProposedAllocs(option.ID)
+	if err != nil {
+		iter.ctx.Logger().Printf(
+			"[ERR] scheduler.dynamic-constraint: failed to get proposed allocations: %v", err)
+		return false
+	}
+
+	// Skip the node if the task group has already been allocated on it.
+	for _, alloc := range proposed {
+		jobCollision := alloc.JobID == iter.job.ID
+		taskCollision := alloc.TaskGroup == iter.tg.Name
+
+		// If the job has a distinctHosts constraint we only need an alloc
+		// collision on the JobID but if the constraint is on the TaskGroup then
+		// we need both a job and TaskGroup collision.
+		jobInvalid := job && jobCollision
+		tgInvalid := !job && jobCollision && taskCollision
+		if jobInvalid || tgInvalid {
+			return false
+		}
+	}
+
+	return true
+}
+
+func (iter *DynamicConstraintIterator) Reset() {
+	iter.source.Reset()
+}
+
 // 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.
@@ -257,6 +368,14 @@ func resolveConstraintTarget(target string, node *structs.Node) (interface{}, bo
 
 // checkConstraint checks if a constraint is satisfied
 func checkConstraint(ctx Context, operand string, lVal, rVal interface{}) bool {
+	// Check for constraints not handled by this iterator.
+	switch operand {
+	case "distinctHosts":
+		return true
+	default:
+		break
+	}
+
 	switch operand {
 	case "=", "==", "is":
 		return reflect.DeepEqual(lVal, rVal)

scheduler/feasible_test.go

@@ -382,6 +382,180 @@ func TestCheckRegexpConstraint(t *testing.T) {
 	}
 }
 
+func TestDynamicConstraint_JobDistinctHosts(t *testing.T) {
+	_, ctx := testContext(t)
+	nodes := []*structs.Node{
+		mock.Node(),
+		mock.Node(),
+		mock.Node(),
+		mock.Node(),
+	}
+	static := NewStaticIterator(ctx, nodes)
+
+	// Create a job with a distinctHosts constraint and two task groups.
+	tg1 := &structs.TaskGroup{Name: "bar"}
+	tg2 := &structs.TaskGroup{Name: "baz"}
+
+	job := &structs.Job{
+		ID:          "foo",
+		Constraints: []*structs.Constraint{{Operand: "distinctHosts"}},
+		TaskGroups:  []*structs.TaskGroup{tg1, tg2},
+	}
+
+	dynamic := NewDynamicConstraintIterator(ctx, static)
+	dynamic.SetTaskGroup(tg1)
+	dynamic.SetJob(job)
+
+	out := collectFeasible(dynamic)
+	if len(out) != 4 {
+		t.Fatalf("Bad: %#v", out)
+	}
+
+	selected := make(map[string]struct{}, 4)
+	for _, option := range out {
+		if _, ok := selected[option.ID]; ok {
+			t.Fatalf("selected node %v for more than one alloc", option)
+		}
+		selected[option.ID] = struct{}{}
+	}
+}
+
+func TestDynamicConstraint_JobDistinctHosts_Infeasible(t *testing.T) {
+	_, ctx := testContext(t)
+	nodes := []*structs.Node{
+		mock.Node(),
+		mock.Node(),
+	}
+	static := NewStaticIterator(ctx, nodes)
+
+	// Create a job with a distinctHosts constraint and two task groups.
+	tg1 := &structs.TaskGroup{Name: "bar"}
+	tg2 := &structs.TaskGroup{Name: "baz"}
+
+	job := &structs.Job{
+		ID:          "foo",
+		Constraints: []*structs.Constraint{{Operand: "distinctHosts"}},
+		TaskGroups:  []*structs.TaskGroup{tg1, tg2},
+	}
+
+	// Add allocs placing tg1 on node1 and tg2 on node2. This should make the
+	// job unsatisfiable.
+	plan := ctx.Plan()
+	plan.NodeAllocation[nodes[0].ID] = []*structs.Allocation{
+		&structs.Allocation{
+			TaskGroup: tg1.Name,
+			JobID:     job.ID,
+		},
+
+		// Should be ignored as it is a different job.
+		&structs.Allocation{
+			TaskGroup: tg2.Name,
+			JobID:     "ignore 2",
+		},
+	}
+	plan.NodeAllocation[nodes[1].ID] = []*structs.Allocation{
+		&structs.Allocation{
+			TaskGroup: tg2.Name,
+			JobID:     job.ID,
+		},
+
+		// Should be ignored as it is a different job.
+		&structs.Allocation{
+			TaskGroup: tg1.Name,
+			JobID:     "ignore 2",
+		},
+	}
+
+	dynamic := NewDynamicConstraintIterator(ctx, static)
+	dynamic.SetTaskGroup(tg1)
+	dynamic.SetJob(job)
+
+	out := collectFeasible(dynamic)
+	if len(out) != 0 {
+		t.Fatalf("Bad: %#v", out)
+	}
+}
+
+func TestDynamicConstraint_JobDistinctHosts_InfeasibleCount(t *testing.T) {
+	_, ctx := testContext(t)
+	nodes := []*structs.Node{
+		mock.Node(),
+		mock.Node(),
+	}
+	static := NewStaticIterator(ctx, nodes)
+
+	// Create a job with a distinctHosts constraint and three task groups.
+	tg1 := &structs.TaskGroup{Name: "bar"}
+	tg2 := &structs.TaskGroup{Name: "baz"}
+	tg3 := &structs.TaskGroup{Name: "bam"}
+
+	job := &structs.Job{
+		ID:          "foo",
+		Constraints: []*structs.Constraint{{Operand: "distinctHosts"}},
+		TaskGroups:  []*structs.TaskGroup{tg1, tg2, tg3},
+	}
+
+	dynamic := NewDynamicConstraintIterator(ctx, static)
+	dynamic.SetTaskGroup(tg1)
+	dynamic.SetJob(job)
+
+	// It should not be able to place 3 tasks with only two nodes.
+	out := collectFeasible(dynamic)
+	if len(out) != 2 {
+		t.Fatalf("Bad: %#v", out)
+	}
+}
+
+func TestDynamicConstraint_TaskGroupDistinctHosts(t *testing.T) {
+	_, ctx := testContext(t)
+	nodes := []*structs.Node{
+		mock.Node(),
+		mock.Node(),
+	}
+	static := NewStaticIterator(ctx, nodes)
+
+	// Create a task group with a distinctHosts constraint.
+	taskGroup := &structs.TaskGroup{
+		Name: "example",
+		Constraints: []*structs.Constraint{
+			{Operand: "distinctHosts"},
+		},
+	}
+
+	// Add a planned alloc to node1.
+	plan := ctx.Plan()
+	plan.NodeAllocation[nodes[0].ID] = []*structs.Allocation{
+		&structs.Allocation{
+			TaskGroup: taskGroup.Name,
+			JobID:     "foo",
+		},
+	}
+
+	// Add a planned alloc to node2 with the same task group name but a
+	// different job.
+	plan.NodeAllocation[nodes[1].ID] = []*structs.Allocation{
+		&structs.Allocation{
+			TaskGroup: taskGroup.Name,
+			JobID:     "bar",
+		},
+	}
+
+	dynamic := NewDynamicConstraintIterator(ctx, static)
+	dynamic.SetTaskGroup(taskGroup)
+	dynamic.SetJob(&structs.Job{ID: "foo"})
+
+	out := collectFeasible(dynamic)
+	if len(out) != 1 {
+		t.Fatalf("Bad: %#v", out)
+	}
+
+	// Expect it to skip the first node as there is a previous alloc on it for
+	// the same task group.
+	if out[0] != nodes[1] {
+		t.Fatalf("Bad: %v", out)
+	}
+}
+
 func collectFeasible(iter FeasibleIterator) (out []*structs.Node) {
 	for {
 		next := iter.Next()