Add PVC retention policy to WorkerPool

api/v1alpha2/workerpool_types.go

@@ -17,6 +17,7 @@ limitations under the License.
 package v1alpha2
 
 import (
+	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
@@ -74,6 +75,16 @@ type WorkerPoolSpec struct {
 	// +kubebuilder:validation:Required
 	// +kubebuilder:validation:XValidation:rule="self == oldSelf",message="volumeClaimTemplate is immutable"
 	VolumeClaimTemplate corev1.PersistentVolumeClaimSpec `json:"volumeClaimTemplate"`
+
+	// PersistentVolumeClaimRetentionPolicy controls whether the worker PVCs
+	// created from volumeClaimTemplate are deleted when the WorkerPool is
+	// deleted (whenDeleted) or scaled down (whenScaled). Each field accepts
+	// Retain or Delete. When unset, or when a field is omitted, PVCs are
+	// retained.
+	// +optional
+	// +kubebuilder:validation:XValidation:rule="!has(self.whenDeleted) || self.whenDeleted in ['Retain', 'Delete']",message="whenDeleted must be Retain or Delete"
+	// +kubebuilder:validation:XValidation:rule="!has(self.whenScaled) || self.whenScaled in ['Retain', 'Delete']",message="whenScaled must be Retain or Delete"
+	PersistentVolumeClaimRetentionPolicy *appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy `json:"persistentVolumeClaimRetentionPolicy,omitempty"`
 }
 
 // WorkerPoolStatus defines the observed state of WorkerPool.

docs/reference.md

@@ -1134,6 +1134,8 @@ keep repository access without restarting (see [Workspace authentication](#works
 | `spec.worker.podOverrides` | Pod customization for worker pods | No |
 | `spec.replicas` | Number of persistent worker pods (defaults to 1) | No |
 | `spec.volumeClaimTemplate` | PersistentVolumeClaimSpec for each worker pod's storage | Yes |
+| `spec.persistentVolumeClaimRetentionPolicy.whenDeleted` | `Retain` or `Delete`. Whether worker PVCs are deleted when the WorkerPool is deleted (defaults to `Retain`) | No |
+| `spec.persistentVolumeClaimRetentionPolicy.whenScaled` | `Retain` or `Delete`. Whether the PVCs of removed worker pods are deleted when `spec.replicas` is reduced (defaults to `Retain`) | No |
 
 ## Workspace
 

internal/controller/workerpool_controller.go

@@ -492,7 +492,7 @@ func (r *WorkerPoolReconciler) reconcileStatefulSet(ctx context.Context, pool *k
 		return ctrl.Result{RequeueAfter: 5 * time.Second}, nil
 	}
 
-	// Update StatefulSet if replicas or pod template changed
+	// Update StatefulSet if replicas, pod template, or PVC retention policy changed
 	needsUpdate := false
 	if sts.Spec.Replicas == nil || *sts.Spec.Replicas != ptr.Deref(pool.Spec.Replicas, 1) {
 		sts.Spec.Replicas = pool.Spec.Replicas
@@ -506,6 +506,11 @@ func (r *WorkerPoolReconciler) reconcileStatefulSet(ctx context.Context, pool *k
 		needsUpdate = true
 	}
 
+	if !pvcRetentionPolicyEqual(sts.Spec.PersistentVolumeClaimRetentionPolicy, desired.Spec.PersistentVolumeClaimRetentionPolicy) {
+		sts.Spec.PersistentVolumeClaimRetentionPolicy = desired.Spec.PersistentVolumeClaimRetentionPolicy
+		needsUpdate = true
+	}
+
 	if needsUpdate {
 		if err := r.Update(ctx, &sts); err != nil {
 			return ctrl.Result{}, fmt.Errorf("workerpool %s: updating StatefulSet: %w", pool.Name, err)
@@ -1036,9 +1041,10 @@ func (r *WorkerPoolReconciler) buildStatefulSet(pool *kelos.WorkerPool, stsName,
 			Labels:    labels,
 		},
 		Spec: appsv1.StatefulSetSpec{
-			Replicas:            pool.Spec.Replicas,
-			ServiceName:         svcName,
-			PodManagementPolicy: parallel,
+			Replicas:                             pool.Spec.Replicas,
+			ServiceName:                          svcName,
+			PodManagementPolicy:                  parallel,
+			PersistentVolumeClaimRetentionPolicy: pool.Spec.PersistentVolumeClaimRetentionPolicy.DeepCopy(),
 			Selector: &metav1.LabelSelector{
 				MatchLabels: labels,
 			},
@@ -1889,6 +1895,27 @@ func podTemplateSpecEqual(a, b corev1.PodSpec) bool {
 		apiequality.Semantic.DeepEqual(a.ImagePullSecrets, b.ImagePullSecrets)
 }
 
+// pvcRetentionPolicyEqual reports whether two StatefulSet PVC retention
+// policies are equivalent, treating a nil policy or an empty field as Retain
+// to match the API server default.
+func pvcRetentionPolicyEqual(a, b *appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy) bool {
+	effective := func(p *appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy) (appsv1.PersistentVolumeClaimRetentionPolicyType, appsv1.PersistentVolumeClaimRetentionPolicyType) {
+		whenDeleted, whenScaled := appsv1.RetainPersistentVolumeClaimRetentionPolicyType, appsv1.RetainPersistentVolumeClaimRetentionPolicyType
+		if p != nil {
+			if p.WhenDeleted != "" {
+				whenDeleted = p.WhenDeleted
+			}
+			if p.WhenScaled != "" {
+				whenScaled = p.WhenScaled
+			}
+		}
+		return whenDeleted, whenScaled
+	}
+	aDeleted, aScaled := effective(a)
+	bDeleted, bScaled := effective(b)
+	return aDeleted == bDeleted && aScaled == bScaled
+}
+
 func mergeReconcileResults(a, b ctrl.Result) ctrl.Result {
 	a.Requeue = a.Requeue || b.Requeue
 	if b.RequeueAfter != 0 && (a.RequeueAfter == 0 || b.RequeueAfter < a.RequeueAfter) {

internal/controller/workerpool_controller_test.go

@@ -409,6 +409,92 @@ func TestWorkerPoolReconciler_UpdatesReplicas(t *testing.T) {
 	assert.Equal(t, int32(3), *sts.Spec.Replicas)
 }
 
+func reconcileWorkerPoolStatefulSet(t *testing.T, objs ...client.Object) appsv1.StatefulSet {
+	t.Helper()
+	scheme := newWorkerPoolTestScheme()
+	cl := fake.NewClientBuilder().
+		WithScheme(scheme).
+		WithStatusSubresource(&kelos.WorkerPool{}).
+		WithObjects(objs...).
+		Build()
+
+	r := newWorkerPoolReconciler(cl, scheme)
+	_, err := r.Reconcile(context.Background(), ctrl.Request{
+		NamespacedName: types.NamespacedName{Name: "my-pool", Namespace: "default"},
+	})
+	require.NoError(t, err)
+
+	var sts appsv1.StatefulSet
+	require.NoError(t, cl.Get(context.Background(), types.NamespacedName{Name: "wp-my-pool", Namespace: "default"}, &sts))
+	return sts
+}
+
+func TestWorkerPoolReconciler_PVCRetentionPolicyUnsetOnCreate(t *testing.T) {
+	pool := newTestWorkerPool("my-pool", "default", 1)
+
+	sts := reconcileWorkerPoolStatefulSet(t, pool, newTestWorkspace("default"))
+
+	assert.Nil(t, sts.Spec.PersistentVolumeClaimRetentionPolicy)
+}
+
+func TestWorkerPoolReconciler_PVCRetentionPolicyPropagatedOnCreate(t *testing.T) {
+	pool := newTestWorkerPool("my-pool", "default", 1)
+	pool.Spec.PersistentVolumeClaimRetentionPolicy = &appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+		WhenDeleted: appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+		WhenScaled:  appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+	}
+
+	sts := reconcileWorkerPoolStatefulSet(t, pool, newTestWorkspace("default"))
+
+	assert.Equal(t, &appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+		WhenDeleted: appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+		WhenScaled:  appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+	}, sts.Spec.PersistentVolumeClaimRetentionPolicy)
+}
+
+func TestWorkerPoolReconciler_PVCRetentionPolicyUpdatesExistingStatefulSet(t *testing.T) {
+	pool := newTestWorkerPool("my-pool", "default", 1)
+	pool.Spec.PersistentVolumeClaimRetentionPolicy = &appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+		WhenDeleted: appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+		WhenScaled:  appsv1.RetainPersistentVolumeClaimRetentionPolicyType,
+	}
+
+	// Build the StatefulSet as the controller would for the same pool, but
+	// with the API server's Retain/Retain default, so the retention policy is
+	// the only difference.
+	scheme := newWorkerPoolTestScheme()
+	r := newWorkerPoolReconciler(fake.NewClientBuilder().WithScheme(scheme).Build(), scheme)
+	ws := newTestWorkspace("default")
+	existingSTS, err := r.buildStatefulSet(pool, "wp-my-pool", "wp-my-pool", &ws.Spec, nil)
+	require.NoError(t, err)
+	existingSTS.Spec.PersistentVolumeClaimRetentionPolicy = &appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+		WhenDeleted: appsv1.RetainPersistentVolumeClaimRetentionPolicyType,
+		WhenScaled:  appsv1.RetainPersistentVolumeClaimRetentionPolicyType,
+	}
+
+	sts := reconcileWorkerPoolStatefulSet(t, pool, ws, existingSTS)
+
+	assert.Equal(t, &appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+		WhenDeleted: appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+		WhenScaled:  appsv1.RetainPersistentVolumeClaimRetentionPolicyType,
+	}, sts.Spec.PersistentVolumeClaimRetentionPolicy)
+}
+
+func TestPVCRetentionPolicyEqual_TreatsUnsetAsRetain(t *testing.T) {
+	retainAll := &appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+		WhenDeleted: appsv1.RetainPersistentVolumeClaimRetentionPolicyType,
+		WhenScaled:  appsv1.RetainPersistentVolumeClaimRetentionPolicyType,
+	}
+	deleteOnScale := &appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+		WhenScaled: appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+	}
+
+	assert.True(t, pvcRetentionPolicyEqual(nil, retainAll))
+	assert.True(t, pvcRetentionPolicyEqual(&appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{}, retainAll))
+	assert.False(t, pvcRetentionPolicyEqual(nil, deleteOnScale))
+	assert.False(t, pvcRetentionPolicyEqual(retainAll, deleteOnScale))
+}
+
 func TestWorkerPoolReconciler_StatusPhases(t *testing.T) {
 	tests := []struct {
 		name          string

internal/manifests/charts/kelos/charts/kelos-crds/templates/workerpool-crd.yaml

@@ -66,6 +66,37 @@ spec:
           spec:
             description: WorkerPoolSpec defines the desired state of WorkerPool.
             properties:
+              persistentVolumeClaimRetentionPolicy:
+                description: |-
+                  PersistentVolumeClaimRetentionPolicy controls whether the worker PVCs
+                  created from volumeClaimTemplate are deleted when the WorkerPool is
+                  deleted (whenDeleted) or scaled down (whenScaled). Each field accepts
+                  Retain or Delete. When unset, or when a field is omitted, PVCs are
+                  retained.
+                properties:
+                  whenDeleted:
+                    description: |-
+                      WhenDeleted specifies what happens to PVCs created from StatefulSet
+                      VolumeClaimTemplates when the StatefulSet is deleted. The default policy
+                      of `Retain` causes PVCs to not be affected by StatefulSet deletion. The
+                      `Delete` policy causes those PVCs to be deleted.
+                    type: string
+                  whenScaled:
+                    description: |-
+                      WhenScaled specifies what happens to PVCs created from StatefulSet
+                      VolumeClaimTemplates when the StatefulSet is scaled down. The default
+                      policy of `Retain` causes PVCs to not be affected by a scaledown. The
+                      `Delete` policy causes the associated PVCs for any excess pods above
+                      the replica count to be deleted.
+                    type: string
+                type: object
+                x-kubernetes-validations:
+                - message: whenDeleted must be Retain or Delete
+                  rule: '!has(self.whenDeleted) || self.whenDeleted in [''Retain'',
+                    ''Delete'']'
+                - message: whenScaled must be Retain or Delete
+                  rule: '!has(self.whenScaled) || self.whenScaled in [''Retain'',
+                    ''Delete'']'
               replicas:
                 default: 1
                 description: |-

internal/manifests/install-crd.yaml

@@ -69017,6 +69017,37 @@ spec:
           spec:
             description: WorkerPoolSpec defines the desired state of WorkerPool.
             properties:
+              persistentVolumeClaimRetentionPolicy:
+                description: |-
+                  PersistentVolumeClaimRetentionPolicy controls whether the worker PVCs
+                  created from volumeClaimTemplate are deleted when the WorkerPool is
+                  deleted (whenDeleted) or scaled down (whenScaled). Each field accepts
+                  Retain or Delete. When unset, or when a field is omitted, PVCs are
+                  retained.
+                properties:
+                  whenDeleted:
+                    description: |-
+                      WhenDeleted specifies what happens to PVCs created from StatefulSet
+                      VolumeClaimTemplates when the StatefulSet is deleted. The default policy
+                      of `Retain` causes PVCs to not be affected by StatefulSet deletion. The
+                      `Delete` policy causes those PVCs to be deleted.
+                    type: string
+                  whenScaled:
+                    description: |-
+                      WhenScaled specifies what happens to PVCs created from StatefulSet
+                      VolumeClaimTemplates when the StatefulSet is scaled down. The default
+                      policy of `Retain` causes PVCs to not be affected by a scaledown. The
+                      `Delete` policy causes the associated PVCs for any excess pods above
+                      the replica count to be deleted.
+                    type: string
+                type: object
+                x-kubernetes-validations:
+                - message: whenDeleted must be Retain or Delete
+                  rule: '!has(self.whenDeleted) || self.whenDeleted in [''Retain'',
+                    ''Delete'']'
+                - message: whenScaled must be Retain or Delete
+                  rule: '!has(self.whenScaled) || self.whenScaled in [''Retain'',
+                    ''Delete'']'
               replicas:
                 default: 1
                 description: |-

