Reserve capacity while a console waits for authorisation

cmd/workloads-manager/main.go

@@ -27,17 +27,20 @@ import (
 var (
 	scheme = runtime.NewScheme()
 
-	app                        = kingpin.New("workloads-manager", "Manages workloads.crd.gocardless.com resources").Version(cmd.VersionStanza())
-	contextName                = app.Flag("context-name", "Distinct name for the context this controller runs within. Usually the user-facing name of the kubernetes context for the cluster").Envar("CONTEXT_NAME").String()
-	pubsubProjectId            = app.Flag("pubsub-project-id", "ID for the project containing the Pub/Sub topic for console event publishing").Envar("PUBSUB_PROJECT_ID").String()
-	pubsubTopicId              = app.Flag("pubsub-topic-id", "ID of the topic to publish lifecycle event messages").Envar("PUBSUB_TOPIC_ID").String()
-	enableDirectoryRoleBinding = app.Flag("directory-role-binding", "Use DirectoryRoleBinding for provisioning RBAC against console objects").Envar("ENABLE_DIRECTORY_ROLE_BINDING").Default("true").Bool()
-	restrictApproverRole       = app.Flag("restrict-approver-role", "Give console approvers attach and logs but not exec, instead of the creator's permissions").Envar("RESTRICT_APPROVER_ROLE").Default("false").Bool()
-	subjectsWithoutAccess      = app.Flag("subjects-without-access", "Usernames that never get access to consoles, such as an authoriser's service account (repeatable)").Strings()
-	enableSessionRecording     = app.Flag("session-recording", "Enable session recording features").Envar("ENABLE_SESSION_RECORDING").Default("false").Bool()
-	sessionSidecarImage        = app.Flag("session-sidecar-image", "Container image to use for the session recording sidecar container").Envar("SESSION_SIDECAR_IMAGE").Default("").String()
-	sessionPubsubProjectId     = app.Flag("session-pubsub-project-id", "ID for the project containing the Pub/Sub topic for session recording").Envar("SESSION_PUBSUB_PROJECT_ID").Default("").String()
-	sessionPubsubTopicId       = app.Flag("session-pubsub-topic-id", "ID of the topic to publish session recording data to").Envar("SESSION_PUBSUB_TOPIC_ID").Default("").String()
+	app                          = kingpin.New("workloads-manager", "Manages workloads.crd.gocardless.com resources").Version(cmd.VersionStanza())
+	contextName                  = app.Flag("context-name", "Distinct name for the context this controller runs within. Usually the user-facing name of the kubernetes context for the cluster").Envar("CONTEXT_NAME").String()
+	pubsubProjectId              = app.Flag("pubsub-project-id", "ID for the project containing the Pub/Sub topic for console event publishing").Envar("PUBSUB_PROJECT_ID").String()
+	pubsubTopicId                = app.Flag("pubsub-topic-id", "ID of the topic to publish lifecycle event messages").Envar("PUBSUB_TOPIC_ID").String()
+	enableDirectoryRoleBinding   = app.Flag("directory-role-binding", "Use DirectoryRoleBinding for provisioning RBAC against console objects").Envar("ENABLE_DIRECTORY_ROLE_BINDING").Default("true").Bool()
+	restrictApproverRole         = app.Flag("restrict-approver-role", "Give console approvers attach and logs but not exec, instead of the creator's permissions").Envar("RESTRICT_APPROVER_ROLE").Default("false").Bool()
+	subjectsWithoutAccess        = app.Flag("subjects-without-access", "Usernames that never get access to consoles, such as an authoriser's service account (repeatable)").Strings()
+	enableSessionRecording       = app.Flag("session-recording", "Enable session recording features").Envar("ENABLE_SESSION_RECORDING").Default("false").Bool()
+	enablePendingCapacity        = app.Flag("pending-capacity", "Reserve a low-priority Job while a console waits for authorisation").Envar("ENABLE_PENDING_CAPACITY").Default("false").Bool()
+	pendingCapacityPriorityClass = app.Flag("pending-capacity-priority-class", "PriorityClass for the pending-authorisation capacity Job; must be preemptible by a normal pod").Envar("PENDING_CAPACITY_PRIORITY_CLASS").String()
+	pendingCapacityTTL           = app.Flag("pending-capacity-ttl", "Active deadline for the pending-authorisation capacity Job").Envar("PENDING_CAPACITY_TTL").Default("15m").Duration()
+	sessionSidecarImage          = app.Flag("session-sidecar-image", "Container image to use for the session recording sidecar container").Envar("SESSION_SIDECAR_IMAGE").Default("").String()
+	sessionPubsubProjectId       = app.Flag("session-pubsub-project-id", "ID for the project containing the Pub/Sub topic for session recording").Envar("SESSION_PUBSUB_PROJECT_ID").Default("").String()
+	sessionPubsubTopicId         = app.Flag("session-pubsub-topic-id", "ID of the topic to publish session recording data to").Envar("SESSION_PUBSUB_TOPIC_ID").Default("").String()
 
 	commonOpts = cmd.NewCommonOptions(app).WithMetrics(app)
 )
@@ -70,6 +73,9 @@ func main() {
 			app.Fatalf("Session recording Google pubsub ID must be set")
 		}
 	}
