fix: don't rank slots when sticky == NONE

hatchet-dev · Jan 28, 2025 · 77d69a0 · 77d69a0
1 parent 2765abd
commit 77d69a0
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 20 deletions.
diff --git a/pkg/repository/v2/sqlcv2/models.go b/pkg/repository/v2/sqlcv2/models.go
diff --git a/pkg/scheduling/v2/scheduler.go b/pkg/scheduling/v2/scheduler.go
@@ -585,7 +585,7 @@ func (s *Scheduler) tryAssignSingleton(
 	ctx, span := telemetry.NewSpan(ctx, "try-assign-singleton") // nolint: ineffassign
 	defer span.End()
 
-	if qi.Sticky.Valid || len(labels) > 0 {
+	if (qi.Sticky != sqlcv2.V2StickyStrategyNONE) || len(labels) > 0 {
 		candidateSlots = getRankedSlots(qi, labels, candidateSlots)
 	}
 

diff --git a/pkg/scheduling/v2/slot.go b/pkg/scheduling/v2/slot.go
@@ -188,7 +188,7 @@ func getRankedSlots(
 
 		// if this is a HARD sticky strategy, and there's a desired worker id, it can only be assigned to that
 		// worker. if there's no desired worker id, we assign to any worker.
-		if qi.Sticky.Valid && qi.Sticky.StickyStrategy == sqlcv2.StickyStrategyHARD {
+		if qi.Sticky == sqlcv2.V2StickyStrategyHARD {
 			if qi.DesiredWorkerID.Valid && workerId == sqlchelpers.UUIDToStr(qi.DesiredWorkerID) {
 				validSlots.addSlot(slot, 0)
 			} else if !qi.DesiredWorkerID.Valid {
@@ -200,7 +200,7 @@ func getRankedSlots(
 
 		// if this is a SOFT sticky strategy, we should prefer the desired worker, but if it is not
 		// available, we can assign to any worker.
-		if qi.Sticky.Valid && qi.Sticky.StickyStrategy == sqlcv2.StickyStrategySOFT {
+		if qi.Sticky == sqlcv2.V2StickyStrategySOFT {
 			if qi.DesiredWorkerID.Valid && workerId == sqlchelpers.UUIDToStr(qi.DesiredWorkerID) {
 				validSlots.addSlot(slot, 1)
 			} else {

diff --git a/pkg/scheduling/v2/slot_test.go b/pkg/scheduling/v2/slot_test.go
@@ -27,7 +27,7 @@ func TestGetRankedSlots(t *testing.T) {
 		{
 			name: "HARD sticky strategy with desired worker available",
 			qi: &sqlcv2.V2QueueItem{
-				Sticky:          sqlcv2.NullStickyStrategy{Valid: true, StickyStrategy: sqlcv2.StickyStrategyHARD},
+				Sticky:          sqlcv2.V2StickyStrategyHARD,
 				DesiredWorkerID: sqlchelpers.UUIDFromStr(stableWorkerId1),
 			},
 			slots: []*slot{
@@ -39,7 +39,7 @@ func TestGetRankedSlots(t *testing.T) {
 		{
 			name: "HARD sticky strategy without desired worker",
 			qi: &sqlcv2.V2QueueItem{
-				Sticky:          sqlcv2.NullStickyStrategy{Valid: true, StickyStrategy: sqlcv2.StickyStrategyHARD},
+				Sticky:          sqlcv2.V2StickyStrategyHARD,
 				DesiredWorkerID: sqlchelpers.UUIDFromStr(uuid.New().String()),
 			},
 			slots: []*slot{
@@ -51,7 +51,7 @@ func TestGetRankedSlots(t *testing.T) {
 		{
 			name: "SOFT sticky strategy with desired worker available",
 			qi: &sqlcv2.V2QueueItem{
-				Sticky:          sqlcv2.NullStickyStrategy{Valid: true, StickyStrategy: sqlcv2.StickyStrategySOFT},
+				Sticky:          sqlcv2.V2StickyStrategySOFT,
 				DesiredWorkerID: sqlchelpers.UUIDFromStr(stableWorkerId1),
 			},
 			slots: []*slot{