Merge pull request #4513 from hashicorp/f-affinities-backend

Implement affinity support in the scheduler (batch and service jobs)

scheduler/context.go

@@ -34,8 +34,8 @@ type Context interface {
 	// RegexpCache is a cache of regular expressions
 	RegexpCache() map[string]*regexp.Regexp
 
-	// ConstraintCache is a cache of version constraints
-	ConstraintCache() map[string]version.Constraints
+	// VersionConstraintCache is a cache of version constraints
+	VersionConstraintCache() map[string]version.Constraints
 
 	// Eligibility returns a tracker for node eligibility in the context of the
 	// eval.
@@ -54,7 +54,8 @@ func (e *EvalCache) RegexpCache() map[string]*regexp.Regexp {
 	}
 	return e.reCache
 }
-func (e *EvalCache) ConstraintCache() map[string]version.Constraints {
+
+func (e *EvalCache) VersionConstraintCache() map[string]version.Constraints {
 	if e.constraintCache == nil {
 		e.constraintCache = make(map[string]version.Constraints)
 	}

scheduler/feasible.go

@@ -403,11 +403,11 @@ func (c *ConstraintChecker) Feasible(option *structs.Node) bool {
 
 func (c *ConstraintChecker) meetsConstraint(constraint *structs.Constraint, option *structs.Node) bool {
 	// Resolve the targets
-	lVal, ok := resolveConstraintTarget(constraint.LTarget, option)
+	lVal, ok := resolveTarget(constraint.LTarget, option)
 	if !ok {
 		return false
 	}
-	rVal, ok := resolveConstraintTarget(constraint.RTarget, option)
+	rVal, ok := resolveTarget(constraint.RTarget, option)
 	if !ok {
 		return false
 	}
@@ -416,8 +416,8 @@ func (c *ConstraintChecker) meetsConstraint(constraint *structs.Constraint, opti
 	return checkConstraint(c.ctx, constraint.Operand, lVal, rVal)
 }
 
-// resolveConstraintTarget is used to resolve the LTarget and RTarget of a Constraint
-func resolveConstraintTarget(target string, node *structs.Node) (interface{}, bool) {
+// resolveTarget is used to resolve the LTarget and RTarget of a Constraint
+func resolveTarget(target string, node *structs.Node) (interface{}, bool) {
 	// If no prefix, this must be a literal value
 	if !strings.HasPrefix(target, "${") {
 		return target, true
@@ -470,16 +470,28 @@ func checkConstraint(ctx Context, operand string, lVal, rVal interface{}) bool {
 	case "<", "<=", ">", ">=":
 		return checkLexicalOrder(operand, lVal, rVal)
 	case structs.ConstraintVersion:
-		return checkVersionConstraint(ctx, lVal, rVal)
+		return checkVersionMatch(ctx, lVal, rVal)
 	case structs.ConstraintRegex:
-		return checkRegexpConstraint(ctx, lVal, rVal)
+		return checkRegexpMatch(ctx, lVal, rVal)
 	case structs.ConstraintSetContains:
-		return checkSetContainsConstraint(ctx, lVal, rVal)
+		return checkSetContainsAll(ctx, lVal, rVal)
 	default:
 		return false
 	}
 }
 
+// checkAffinity checks if a specific affinity is satisfied
+func checkAffinity(ctx Context, operand string, lVal, rVal interface{}) bool {
+	switch operand {
+	case structs.ConstraintSetContaintsAny:
+		return checkSetContainsAny(lVal, rVal)
+	case structs.ConstraintSetContainsAll, structs.ConstraintSetContains:
+		return checkSetContainsAll(ctx, lVal, rVal)
+	default:
+		return checkConstraint(ctx, operand, lVal, rVal)
+	}
+}
+
 // checkLexicalOrder is used to check for lexical ordering
 func checkLexicalOrder(op string, lVal, rVal interface{}) bool {
 	// Ensure the values are strings
@@ -506,9 +518,9 @@ func checkLexicalOrder(op string, lVal, rVal interface{}) bool {
 	}
 }
 
-// checkVersionConstraint is used to compare a version on the
+// checkVersionMatch is used to compare a version on the
 // left hand side with a set of constraints on the right hand side
-func checkVersionConstraint(ctx Context, lVal, rVal interface{}) bool {
+func checkVersionMatch(ctx Context, lVal, rVal interface{}) bool {
 	// Parse the version
 	var versionStr string
 	switch v := lVal.(type) {
@@ -533,7 +545,7 @@ func checkVersionConstraint(ctx Context, lVal, rVal interface{}) bool {
 	}
 
 	// Check the cache for a match
-	cache := ctx.ConstraintCache()
+	cache := ctx.VersionConstraintCache()
 	constraints := cache[constraintStr]
 
 	// Parse the constraints
@@ -549,9 +561,9 @@ func checkVersionConstraint(ctx Context, lVal, rVal interface{}) bool {
 	return constraints.Check(vers)
 }
 
-// checkRegexpConstraint is used to compare a value on the
+// checkRegexpMatch is used to compare a value on the
 // left hand side with a regexp on the right hand side
-func checkRegexpConstraint(ctx Context, lVal, rVal interface{}) bool {
+func checkRegexpMatch(ctx Context, lVal, rVal interface{}) bool {
 	// Ensure left-hand is string
 	lStr, ok := lVal.(string)
 	if !ok {
@@ -582,9 +594,9 @@ func checkRegexpConstraint(ctx Context, lVal, rVal interface{}) bool {
 	return re.MatchString(lStr)
 }
 
-// checkSetContainsConstraint is used to see if the left hand side contains the
+// checkSetContainsAll is used to see if the left hand side contains the
 // string on the right hand side
-func checkSetContainsConstraint(ctx Context, lVal, rVal interface{}) bool {
+func checkSetContainsAll(ctx Context, lVal, rVal interface{}) bool {
 	// Ensure left-hand is string
 	lStr, ok := lVal.(string)
 	if !ok {
@@ -614,6 +626,38 @@ func checkSetContainsConstraint(ctx Context, lVal, rVal interface{}) bool {
 	return true
 }
 
+// checkSetContainsAny is used to see if the left hand side contains any
+// values on the right hand side
+func checkSetContainsAny(lVal, rVal interface{}) bool {
+	// Ensure left-hand is string
+	lStr, ok := lVal.(string)
+	if !ok {
+		return false
+	}
+
+	// RHS must be a string
+	rStr, ok := rVal.(string)
+	if !ok {
+		return false
+	}
+
+	input := strings.Split(lStr, ",")
+	lookup := make(map[string]struct{}, len(input))
+	for _, in := range input {
+		cleaned := strings.TrimSpace(in)
+		lookup[cleaned] = struct{}{}
+	}
+
+	for _, r := range strings.Split(rStr, ",") {
+		cleaned := strings.TrimSpace(r)
+		if _, ok := lookup[cleaned]; ok {
+			return true
+		}
+	}
+
+	return false
+}
+
 // FeasibilityWrapper is a FeasibleIterator which wraps both job and task group
 // FeasibilityCheckers in which feasibility checking can be skipped if the
 // computed node class has previously been marked as eligible or ineligible.

scheduler/feasible_test.go

@@ -309,7 +309,7 @@ func TestResolveConstraintTarget(t *testing.T) {
 	}
 
 	for _, tc := range cases {
-		res, ok := resolveConstraintTarget(tc.target, tc.node)
+		res, ok := resolveTarget(tc.target, tc.node)
 		if ok != tc.result {
 			t.Fatalf("TC: %#v, Result: %v %v", tc, res, ok)
 		}
@@ -460,7 +460,7 @@ func TestCheckVersionConstraint(t *testing.T) {
 	}
 	for _, tc := range cases {
 		_, ctx := testContext(t)
-		if res := checkVersionConstraint(ctx, tc.lVal, tc.rVal); res != tc.result {
+		if res := checkVersionMatch(ctx, tc.lVal, tc.rVal); res != tc.result {
 			t.Fatalf("TC: %#v, Result: %v", tc, res)
 		}
 	}
@@ -495,7 +495,7 @@ func TestCheckRegexpConstraint(t *testing.T) {
 	}
 	for _, tc := range cases {
 		_, ctx := testContext(t)
-		if res := checkRegexpConstraint(ctx, tc.lVal, tc.rVal); res != tc.result {
+		if res := checkRegexpMatch(ctx, tc.lVal, tc.rVal); res != tc.result {
 			t.Fatalf("TC: %#v, Result: %v", tc, res)
 		}
 	}
@@ -1611,3 +1611,11 @@ func TestFeasibilityWrapper_JobEligible_TgEscaped(t *testing.T) {
 		t.Fatalf("bad: %v %v", e, ok)
 	}
 }
+
+func TestSetContainsAny(t *testing.T) {
+	require.True(t, checkSetContainsAny("a", "a"))
+	require.True(t, checkSetContainsAny("a,b", "a"))
+	require.True(t, checkSetContainsAny("  a,b  ", "a "))
+	require.True(t, checkSetContainsAny("a", "a"))
+	require.False(t, checkSetContainsAny("b", "a"))
+}

scheduler/propertyset.go

@@ -313,7 +313,7 @@ func getProperty(n *structs.Node, property string) (string, bool) {
 		return "", false
 	}
 
-	val, ok := resolveConstraintTarget(property, n)
+	val, ok := resolveTarget(property, n)
 	if !ok {
 		return "", false
 	}