Cordon queue (#183)

* Dropping queues table

* Plumbing in cordoning mechanism

* Adding tests

* Generating proto with the correct version

* Removing method name from queue_service error messages

* Moving rate-limit comment

* Removing swagger diffs

internal/armada/permissions/permissions.go

@@ -13,5 +13,6 @@ const (
 	WatchAllEvents                            = "watch_all_events"
 	CreateQueue                               = "create_queue"
 	DeleteQueue                               = "delete_queue"
+	CordonQueue                               = "cordon_queue"
 	CordonNodes                               = "cordon_nodes"
 )

internal/armada/queue/queue_repository.go

@@ -35,6 +35,8 @@ type QueueRepository interface {
 	CreateQueue(*armadacontext.Context, queue.Queue) error
 	UpdateQueue(*armadacontext.Context, queue.Queue) error
 	DeleteQueue(ctx *armadacontext.Context, name string) error
+	CordonQueue(ctx *armadacontext.Context, name string) error
+	UncordonQueue(ctx *armadacontext.Context, name string) error
 }
 
 type ReadOnlyQueueRepository interface {
@@ -116,6 +118,24 @@ func (r *PostgresQueueRepository) DeleteQueue(ctx *armadacontext.Context, name s
 	return nil
 }
 
+func (r *PostgresQueueRepository) CordonQueue(ctx *armadacontext.Context, name string) error {
+	queueToCordon, err := r.GetQueue(ctx, name)
+	if err != nil {
+		return err
+	}
+	queueToCordon.Cordoned = true
+	return r.upsertQueue(ctx, queueToCordon)
+}
+
+func (r *PostgresQueueRepository) UncordonQueue(ctx *armadacontext.Context, name string) error {
+	queueToUncordon, err := r.GetQueue(ctx, name)
+	if err != nil {
+		return err
+	}
+	queueToUncordon.Cordoned = false
+	return r.upsertQueue(ctx, queueToUncordon)
+}
+
 func (r *PostgresQueueRepository) upsertQueue(ctx *armadacontext.Context, queue queue.Queue) error {
 	data, err := proto.Marshal(queue.ToAPI())
 	if err != nil {

internal/armada/queue/queue_service.go

@@ -2,6 +2,7 @@ package queue
 
 import (
 	"context"
+	"fmt"
 	"math"
 
 	"github.com/gogo/protobuf/types"
@@ -37,9 +38,9 @@ func (s *Server) CreateQueue(grpcCtx context.Context, req *api.Queue) (*types.Em
 	err := s.authorizer.AuthorizeAction(ctx, permissions.CreateQueue)
 	var ep *armadaerrors.ErrUnauthorized
 	if errors.As(err, &ep) {
-		return nil, status.Errorf(codes.PermissionDenied, "[CreateQueue] error creating queue %s: %s", req.Name, ep)
+		return nil, status.Errorf(codes.PermissionDenied, "error creating queue %s: %s", req.Name, ep)
 	} else if err != nil {
-		return nil, status.Errorf(codes.Unavailable, "[CreateQueue] error checking permissions: %s", err)
+		return nil, status.Errorf(codes.Unavailable, "error checking permissions: %s", err)
 	}
 
 	if len(req.UserOwners) == 0 {
@@ -49,15 +50,15 @@ func (s *Server) CreateQueue(grpcCtx context.Context, req *api.Queue) (*types.Em
 
 	queue, err := queue.NewQueue(req)
 	if err != nil {
-		return nil, status.Errorf(codes.InvalidArgument, "[CreateQueue] error validating queue: %s", err)
+		return nil, status.Errorf(codes.InvalidArgument, "error validating queue: %s", err)
 	}
 
 	err = s.queueRepository.CreateQueue(ctx, queue)
 	var eq *ErrQueueAlreadyExists
 	if errors.As(err, &eq) {
-		return nil, status.Errorf(codes.AlreadyExists, "[CreateQueue] error creating queue: %s", err)
+		return nil, status.Errorf(codes.AlreadyExists, "error creating queue: %s", err)
 	} else if err != nil {
-		return nil, status.Errorf(codes.Unavailable, "[CreateQueue] error creating queue: %s", err)
+		return nil, status.Errorf(codes.Unavailable, "error creating queue: %s", err)
 	}
 
 	return &types.Empty{}, nil
@@ -87,22 +88,22 @@ func (s *Server) UpdateQueue(grpcCtx context.Context, req *api.Queue) (*types.Em
 	err := s.authorizer.AuthorizeAction(ctx, permissions.CreateQueue)
 	var ep *armadaerrors.ErrUnauthorized
 	if errors.As(err, &ep) {
-		return nil, status.Errorf(codes.PermissionDenied, "[UpdateQueue] error updating queue %s: %s", req.Name, ep)
+		return nil, status.Errorf(codes.PermissionDenied, "error updating queue %s: %s", req.Name, ep)
 	} else if err != nil {
-		return nil, status.Errorf(codes.Unavailable, "[UpdateQueue] error checking permissions: %s", err)
+		return nil, status.Errorf(codes.Unavailable, "error checking permissions: %s", err)
 	}
 
 	queue, err := queue.NewQueue(req)
 	if err != nil {
-		return nil, status.Errorf(codes.InvalidArgument, "[UpdateQueue] error: %s", err)
+		return nil, status.Errorf(codes.InvalidArgument, "error: %s", err)
 	}
 
 	err = s.queueRepository.UpdateQueue(ctx, queue)
 	var e *ErrQueueNotFound
 	if errors.As(err, &e) {
-		return nil, status.Errorf(codes.NotFound, "[UpdateQueue] error: %s", err)
+		return nil, status.Errorf(codes.NotFound, "error: %s", err)
 	} else if err != nil {
-		return nil, status.Errorf(codes.Unavailable, "[UpdateQueue] error getting queue %q: %s", queue.Name, err)
+		return nil, status.Errorf(codes.Unavailable, "error getting queue %q: %s", queue.Name, err)
 	}
 
 	return &types.Empty{}, nil
@@ -133,13 +134,13 @@ func (s *Server) DeleteQueue(grpcCtx context.Context, req *api.QueueDeleteReques
 	err := s.authorizer.AuthorizeAction(ctx, permissions.DeleteQueue)
 	var ep *armadaerrors.ErrUnauthorized
 	if errors.As(err, &ep) {
-		return nil, status.Errorf(codes.PermissionDenied, "[DeleteQueue] error deleting queue %s: %s", req.Name, ep)
+		return nil, status.Errorf(codes.PermissionDenied, "error deleting queue %s: %s", req.Name, ep)
 	} else if err != nil {
-		return nil, status.Errorf(codes.Unavailable, "[DeleteQueue] error checking permissions: %s", err)
+		return nil, status.Errorf(codes.Unavailable, "error checking permissions: %s", err)
 	}
 	err = s.queueRepository.DeleteQueue(ctx, req.Name)
 	if err != nil {
-		return nil, status.Errorf(codes.InvalidArgument, "[DeleteQueue] error deleting queue %s: %s", req.Name, err)
+		return nil, status.Errorf(codes.InvalidArgument, "error deleting queue %s: %s", req.Name, err)
 	}
 	return &types.Empty{}, nil
 }
@@ -149,9 +150,9 @@ func (s *Server) GetQueue(grpcCtx context.Context, req *api.QueueGetRequest) (*a
 	queue, err := s.queueRepository.GetQueue(ctx, req.Name)
 	var e *ErrQueueNotFound
 	if errors.As(err, &e) {
-		return nil, status.Errorf(codes.NotFound, "[GetQueue] error: %s", err)
+		return nil, status.Errorf(codes.NotFound, "error: %s", err)
 	} else if err != nil {
-		return nil, status.Errorf(codes.Unavailable, "[GetQueue] error getting queue %q: %s", req.Name, err)
+		return nil, status.Errorf(codes.Unavailable, "error getting queue %q: %s", req.Name, err)
 	}
 	return queue.ToAPI(), nil
 }
@@ -189,3 +190,41 @@ func (s *Server) GetQueues(req *api.StreamingQueueGetRequest, stream api.QueueSe
 	}
 	return nil
 }
+
+func (s *Server) CordonQueue(grpcCtx context.Context, req *api.QueueCordonRequest) (*types.Empty, error) {
+	ctx := armadacontext.FromGrpcCtx(grpcCtx)
+
+	err := s.authorizer.AuthorizeAction(ctx, permissions.CordonQueue)
+	var ep *armadaerrors.ErrUnauthorized
+	if errors.As(err, &ep) {
+		return nil, status.Errorf(codes.PermissionDenied, "error cordoning queue %s: %s", req.Name, ep)
+	} else if err != nil {
+		return nil, status.Errorf(codes.Unavailable, "error checking permissions: %s", err)
+	}
+
+	queueName := req.Name
+	if queueName == "" {
+		return nil, fmt.Errorf("cannot cordon queue with empty name")
+	}
+
+	return &types.Empty{}, s.queueRepository.CordonQueue(ctx, queueName)
+}
+
+func (s *Server) UncordonQueue(grpcCtx context.Context, req *api.QueueUncordonRequest) (*types.Empty, error) {
+	ctx := armadacontext.FromGrpcCtx(grpcCtx)
+
+	err := s.authorizer.AuthorizeAction(ctx, permissions.CordonQueue)
+	var ep *armadaerrors.ErrUnauthorized
+	if errors.As(err, &ep) {
+		return nil, status.Errorf(codes.PermissionDenied, "error uncordoning queue %s: %s", req.Name, ep)
+	} else if err != nil {
+		return nil, status.Errorf(codes.Unavailable, "error checking permissions: %s", err)
+	}
+
+	queueName := req.Name
+	if queueName == "" {
+		return nil, fmt.Errorf("cannot uncordon queue with empty name")
+	}
+
+	return &types.Empty{}, s.queueRepository.UncordonQueue(ctx, queueName)
+}

internal/armadactl/app.go

@@ -47,10 +47,13 @@ type Params struct {
 // However, they are user-replaceable to facilitate testing.
 // TODO Consider replacing with an interface
 type QueueAPI struct {
-	Create queue.CreateAPI
-	Delete queue.DeleteAPI
-	Get    queue.GetAPI
-	Update queue.UpdateAPI
+	Create   queue.CreateAPI
+	Delete   queue.DeleteAPI
+	Get      queue.GetAPI
+	GetAll   queue.GetAllAPI
+	Update   queue.UpdateAPI
+	Cordon   queue.CordonAPI
+	Uncordon queue.UncordonAPI
 }
 
 // New instantiates an App with default parameters, including standard output

internal/scheduler/constraints/constraints.go

@@ -23,6 +23,7 @@ const (
 	// Indicates that the scheduling rate limit has been exceeded.
 	GlobalRateLimitExceededUnschedulableReason = "global scheduling rate limit exceeded"
 	QueueRateLimitExceededUnschedulableReason  = "queue scheduling rate limit exceeded"
+	SchedulingPausedOnQueueUnschedulableReason = "scheduling paused on queue"
 
 	// Indicates that scheduling a gang would exceed the rate limit.
 	GlobalRateLimitExceededByGangUnschedulableReason = "gang would exceed global scheduling rate limit"
@@ -54,7 +55,7 @@ func IsTerminalUnschedulableReason(reason string) bool {
 // IsTerminalQueueUnschedulableReason returns true if reason indicates
 // it's not possible to schedule any more jobs from this queue in this round.
 func IsTerminalQueueUnschedulableReason(reason string) bool {
-	return reason == QueueRateLimitExceededUnschedulableReason
+	return reason == QueueRateLimitExceededUnschedulableReason || reason == SchedulingPausedOnQueueUnschedulableReason
 }
 
 // SchedulingConstraints contains scheduling constraints, e.g., per-queue resource limits.
@@ -75,6 +76,8 @@ type SchedulingConstraints struct {
 type queueSchedulingConstraints struct {
 	// Scheduling constraints by priority class.
 	PriorityClassSchedulingConstraintsByPriorityClassName map[string]priorityClassSchedulingConstraints
+	// Determines whether scheduling has been paused for this queue
+	Cordoned bool
 }
 
 // priorityClassSchedulingConstraints contains scheduling constraints that apply to jobs of a specific priority class.
@@ -84,12 +87,7 @@ type priorityClassSchedulingConstraints struct {
 	MaximumResourcesPerQueue map[string]resource.Quantity
 }
 
-func NewSchedulingConstraints(
-	pool string,
-	totalResources schedulerobjects.ResourceList,
-	config configuration.SchedulingConfig,
-	queues []*api.Queue,
-) SchedulingConstraints {
+func NewSchedulingConstraints(pool string, totalResources schedulerobjects.ResourceList, config configuration.SchedulingConfig, queues []*api.Queue, cordonStatusByQueue map[string]bool) SchedulingConstraints {
 	priorityClassSchedulingConstraintsByPriorityClassName := make(map[string]priorityClassSchedulingConstraints, len(config.PriorityClasses))
 	for name, priorityClass := range config.PriorityClasses {
 		maximumResourceFractionPerQueue := priorityClass.MaximumResourceFractionPerQueue
@@ -117,10 +115,9 @@ func NewSchedulingConstraints(
 				MaximumResourcesPerQueue: absoluteFromRelativeLimits(totalResources.Resources, maximumResourceFraction),
 			}
 		}
-		if len(priorityClassSchedulingConstraintsByPriorityClassNameForQueue) > 0 {
-			queueSchedulingConstraintsByQueueName[queue.Name] = queueSchedulingConstraints{
-				PriorityClassSchedulingConstraintsByPriorityClassName: priorityClassSchedulingConstraintsByPriorityClassNameForQueue,
-			}
+		queueSchedulingConstraintsByQueueName[queue.Name] = queueSchedulingConstraints{
+			PriorityClassSchedulingConstraintsByPriorityClassName: priorityClassSchedulingConstraintsByPriorityClassNameForQueue,
+			Cordoned: cordonStatusByQueue[queue.Name],
 		}
 	}
 
@@ -182,6 +179,9 @@ func (constraints *SchedulingConstraints) CheckConstraints(
 		return false, GlobalRateLimitExceededByGangUnschedulableReason, nil
 	}
 
+	if queueConstraints, ok := constraints.queueSchedulingConstraintsByQueueName[qctx.Queue]; ok && queueConstraints.Cordoned {
+		return false, SchedulingPausedOnQueueUnschedulableReason, nil
+	}
 	// Per-queue rate limiter check.
 	tokens = qctx.Limiter.TokensAt(sctx.Started)
 	if tokens <= 0 {