[YUNIKORN-2500] Use GetPreemptableResource, GetRemainingGuaranteedRes…

…ource in Preemption flow (apache#830) Closes: apache#830 Signed-off-by: Manikandan R <[email protected]>
manirajv06 · Jun 25, 2024 · cf73072 · cf73072
1 parent ff2a881
commit cf73072
Show file tree

Hide file tree

Showing 9 changed files with 578 additions and 247 deletions.
diff --git a/pkg/common/resources/resources.go b/pkg/common/resources/resources.go
@@ -645,6 +645,25 @@ func Equals(left, right *Resource) bool {
 	return true
 }
 
+// MatchAny returns true if at least one type in the defined resource exists in the other resource.
+// False if none of the types exist in the other resource.
+// A nil resource is treated as an empty resource (no types defined) and returns false
+// Values are not considered during the checks
+func (r *Resource) MatchAny(other *Resource) bool {
+	if r == nil || other == nil {
+		return false
+	}
+	if r == other {
+		return true
+	}
+	for k := range r.Resources {
+		if _, ok := other.Resources[k]; ok {
+			return true
+		}
+	}
+	return false
+}
+
 // Compare the resources equal returns the specific values for following cases:
 // left  right  return
 // nil   nil    true
@@ -751,6 +770,55 @@ func StrictlyGreaterThanOrEquals(larger, smaller *Resource) bool {
 	return true
 }
 
+// StrictlyGreaterThanOnlyExisting returns true if all quantities for types in the defined resource are greater than
+// the quantity for the same type in smaller.
+// Types defined in smaller that are not in the defined resource are ignored.
+// Two resources that are equal are not considered strictly larger than each other.
+func (r *Resource) StrictlyGreaterThanOnlyExisting(smaller *Resource) bool {
+	if r == nil {
+		r = Zero
+	}
+	if smaller == nil {
+		smaller = Zero
+	}
+
+	// keep track of the number of not equal values
+	notEqual := false
+
+	// Is larger and smaller completely disjoint?
+	atleastOneResourcePresent := false
+	// Is all resource in larger greater than zero?
+	isAllPositiveInLarger := true
+
+	for k, v := range r.Resources {
+		// even when smaller is empty, all resource type in larger should be greater than zero
+		if smaller.IsEmpty() && v <= 0 {
+			isAllPositiveInLarger = false
+		}
+		// when smaller is not empty
+		if val, ok := smaller.Resources[k]; ok {
+			// at least one common resource type is there
+			atleastOneResourcePresent = true
+			if val > v {
+				return false
+			}
+			if val != v {
+				notEqual = true
+			}
+		}
+	}
+
+	switch {
+	case smaller.IsEmpty() && !r.IsEmpty():
+		return isAllPositiveInLarger
+	case atleastOneResourcePresent:
+		return notEqual
+	default:
+		// larger and smaller is completely disjoint. none of the resource match.
+		return !r.IsEmpty() && !smaller.IsEmpty()
+	}
+}
+
 // Have at least one quantity > 0, and no quantities < 0
 // A nil resource is not strictly greater than zero.
 func StrictlyGreaterThanZero(larger *Resource) bool {
@@ -816,6 +884,29 @@ func ComponentWiseMinPermissive(left, right *Resource) *Resource {
 	return out
 }
 
+// ComponentWiseMinOnlyExisting Returns a new Resource with the smallest value for resource type
+// existing only in left but not vice versa.
+func ComponentWiseMinOnlyExisting(left, right *Resource) *Resource {
+	out := NewResource()
+	if right == nil && left == nil {
+		return nil
+	}
+	if left == nil {
+		return nil
+	}
+	if right == nil {
+		return left.Clone()
+	}
+	for k, v := range left.Resources {
+		if val, ok := right.Resources[k]; ok {
+			out.Resources[k] = min(v, val)
+		} else {
+			out.Resources[k] = v
+		}
+	}
+	return out
+}
+
 func (r *Resource) HasNegativeValue() bool {
 	if r == nil {
 		return false

diff --git a/pkg/common/resources/resources_test.go b/pkg/common/resources/resources_test.go
@@ -319,6 +319,111 @@ func TestStrictlyGreaterThan(t *testing.T) {
 	}
 }
 
+func TestMatchAnyOnlyExisting(t *testing.T) {
+	var tests = []struct {
+		caseName string
+		left     map[string]Quantity
+		right    map[string]Quantity
+		expected bool
+	}{
+		{"nil resource should not match", nil, nil, false},
+		{"empty resource should not match", map[string]Quantity{}, map[string]Quantity{}, false},
+		{"equal positive resources should match", map[string]Quantity{"first": 1}, map[string]Quantity{"first": 1}, true},
+		{"equal positive resources with different values should match", map[string]Quantity{"first": 1}, map[string]Quantity{"first": 2}, true},
+		{"positive resource and nil resource should not match", map[string]Quantity{"first": 1}, nil, false},
+		{"nil resource and positive resource should not match", nil, map[string]Quantity{"first": 1}, false},
+		{"positive resource and empty resource should not match", map[string]Quantity{"first": 1}, map[string]Quantity{}, false},
+		{"empty resource and positive resource should not match", map[string]Quantity{}, map[string]Quantity{"first": 1}, false},
+		{"positive resource should match even though value is zero", map[string]Quantity{"zero": 0}, map[string]Quantity{"zero": 0}, true},
+		{"positive resource should match even though extra resource type is there", map[string]Quantity{"first": 10}, map[string]Quantity{"first": 10, "second": 1}, true},
+		{"positive resource should match even though extra resource type is there", map[string]Quantity{"first": 10, "second": 1}, map[string]Quantity{"first": 10}, true},
+		{"resource should not match", map[string]Quantity{"first": 10}, map[string]Quantity{"second": 10}, false},
+		{"resource should not match", map[string]Quantity{"second": 10}, map[string]Quantity{"first": 10}, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.caseName, func(t *testing.T) {
+			var left *Resource
+			var right *Resource
+			if tt.left != nil {
+				left = NewResourceFromMap(tt.left)
+			}
+			if tt.right != nil {
+				right = NewResourceFromMap(tt.right)
+			}
+			if result := left.MatchAny(right); result != tt.expected {
+				t.Errorf("MatchAny: got %v, expected %v", result, tt.expected)
+			}
+		})
+	}
+}
+
+func TestStrictlyGreaterThanOnlyExisting(t *testing.T) {
+	type inputs struct {
+		larger  map[string]Quantity
+		smaller map[string]Quantity
+		sameRef bool
+	}
+	type outputs struct {
+		larger  bool
+		smaller bool
+	}
+	var tests = []struct {
+		caseName string
+		input    inputs
+		expected outputs
+	}{
+		{"Nil check", inputs{nil, nil, false}, outputs{false, false}},
+		{"Positive resource and empty resources", inputs{map[string]Quantity{"first": 10}, map[string]Quantity{}, false}, outputs{true, false}},
+		{"Positive resource and nil resources", inputs{map[string]Quantity{"first": 10}, nil, false}, outputs{true, false}},
+
+		{"Equal Positive resources", inputs{map[string]Quantity{"first": 10}, map[string]Quantity{"first": 10}, false}, outputs{false, false}},
+		{"Different Positive resources", inputs{map[string]Quantity{"first": 10}, map[string]Quantity{"first": 20}, false}, outputs{false, true}},
+		{"Different Positive resources", inputs{map[string]Quantity{"first": 10}, map[string]Quantity{"first": 5}, false}, outputs{true, false}},
+
+		{"Equal Positive resources with extra resource types", inputs{map[string]Quantity{"first": 10}, map[string]Quantity{"first": 10, "sec": 10}, false}, outputs{false, false}},
+		{"Different Positive resources with extra resource types", inputs{map[string]Quantity{"first": 10}, map[string]Quantity{"first": 20, "sec": 10}, false}, outputs{false, true}},
+		{"Different Positive resources with extra resource types", inputs{map[string]Quantity{"first": 10}, map[string]Quantity{"first": 5, "sec": 10}, false}, outputs{true, false}},
+
+		{"Equal Positive resources but completely disjoint", inputs{map[string]Quantity{"first": 10}, map[string]Quantity{"sec": 10}, false}, outputs{true, true}},
+		{"Zero resource and empty resources", inputs{map[string]Quantity{"first": 0}, map[string]Quantity{}, false}, outputs{false, false}},
+		{"Zero resource and nil resources", inputs{map[string]Quantity{"first": 0}, nil, false}, outputs{false, false}},
+
+		{"Negative resource and empty resources", inputs{map[string]Quantity{"first": -10}, map[string]Quantity{}, false}, outputs{false, false}},
+		{"Negative resource and nil resources", inputs{map[string]Quantity{"first": -10}, nil, false}, outputs{false, false}},
+		{"Negative resource and empty resources", inputs{map[string]Quantity{"first": -10, "sec": 10}, map[string]Quantity{}, false}, outputs{false, false}},
+		{"Negative resource and nil resources", inputs{map[string]Quantity{"first": -10, "sec": 10}, nil, false}, outputs{false, false}},
+
+		{"Equal Negative resources", inputs{map[string]Quantity{"first": -10}, map[string]Quantity{"first": -10}, false}, outputs{false, false}},
+		{"Different Negative resources", inputs{map[string]Quantity{"first": -10}, map[string]Quantity{"first": -20}, false}, outputs{true, false}},
+		{"Different Negative resources", inputs{map[string]Quantity{"first": -10}, map[string]Quantity{"first": -5}, false}, outputs{false, true}},
+
+		{"Equal Negative resources with extra resource types", inputs{map[string]Quantity{"first": -10}, map[string]Quantity{"first": -10, "sec": 10}, false}, outputs{false, false}},
+		{"Different Negative resources with extra resource types", inputs{map[string]Quantity{"first": -10}, map[string]Quantity{"first": -20, "sec": 10}, false}, outputs{true, false}},
+		{"Different Negative resources with extra resource types", inputs{map[string]Quantity{"first": -10}, map[string]Quantity{"first": -5, "sec": 10}, false}, outputs{false, true}},
+
+		{"Equal Negative resources but completely disjoint", inputs{map[string]Quantity{"first": -10}, map[string]Quantity{"sec": -10}, false}, outputs{true, true}},
+	}
+	for _, tt := range tests {
+		t.Run(tt.caseName, func(t *testing.T) {
+			var compare, base *Resource
+			if tt.input.larger != nil {
+				compare = NewResourceFromMap(tt.input.larger)
+			}
+			if tt.input.sameRef {
+				base = compare
+			} else {
+				base = NewResourceFromMap(tt.input.smaller)
+			}
+			if result := compare.StrictlyGreaterThanOnlyExisting(base); result != tt.expected.larger {
+				t.Errorf("comapre %v, base %v, got %v, expeceted %v", compare, base, result, tt.expected.larger)
+			}
+			if result := base.StrictlyGreaterThanOnlyExisting(compare); result != tt.expected.smaller {
+				t.Errorf("base %v, compare %v, got %v, expeceted %v", base, compare, result, tt.expected.smaller)
+			}
+		})
+	}
+}
+
 func TestStrictlyGreaterThanOrEquals(t *testing.T) {
 	type inputs struct {
 		larger  map[string]Quantity
@@ -443,6 +548,46 @@ func TestComponentWiseMinPermissive(t *testing.T) {
 	}
 }
 
+func TestComponentWiseMinOnlyExisting(t *testing.T) {
+	testCases := []struct {
+		name     string
+		left     map[string]Quantity
+		right    map[string]Quantity
+		expected map[string]Quantity
+	}{
+		{"Min of nil resources should be nil", nil, nil, nil},
+		{"Min of empty resources should be empty resource ", map[string]Quantity{}, map[string]Quantity{}, map[string]Quantity{}},
+		{"Min of positive resource and nil resource", map[string]Quantity{"first": 5}, nil, map[string]Quantity{"first": 5}},
+		{"Min of nil resource and positive resource", nil, map[string]Quantity{"first": 5}, nil},
+		{"Min of two positive resources", map[string]Quantity{"first": 5}, map[string]Quantity{"first": 10}, map[string]Quantity{"first": 5}},
+		{"Min of two positive resources", map[string]Quantity{"first": 10}, map[string]Quantity{"first": 5}, map[string]Quantity{"first": 5}},
+		{"Min of positive resource and negative resource", map[string]Quantity{"first": 5}, map[string]Quantity{"first": -5}, map[string]Quantity{"first": -5}},
+		{"Min of positive resource and negative resource", map[string]Quantity{"first": -5}, map[string]Quantity{"first": 5}, map[string]Quantity{"first": -5}},
+		{"Min of two positive resources with extra resource types", map[string]Quantity{"first": 10}, map[string]Quantity{"first": 5, "second": 15}, map[string]Quantity{"first": 5}},
+		{"Min of two positive resources with extra resource types", map[string]Quantity{"first": 5, "second": 15}, map[string]Quantity{"first": 10}, map[string]Quantity{"first": 5, "second": 15}},
+		{"Min of positive resource and negative resource with extra resource types", map[string]Quantity{"first": 10}, map[string]Quantity{"first": -5, "second": 15}, map[string]Quantity{"first": -5}},
+		{"Min of positive resource and negative resource with extra resource types", map[string]Quantity{"first": -5, "second": 15}, map[string]Quantity{"first": 10}, map[string]Quantity{"first": -5, "second": 15}},
+	}
+	for _, tc := range testCases {
+		var left *Resource
+		var right *Resource
+		var expected *Resource
+		if tc.left != nil {
+			left = NewResourceFromMap(tc.left)
+		}
+		if tc.right != nil {
+			right = NewResourceFromMap(tc.right)
+		}
+		if tc.expected != nil {
+			expected = NewResourceFromMap(tc.expected)
+		}
+		t.Run(tc.name, func(t *testing.T) {
+			result := ComponentWiseMinOnlyExisting(left, right)
+			assert.DeepEqual(t, result, expected)
+		})
+	}
+}
+
 func TestComponentWiseMax(t *testing.T) {
 	type inputs struct {
 		res1    map[string]Quantity

diff --git a/pkg/mock/preemption_predicate_plugin.go b/pkg/mock/preemption_predicate_plugin.go
@@ -22,6 +22,8 @@ import (
 	"errors"
 	"fmt"
 
+	"github.com/apache/yunikorn-core/pkg/locking"
+
 	"github.com/apache/yunikorn-scheduler-interface/lib/go/si"
 )
 
@@ -31,6 +33,8 @@ type PreemptionPredicatePlugin struct {
 	allocations  map[string]string
 	preemptions  []Preemption
 	errHolder    *errHolder
+
+	locking.RWMutex
 }
 
 type Preemption struct {
@@ -47,6 +51,8 @@ type errHolder struct {
 }
 
 func (m *PreemptionPredicatePlugin) Predicates(args *si.PredicatesArgs) error {
+	m.RLock()
+	defer m.RUnlock()
 	if args.Allocate {
 		nodeID, ok := m.allocations[args.AllocationKey]
 		if !ok {
@@ -69,6 +75,8 @@ func (m *PreemptionPredicatePlugin) Predicates(args *si.PredicatesArgs) error {
 }
 
 func (m *PreemptionPredicatePlugin) PreemptionPredicates(args *si.PreemptionPredicatesArgs) *si.PreemptionPredicatesResponse {
+	m.Lock()
+	defer m.Unlock()
 	result := &si.PreemptionPredicatesResponse{
 		Success: false,
 		Index:   -1,
@@ -109,6 +117,8 @@ func (m *PreemptionPredicatePlugin) PreemptionPredicates(args *si.PreemptionPred
 // GetPredicateError returns the error set by the preemption predicate check that failed.
 // Returns a nil error on success.
 func (m *PreemptionPredicatePlugin) GetPredicateError() error {
+	m.RLock()
+	defer m.RUnlock()
 	return m.errHolder.err
 }