This is an automated email from the ASF dual-hosted git repository.
manirajv06 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/yunikorn-core.git
The following commit(s) were added to refs/heads/master by this push:
new 784bdc46 [YUNIKORN-3325] Deduplicate cancel reservations code flow
(#1116)
784bdc46 is described below
commit 784bdc46d96ad7700cf33d359e7679c1f403c7d1
Author: kaijaytu <[email protected]>
AuthorDate: Tue Aug 4 11:53:20 2026 +0530
[YUNIKORN-3325] Deduplicate cancel reservations code flow (#1116)
Extract unreserveForApp() and cancelMatchingReservations() from the
duplicated reservation cancellation logic in application.go and
preemption.go.
Both cancelReservations (tryRequiredNode path) and initWorkingState
(preemption path) shared the same lock-aware unreserve pattern that
branches on whether the reservation belongs to the calling app.
Consolidate into a single predicate-based method that both call sites
now use with their own filter logic.
Closes: #1116
Signed-off-by: mani <[email protected]>
---
pkg/scheduler/objects/application.go | 60 ++++++++++++-------
pkg/scheduler/objects/application_test.go | 98 +++++++++++++++++++++++++++++++
pkg/scheduler/objects/preemption.go | 39 ++++--------
3 files changed, 146 insertions(+), 51 deletions(-)
diff --git a/pkg/scheduler/objects/application.go
b/pkg/scheduler/objects/application.go
index d5b408cf..a125e0ac 100644
--- a/pkg/scheduler/objects/application.go
+++ b/pkg/scheduler/objects/application.go
@@ -1227,33 +1227,49 @@ func (sa *Application) tryRequiredNode(request
*Allocation, getNodeFn func(strin
return result
}
-// cancelReservations will cancel all non required node reservations for a
node. The list of reservations passed in is
-// a copy of all reservations of a single node. This is called during the
required node allocation cycle only.
-// The returned int value is used to update the partition counter of active
reservations.
-func (sa *Application) cancelReservations(reservations []*reservation) int {
- var released, num int
- // un reserve all the apps that were reserved on the node
+// unreserveForApp handles the lock-aware unreserve for a single reservation.
+// Uses the internal unlocked path when the reservation belongs to this app.
+func (sa *Application) unreserveForApp(res *reservation) int {
+ var num int
+ if sa.ApplicationID == res.appID {
+ num = sa.unReserveInternal(res)
+ sa.queue.UnReserve(sa.ApplicationID, num)
+ } else {
+ num = res.app.UnReserve(res.node, res.alloc)
+ res.app.GetQueue().UnReserve(res.app.ApplicationID, num)
+ }
+ if num > 0 {
+ log.Log(log.SchedApplication).Info("Reservation cancelled",
+ zap.String("triggeringAppID", sa.ApplicationID),
+ zap.String("reservingAppID", res.appID),
+ zap.String("reservingAllocationKey", res.allocKey),
+ zap.String("node", res.nodeID))
+ }
+ return num
+}
+
+// cancelMatchingReservations cancels reservations that match the predicate.
+// Returns the number of reservations released and the number remaining.
+func (sa *Application) cancelMatchingReservations(reservations []*reservation,
shouldCancel func(*reservation) bool) (released, remaining int) {
+ remaining = len(reservations)
for _, res := range reservations {
- // cleanup if the reservation does not have this node as a
requirement
- if res.alloc.requiredNode != "" {
+ if !shouldCancel(res) {
continue
}
- thisApp := res.app.ApplicationID == sa.ApplicationID
- if thisApp {
- num = sa.unReserveInternal(res)
- sa.queue.UnReserve(sa.ApplicationID, num)
- } else {
- num = res.app.UnReserve(res.node, res.alloc)
- res.app.GetQueue().UnReserve(res.app.ApplicationID, num)
- }
- log.Log(log.SchedApplication).Info("Cancelled reservation for
required node allocation",
- zap.String("triggered by appID", sa.ApplicationID),
- zap.String("affected application ID", res.appID),
- zap.String("affected allocationKey", res.allocKey),
- zap.String("required node", res.nodeID),
- zap.Int("reservations count", num))
+ num := sa.unreserveForApp(res)
released += num
+ remaining -= num
}
+ return
+}
+
+// cancelReservations will cancel all non required node reservations for a
node. The list of reservations passed in is
+// a copy of all reservations of a single node. This is called during the
required node allocation cycle only.
+// The returned int value is used to update the partition counter of active
reservations.
+func (sa *Application) cancelReservations(reservations []*reservation) int {
+ released, _ := sa.cancelMatchingReservations(reservations, func(res
*reservation) bool {
+ return res.alloc.requiredNode == ""
+ })
return released
}
diff --git a/pkg/scheduler/objects/application_test.go
b/pkg/scheduler/objects/application_test.go
index fe65d485..0f237df6 100644
--- a/pkg/scheduler/objects/application_test.go
+++ b/pkg/scheduler/objects/application_test.go
@@ -4332,3 +4332,101 @@ func TestReservationReleasedCallback(t *testing.T) {
t.Fatal("callback was not called")
}
}
+
+func TestUnreserveForApp(t *testing.T) {
+ resMap := map[string]string{"first": "10"}
+ rootQ, err := createRootQueue(resMap)
+ assert.NilError(t, err)
+ childQ, err := createManagedQueue(rootQ, "child", false, resMap)
+ assert.NilError(t, err)
+
+ node := newNode(nodeID1, map[string]resources.Quantity{"first": 10})
+ allocRes :=
resources.NewResourceFromMap(map[string]resources.Quantity{"first": 3})
+
+ // app1 holds the reservation, app2 is the triggering app
+ app1 := newApplication(appID1, "default", childQ.QueuePath)
+ app1.SetQueue(childQ)
+ childQ.applications[appID1] = app1
+ ask1 := newAllocationAsk(aKey, appID1, allocRes)
+ err = app1.AddAllocationAsk(ask1)
+ assert.NilError(t, err)
+ err = app1.reserveInternal(node, ask1)
+ assert.NilError(t, err)
+
+ app2 := newApplication(appID2, "default", childQ.QueuePath)
+ app2.SetQueue(childQ)
+ childQ.applications[appID2] = app2
+
+ // different app: unreserveForApp uses the locked UnReserve path
+ res := app1.reservations[aKey]
+ assert.Assert(t, res != nil)
+ num := app2.unreserveForApp(res)
+ assert.Equal(t, num, 1, "expected reservation to be released")
+ assert.Equal(t, app1.NodeReservedForAsk(aKey), "", "reservation should
be removed from app1")
+
+ // same app: re-reserve and use same-app path
+ err = app1.reserveInternal(node, ask1)
+ assert.NilError(t, err)
+ res = app1.reservations[aKey]
+ assert.Assert(t, res != nil)
+ num = app1.unreserveForApp(res)
+ assert.Equal(t, num, 1, "expected reservation to be released via
internal path")
+ assert.Equal(t, app1.NodeReservedForAsk(aKey), "", "reservation should
be removed")
+}
+
+func TestCancelMatchingReservations(t *testing.T) {
+ resMap := map[string]string{"first": "10"}
+ rootQ, err := createRootQueue(resMap)
+ assert.NilError(t, err)
+ childQ, err := createManagedQueue(rootQ, "child", false, resMap)
+ assert.NilError(t, err)
+
+ node1 := newNode(nodeID1, map[string]resources.Quantity{"first": 10})
+ node2 := newNode(nodeID2, map[string]resources.Quantity{"first": 10})
+ allocRes :=
resources.NewResourceFromMap(map[string]resources.Quantity{"first": 2})
+
+ app := newApplication(appID1, "default", childQ.QueuePath)
+ app.SetQueue(childQ)
+ childQ.applications[appID1] = app
+
+ // reserve ask1 on node1, ask2 on node2
+ ask1 := newAllocationAsk(aKey, appID1, allocRes)
+ err = app.AddAllocationAsk(ask1)
+ assert.NilError(t, err)
+ err = app.reserveInternal(node1, ask1)
+ assert.NilError(t, err)
+
+ ask2 := newAllocationAsk(aKey2, appID1, allocRes)
+ err = app.AddAllocationAsk(ask2)
+ assert.NilError(t, err)
+ err = app.reserveInternal(node2, ask2)
+ assert.NilError(t, err)
+
+ // gather both reservations
+ reservations := append(node1.GetReservations(),
node2.GetReservations()...)
+ assert.Equal(t, len(reservations), 2, "expected 2 reservations total")
+
+ // cancel only the first ask's reservation using predicate on allocKey
+ released, remaining := app.cancelMatchingReservations(reservations,
func(res *reservation) bool {
+ return res.allocKey == aKey
+ })
+ assert.Equal(t, released, 1, "expected 1 reservation released")
+ assert.Equal(t, remaining, 1, "expected 1 reservation remaining")
+ assert.Equal(t, app.NodeReservedForAsk(aKey), "", "first reservation
should be gone")
+ assert.Equal(t, app.NodeReservedForAsk(aKey2), nodeID2, "second
reservation should remain")
+
+ // cancel all remaining
+ reservations = node2.GetReservations()
+ released, remaining = app.cancelMatchingReservations(reservations,
func(res *reservation) bool {
+ return true
+ })
+ assert.Equal(t, released, 1, "expected 1 reservation released")
+ assert.Equal(t, remaining, 0, "expected no reservations remaining")
+
+ // empty list returns zero
+ released, remaining = app.cancelMatchingReservations(nil, func(res
*reservation) bool {
+ return true
+ })
+ assert.Equal(t, released, 0)
+ assert.Equal(t, remaining, 0)
+}
diff --git a/pkg/scheduler/objects/preemption.go
b/pkg/scheduler/objects/preemption.go
index dc791184..211e883b 100644
--- a/pkg/scheduler/objects/preemption.go
+++ b/pkg/scheduler/objects/preemption.go
@@ -165,40 +165,21 @@ func (p *Preemptor) initWorkingState() int {
p.iterator.ForEachNode(func(node *Node) bool {
isReserved := false
if node.IsReserved() &&
!node.isReservedForAllocation(p.ask.GetAllocationKey()) {
- leftCount := 0
- for _, res := range node.GetReservations() {
- leftCount++
- // Is Allocation daemon set?
- // Has this allocation already triggered
preemption?
+ askPriority := p.ask.priority
+ released, remaining :=
p.application.cancelMatchingReservations(node.GetReservations(), func(res
*reservation) bool {
if res.alloc.requiredNode != "" ||
res.alloc.HasTriggeredPreemption() {
- continue
+ return false
}
- // Cancel reservation based on its priority and
waiting time in reservation queue
- if res.alloc.GetPriority() < p.ask.priority &&
time.Since(res.createTime) > reservationWaitTimeout {
-
log.Log(log.SchedPreemption).Info("Cancelling reservation to consider node for
preemption",
- zap.String("triggeringAppID",
p.application.ApplicationID),
-
zap.String("triggeringAllocationKey", p.ask.allocationKey),
- zap.String("reservingAppID",
res.appID),
-
zap.String("reservingAllocationKey", res.allocKey),
- zap.String("node", node.NodeID))
- num := 0
- if p.application.ApplicationID ==
res.appID {
- num =
res.app.unReserveInternal(res)
-
res.app.queue.UnReserve(res.app.ApplicationID, num)
- } else {
- num =
res.app.UnReserve(res.node, res.alloc)
-
res.app.GetQueue().UnReserve(res.app.ApplicationID, num)
- }
- totalReservationCancel += num
- leftCount -= num
- }
- }
- log.Log(log.SchedPreemption).Debug("Reservations left
on node are cleanup",
+ return res.alloc.GetPriority() < askPriority &&
time.Since(res.createTime) > reservationWaitTimeout
+ })
+ totalReservationCancel += released
+ log.Log(log.SchedPreemption).Debug("Reservations left
on node after cleanup",
zap.String("triggeringAppID",
p.application.ApplicationID),
zap.String("triggeringAllocationKey",
p.ask.allocationKey),
zap.String("node", node.NodeID),
- zap.Int("leftCount", leftCount))
- isReserved = leftCount > 0
+ zap.Int("released", released),
+ zap.Int("remaining", remaining))
+ isReserved = remaining > 0
}
if !node.IsSchedulable() || isReserved ||
!node.FitInNode(p.ask.GetAllocatedResource()) {
// node is not available, remove any potential victims
from consideration
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