test/integration/workerpool_test.go

@@ -116,4 +116,90 @@ var _ = Describe("WorkerPool Controller", func() {
 			Expect(updatedPool.Status.Phase).To(Equal(kelos.WorkerPoolPhasePending))
 		})
 	})
+
+	Context("When setting a PVC retention policy", func() {
+		newPool := func(name, namespace string, policy *appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy) *kelos.WorkerPool {
+			return &kelos.WorkerPool{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      name,
+					Namespace: namespace,
+				},
+				Spec: kelos.WorkerPoolSpec{
+					Worker: kelos.WorkerSpec{
+						Type: "claude-code",
+						Credentials: &kelos.Credentials{
+							Type: kelos.CredentialTypeNone,
+						},
+						WorkspaceRef: &kelos.WorkspaceReference{
+							Name: "pool-workspace",
+						},
+					},
+					Replicas: ptr.To(int32(1)),
+					VolumeClaimTemplate: corev1.PersistentVolumeClaimSpec{
+						AccessModes: []corev1.PersistentVolumeAccessMode{corev1.ReadWriteOnce},
+						Resources: corev1.VolumeResourceRequirements{
+							Requests: corev1.ResourceList{
+								corev1.ResourceStorage: resource.MustParse("1Gi"),
+							},
+						},
+					},
+					PersistentVolumeClaimRetentionPolicy: policy,
+				},
+			}
+		}
+
+		It("Should propagate the policy to the StatefulSet and update it when changed", func() {
+			ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: "test-workerpool-pvc-retention"}}
+			Expect(k8sClient.Create(ctx, ns)).Should(Succeed())
+
+			ws := &kelos.Workspace{
+				ObjectMeta: metav1.ObjectMeta{Name: "pool-workspace", Namespace: ns.Name},
+				Spec:       kelos.WorkspaceSpec{Repo: "https://github.com/example/repo.git", Ref: "main"},
+			}
+			Expect(k8sClient.Create(ctx, ws)).Should(Succeed())
+
+			By("Creating a WorkerPool that deletes PVCs on scale down")
+			pool := newPool("retention-pool", ns.Name, &appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+				WhenScaled: appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+			})
+			Expect(k8sClient.Create(ctx, pool)).Should(Succeed())
+
+			stsKey := types.NamespacedName{Name: "wp-retention-pool", Namespace: ns.Name}
+			Eventually(func(g Gomega) {
+				sts := &appsv1.StatefulSet{}
+				g.Expect(k8sClient.Get(ctx, stsKey, sts)).To(Succeed())
+				g.Expect(sts.Spec.PersistentVolumeClaimRetentionPolicy).NotTo(BeNil())
+				g.Expect(sts.Spec.PersistentVolumeClaimRetentionPolicy.WhenScaled).To(Equal(appsv1.DeletePersistentVolumeClaimRetentionPolicyType))
+			}, timeout, interval).Should(Succeed())
+
+			By("Changing the WorkerPool to also delete PVCs on deletion")
+			poolKey := types.NamespacedName{Name: "retention-pool", Namespace: ns.Name}
+			Eventually(func(g Gomega) {
+				latest := &kelos.WorkerPool{}
+				g.Expect(k8sClient.Get(ctx, poolKey, latest)).To(Succeed())
+				latest.Spec.PersistentVolumeClaimRetentionPolicy = &appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+					WhenDeleted: appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+					WhenScaled:  appsv1.DeletePersistentVolumeClaimRetentionPolicyType,
+				}
+				g.Expect(k8sClient.Update(ctx, latest)).To(Succeed())
+			}, timeout, interval).Should(Succeed())
+
+			Eventually(func(g Gomega) {
+				sts := &appsv1.StatefulSet{}
+				g.Expect(k8sClient.Get(ctx, stsKey, sts)).To(Succeed())
+				g.Expect(sts.Spec.PersistentVolumeClaimRetentionPolicy).NotTo(BeNil())
+				g.Expect(sts.Spec.PersistentVolumeClaimRetentionPolicy.WhenDeleted).To(Equal(appsv1.DeletePersistentVolumeClaimRetentionPolicyType))
+				g.Expect(sts.Spec.PersistentVolumeClaimRetentionPolicy.WhenScaled).To(Equal(appsv1.DeletePersistentVolumeClaimRetentionPolicyType))
+			}, timeout, interval).Should(Succeed())
+		})
+
+		It("Should reject an unknown policy value", func() {
+			pool := newPool("invalid-retention-pool", "default", &appsv1.StatefulSetPersistentVolumeClaimRetentionPolicy{
+				WhenScaled: "Archive",
+			})
+			err := k8sClient.Create(ctx, pool)
+			Expect(err).To(HaveOccurred())
+			Expect(err.Error()).To(ContainSubstring("whenScaled must be Retain or Delete"))
+		})
+	})
 })