diff --git a/controllers/tau-core/internal/controller/controller_test.go b/controllers/tau-core/internal/controller/controller_test.go index a688a007..3154318f 100644 --- a/controllers/tau-core/internal/controller/controller_test.go +++ b/controllers/tau-core/internal/controller/controller_test.go @@ -1257,9 +1257,14 @@ func TestWorkspaceDeleteCleansWorkspaceAccess(t *testing.T) { now := metav1.Now() workspace.Finalizers = []string{workspaceFinalizer} workspace.DeletionTimestamp = &now + workspace.Spec.Target.Namespace = "new-aurora" workspace.Status.Target.ResolvedNamespace = "aurora" targetNamespace := testNamespace("aurora") - targetNamespace.Labels = map[string]string{labelWorkspace: "aurora"} + targetNamespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + targetNamespace.Labels[labelTeam] = "vision" + targetNamespace.Labels["preserved"] = "true" + targetNamespace.Annotations[annotationResultScope] = workspace.Spec.Defaults.OutputRoot + targetNamespace.Annotations["preserved"] = "true" targetRole := testRole("aurora", defaultRoleName, "aurora") targetBinding := testRoleBinding("aurora", defaultRoleName, "aurora") targetServiceAccount := testServiceAccount("aurora", "tau-workload", "aurora") @@ -1298,6 +1303,35 @@ func TestWorkspaceDeleteCleansWorkspaceAccess(t *testing.T) { if err := c.Get(ctx, client.ObjectKey{Name: "aurora"}, targetNamespace); err != nil { t.Fatalf("target namespace was deleted: %v", err) } + for _, key := range []string{ + labelManagedBy, + labelTeam, + labelWorkspace, + labelWorkspaceLocalQueue, + labelKueueDefaultLocalQueue, + } { + if _, ok := targetNamespace.Labels[key]; ok { + t.Fatalf("target namespace retained workspace label %q: %#v", key, targetNamespace.Labels) + } + } + for _, key := range []string{annotationOwnerUID, annotationResultScope} { + if _, ok := targetNamespace.Annotations[key]; ok { + t.Fatalf("target namespace retained workspace annotation %q: %#v", key, targetNamespace.Annotations) + } + } + for key, want := range map[string]string{ + labelPSAEnforce: "baseline", + labelPSAAudit: "restricted", + labelPSAWarn: "restricted", + "preserved": "true", + } { + if got := targetNamespace.Labels[key]; got != want { + t.Fatalf("target namespace label %q = %q, want %q", key, got, want) + } + } + if got := targetNamespace.Annotations["preserved"]; got != "true" { + t.Fatalf("target namespace preserved annotation = %q, want true", got) + } if err := c.Get(ctx, client.ObjectKey{Name: clusterQueueReaderBindingName("aurora")}, &rbacv1.ClusterRoleBinding{}); !apierrors.IsNotFound(err) { t.Fatalf("expected ClusterQueue reader ClusterRoleBinding deleted, got err=%v", err) } @@ -1308,6 +1342,500 @@ func TestWorkspaceDeleteCleansWorkspaceAccess(t *testing.T) { } } +func TestWorkspaceDeletePreservesNamespaceMetadataForDifferentOwnerUID(t *testing.T) { + ctx := context.Background() + scheme := testScheme(t) + workspace := testWorkspace("aurora") + now := metav1.Now() + workspace.Finalizers = []string{workspaceFinalizer} + workspace.DeletionTimestamp = &now + workspace.Status.Target.ResolvedNamespace = "aurora" + targetNamespace := testNamespace("aurora") + targetNamespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + targetNamespace.Labels[labelTeam] = "vision" + targetNamespace.Labels["preserved"] = "true" + targetNamespace.Annotations[annotationOwnerUID] = "different-workspace-uid" + targetNamespace.Annotations[annotationResultScope] = workspace.Spec.Defaults.OutputRoot + targetNamespace.Annotations["preserved"] = "true" + c := fake.NewClientBuilder(). + WithScheme(scheme). + WithObjects(workspace, targetNamespace). + Build() + reconciler := newTestWorkspaceReconciler(c) + + if _, err := reconciler.Reconcile(ctx, ctrl.Request{NamespacedName: types.NamespacedName{Name: "aurora", Namespace: tauv1alpha1.SystemNamespace}}); err != nil { + t.Fatalf("Reconcile delete error = %v", err) + } + if err := c.Get(ctx, client.ObjectKey{Name: "aurora"}, targetNamespace); err != nil { + t.Fatalf("target namespace was deleted: %v", err) + } + for key, want := range map[string]string{ + labelManagedBy: labelManagedByValue, + labelTeam: "vision", + labelWorkspace: workspace.Name, + labelWorkspaceLocalQueue: workspace.Spec.Queue, + labelKueueDefaultLocalQueue: workspace.Spec.Queue, + labelPSAEnforce: "baseline", + labelPSAAudit: "restricted", + labelPSAWarn: "restricted", + "preserved": "true", + } { + if got := targetNamespace.Labels[key]; got != want { + t.Fatalf("target namespace label %q = %q, want %q", key, got, want) + } + } + for key, want := range map[string]string{ + annotationOwnerUID: "different-workspace-uid", + annotationResultScope: workspace.Spec.Defaults.OutputRoot, + "preserved": "true", + } { + if got := targetNamespace.Annotations[key]; got != want { + t.Fatalf("target namespace annotation %q = %q, want %q", key, got, want) + } + } +} + +func TestWorkspaceDeletePreservesNamespaceRoutingMetadataForExternalQueueWorkloads(t *testing.T) { + for _, test := range []struct { + name string + workloadStatus map[string]any + queueStatus tauv1alpha1.WorkspaceQueueStatus + teamRef *tauv1alpha1.TauClusterObjectReference + wantBlock bool + }{ + { + name: "pending", + queueStatus: tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + }, + wantBlock: true, + }, + { + name: "admitted", + workloadStatus: map[string]any{ + "admission": map[string]any{"clusterQueue": "jobqueue"}, + }, + queueStatus: tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + }, + wantBlock: true, + }, + {name: "pending with empty queue status", wantBlock: true}, + { + name: "pending on compatibility alias", + teamRef: &tauv1alpha1.TauClusterObjectReference{Name: "vision"}, + queueStatus: tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: workspaceClusterQueueName("aurora"), + ClusterQueue: workspaceClusterQueueName("aurora"), + }, + wantBlock: true, + }, + { + name: "finished", + workloadStatus: map[string]any{ + "conditions": []any{map[string]any{"type": "Finished", "status": "True"}}, + }, + queueStatus: tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + now := metav1.Now() + workspace.DeletionTimestamp = &now + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Spec.TeamRef = test.teamRef + workspace.Status.Target.ResolvedNamespace = workspace.Name + workspace.Status.Queue = test.queueStatus + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{ + "key": labelWorkspace, + "operator": "Exists", + }}, + }, + } + workload := newQueueObject(workloadGVK) + workload.SetName("training") + workload.SetNamespace(namespace.Name) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + if test.workloadStatus != nil { + workload.Object["status"] = test.workloadStatus + } + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, namespace, localQueue, clusterQueue, workload). + Build() + reconciler := newTestWorkspaceReconciler(c) + + _, err := reconciler.Reconcile(ctx, ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)}) + if test.wantBlock { + if err == nil || !strings.Contains(err.Error(), "before removing namespace routing metadata") { + t.Fatalf("Reconcile() error = %v, want routing metadata retention", err) + } + } else if err != nil { + t.Fatalf("Reconcile() error = %v", err) + } + + var gotNamespace corev1.Namespace + if err := c.Get(ctx, client.ObjectKeyFromObject(namespace), &gotNamespace); err != nil { + t.Fatal(err) + } + if test.wantBlock { + if gotNamespace.Labels[labelWorkspace] != workspace.Name || + gotNamespace.Annotations[annotationOwnerUID] != string(workspace.UID) { + t.Fatalf("namespace routing metadata was removed while workload remained: labels=%#v annotations=%#v", gotNamespace.Labels, gotNamespace.Annotations) + } + var remaining tauv1alpha1.TauWorkspace + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &remaining); err != nil { + t.Fatal(err) + } + if !controllerutil.ContainsFinalizer(&remaining, workspaceFinalizer) { + t.Fatal("workspace finalizer was removed while routing metadata remained in use") + } + } else { + if gotNamespace.Labels[labelWorkspace] != "" || + gotNamespace.Annotations[annotationOwnerUID] != "" { + t.Fatalf("drained namespace retained stale metadata: labels=%#v annotations=%#v", gotNamespace.Labels, gotNamespace.Annotations) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &tauv1alpha1.TauWorkspace{}); !apierrors.IsNotFound(err) { + t.Fatalf("workspace finalizer was not removed after workload drain: %v", err) + } + } + for _, queue := range []*unstructured.Unstructured{localQueue, clusterQueue} { + got := newQueueObject(queue.GroupVersionKind()) + if err := c.Get(ctx, client.ObjectKeyFromObject(queue), got); err != nil { + t.Fatalf("external %s was mutated or deleted: %v", queue.GetKind(), err) + } + } + }) + } +} + +func TestWorkspaceDeletePreservesWorkloadIdentityWhileExternalQueueWorkloadDrains(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + now := metav1.Now() + workspace.DeletionTimestamp = &now + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Target.ResolvedNamespace = workspace.Name + workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + } + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + account := testServiceAccount(namespace.Name, workspace.Spec.WorkloadIdentity.ServiceAccountName, workspace.Name) + binding := testRoleBinding(namespace.Name, defaultRoleName, workspace.Name) + job := &batchv1.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: "training", + Namespace: namespace.Name, + UID: types.UID("training-job-uid"), + }, + Spec: batchv1.JobSpec{ + Template: corev1.PodTemplateSpec{ + Spec: corev1.PodSpec{ + ServiceAccountName: account.Name, + RestartPolicy: corev1.RestartPolicyNever, + Containers: []corev1.Container{{ + Name: "trainer", + Image: "example.invalid/trainer", + Command: []string{"train"}, + }}, + }, + }, + }, + } + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{ + "key": labelWorkspace, + "operator": "Exists", + }}, + }, + } + workload := newQueueObject(workloadGVK) + workload.SetName("job-training") + workload.SetNamespace(namespace.Name) + workload.SetOwnerReferences([]metav1.OwnerReference{{ + APIVersion: batchv1.SchemeGroupVersion.String(), + Kind: "Job", + Name: job.Name, + UID: job.UID, + }}) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, namespace, account, binding, job, localQueue, clusterQueue, workload). + Build() + reconciler := newTestWorkspaceReconciler(c) + request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} + + if _, err := reconciler.Reconcile(ctx, request); err == nil || + !strings.Contains(err.Error(), "before removing namespace routing metadata") { + t.Fatalf("Reconcile() error = %v, want workload drain wait", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { + t.Fatalf("workspace workload identity was removed before Workload drain: %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(binding), &rbacv1.RoleBinding{}); !apierrors.IsNotFound(err) { + t.Fatalf("researcher RoleBinding survived deletion drain: %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(job), &batchv1.Job{}); err != nil { + t.Fatalf("pending Job was removed during Workload drain: %v", err) + } + + workload.Object["status"] = map[string]any{"conditions": []any{ + map[string]any{"type": "Finished", "status": "True"}, + }} + if err := c.Update(ctx, workload); err != nil { + t.Fatal(err) + } + if _, err := reconciler.Reconcile(ctx, request); err != nil { + t.Fatalf("Reconcile() after Workload drain = %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { + t.Fatalf("workspace workload identity survived after Workload drain: %v", err) + } +} + +func TestLegacyWorkspaceDeletePreservesRoutingMetadataWithEmptyQueueStatus(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationV0Primary: "true"} + now := metav1.Now() + workspace.DeletionTimestamp = &now + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + namespace.Annotations = nil + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{ + "key": labelWorkspace, + "operator": "Exists", + }}, + }, + } + workload := newQueueObject(workloadGVK) + workload.SetName("training") + workload.SetNamespace(namespace.Name) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, namespace, localQueue, clusterQueue, workload). + WithStatusSubresource(&tauv1alpha1.TauWorkspace{}). + Build() + reconciler := newTestWorkspaceReconciler(c) + request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} + + var cleanupErr error + for range 4 { + if _, err := reconciler.Reconcile(ctx, request); err != nil { + cleanupErr = err + break + } + } + if cleanupErr == nil || !strings.Contains(cleanupErr.Error(), "before removing namespace routing metadata") { + t.Fatalf("Reconcile() error = %v, want routing metadata retention after ownership migration", cleanupErr) + } + var gotNamespace corev1.Namespace + if err := c.Get(ctx, client.ObjectKeyFromObject(namespace), &gotNamespace); err != nil { + t.Fatal(err) + } + if gotNamespace.Labels[labelWorkspace] != workspace.Name || + gotNamespace.Annotations[annotationOwnerUID] != string(workspace.UID) { + t.Fatalf("legacy deletion removed migrated routing metadata: labels=%#v annotations=%#v", gotNamespace.Labels, gotNamespace.Annotations) + } + var remaining tauv1alpha1.TauWorkspace + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &remaining); err != nil { + t.Fatal(err) + } + if remaining.Status.Queue != (tauv1alpha1.WorkspaceQueueStatus{}) { + t.Fatalf("legacy deletion unexpectedly populated queue status: %#v", remaining.Status.Queue) + } + if !controllerutil.ContainsFinalizer(&remaining, workspaceFinalizer) { + t.Fatal("workspace finalizer was removed while migrated routing metadata remained in use") + } +} + +func TestWorkspaceDeleteIgnoresUnrelatedWorkloadsWhenCleaningNamespaceMetadata(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + now := metav1.Now() + workspace.DeletionTimestamp = &now + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Target.ResolvedNamespace = workspace.Name + workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + } + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + namespace.Labels["queue-scope"] = "shared" + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := testClusterQueue("jobqueue") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchLabels": map[string]any{"queue-scope": "shared"}, + }, + } + selectorIndependent := newQueueObject(workloadGVK) + selectorIndependent.SetName("selector-independent") + selectorIndependent.SetNamespace(namespace.Name) + selectorIndependent.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + unrelated := newQueueObject(workloadGVK) + unrelated.SetName("unrelated") + unrelated.SetNamespace(namespace.Name) + unrelated.Object["spec"] = map[string]any{"queueName": "missing-queue"} + foreign := newQueueObject(workloadGVK) + foreign.SetName("foreign") + foreign.SetNamespace("other") + foreign.Object["spec"] = map[string]any{"queueName": "jobqueue"} + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, namespace, localQueue, clusterQueue, selectorIndependent, unrelated, foreign). + Build() + + if _, err := newTestWorkspaceReconciler(c).Reconcile(ctx, ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)}); err != nil { + t.Fatalf("Reconcile() error = %v", err) + } + var gotNamespace corev1.Namespace + if err := c.Get(ctx, client.ObjectKeyFromObject(namespace), &gotNamespace); err != nil { + t.Fatal(err) + } + if gotNamespace.Labels[labelWorkspace] != "" || + gotNamespace.Annotations[annotationOwnerUID] != "" { + t.Fatalf("unrelated workloads retained stale namespace metadata: labels=%#v annotations=%#v", gotNamespace.Labels, gotNamespace.Annotations) + } +} + +func TestWorkspaceDeleteFailsClosedWhenExternalQueueRoutingCannotBeVerified(t *testing.T) { + for _, test := range []struct { + name string + failGVK schema.GroupVersionKind + failKey client.ObjectKey + mutateQueue func(*unstructured.Unstructured) + wantError string + }{ + { + name: "local queue read", + failGVK: localQueueGVK, + failKey: client.ObjectKey{Name: "jobqueue", Namespace: "aurora"}, + wantError: "read LocalQueue aurora/jobqueue", + }, + { + name: "cluster queue read", + failGVK: clusterQueueGVK, + failKey: client.ObjectKey{Name: "jobqueue"}, + wantError: "read ClusterQueue jobqueue", + }, + { + name: "malformed selector", + mutateQueue: func(queue *unstructured.Unstructured) { + queue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchExpressions": "invalid", + }, + } + }, + wantError: "read ClusterQueue \"jobqueue\" namespaceSelector", + }, + } { + t.Run(test.name, func(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + now := metav1.Now() + workspace.DeletionTimestamp = &now + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Target.ResolvedNamespace = workspace.Name + workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + } + namespace := testNamespace(workspace.Name) + namespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + localQueue := testLocalQueue(namespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := testClusterQueue("jobqueue") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchLabels": map[string]any{labelWorkspace: workspace.Name}, + }, + } + if test.mutateQueue != nil { + test.mutateQueue(clusterQueue) + } + workload := newQueueObject(workloadGVK) + workload.SetName("training") + workload.SetNamespace(namespace.Name) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + base := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, namespace, localQueue, clusterQueue, workload). + Build() + reconciler := newTestWorkspaceReconciler(base) + if !test.failGVK.Empty() { + reconciler.APIReader = &objectReadFailingClient{ + Client: base, + gvk: test.failGVK, + key: test.failKey, + err: errors.New("access denied"), + } + } + + _, err := reconciler.Reconcile(ctx, ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)}) + if err == nil || !strings.Contains(err.Error(), test.wantError) { + t.Fatalf("Reconcile() error = %v, want %q", err, test.wantError) + } + var gotNamespace corev1.Namespace + if err := base.Get(ctx, client.ObjectKeyFromObject(namespace), &gotNamespace); err != nil { + t.Fatal(err) + } + if gotNamespace.Labels[labelWorkspace] != workspace.Name || + gotNamespace.Annotations[annotationOwnerUID] != string(workspace.UID) { + t.Fatalf("failed routing verification removed namespace metadata: labels=%#v annotations=%#v", gotNamespace.Labels, gotNamespace.Annotations) + } + }) + } +} + func TestWorkspaceDeleteWaitsForClusterQueueFinalizer(t *testing.T) { ctx := context.Background() scheme := testScheme(t) @@ -2288,9 +2816,126 @@ func TestUnmigratedWorkspaceRetargetMigratesHistoryBeforeDeletion(t *testing.T) if err := deletionClient.Get(ctx, client.ObjectKeyFromObject(newNamespace), &gotNew); err != nil { t.Fatalf("Get new desired Namespace: %v", err) } - if gotNew.Annotations[annotationOwnerUID] != string(workspace.UID) || - gotNew.Labels[labelWorkspace] != workspace.Name { - t.Fatalf("normal retarget deletion did not preserve adopted Namespace ownership: labels=%#v annotations=%#v", gotNew.Labels, gotNew.Annotations) + if gotNew.Annotations[annotationOwnerUID] != "" || + gotNew.Labels[labelWorkspace] != "" { + t.Fatalf("normal retarget deletion retained adopted Namespace ownership: labels=%#v annotations=%#v", gotNew.Labels, gotNew.Annotations) + } +} + +func TestWorkspaceRetargetPreservesOldWorkloadIdentityWhileExternalQueueWorkloadDrains(t *testing.T) { + ctx := context.Background() + workspace := testWorkspace("aurora") + workspace.Finalizers = []string{workspaceFinalizer} + workspace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + workspace.Spec.Target.Namespace = "new-target" + workspace.Spec.Target.CreateNamespace = false + workspace.Spec.Queue = "jobqueue" + workspace.Status.Phase = tauv1alpha1.WorkspacePhaseReady + workspace.Status.Target.ResolvedNamespace = "old-target" + workspace.Status.Queue = tauv1alpha1.WorkspaceQueueStatus{ + LocalQueue: "jobqueue", + ClusterQueue: "jobqueue", + ClusterQueueUID: "jobqueue-uid", + } + oldNamespace := testNamespace("old-target") + oldNamespace.Labels = workspaceNamespaceLabels(workspace.Name, workspace.Spec.Queue) + oldNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + newNamespace := testNamespace("new-target") + newNamespace.Labels = map[string]string{labelWorkspace: workspace.Name} + newNamespace.Annotations = map[string]string{annotationOwnerUID: string(workspace.UID)} + account := testServiceAccount(oldNamespace.Name, workspace.Spec.WorkloadIdentity.ServiceAccountName, workspace.Name) + binding := testRoleBinding(oldNamespace.Name, defaultRoleName, workspace.Name) + binding.Subjects = []rbacv1.Subject{{ + Kind: rbacv1.GroupKind, + APIGroup: rbacv1.GroupName, + Name: "former-researchers", + }} + job := &batchv1.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: "training", + Namespace: oldNamespace.Name, + UID: types.UID("training-job-uid"), + }, + Spec: batchv1.JobSpec{ + Template: corev1.PodTemplateSpec{ + Spec: corev1.PodSpec{ + ServiceAccountName: account.Name, + RestartPolicy: corev1.RestartPolicyNever, + Containers: []corev1.Container{{ + Name: "trainer", + Image: "example.invalid/trainer", + Command: []string{"train"}, + }}, + }, + }, + }, + } + localQueue := testLocalQueue(oldNamespace.Name, "jobqueue", "jobqueue") + localQueue.SetAnnotations(nil) + clusterQueue := clusterQueueWithUID("jobqueue", "jobqueue-uid") + clusterQueue.Object["spec"] = map[string]any{ + "namespaceSelector": map[string]any{ + "matchExpressions": []any{map[string]any{ + "key": labelWorkspace, + "operator": "Exists", + }}, + }, + } + workload := newQueueObject(workloadGVK) + workload.SetName("job-training") + workload.SetNamespace(oldNamespace.Name) + workload.SetOwnerReferences([]metav1.OwnerReference{{ + APIVersion: batchv1.SchemeGroupVersion.String(), + Kind: "Job", + Name: job.Name, + UID: job.UID, + }}) + workload.Object["spec"] = map[string]any{"queueName": localQueue.GetName()} + c := fake.NewClientBuilder(). + WithScheme(testScheme(t)). + WithObjects(workspace, oldNamespace, newNamespace, account, binding, job, localQueue, clusterQueue, workload). + WithStatusSubresource(&tauv1alpha1.TauWorkspace{}). + Build() + reconciler := newTestWorkspaceReconciler(c) + request := ctrl.Request{NamespacedName: client.ObjectKeyFromObject(workspace)} + + if _, err := reconciler.Reconcile(ctx, request); err != nil { + t.Fatalf("Reconcile() during retarget drain = %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); err != nil { + t.Fatalf("old workload identity was removed before Workload drain: %v", err) + } + var gotBinding rbacv1.RoleBinding + if err := c.Get(ctx, client.ObjectKeyFromObject(binding), &gotBinding); err != nil { + t.Fatalf("Get old researcher RoleBinding: %v", err) + } + wantSubjects := []rbacv1.Subject{rbacSubject(*workspace.Spec.KubernetesSubject, oldNamespace.Name)} + if !reflect.DeepEqual(gotBinding.Subjects, wantSubjects) { + t.Fatalf("old researcher subjects = %#v, want current subject %#v", gotBinding.Subjects, wantSubjects) + } + var waiting tauv1alpha1.TauWorkspace + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { + t.Fatal(err) + } + if waiting.Status.Target.ResolvedNamespace != oldNamespace.Name { + t.Fatalf("resolved namespace = %q, want old target retained during drain", waiting.Status.Target.ResolvedNamespace) + } + + workload.Object["status"] = map[string]any{"conditions": []any{ + map[string]any{"type": "Finished", "status": "True"}, + }} + if err := c.Update(ctx, workload); err != nil { + t.Fatal(err) + } + reconcileWorkspaceTimes(t, reconciler, ctx, workspace.Name, 3) + if err := c.Get(ctx, client.ObjectKeyFromObject(account), &corev1.ServiceAccount{}); !apierrors.IsNotFound(err) { + t.Fatalf("old workload identity survived completed retarget: %v", err) + } + if err := c.Get(ctx, client.ObjectKeyFromObject(workspace), &waiting); err != nil { + t.Fatal(err) + } + if waiting.Status.Target.ResolvedNamespace != newNamespace.Name { + t.Fatalf("resolved namespace = %q, want completed retarget %q", waiting.Status.Target.ResolvedNamespace, newNamespace.Name) } } @@ -4613,6 +5258,25 @@ type statusUpdateFailingClient struct { updateErr error } +type objectReadFailingClient struct { + client.Client + gvk schema.GroupVersionKind + key client.ObjectKey + err error +} + +func (c *objectReadFailingClient) Get( + ctx context.Context, + key client.ObjectKey, + obj client.Object, + opts ...client.GetOption, +) error { + if obj.GetObjectKind().GroupVersionKind() == c.gvk && key == c.key { + return c.err + } + return c.Client.Get(ctx, key, obj, opts...) +} + func (c *statusUpdateFailingClient) Status() client.SubResourceWriter { return &statusUpdateFailingWriter{SubResourceWriter: c.Client.Status(), client: c} } diff --git a/controllers/tau-core/internal/controller/workspace_cleanup.go b/controllers/tau-core/internal/controller/workspace_cleanup.go index 6009f4ba..6fa12599 100644 --- a/controllers/tau-core/internal/controller/workspace_cleanup.go +++ b/controllers/tau-core/internal/controller/workspace_cleanup.go @@ -5,6 +5,7 @@ package controller import ( "context" + "errors" "fmt" "strings" @@ -15,6 +16,8 @@ import ( "k8s.io/apimachinery/pkg/api/meta" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "k8s.io/apimachinery/pkg/labels" + "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/runtime/schema" "sigs.k8s.io/controller-runtime/pkg/client" ) @@ -77,20 +80,90 @@ func (r *TauWorkspaceReconciler) cleanupStaleNamespaceMetadata( workspace *tauv1alpha1.TauWorkspace, namespaceName string, ) error { + namespace, err := r.workspaceNamespaceForMetadataCleanup(ctx, workspace, namespaceName) + if err != nil || namespace == nil { + return err + } + inUse, message, err := r.namespaceRoutingMetadataInUse(ctx, workspace, namespace) + if err != nil { + return err + } + if inUse { + return errors.New(message) + } + for _, key := range []string{ + labelManagedBy, + labelTeam, + labelWorkspace, + labelWorkspaceLocalQueue, + labelKueueDefaultLocalQueue, + } { + delete(namespace.Labels, key) + } + if namespace.Annotations != nil { + delete(namespace.Annotations, annotationResultScope) + delete(namespace.Annotations, annotationOwnerUID) + } + return r.Update(ctx, namespace) +} + +func (r *TauWorkspaceReconciler) workspaceNamespaceForMetadataCleanup( + ctx context.Context, + workspace *tauv1alpha1.TauWorkspace, + namespaceName string, +) (*corev1.Namespace, error) { if reservedNamespaceReason(namespaceName, systemNamespace(r.SystemNamespace)) != "" { - return nil + return nil, nil } var namespace corev1.Namespace - if err := r.Get(ctx, client.ObjectKey{Name: namespaceName}, &namespace); err != nil { + if err := r.apiReader().Get(ctx, client.ObjectKey{Name: namespaceName}, &namespace); err != nil { if apierrors.IsNotFound(err) { - return nil + return nil, nil } - return err + return nil, err } if namespace.Labels[labelWorkspace] != workspace.Name || namespace.Annotations[annotationOwnerUID] != string(workspace.UID) || workspace.UID == "" { - return nil + return nil, nil + } + return &namespace, nil +} + +func (r *TauWorkspaceReconciler) waitForNamespaceRoutingMetadataDrain( + ctx context.Context, + workspace *tauv1alpha1.TauWorkspace, + namespaceName string, +) error { + namespace, err := r.workspaceNamespaceForMetadataCleanup(ctx, workspace, namespaceName) + if err != nil || namespace == nil { + return err + } + inUse, message, err := r.namespaceRoutingMetadataInUse(ctx, workspace, namespace) + if err != nil { + return err + } + if inUse { + return errors.New(message) + } + return nil +} + +func (r *TauWorkspaceReconciler) namespaceRoutingMetadataInUse( + ctx context.Context, + workspace *tauv1alpha1.TauWorkspace, + namespace *corev1.Namespace, +) (bool, string, error) { + workloads := &unstructured.UnstructuredList{} + workloads.SetGroupVersionKind(schema.GroupVersionKind{ + Group: workloadGVK.Group, Version: workloadGVK.Version, Kind: workloadGVK.Kind + "List", + }) + if err := r.apiReader().List(ctx, workloads, client.InNamespace(namespace.Name)); err != nil { + return false, "", fmt.Errorf("list Workloads in namespace %q before routing metadata cleanup: %w", namespace.Name, err) + } + cleanedLabels := make(map[string]string, len(namespace.Labels)) + for key, value := range namespace.Labels { + cleanedLabels[key] = value } for _, key := range []string{ labelManagedBy, @@ -99,16 +172,103 @@ func (r *TauWorkspaceReconciler) cleanupStaleNamespaceMetadata( labelWorkspaceLocalQueue, labelKueueDefaultLocalQueue, } { - delete(namespace.Labels, key) + delete(cleanedLabels, key) } - if namespace.Annotations != nil { - delete(namespace.Annotations, annotationResultScope) - delete(namespace.Annotations, annotationOwnerUID) - } - return r.Update(ctx, &namespace) + for i := range workloads.Items { + workload := &workloads.Items[i] + if workloadFinished(workload) { + continue + } + queueName, found, err := unstructured.NestedString(workload.Object, "spec", "queueName") + if err != nil { + return false, "", fmt.Errorf("read Workload %s queueName: %w", client.ObjectKeyFromObject(workload), err) + } + queueName = strings.TrimSpace(queueName) + if !found || queueName == "" { + continue + } + localQueue := newQueueObject(localQueueGVK) + localQueueKey := client.ObjectKey{Name: queueName, Namespace: namespace.Name} + if err := r.apiReader().Get(ctx, localQueueKey, localQueue); err != nil { + if apierrors.IsNotFound(err) { + continue + } + return false, "", fmt.Errorf("read LocalQueue %s before routing metadata cleanup: %w", localQueueKey, err) + } + clusterQueueName, found, err := unstructured.NestedString(localQueue.Object, "spec", "clusterQueue") + if err != nil { + return false, "", fmt.Errorf("read LocalQueue %s backing ClusterQueue: %w", localQueueKey, err) + } + if !found || strings.TrimSpace(clusterQueueName) == "" { + return false, "", fmt.Errorf("LocalQueue %s has no backing ClusterQueue", localQueueKey) + } + admittedClusterQueue, admitted, err := unstructured.NestedString(workload.Object, "status", "admission", "clusterQueue") + if err != nil { + return false, "", fmt.Errorf("read Workload %s admitted ClusterQueue: %w", client.ObjectKeyFromObject(workload), err) + } + if admitted && strings.TrimSpace(admittedClusterQueue) != "" { + clusterQueueName = admittedClusterQueue + } + clusterQueue := newQueueObject(clusterQueueGVK) + if err := r.apiReader().Get(ctx, client.ObjectKey{Name: clusterQueueName}, clusterQueue); err != nil { + return false, "", fmt.Errorf("read ClusterQueue %s before routing metadata cleanup: %w", clusterQueueName, err) + } + if expectedUID := strings.TrimSpace(workspace.Status.Queue.ClusterQueueUID); expectedUID != "" && + clusterQueueName == strings.TrimSpace(workspace.Status.Queue.ClusterQueue) && + string(clusterQueue.GetUID()) != expectedUID { + return false, "", fmt.Errorf( + "ClusterQueue %q UID changed from %q to %q before routing metadata cleanup", + clusterQueueName, + expectedUID, + clusterQueue.GetUID(), + ) + } + selectorMap, found, err := unstructured.NestedMap(clusterQueue.Object, "spec", "namespaceSelector") + if err != nil { + return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) + } + labelSelector := metav1.LabelSelector{} + if found { + if err := runtime.DefaultUnstructuredConverter.FromUnstructured(selectorMap, &labelSelector); err != nil { + return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) + } + } + selector, err := metav1.LabelSelectorAsSelector(&labelSelector) + if err != nil { + return false, "", fmt.Errorf("read ClusterQueue %q namespaceSelector: %w", clusterQueueName, err) + } + if !selector.Matches(labels.Set(namespace.Labels)) { + return false, "", fmt.Errorf( + "ClusterQueue %q namespaceSelector does not match namespace %q for unfinished Workload %s", + clusterQueueName, + namespace.Name, + client.ObjectKeyFromObject(workload), + ) + } + if selector.Matches(labels.Set(cleanedLabels)) { + continue + } + return true, fmt.Sprintf( + "waiting for Workload %s using retained LocalQueue %q and ClusterQueue %q before removing namespace routing metadata", + client.ObjectKeyFromObject(workload), + queueName, + clusterQueueName, + ), nil + } + return false, "", nil } func (r *TauWorkspaceReconciler) cleanupWorkspaceAccess(ctx context.Context, workspace *tauv1alpha1.TauWorkspace) error { + targetNamespace := strings.TrimSpace(workspace.Status.Target.ResolvedNamespace) + if targetNamespace == "" { + targetNamespace = resolvedNamespace(workspace) + } + if err := r.cleanupOwnedResearcherRoleBindings(ctx, workspace); err != nil { + return err + } + if err := r.waitForNamespaceRoutingMetadataDrain(ctx, workspace, targetNamespace); err != nil { + return err + } if err := r.cleanupStaleTargetRBAC(ctx, workspace, "", "", false); err != nil { return err } @@ -121,7 +281,10 @@ func (r *TauWorkspaceReconciler) cleanupWorkspaceAccess(ctx context.Context, wor if err := r.cleanupWorkspaceClusterQueue(ctx, workspace); err != nil { return err } - return r.cleanupSystemReaderRBAC(ctx, workspace) + if err := r.cleanupSystemReaderRBAC(ctx, workspace); err != nil { + return err + } + return r.cleanupStaleNamespaceMetadata(ctx, workspace, targetNamespace) } func (r *TauWorkspaceReconciler) cleanupWorkspaceClusterQueue(ctx context.Context, workspace *tauv1alpha1.TauWorkspace) error { @@ -234,10 +397,10 @@ func (r *TauWorkspaceReconciler) cleanupStaleTargetRBAC( return nil } -func (r *TauWorkspaceReconciler) cleanupOwnedNamespaceAccess( +func (r *TauWorkspaceReconciler) cleanupOwnedNamespaceWorkloadIdentity( ctx context.Context, workspace *tauv1alpha1.TauWorkspace, - targetNamespace, keepServiceAccount string, + keepServiceAccount string, ) error { var namespaces corev1.NamespaceList if err := r.List(ctx, &namespaces, client.MatchingLabels{ @@ -260,7 +423,7 @@ func (r *TauWorkspaceReconciler) cleanupOwnedNamespaceAccess( return err } } - return r.reconcileOwnedNamespaceResearcherAccess(ctx, workspace, targetNamespace) + return nil } func canDeleteLegacyWorkspaceRole(role *rbacv1.Role, workspace *tauv1alpha1.TauWorkspace) bool { diff --git a/controllers/tau-core/internal/controller/workspace_rbac.go b/controllers/tau-core/internal/controller/workspace_rbac.go index 75ce902f..14c42cda 100644 --- a/controllers/tau-core/internal/controller/workspace_rbac.go +++ b/controllers/tau-core/internal/controller/workspace_rbac.go @@ -126,6 +126,29 @@ func (r *TauWorkspaceReconciler) cleanupResearcherRBAC( return r.deleteOwnedObject(ctx, binding, workspace) } +func (r *TauWorkspaceReconciler) cleanupOwnedResearcherRoleBindings( + ctx context.Context, + workspace *tauv1alpha1.TauWorkspace, +) error { + var bindings rbacv1.RoleBindingList + if err := r.List(ctx, &bindings, client.MatchingLabels{ + labelManagedBy: labelManagedByValue, + labelWorkspace: workspace.Name, + }); err != nil { + return err + } + for i := range bindings.Items { + binding := &bindings.Items[i] + if binding.Name != defaultRoleName { + continue + } + if err := r.deleteOwnedObject(ctx, binding, workspace); err != nil { + return err + } + } + return nil +} + func clusterQueueReaderBindingName(workspaceName string) string { return "tau-clusterqueue-reader-" + workspaceName } diff --git a/controllers/tau-core/internal/controller/workspace_sync.go b/controllers/tau-core/internal/controller/workspace_sync.go index d8e1a405..c14d58ae 100644 --- a/controllers/tau-core/internal/controller/workspace_sync.go +++ b/controllers/tau-core/internal/controller/workspace_sync.go @@ -103,10 +103,22 @@ func (r *TauWorkspaceReconciler) syncWorkspace(ctx context.Context, workspace *t } previousNamespace := workspace.Status.Target.ResolvedNamespace previousAccessReady := true - if err := r.cleanupOwnedNamespaceAccess(ctx, workspace, targetNamespace, keepServiceAccount); err != nil { + if previousNamespace != "" && previousNamespace != targetNamespace { + if err := r.waitForNamespaceRoutingMetadataDrain(ctx, workspace, previousNamespace); err != nil { + previousAccessReady = false + conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "NamespaceCleanupFailed", err.Error(), workspace.Generation)) + } + } + if err := r.reconcileOwnedNamespaceResearcherAccess(ctx, workspace, targetNamespace); err != nil { previousAccessReady = false conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "RBACCleanupFailed", err.Error(), workspace.Generation)) } + if previousAccessReady { + if err := r.cleanupOwnedNamespaceWorkloadIdentity(ctx, workspace, keepServiceAccount); err != nil { + previousAccessReady = false + conditions = append(conditions, boolCondition(tauv1alpha1.ConditionDriftDetected, true, "RBACCleanupFailed", err.Error(), workspace.Generation)) + } + } currentTargetReady := namespaceReady && rbacReady && rbacErr == nil && queueReady && previousAccessReady if workspace.Spec.WorkloadIdentity != nil {