feat: adds pdbs to workbench sessions (#67)

Author: amdove

api/product/disruption_budget.go

@@ -1,6 +1,8 @@
 package product
 
 import (
+	"fmt"
+
 	policyv1 "k8s.io/api/policy/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
@@ -42,3 +44,29 @@ func DefineDisruptionBudget(p NamerAndOwnerProvider, req ctrl.Request, minAvaila
 
 	return pdb
 }
+
+// DefineSessionDisruptionBudget creates a PodDisruptionBudget for workbench session pods.
+// This PDB uses maxUnavailable=0 to prevent any session pods from being evicted during
+// node drains or cluster maintenance, ensuring session persistence.
+func DefineSessionDisruptionBudget(p NamerAndOwnerProvider, req ctrl.Request) *policyv1.PodDisruptionBudget {
+	return &policyv1.PodDisruptionBudget{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:            fmt.Sprintf("%s-sessions", p.ComponentName()),
+			Namespace:       req.Namespace,
+			OwnerReferences: p.OwnerReferencesForChildren(),
+			Labels:          p.KubernetesLabels(),
+		},
+		Spec: policyv1.PodDisruptionBudgetSpec{
+			Selector: &metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					// This label is set by the launcher on all session pods
+					"launcher-instance-id": p.ComponentName(),
+				},
+			},
+			MaxUnavailable: &intstr.IntOrString{
+				Type:   intstr.Int,
+				IntVal: 0,
+			},
+		},
+	}
+}

client-go/applyconfiguration/core/v1beta1/connectruntimeimagespec.go

@@ -822,9 +822,8 @@ func (r *ConnectReconciler) ensureDeployedService(ctx context.Context, req ctrl.
 	}
 
 	// POD DISRUPTION BUDGET
