Implementing cancel/preempt rpcs on queue service

internal/scheduleringester/dbops.go

@@ -7,6 +7,7 @@ import (
 	"golang.org/x/exp/maps"
 
 	schedulerdb "github.com/armadaproject/armada/internal/scheduler/database"
+	"github.com/armadaproject/armada/pkg/controlplaneevents"
 )
 
 type Operation int
@@ -87,6 +88,18 @@ type CancelOnExecutor struct {
 	PriorityClasses []string
 }
 
+type PreemptOnQueue struct {
+	Name            string
+	PriorityClasses []string
+	JobStates       []controlplaneevents.ActiveJobState
+}
+
+type CancelOnQueue struct {
+	Name            string
+	PriorityClasses []string
+	JobStates       []controlplaneevents.ActiveJobState
+}
+
 // DbOperation captures a generic batch database operation.
 //
 // There are 5 types of operations:
@@ -185,6 +198,8 @@ type (
 	DeleteExecutorSettings map[string]*ExecutorSettingsDelete
 	PreemptExecutor        map[string]*PreemptOnExecutor
 	CancelExecutor         map[string]*CancelOnExecutor
+	PreemptQueue           map[string]*PreemptOnQueue
+	CancelQueue            map[string]*CancelOnQueue
 )
 
 type jobSetOperation interface {
@@ -337,6 +352,14 @@ func (ce CancelExecutor) Merge(_ DbOperation) bool {
 	return false
 }
 
+func (pq PreemptQueue) Merge(_ DbOperation) bool {
+	return false
+}
+
+func (cq CancelQueue) Merge(_ DbOperation) bool {
+	return false
+}
+
 // MergeInMap merges an op b into a, provided that b is of the same type as a.
 // For example, if a is of type MarkJobSetsCancelRequested, b is only merged if also of type MarkJobSetsCancelRequested.
 // Returns true if the ops were merged and false otherwise.
@@ -536,6 +559,30 @@ func (ce CancelExecutor) CanBeAppliedBefore(b DbOperation) bool {
 	return true
 }
 
+func (pq PreemptQueue) CanBeAppliedBefore(b DbOperation) bool {
+	switch op := b.(type) {
+	case queueOperation:
+		for queue := range pq {
+			if affectsQueue := op.affectsQueue(queue); affectsQueue {
+				return false
+			}
+		}
+	}
+	return true
+}
+
+func (cq CancelQueue) CanBeAppliedBefore(b DbOperation) bool {
+	switch op := b.(type) {
+	case queueOperation:
+		for queue := range cq {
+			if affectsQueue := op.affectsQueue(queue); affectsQueue {
+				return false
+			}
+		}
+	}
+	return true
+}
+
 // definesJobInSet returns true if b is an InsertJobs operation
 // that inserts at least one job in any of the job sets that make
 // up the keys of a.
@@ -684,6 +731,14 @@ func (ce CancelExecutor) GetOperation() Operation {
 	return ControlPlaneOperation
 }
 
+func (pq PreemptQueue) GetOperation() Operation {
+	return ControlPlaneOperation
+}
+
+func (cq CancelQueue) GetOperation() Operation {
+	return ControlPlaneOperation
+}
+
 type executorOperation interface {
 	affectsExecutor(string) bool
 }
@@ -707,3 +762,17 @@ func (ce CancelExecutor) affectsExecutor(executor string) bool {
 	_, ok := ce[executor]
 	return ok
 }
+
+type queueOperation interface {
+	affectsQueue(string) bool
+}
+
+func (pq PreemptQueue) affectsQueue(queue string) bool {
+	_, ok := pq[queue]
+	return ok
+}
+
+func (cq CancelQueue) affectsQueue(queue string) bool {
+	_, ok := cq[queue]
+	return ok
+}

internal/scheduleringester/instructions.go

@@ -414,6 +414,10 @@ func (c *ControlPlaneEventsInstructionConverter) dbOperationFromControlPlaneEven
 		operations, err = c.handlePreemptOnExecutor(event.GetPreemptOnExecutor())
 	case *controlplaneevents.Event_CancelOnExecutor:
 		operations, err = c.handleCancelOnExecutor(event.GetCancelOnExecutor())
+	case *controlplaneevents.Event_PreemptOnQueue:
+		operations, err = c.handlePreemptOnQueue(event.GetPreemptOnQueue())
+	case *controlplaneevents.Event_CancelOnQueue:
+		operations, err = c.handleCancelOnQueue(event.GetCancelOnQueue())
 	default:
 		log.Errorf("Unknown event of type %T", ev)
 	}
