Repository: armadaproject/armada
Pull request #3540 (2024-04-25 16:53:08)

Improve correctness of cancel + reprioritise by jobIds
Now only cancel/reprioritise jobs that belong to the queue/jobset of the request





Patch:
diff --git a/internal/scheduler/database/query.sql.go b/internal/scheduler/database/query.sql.go
index d9675115..593e52d6 100644
--- a/internal/scheduler/database/query.sql.go
+++ b/internal/scheduler/database/query.sql.go
@@ -133,11 +133,17 @@ func (q *Queries) MarkJobRunsSucceededById(ctx context.Context, runIds []uuid.UU
 }
 
 const markJobsCancelRequestedById = `-- name: MarkJobsCancelRequestedById :exec
-UPDATE jobs SET cancel_requested = true WHERE job_id = ANY($1::text[])
+UPDATE jobs SET cancel_requested = true WHERE queue = $1 and job_set = $2 and job_id = ANY($3::text[])
 `
 
-func (q *Queries) MarkJobsCancelRequestedById(ctx context.Context, jobIds []string) error {
-	_, err := q.db.Exec(ctx, markJobsCancelRequestedById, jobIds)
+type MarkJobsCancelRequestedByIdParams struct {
+	Queue  string   `db:"queue"`
+	JobSet string   `db:"job_set"`
+	JobIds []string `db:"job_ids"`
+}
+
+func (q *Queries) MarkJobsCancelRequestedById(ctx context.Context, arg MarkJobsCancelRequestedByIdParams) error {
+	_, err := q.db.Exec(ctx, markJobsCancelRequestedById, arg.Queue, arg.JobSet, arg.JobIds)
 	return err
 }
 
@@ -708,16 +714,23 @@ func (q *Queries) SetTerminatedTime(ctx context.Context, arg SetTerminatedTimePa
 }
 
 const updateJobPriorityById = `-- name: UpdateJobPriorityById :exec
-UPDATE jobs SET priority = $1 WHERE job_id = $2
+UPDATE jobs SET priority = $1 WHERE queue = $2 and job_set = $3 and job_id = ANY($4::text[])
 `
 
 type UpdateJobPriorityByIdParams struct {
-	Priority int64  `db:"priority"`
-	JobID    string `db:"job_id"`
+	Priority int64    `db:"priority"`
+	Queue    string   `db:"queue"`
+	JobSet   string   `db:"job_set"`
+	JobIds   []string `db:"job_ids"`
 }
 
 func (q *Queries) UpdateJobPriorityById(ctx context.Context, arg UpdateJobPriorityByIdParams) error {
-	_, err := q.db.Exec(ctx, updateJobPriorityById, arg.Priority, arg.JobID)
+	_, err := q.db.Exec(ctx, updateJobPriorityById,
+		arg.Priority,
+		arg.Queue,
+		arg.JobSet,
+		arg.JobIds,
+	)
 	return err
 }
 
diff --git a/internal/scheduler/database/query/query.sql b/internal/scheduler/database/query/query.sql
index 6b582ccb..f19bd5e7 100644
--- a/internal/scheduler/database/query/query.sql
+++ b/internal/scheduler/database/query/query.sql
@@ -17,7 +17,7 @@ UPDATE jobs SET cancel_by_jobset_requested = true WHERE job_set = sqlc.arg(job_s
 UPDATE jobs SET succeeded = true WHERE job_id = ANY(sqlc.arg(job_ids)::text[]);
 
 -- name: MarkJobsCancelRequestedById :exec
-UPDATE jobs SET cancel_requested = true WHERE job_id = ANY(sqlc.arg(job_ids)::text[]);
+UPDATE jobs SET cancel_requested = true WHERE queue = sqlc.arg(queue) and job_set = sqlc.arg(job_set) and job_id = ANY(sqlc.arg(job_ids)::text[]);
 
 -- name: MarkJobsCancelledById :exec
 UPDATE jobs SET cancelled = true WHERE job_id = ANY(sqlc.arg(job_ids)::text[]);
@@ -26,7 +26,7 @@ UPDATE jobs SET cancelled = true WHERE job_id = ANY(sqlc.arg(job_ids)::text[]);
 UPDATE jobs SET failed = true WHERE job_id = ANY(sqlc.arg(job_ids)::text[]);
 
 -- name: UpdateJobPriorityById :exec
-UPDATE jobs SET priority = $1 WHERE job_id = $2;
+UPDATE jobs SET priority = $1 WHERE queue = sqlc.arg(queue) and job_set = sqlc.arg(job_set) and job_id = ANY(sqlc.arg(job_ids)::text[]);
 
 -- name: SelectNewRuns :many
 SELECT * FROM runs WHERE serial > $1 ORDER BY serial LIMIT $2;
diff --git a/internal/scheduleringester/dbops.go b/internal/scheduleringester/dbops.go
index e5330c0f..6f6d00e6 100644
--- a/internal/scheduleringester/dbops.go
+++ b/internal/scheduleringester/dbops.go
@@ -7,6 +7,7 @@ import (
 	"github.com/google/uuid"
 	"golang.org/x/exp/maps"
 
+	"github.com/armadaproject/armada/internal/common/slices"
 	schedulerdb "github.com/armadaproject/armada/internal/scheduler/database"
 )
 
@@ -42,6 +43,11 @@ type JobSetKey struct {
 	jobSet string
 }
 
+type JobReprioritiseKey struct {
+	JobSetKey
+	Priority int64
+}
+
 type JobRunDetails struct {
 	Queue string
 	DbRun *schedulerdb.Run
@@ -120,11 +126,10 @@ type (
 	InsertRuns                     map[uuid.UUID]*JobRunDetails
 	UpdateJobSetPriorities         map[JobSetKey]int64
 	MarkJobSetsCancelRequested     map[JobSetKey]*JobSetCancelAction
-	MarkJobsCancelRequested        map[string]bool
+	MarkJobsCancelRequested        map[JobSetKey][]string
 	MarkJobsCancelled              map[string]time.Time
 	MarkJobsSucceeded              map[string]bool
 	MarkJobsFailed                 map[string]bool
-	UpdateJobPriorities            map[string]int64
 	UpdateJobSchedulingInfo        map[string]*JobSchedulingInfoUpdate
 	UpdateJobQueuedState           map[string]*JobQueuedStateUpdate
 	MarkRunsSucceeded              map[uuid.UUID]time.Time
@@ -134,8 +139,12 @@ type (
 	MarkRunsPending                map[uuid.UUID]time.Time
 	MarkRunsPreempted              map[uuid.UUID]time.Time
 	InsertJobRunErrors             map[uuid.UUID]*schedulerdb.JobRunError
-	MarkJobsValidated              map[string]bool
-	InsertPartitionMarker          struct {
+	UpdateJobPriorities            struct {
+		key    JobReprioritiseKey
+		jobIds []string
+	}
+	MarkJobsValidated     map[string]bool
+	InsertPartitionMarker struct {
 		markers []*schedulerdb.Marker
 	}
 )
@@ -171,7 +180,7 @@ func (a MarkJobSetsCancelRequested) Merge(b DbOperation) bool {
 }
 
 func (a MarkJobsCancelRequested) Merge(b DbOperation) bool {
-	return mergeInMap(a, b)
+	return mergeListMaps(a, b)
 }
 
 func (a MarkRunsForJobPreemptRequested) Merge(b DbOperation) bool {
@@ -227,7 +236,14 @@ func (a MarkJobsFailed) Merge(b DbOperation) bool {
 }
 
 func (a UpdateJobPriorities) Merge(b DbOperation) bool {
-	return mergeInMap(a, b)
+	switch op := b.(type) {
+	case UpdateJobPriorities:
+		if a.key == op.key {
+			a.jobIds = slices.Unique(append(a.jobIds, op.jobIds...))
+			return true
+		}
+	}
+	return false
 }
 
 func (a MarkRunsSucceeded) Merge(b DbOperation) bool {
@@ -346,7 +362,7 @@ func (a MarkJobSetsCancelRequested) CanBeAppliedBefore(b DbOperation) bool {
 }
 
 func (a MarkJobsCancelRequested) CanBeAppliedBefore(b DbOperation) bool {
-	return !definesJob(a, b) && !definesRunForJob(a, b)
+	return !definesJobInSet(a, b) && !definesRunInSet(a, b)
 }
 
 func (a MarkRunsForJobPreemptRequested) CanBeAppliedBefore(b DbOperation) bool {
@@ -375,7 +391,10 @@ func (a UpdateJobQueuedState) CanBeAppliedBefore(b DbOperation) bool {
 
 func (a UpdateJobPriorities) CanBeAppliedBefore(b DbOperation) bool {
 	_, isUpdateJobSetPriorities := b.(UpdateJobSetPriorities)
-	return !isUpdateJobSetPriorities && !definesJob(a, b)
+	_, isUpdateJobPriorities := b.(UpdateJobPriorities)
+	return !isUpdateJobPriorities && !isUpdateJobSetPriorities &&
+		!definesJobInSet(map[JobSetKey]bool{a.key.JobSetKey: true}, b) &&
+		!definesRunInSet(map[JobSetKey]bool{a.key.JobSetKey: true}, b)
 }
 
 func (a MarkRunsSucceeded) CanBeAppliedBefore(b DbOperation) bool {
@@ -464,16 +483,3 @@ func definesRun[M ~map[uuid.UUID]V, V any](a M, b DbOperation) bool {
 	}
 	return false
 }
-
-// definesRunForJob returns true if b is an InsertRuns operation
-// that inserts at least one run with job id equal to any of the keys of a.
-func definesRunForJob[M ~map[string]V, V any](a M, b DbOperation) bool {
-	if op, ok := b.(InsertRuns); ok {
-		for _, run := range op {
-			if _, ok := a[run.DbRun.JobID]; ok {
-				return true
-			}
-		}
-	}
-	return false
-}
diff --git a/internal/scheduleringester/instructions.go b/internal/scheduleringester/instructions.go
index b732235f..431c2dc9 100644
--- a/internal/scheduleringester/instructions.go
+++ b/internal/scheduleringester/instructions.go
@@ -96,11 +96,11 @@ func (c *InstructionConverter) dbOperationsFromEventSequence(es *armadaevents.Ev
 		case *armadaevents.EventSequence_Event_JobPreemptionRequested:
 			operationsFromEvent, err = c.handleJobPreemptionRequested(event.GetJobPreemptionRequested(), meta)
 		case *armadaevents.EventSequence_Event_ReprioritiseJob:
-			operationsFromEvent, err = c.handleReprioritiseJob(event.GetReprioritiseJob())
+			operationsFromEvent, err = c.handleReprioritiseJob(event.GetReprioritiseJob(), meta)
 		case *armadaevents.EventSequence_Event_ReprioritiseJobSet:
 			operationsFromEvent, err = c.handleReprioritiseJobSet(event.GetReprioritiseJobSet(), meta)
 		case *armadaevents.EventSequence_Event_CancelJob:
-			operationsFromEvent, err = c.handleCancelJob(event.GetCancelJob())
+			operationsFromEvent, err = c.handleCancelJob(event.GetCancelJob(), meta)
 		case *armadaevents.EventSequence_Event_CancelJobSet:
 			operationsFromEvent, err = c.handleCancelJobSet(event.GetCancelJobSet(), meta)
 		case *armadaevents.EventSequence_Event_CancelledJob:
@@ -339,13 +339,20 @@ func (c *InstructionConverter) handleJobPreemptionRequested(preemptionRequested
 	}}, nil
 }
 
-func (c *InstructionConverter) handleReprioritiseJob(reprioritiseJob *armadaevents.ReprioritiseJob) ([]DbOperation, error) {
+func (c *InstructionConverter) handleReprioritiseJob(reprioritiseJob *armadaevents.ReprioritiseJob, meta eventSequenceCommon) ([]DbOperation, error) {
 	jobId, err := armadaevents.UlidStringFromProtoUuid(reprioritiseJob.GetJobId())
 	if err != nil {
 		return nil, err
 	}
 	return []DbOperation{UpdateJobPriorities{
-		jobId: int64(reprioritiseJob.Priority),
+		key: JobReprioritiseKey{
+			JobSetKey: JobSetKey{
+				queue:  meta.queue,
+				jobSet: meta.jobset,
+			},
+			Priority: int64(reprioritiseJob.Priority),
+		},
+		jobIds: []string{jobId},
 	}}, nil
 }
 
@@ -355,13 +362,16 @@ func (c *InstructionConverter) handleReprioritiseJobSet(reprioritiseJobSet *arma
 	}}, nil
 }
 
-func (c *InstructionConverter) handleCancelJob(cancelJob *armadaevents.CancelJob) ([]DbOperation, error) {
+func (c *InstructionConverter) handleCancelJob(cancelJob *armadaevents.CancelJob, meta eventSequenceCommon) ([]DbOperation, error) {
 	jobId, err := armadaevents.UlidStringFromProtoUuid(cancelJob.GetJobId())
 	if err != nil {
 		return nil, err
 	}
 	return []DbOperation{MarkJobsCancelRequested{
-		jobId: true,
+		JobSetKey{
+			queue:  meta.queue,
+			jobSet: meta.jobset,
+		}: []string{jobId},
 	}}, nil
 }
 
diff --git a/internal/scheduleringester/schedulerdb.go b/internal/scheduleringester/schedulerdb.go
index 3573304b..bbe868de 100644
--- a/internal/scheduleringester/schedulerdb.go
+++ b/internal/scheduleringester/schedulerdb.go
@@ -174,10 +174,16 @@ func (s *SchedulerDb) WriteDbOp(ctx *armadacontext.Context, tx pgx.Tx, op DbOper
 			}
 		}
 	case MarkJobsCancelRequested:
-		jobIds := maps.Keys(o)
-		err := queries.MarkJobsCancelRequestedById(ctx, jobIds)
-		if err != nil {
-			return errors.WithStack(err)
+		for key, value := range o {
+			params := schedulerdb.MarkJobsCancelRequestedByIdParams{
+				Queue:  key.queue,
+				JobSet: key.jobSet,
+				JobIds: value,
+			}
+			err := queries.MarkJobsCancelRequestedById(ctx, params)
+			if err != nil {
+				return errors.WithStack(err)
+			}
 		}
 	case MarkJobsCancelled:
 		jobIds := maps.Keys(o)
@@ -220,16 +226,14 @@ func (s *SchedulerDb) WriteDbOp(ctx *armadacontext.Context, tx pgx.Tx, op DbOper
 			return errors.WithStack(err)
 		}
 	case UpdateJobPriorities:
-		// TODO: This will be slow if there's a large number of ids.
-		// Could be addressed by using a separate table for priority + upsert.
-		for jobId, priority := range o {
-			err := queries.UpdateJobPriorityById(ctx, schedulerdb.UpdateJobPriorityByIdParams{
-				JobID:    jobId,
-				Priority: priority,
-			})
-			if err != nil {
-				return errors.WithStack(err)
-			}
+		err := queries.UpdateJobPriorityById(ctx, schedulerdb.UpdateJobPriorityByIdParams{
+			Queue:    o.key.queue,
+			JobSet:   o.key.jobSet,
+			Priority: o.key.Priority,
+			JobIds:   o.jobIds,
+		})
+		if err != nil {
+			return errors.WithStack(err)
 		}
 	case MarkRunsSucceeded:
 		successTimes := make([]interface{}, 0, len(o))


Test patch:
diff --git a/internal/scheduleringester/dbops_test.go b/internal/scheduleringester/dbops_test.go
index 552f24a2..a666d3f9 100644
--- a/internal/scheduleringester/dbops_test.go
+++ b/internal/scheduleringester/dbops_test.go
@@ -19,17 +19,17 @@ func TestMerge(t *testing.T) {
 	jobId1 := util.NewULID()
 	jobId2 := util.NewULID()
 	jobId3 := util.NewULID()
-	markJobsCancelled1 := MarkJobsCancelRequested{jobId1: false, jobId2: false}
-	markJobsCancelled2 := MarkJobsCancelRequested{jobId2: true, jobId3: true}
-	ok := markJobsCancelled1.Merge(markJobsCancelled2)
+	markJobsFailed1 := MarkJobsFailed{jobId1: false, jobId2: false}
+	markJobsFailed2 := MarkJobsFailed{jobId2: true, jobId3: true}
+	ok := markJobsFailed1.Merge(markJobsFailed2)
 	assert.True(t, ok)
-	assert.Equal(t, MarkJobsCancelRequested{jobId1: false, jobId2: true, jobId3: true}, markJobsCancelled1)
+	assert.Equal(t, MarkJobsFailed{jobId1: false, jobId2: true, jobId3: true}, markJobsFailed1)
 
 	jobId4 := util.NewULID()
 	markJobsSucceeded1 := MarkJobsSucceeded{jobId1: true, jobId4: true}
-	ok = markJobsCancelled1.Merge(markJobsSucceeded1)
+	ok = markJobsFailed1.Merge(markJobsSucceeded1)
 	assert.False(t, ok)
-	assert.Equal(t, MarkJobsCancelRequested{jobId1: false, jobId2: true, jobId3: true}, markJobsCancelled1)
+	assert.Equal(t, MarkJobsFailed{jobId1: false, jobId2: true, jobId3: true}, markJobsFailed1)
 }
 
 func TestMerge_MarkRunsForJobPreemptRequested(t *testing.T) {
@@ -159,15 +159,15 @@ func TestDbOperationOptimisation(t *testing.T) {
 			UpdateJobSetPriorities{JobSetKey{queue: testQueueName, jobSet: "set2"}: 2},                      // 3
 			InsertJobs{jobIds[2]: &schedulerdb.Job{JobID: jobIds[2], Queue: testQueueName, JobSet: "set1"}}, // 3
 		}},
-		"UpdateJobSetPriorities, UpdateJobPriorities": {N: 4, Ops: []DbOperation{
-			InsertJobs{jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], Queue: testQueueName, JobSet: "set1"}}, // 1
-			InsertJobs{jobIds[1]: &schedulerdb.Job{JobID: jobIds[1], Queue: testQueueName, JobSet: "set1"}}, // 1
-			UpdateJobPriorities{jobIds[0]: 1},                                                               // 2
-			UpdateJobSetPriorities{JobSetKey{queue: testQueueName, jobSet: "set1"}: 2},                      // 3
-			UpdateJobPriorities{jobIds[1]: 3},                                                               // 4
-			InsertJobs{jobIds[2]: &schedulerdb.Job{JobID: jobIds[2], Queue: testQueueName, JobSet: "set2"}}, // 4
-			UpdateJobPriorities{jobIds[1]: 4},                                                               // 4
-			UpdateJobPriorities{jobIds[2]: 5},                                                               // 4
+		"UpdateJobSetPriorities, UpdateJobPriorities": {N: 5, Ops: []DbOperation{
+			InsertJobs{jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], Queue: testQueueName, JobSet: "set1"}},                  // 1
+			InsertJobs{jobIds[1]: &schedulerdb.Job{JobID: jobIds[1], Queue: testQueueName, JobSet: "set1"}},                  // 1
+			UpdateJobPriorities{JobReprioritiseKey{JobSetKey{queue: testQueueName, jobSet: "set1"}, 1}, []string{jobIds[0]}}, // 2                                                            // 2
+			UpdateJobSetPriorities{JobSetKey{queue: testQueueName, jobSet: "set1"}: 2},                                       // 3
+			UpdateJobPriorities{JobReprioritiseKey{JobSetKey{queue: testQueueName, jobSet: "set1"}, 3}, []string{jobIds[1]}}, // 4
+			InsertJobs{jobIds[2]: &schedulerdb.Job{JobID: jobIds[2], Queue: testQueueName, JobSet: "set2"}},                  // 1
+			UpdateJobPriorities{JobReprioritiseKey{JobSetKey{queue: testQueueName, jobSet: "set1"}, 4}, []string{jobIds[1]}}, // 5
+			UpdateJobPriorities{JobReprioritiseKey{JobSetKey{queue: testQueueName, jobSet: "set1"}, 4}, []string{jobIds[2]}}, // 5
 		}},
 		"MarkJobSetsCancelRequested": {N: 3, Ops: []DbOperation{
 			InsertJobs{jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], Queue: testQueueName, JobSet: "set1"}},                                          // 1
@@ -179,11 +179,11 @@ func TestDbOperationOptimisation(t *testing.T) {
 		"MarkJobSetsCancelRequested, MarkJobsCancelRequested": {N: 4, Ops: []DbOperation{
 			InsertJobs{jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], Queue: testQueueName, JobSet: "set1"}},                                          // 1
 			InsertJobs{jobIds[1]: &schedulerdb.Job{JobID: jobIds[1], Queue: testQueueName, JobSet: "set1"}},                                          // 1
-			MarkJobsCancelRequested{jobIds[0]: true},                                                                                                 // 2
+			MarkJobsCancelRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: []string{jobIds[0]}},                                            // 2                                                     // 2
 			MarkJobSetsCancelRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: &JobSetCancelAction{cancelQueued: true, cancelLeased: true}}, // 3
 			InsertJobs{jobIds[2]: &schedulerdb.Job{JobID: jobIds[2], Queue: testQueueName, JobSet: "set1"}},                                          // 4
-			MarkJobsCancelRequested{jobIds[1]: true},                                                                                                 // 4
-			MarkJobsCancelRequested{jobIds[2]: true},                                                                                                 // 4
+			MarkJobsCancelRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: []string{jobIds[1]}},                                            // 4                                              // 4
+			MarkJobsCancelRequested{JobSetKey{queue: testQueueName, jobSet: "set1"}: []string{jobIds[2]}},                                            // 4                                                 // 4
 		}},
 		"MarkRunsForJobPreemptRequested": {N: 2, Ops: []DbOperation{
 			InsertJobs{jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], Queue: testQueueName, JobSet: "set1"}},      // 1
@@ -421,11 +421,15 @@ func (db *mockDb) apply(op DbOperation) error {
 			}
 		}
 	case MarkJobsCancelRequested:
-		for jobId := range o {
-			if job, ok := db.Jobs[jobId]; ok {
-				job.CancelRequested = true
-			} else {
-				return errors.Errorf("job %s not in db", jobId)
+		for jobSetKey, jobIds := range o {
+			for _, jobId := range jobIds {
+				job, ok := db.Jobs[jobId]
+				if !ok {
+					return errors.Errorf("job %s not in db", jobId)
+				}
+				if job.JobSet == jobSetKey.jobSet && job.Queue == jobSetKey.queue {
+					job.CancelRequested = true
+				}
 			}
 		}
 	case MarkJobsSucceeded:
@@ -445,12 +449,14 @@ func (db *mockDb) apply(op DbOperation) error {
 			}
 		}
 	case UpdateJobPriorities:
-		for jobId, priority := range o {
-			if job, ok := db.Jobs[jobId]; ok {
-				job.Priority = priority
-			} else {
+		for _, jobId := range o.jobIds {
+			job, ok := db.Jobs[jobId]
+			if !ok {
 				return errors.Errorf("job %s not in db", jobId)
 			}
+			if job.JobSet == o.key.jobSet && job.Queue == o.key.queue {
+				job.Priority = o.key.Priority
+			}
 		}
 	case MarkRunsSucceeded:
 		for runId := range o {
diff --git a/internal/scheduleringester/instructions_test.go b/internal/scheduleringester/instructions_test.go
index 8bc783ff..77a2a509 100644
--- a/internal/scheduleringester/instructions_test.go
+++ b/internal/scheduleringester/instructions_test.go
@@ -129,7 +129,13 @@ func TestConvertSequence(t *testing.T) {
 		"reprioritise job": {
 			events: []*armadaevents.EventSequence_Event{f.JobReprioritiseRequested},
 			expected: []DbOperation{
-				UpdateJobPriorities{f.JobIdString: f.NewPriority},
+				UpdateJobPriorities{
+					key: JobReprioritiseKey{
+						JobSetKey: JobSetKey{queue: f.Queue, jobSet: f.JobSetName},
+						Priority:  f.NewPriority,
+					},
+					jobIds: []string{f.JobIdString},
+				},
 			},
 		},
 		"reprioritise jobset": {
@@ -141,7 +147,7 @@ func TestConvertSequence(t *testing.T) {
 		"JobCancelRequested": {
 			events: []*armadaevents.EventSequence_Event{f.JobCancelRequested},
 			expected: []DbOperation{
-				MarkJobsCancelRequested{f.JobIdString: true},
+				MarkJobsCancelRequested{JobSetKey{queue: f.Queue, jobSet: f.JobSetName}: {f.JobIdString}},
 			},
 		},
 		"JobSetCancelRequested": {
diff --git a/internal/scheduleringester/schedulerdb_test.go b/internal/scheduleringester/schedulerdb_test.go
index 4bd28b4c..439d2de8 100644
--- a/internal/scheduleringester/schedulerdb_test.go
+++ b/internal/scheduleringester/schedulerdb_test.go
@@ -32,7 +32,8 @@ func TestWriteOps(t *testing.T) {
 	}
 	scheduledAtPriorities := []int32{5, 10}
 	tests := map[string]struct {
-		Ops []DbOperation
+		Ops             []DbOperation
+		ExpectNoUpdates bool
 	}{
 		"InsertJobs": {Ops: []DbOperation{
 			InsertJobs{
@@ -88,14 +89,20 @@ func TestWriteOps(t *testing.T) {
 		}},
 		"UpdateJobPriorities": {Ops: []DbOperation{
 			InsertJobs{
-				jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], JobSet: "set1"},
-				jobIds[1]: &schedulerdb.Job{JobID: jobIds[1], JobSet: "set2"},
-				jobIds[2]: &schedulerdb.Job{JobID: jobIds[2], JobSet: "set1"},
-				jobIds[3]: &schedulerdb.Job{JobID: jobIds[3], JobSet: "set2"},
+				jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], Queue: testQueueName, JobSet: "set1"},
+				jobIds[1]: &schedulerdb.Job{JobID: jobIds[1], Queue: testQueueName, JobSet: "set2"},
+				jobIds[2]: &schedulerdb.Job{JobID: jobIds[2], Queue: testQueueName, JobSet: "set1"},
+				jobIds[3]: &schedulerdb.Job{JobID: jobIds[3], Queue: testQueueName, JobSet: "set2"},
 			},
 			UpdateJobPriorities{
-				jobIds[0]: 1,
-				jobIds[1]: 3,
+				key: JobReprioritiseKey{
+					JobSetKey: JobSetKey{
+						queue:  testQueueName,
+						jobSet: "set1",
+					},
+					Priority: 3,
+				},
+				jobIds: []string{jobIds[0], jobIds[2]},
 			},
 		}},
 		"MarkRunsForJobPreemptRequested": {Ops: []DbOperation{
@@ -145,14 +152,14 @@ func TestWriteOps(t *testing.T) {
 		}},
 		"MarkJobsCancelRequested": {Ops: []DbOperation{
 			InsertJobs{
-				jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], JobSet: "set1"},
-				jobIds[1]: &schedulerdb.Job{JobID: jobIds[1], JobSet: "set2"},
-				jobIds[2]: &schedulerdb.Job{JobID: jobIds[2], JobSet: "set1"},
-				jobIds[3]: &schedulerdb.Job{JobID: jobIds[3], JobSet: "set2"},
+				jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], Queue: testQueueName, JobSet: "set1"},
+				jobIds[1]: &schedulerdb.Job{JobID: jobIds[1], Queue: testQueueName, JobSet: "set2"},
+				jobIds[2]: &schedulerdb.Job{JobID: jobIds[2], Queue: testQueueName, JobSet: "set1"},
+				jobIds[3]: &schedulerdb.Job{JobID: jobIds[3], Queue: testQueueName, JobSet: "set2"},
 			},
 			MarkJobsCancelRequested{
-				jobIds[0]: true,
-				jobIds[1]: true,
+				JobSetKey{queue: testQueueName, jobSet: "set1"}: []string{jobIds[0]},
+				JobSetKey{queue: testQueueName, jobSet: "set2"}: []string{jobIds[1]},
 			},
 		}},
 		"MarkJobsCancelled": {Ops: []DbOperation{
@@ -500,7 +507,7 @@ func assertOpSuccess(t *testing.T, schedulerDb *SchedulerDb, serials map[string]
 		numChanged := 0
 		jobIds := make([]string, 0)
 		for _, job := range jobs {
-			if _, ok := expected[job.JobID]; ok {
+			if _, ok := expected[JobSetKey{queue: job.Queue, jobSet: job.JobSet}]; ok {
 				assert.True(t, job.CancelRequested)
 				numChanged++
 				jobIds = append(jobIds, job.JobID)
@@ -598,12 +605,13 @@ func assertOpSuccess(t *testing.T, schedulerDb *SchedulerDb, serials map[string]
 		}
 		numChanged := 0
 		for _, job := range jobs {
-			if priority, ok := expected[job.JobID]; ok {
-				assert.Equal(t, priority, job.Priority)
+			jobSetKey := JobSetKey{queue: job.Queue, jobSet: job.JobSet}
+			if expected.key.JobSetKey == jobSetKey {
+				assert.Equal(t, expected.key.Priority, job.Priority)
 				numChanged++
 			}
 		}
-		assert.Equal(t, len(expected), numChanged)
+		assert.Equal(t, len(expected.jobIds), numChanged)
 	case MarkRunsSucceeded:
 		jobs, err := selectNewJobs(ctx, 0)
 		if err != nil {