-	if err := CreateOrUpdateDisruptionBudget(
-		ctx, req, r.Client, r.Scheme, c, c, ptr.To(product.DetermineMinAvailableReplicas(c.Spec.Replicas)), nil,
-	); err != nil {
+	pdb := product.DefineDisruptionBudget(c, req, ptr.To(product.DetermineMinAvailableReplicas(c.Spec.Replicas)), nil)
+	if err := CreateOrUpdateDisruptionBudget(ctx, r.Client, r.Scheme, c, pdb); err != nil {
 		return ctrl.Result{}, err
 	}
 

internal/controller/core/connect.go

@@ -9,13 +9,12 @@ import (
 	policyv1 "k8s.io/api/policy/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
-	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
-func CreateOrUpdateDisruptionBudget(ctx context.Context, req ctrl.Request, c client.Client, scheme *runtime.Scheme, p product.NamerAndOwnerProvider, owner client.Object, minAvailable, maxUnavailable *int) error {
+// CreateOrUpdateDisruptionBudget creates or updates a PodDisruptionBudget from the given spec.
+func CreateOrUpdateDisruptionBudget(ctx context.Context, c client.Client, scheme *runtime.Scheme, owner client.Object, pdbSpec *policyv1.PodDisruptionBudget) error {
 	l := logr.FromContextOrDiscard(ctx)
-	pdbSpec := product.DefineDisruptionBudget(p, req, minAvailable, maxUnavailable)
 
 	pdb := &policyv1.PodDisruptionBudget{
 		ObjectMeta: metav1.ObjectMeta{

internal/controller/core/disruption_budget.go

@@ -686,9 +686,8 @@ func (r *PackageManagerReconciler) ensureDeployedService(ctx context.Context, re
 	}
 
 	// POD DISRUPTION BUDGET
-	if err := CreateOrUpdateDisruptionBudget(
-		ctx, req, r.Client, r.Scheme, pm, pm, ptr.To(product.DetermineMinAvailableReplicas(pm.Spec.Replicas)), nil,
-	); err != nil {
+	pdb := product.DefineDisruptionBudget(pm, req, ptr.To(product.DetermineMinAvailableReplicas(pm.Spec.Replicas)), nil)
+	if err := CreateOrUpdateDisruptionBudget(ctx, r.Client, r.Scheme, pm, pdb); err != nil {
 		return ctrl.Result{}, err
 	}
 

internal/controller/core/package_manager.go

@@ -14,6 +14,7 @@ import (
 	"github.com/traefik/traefik/v3/pkg/provider/kubernetes/crd/traefikio/v1alpha1"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
+	policyv1 "k8s.io/api/policy/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -343,6 +344,13 @@ func getMiddleware(t *testing.T, cli client.Client, siteNamespace, siteName stri
 	return middleware
 }
 
+func getPodDisruptionBudget(t *testing.T, cli client.Client, namespace, name string) *policyv1.PodDisruptionBudget {
+	pdb := &policyv1.PodDisruptionBudget{}
+	err := cli.Get(context.TODO(), client.ObjectKey{Name: name, Namespace: namespace}, pdb, &client.GetOptions{})
+	assert.Nil(t, err)
+	return pdb
+}
+
 func TestSiteReconcileWithTolerations(t *testing.T) {
 	siteName := "tolerations-site"
 	siteNamespace := "posit-team"

internal/controller/core/site_test.go

@@ -997,10 +997,15 @@ func (r *WorkbenchReconciler) ensureDeployedService(ctx context.Context, req ctr
 		return ctrl.Result{}, err
 	}
 
-	// POD DISRUPTION BUDGET
-	if err := CreateOrUpdateDisruptionBudget(
-		ctx, req, r.Client, r.Scheme, w, w, ptr.To(product.DetermineMinAvailableReplicas(w.Spec.Replicas)), nil,
-	); err != nil {
+	// POD DISRUPTION BUDGET for server pods
+	serverPdb := product.DefineDisruptionBudget(w, req, ptr.To(product.DetermineMinAvailableReplicas(w.Spec.Replicas)), nil)
+	if err := CreateOrUpdateDisruptionBudget(ctx, r.Client, r.Scheme, w, serverPdb); err != nil {
+		return ctrl.Result{}, err
+	}
+
+	// POD DISRUPTION BUDGET for session pods (prevents eviction during node drains)
+	sessionPdb := product.DefineSessionDisruptionBudget(w, req)
+	if err := CreateOrUpdateDisruptionBudget(ctx, r.Client, r.Scheme, w, sessionPdb); err != nil {
 		return ctrl.Result{}, err
 	}
 

internal/controller/core/workbench.go

@@ -355,3 +355,37 @@ func TestWorkbenchLoadBalancingDisabled(t *testing.T) {
 		assert.NotEqual(t, "load-balancer-config", v.Name, "Should not have load-balancer-config volume when load balancing is disabled")
 	}
 }
+
+func TestWorkbenchPodDisruptionBudgets(t *testing.T) {
+	ctx := context.Background()
+	ns := "posit-team"
+	name := "workbench-pdb"
+
+	ctx, r, req, cli := initWorkbenchReconciler(t, ctx, ns, name)
+
+	wb := defineDefaultWorkbench(t, ns, name)
+
+	err := internal.BasicCreateOrUpdate(ctx, r, r.GetLogger(ctx), req.NamespacedName, &positcov1beta1.Workbench{}, wb)
+	require.NoError(t, err)
+
+	wb = getWorkbench(t, cli, ns, name)
+
+	res, err := r.ReconcileWorkbench(ctx, req, wb)
+	require.NoError(t, err)
+	require.True(t, res.IsZero())
+
+	// Verify session PDB is created
+	sessionPdb := getPodDisruptionBudget(t, cli, ns, name+"-workbench-sessions")
+	require.NotNil(t, sessionPdb, "Session PDB should be created")
+	assert.Equal(t, name+"-workbench-sessions", sessionPdb.Name)
+
+	// Verify session PDB has correct selector to target session pods
+	require.NotNil(t, sessionPdb.Spec.Selector, "Session PDB should have a selector")
+	assert.Equal(t, wb.ComponentName(), sessionPdb.Spec.Selector.MatchLabels["launcher-instance-id"],
+		"Session PDB should select pods with launcher-instance-id label matching workbench component name")
+
+	// Verify session PDB has maxUnavailable=0 to prevent any evictions
+	require.NotNil(t, sessionPdb.Spec.MaxUnavailable, "Session PDB should have maxUnavailable set")
+	assert.Equal(t, int32(0), sessionPdb.Spec.MaxUnavailable.IntVal,
+		"Session PDB should have maxUnavailable=0 to prevent session evictions")
+}