@@ -473,6 +477,30 @@ func (c *ControlPlaneEventsInstructionConverter) handleCancelOnExecutor(cancel *
 	}, nil
 }
 
+func (c *ControlPlaneEventsInstructionConverter) handlePreemptOnQueue(preempt *controlplaneevents.PreemptOnQueue) ([]DbOperation, error) {
+	return []DbOperation{
+		PreemptQueue{
+			preempt.Name: {
+				Name:            preempt.Name,
+				PriorityClasses: preempt.PriorityClasses,
+				JobStates:       preempt.JobStates,
+			},
+		},
+	}, nil
+}
+
+func (c *ControlPlaneEventsInstructionConverter) handleCancelOnQueue(cancel *controlplaneevents.CancelOnQueue) ([]DbOperation, error) {
+	return []DbOperation{
+		CancelQueue{
+			cancel.Name: {
+				Name:            cancel.Name,
+				PriorityClasses: cancel.PriorityClasses,
+				JobStates:       cancel.JobStates,
+			},
+		},
+	}, nil
+}
+
 // schedulingInfoFromSubmitJob returns a minimal representation of a job containing only the info needed by the scheduler.
 func (c *JobSetEventsInstructionConverter) schedulingInfoFromSubmitJob(submitJob *armadaevents.SubmitJob, submitTime time.Time) (*schedulerobjects.JobSchedulingInfo, error) {
 	return SchedulingInfoFromSubmitJob(submitJob, submitTime)

internal/scheduleringester/schedulerdb.go

@@ -428,6 +428,54 @@ func (s *SchedulerDb) WriteDbOp(ctx *armadacontext.Context, tx pgx.Tx, op DbOper
 				}
 			}
 		}
+	case CancelQueue:
+		for _, cancelRequest := range o {
+			jobs, err := queries.SelectAllJobsByQueueAndJobState(ctx, cancelRequest.Name, cancelRequest.JobStates)
+			if err != nil {
+				return errors.Wrapf(err, "error cancelling jobs by queue, job state and priority class")
+			}
+			inPriorityClasses := jobInPriorityClasses(cancelRequest.PriorityClasses)
+			jobsToCancel := make([]schedulerdb.Job, 0)
+			for _, job := range jobs {
+				ok, err := inPriorityClasses(job)
+				if err != nil {
+					return errors.Wrapf(err, "error cancelling jobs by queue, job state and priority class")
+				}
+				if ok {
+					jobsToCancel = append(jobsToCancel, job)
+				}
+			}
+			for _, requestCancelParams := range createMarkJobsCancelRequestedByIdParams(jobsToCancel) {
+				err = queries.MarkJobsCancelRequestedById(ctx, *requestCancelParams)
+				if err != nil {
+					return errors.Wrapf(err, "error cancelling jobs by queue, job state and priority class")
+				}
+			}
+		}
+	case PreemptQueue:
+		for _, preemptRequest := range o {
+			jobs, err := queries.SelectAllJobsByQueueAndJobState(ctx, preemptRequest.Name, preemptRequest.JobStates)
+			if err != nil {
+				return errors.Wrapf(err, "error preempting jobs by queue, job state and priority class")
+			}
+			inPriorityClasses := jobInPriorityClasses(preemptRequest.PriorityClasses)
+			jobsToPreempt := make([]schedulerdb.Job, 0)
+			for _, job := range jobs {
+				ok, err := inPriorityClasses(job)
+				if err != nil {
+					return errors.Wrapf(err, "error preempting jobs by queue, job state and priority class")
+				}
+				if ok {
+					jobsToPreempt = append(jobsToPreempt, job)
+				}
+			}
+			for _, requestPreemptParams := range createMarkJobRunsPreemptRequestedByJobIdParams(jobsToPreempt) {
+				err = queries.MarkJobRunsPreemptRequestedByJobId(ctx, *requestPreemptParams)
+				if err != nil {
+					return errors.Wrapf(err, "error preempting jobs by queue, job state and priority class")
+				}
+			}
+		}
 	default:
 		return errors.Errorf("received unexpected op %+v", op)
 	}