+	if *enablePendingCapacity && (*pendingCapacityPriorityClass == "" || *pendingCapacityTTL <= 0) {
+		app.Fatalf("pending capacity requires a priority class and a positive ttl")
+	}
 
 	// Create publisher sink for console lifecycle events
 	var publisher events.Publisher
@@ -104,18 +110,21 @@ func main() {
 
 	// controller
 	if err = (&consolecontroller.ConsoleReconciler{
-		Client:                     mgr.GetClient(),
-		LifecycleRecorder:          lifecycleRecorder,
-		ConsoleIdBuilder:           idBuilder,
-		Log:                        ctrl.Log.WithName("controllers").WithName("console"),
-		Scheme:                     mgr.GetScheme(),
-		EnableDirectoryRoleBinding: *enableDirectoryRoleBinding,
-		RestrictApproverRole:       *restrictApproverRole,
-		SubjectsWithoutAccess:      *subjectsWithoutAccess,
-		EnableSessionRecording:     *enableSessionRecording,
-		SessionSidecarImage:        *sessionSidecarImage,
-		SessionPubsubProjectId:     *sessionPubsubProjectId,
-		SessionPubsubTopicId:       *sessionPubsubTopicId,
+		Client:                       mgr.GetClient(),
+		LifecycleRecorder:            lifecycleRecorder,
+		ConsoleIdBuilder:             idBuilder,
+		Log:                          ctrl.Log.WithName("controllers").WithName("console"),
+		Scheme:                       mgr.GetScheme(),
+		EnableDirectoryRoleBinding:   *enableDirectoryRoleBinding,
+		RestrictApproverRole:         *restrictApproverRole,
+		SubjectsWithoutAccess:        *subjectsWithoutAccess,
+		EnableSessionRecording:       *enableSessionRecording,
+		SessionSidecarImage:          *sessionSidecarImage,
+		SessionPubsubProjectId:       *sessionPubsubProjectId,
+		SessionPubsubTopicId:         *sessionPubsubTopicId,
+		EnablePendingCapacity:        *enablePendingCapacity,
+		PendingCapacityPriorityClass: *pendingCapacityPriorityClass,
+		PendingCapacityTTL:           *pendingCapacityTTL,
 	}).SetupWithManager(ctx, mgr); err != nil {
 		app.Fatalf("failed to create controller: %v", err)
 	}

internal/controller/workloads/console_controller.go

@@ -108,6 +108,16 @@ type ConsoleReconciler struct {
 	SessionPubsubProjectId string
 	// The Pub/Sub topic ID that the session recording data should be sent to
 	SessionPubsubTopicId string
+
+	// Reserve a Job while a console waits for authorisation, so the node it
+	// will need is already there. Off unless set.
+	EnablePendingCapacity bool
+	// PriorityClass for that Job. It has to be lower than a normal pod so the
+	// console pod can preempt it. Empty when the feature is off.
+	PendingCapacityPriorityClass string
+	// Active deadline for the Job. The console's own TTL is a separate clock.
+	// Zero uses the 15 minute default.
+	PendingCapacityTTL time.Duration
 }
 
 func (r *ConsoleReconciler) SetupWithManager(ctx context.Context, mgr ctrl.Manager) error {
@@ -453,11 +463,19 @@ func (r *ConsoleReconciler) Reconcile(logger logr.Logger, ctx context.Context, r
 		return ctrl.Result{}, err
 	}
 
+	// Hold or release the capacity job after the phase is known. A failure
+	// here must not undo the status update above: the next reconcile retries
+	// the job, and an authorised console still gets its own job.
+	if err := r.reconcilePendingCapacity(ctx, logger, csl, tpl); err != nil {
+		return ctrl.Result{}, err
+	}
+
 	var res ctrl.Result
 	switch {
 	case csl.PendingAuthorisation():
 		// Requeue for when the console has reached its before-running TTL, so that
 		// it can be deleted if it has not yet been authorised by that point.
+		// The capacity job's own deadline is enforced by the job controller.
 		res = requeueAfterInterval(logger, time.Until(*csl.GetGCTime()))
 	case csl.Pending():
 		// Requeue every second while job has been created but there is not yet a

internal/controller/workloads/integration/integration_test.go

@@ -639,6 +639,15 @@ var _ = Describe("Console", func() {
 				Expect(drb.ObjectMeta.OwnerReferences[0].Name).To(Equal(csl.ObjectMeta.Name))
 			})
 
+			It("does not reserve capacity while authorisation is pending", func() {
+				Consistently(func() bool {
+					err := mgr.GetClient().Get(context.TODO(), client.ObjectKey{
+						Namespace: namespaceName, Name: consoleName + "-capacity",
+					}, &corev1.Pod{})
+					return apierrors.IsNotFound(err)
+				}, 2*time.Second, 200*time.Millisecond).Should(BeTrue())
+			})
+
 			Context("When restrict-approver-role is off", func() {
 				It("binds the approver to the creator's role", func() {
 					user, err := testEnv.AddUser(

internal/controller/workloads/pending_capacity.go

@@ -0,0 +1,253 @@
+package controllers
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/go-logr/logr"
+	"github.com/pkg/errors"
+	batchv1 "k8s.io/api/batch/v1"
+	corev1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/types"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
+
+	workloadsv1alpha1 "github.com/gocardless/theatre/v5/api/workloads/v1alpha1"
+)
+
+const (
+	// app.kubernetes.io/component value. Distinct from a real console pod so
+	// the job-pod lookup does not treat this as the session.
+	pendingCapacityComponent = "console-pending-capacity"
+
+	// cluster-autoscaler.kubernetes.io/safe-to-evict. Copied when the template
+	// sets it, so the placeholder is evicted on the same terms as the console pod.
+	safeToEvictAnnotation = "cluster-autoscaler.kubernetes.io/safe-to-evict"
+
+	// Used when the flag is enabled but the duration was left unset.
+	defaultPendingCapacityTTL = 15 * time.Minute
+)
+
+// reconcilePendingCapacity creates a Job while the console is waiting for
+// authorisation, and deletes it otherwise: once the console Job exists, once
+// the console is stopped, and once the console itself expires.
+//
+// The Job's active deadline ends the placeholder. The Job object is left in
+// place after that, so a later reconcile does not start another one while the
+// console is still waiting. Deleting the console object still relies on the
+// owner reference, because a console that is already being deleted is not
+// reconciled.
+func (r *ConsoleReconciler) reconcilePendingCapacity(ctx context.Context, logger logr.Logger, csl *workloadsv1alpha1.Console, template *workloadsv1alpha1.ConsoleTemplate) error {
+	name := types.NamespacedName{
+		Namespace: csl.Namespace,
+		Name:      pendingCapacityJobName(csl.Name),
+	}
+	hold := r.EnablePendingCapacity && csl.PendingAuthorisation() && !csl.EligibleForGC()
+	if !hold {
+		return r.deletePendingCapacityJob(ctx, logger, name)
+	}
+
+	job, err := r.buildPendingCapacityJob(name, csl, template)
+	if err != nil {
+		return err
+	}
+	if err := controllerutil.SetControllerReference(csl, job, r.Scheme); err != nil {
+		return errors.Wrap(err, "failed to set owner on pending capacity job")
+	}
+
+	existing := &batchv1.Job{}
+	err = r.Get(ctx, name, existing)
+	if apierrors.IsNotFound(err) {
+		if err := r.Create(ctx, job); err != nil {
+			return errors.Wrap(err, "failed to create pending capacity job")
+		}
+		logger.Info("Created pending capacity job", "event", EventSuccessfulCreate, "kind", "job", "job", job.Name)
+		return nil
+	}
+	if err != nil {
+		return errors.Wrap(err, "failed to get pending capacity job")
+	}
+	if !metav1.IsControlledBy(existing, csl) {
+		return errors.Errorf("pending capacity job %s is not owned by this console", name)
+	}
+	// A finished job still occupies the name. Leave it, so the deadline is not
+	// restarted by creating another.
+	return nil
+}
+
+func (r *ConsoleReconciler) deletePendingCapacityJob(ctx context.Context, logger logr.Logger, name types.NamespacedName) error {
+	grace := int64(0)
+	policy := metav1.DeletePropagationBackground
+	job := &batchv1.Job{ObjectMeta: metav1.ObjectMeta{Namespace: name.Namespace, Name: name.Name}}
+	err := r.Delete(ctx, job, &client.DeleteOptions{
+		GracePeriodSeconds: &grace,
+		PropagationPolicy:  &policy,
+	})
+	if apierrors.IsNotFound(err) {
+		return nil
+	}
+	if err != nil {
+		return errors.Wrap(err, "failed to delete pending capacity job")
+	}
+	logger.Info("Deleted pending capacity job", "event", EventDelete, "kind", "job", "job", name.Name)
+	return nil
+}
+
+func (r *ConsoleReconciler) buildPendingCapacityJob(name types.NamespacedName, csl *workloadsv1alpha1.Console, template *workloadsv1alpha1.ConsoleTemplate) (*batchv1.Job, error) {
+	spec := template.Spec.Template.Spec.DeepCopy()
+	if len(spec.Containers) == 0 || spec.Containers[0].Image == "" {
+		return nil, errors.New("console template has no image for the pending capacity job")
+	}
+	// The session sidecar is part of the pod that will actually run, so the
+	// placeholder has to reserve its resources too. The image that is pulled
+	// is still the template's own image.
+	if r.EnableSessionRecording {
+		spec.Containers = append(spec.Containers, r.buildSidecarContainer(r.ConsoleIdBuilder.BuildId(csl)))
+	}
+
+	deadline := int64(r.pendingCapacityTTL().Seconds())
+	backoff := int32(10)
+	completions := int32(1)
+	parallelism := int32(1)
+	grace := int64(0)
+	automount := false
+	labels := map[string]string{
+		"app.kubernetes.io/component": pendingCapacityComponent,
+		"console-name":                sanitiseLabel(csl.Name),
+	}
+	podLabels := copyStringMap(labels)
+	resources := corev1.ResourceRequirements{
+		Requests: effectiveResources(spec, false),
+		Limits:   effectiveResources(spec, true),
+	}
+
+	job := &batchv1.Job{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      name.Name,
+			Namespace: name.Namespace,
+			Labels:    labels,
+		},
+		Spec: batchv1.JobSpec{
+			ActiveDeadlineSeconds: &deadline,
+			// A preempted pod can be replaced. The deadline is counted from
+			// when the job started, so a replacement does not extend it.
+			BackoffLimit: &backoff,
+			Completions:  &completions,
+			Parallelism:  &parallelism,
+			Template: corev1.PodTemplateSpec{
+				ObjectMeta: metav1.ObjectMeta{Labels: podLabels},
+				Spec: corev1.PodSpec{
+					// Not the console service account: that account is the
+					// workload's identity, and this pod only has to sit on a node.
+					AutomountServiceAccountToken:  &automount,
+					PriorityClassName:             r.PendingCapacityPriorityClass,
+					TerminationGracePeriodSeconds: &grace,
+					RestartPolicy:                 corev1.RestartPolicyNever,
+					NodeSelector:                  copyStringMap(spec.NodeSelector),
+					Tolerations:                   spec.Tolerations,
+					Affinity:                      spec.Affinity,
+					Containers: []corev1.Container{{
+						Name:            "sleep",
+						Image:           spec.Containers[0].Image,
+						ImagePullPolicy: corev1.PullIfNotPresent,
+						Command:         []string{"sleep"},
+						Args:            []string{fmt.Sprintf("%d", deadline)},
+						Resources:       resources,
+					}},
+				},
+			},
+		},
+	}
+	if v, ok := template.Spec.Template.Annotations[safeToEvictAnnotation]; ok {
+		job.Spec.Template.Annotations = map[string]string{safeToEvictAnnotation: v}
+	}
+	return job, nil
+}
+
+func (r *ConsoleReconciler) pendingCapacityTTL() time.Duration {
+	if r.PendingCapacityTTL > 0 {
+		return r.PendingCapacityTTL
+	}
+	return defaultPendingCapacityTTL
+}
+
+// pendingCapacityJobName keeps the job name inside the 63 character label
+// limit. The suffix differs from the console job, so its pods are not
+// mistaken for the session.
+func pendingCapacityJobName(consoleName string) string {
+	return fmt.Sprintf("%s-capacity", truncateString(consoleName, 54))
+}
+
+// effectiveResources is the request (or limit) the scheduler uses for the
+// pod: the sum of the app containers, plus sidecar init containers, and at
+// least the largest sequential init container. That is the hole the real
+// console pod needs.
+func effectiveResources(spec *corev1.PodSpec, limits bool) corev1.ResourceList {
+	pick := func(c corev1.Container) corev1.ResourceList {
+		if limits {
+			return c.Resources.Limits
+		}
+		return c.Resources.Requests
+	}
+
+	total := corev1.ResourceList{}
+	for _, c := range spec.Containers {
+		addResourceList(total, pick(c))
+	}
+	initMax := corev1.ResourceList{}
+	for _, c := range spec.InitContainers {
+		if isSidecarInit(c) {
+			addResourceList(total, pick(c))
+			continue
+		}
+		for name, q := range pick(c) {
+			cur, ok := initMax[name]
+			if !ok || q.Cmp(cur) > 0 {
+				initMax[name] = q.DeepCopy()
+			}
+		}
+	}
+	for name, q := range initMax {
+		cur, ok := total[name]
+		if !ok || q.Cmp(cur) > 0 {
+			total[name] = q.DeepCopy()
+		}
+	}
+	if len(total) == 0 {
+		return nil
+	}
+	return total
+}
+
+func isSidecarInit(c corev1.Container) bool {
+	return c.RestartPolicy != nil && *c.RestartPolicy == corev1.ContainerRestartPolicyAlways
+}
+
+func addResourceList(dst, src corev1.ResourceList) {
+	for name, q := range src {
+		if q.IsZero() {
+			continue
+		}
+		q = q.DeepCopy()
+		if cur, ok := dst[name]; ok {
+			cur.Add(q)
+			dst[name] = cur
+			continue
+		}
+		dst[name] = q
+	}
+}
+
+func copyStringMap(in map[string]string) map[string]string {
+	if len(in) == 0 {
+		return nil
+	}
+	out := make(map[string]string, len(in))
+	for k, v := range in {
+		out[k] = v
+	}
+	return out
+}