Fix race conditions and improve finalizer reliability

- Replace stale list iteration with proper polling using fresh queries
- Add timeout and context cancellation handling with configurable constants
- Implement field selector optimization with backward compatibility fallback
- Replace brittle time.Sleep with condition-based polling in integration tests
- Add comprehensive edge case tests for context cancellation and partial cleanup failures

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: jlegrone

internal/controller/finalizer_test.go

@@ -6,6 +6,7 @@ package controller
 
 import (
 	"context"
+	"strings"
 	"testing"
 
 	temporaliov1alpha1 "github.com/temporalio/temporal-worker-controller/api/v1alpha1"
@@ -239,3 +240,163 @@ func TestHandleDeletion(t *testing.T) {
 	// by checking if the finalizer was removed (which we can't easily do since the resource is deleted)
 	// Instead, we'll verify that the deletion handling completed without error, which means cleanup was successful
 }
+
+func TestCleanupWithContextCancellation(t *testing.T) {
+	// Create a context that will be cancelled during cleanup
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// Create a TemporalWorkerDeployment using test helpers
+	workerDeploy := testhelpers.ModifyObj(testhelpers.MakeTWDWithName("test-worker", "default"), func(twd *temporaliov1alpha1.TemporalWorkerDeployment) *temporaliov1alpha1.TemporalWorkerDeployment {
+		twd.UID = "worker-uid-123"
+		return twd
+	})
+
+	// Create fake client using test helpers
+	client := testhelpers.SetupFakeClient()
+
+	reconciler := &TemporalWorkerDeploymentReconciler{
+		Client: client,
+		Scheme: testhelpers.SetupTestScheme(),
+	}
+
+	// Create a test logger using testlogr
+	logger := testlogr.New(t)
+
+	// Cancel the context immediately to simulate cancellation during cleanup
+	cancel()
+
+	// Test cleanup with cancelled context - should handle gracefully
+	err := reconciler.cleanupManagedResources(ctx, logger, workerDeploy)
+	if err == nil {
+		t.Error("Expected error when context is cancelled during cleanup")
+	}
+
+	// Error should indicate context cancellation
+	if ctx.Err() != context.Canceled {
+		t.Error("Context should be cancelled")
+	}
+}
+
+func TestWaitForOwnedDeploymentsTimeout(t *testing.T) {
+	ctx := context.Background()
+
+	// Create a TemporalWorkerDeployment using test helpers
+	workerDeploy := testhelpers.ModifyObj(testhelpers.MakeTWDWithName("test-worker", "default"), func(twd *temporaliov1alpha1.TemporalWorkerDeployment) *temporaliov1alpha1.TemporalWorkerDeployment {
+		twd.UID = "worker-uid-123"
+		return twd
+	})
+
+	// Create a deployment that won't be deleted (simulate stuck deletion)
+	persistentDeployment := &appsv1.Deployment{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "persistent-deployment",
+			Namespace: "default",
+			OwnerReferences: []metav1.OwnerReference{
+				{
+					APIVersion: apiGVStr,
+					Kind:       "TemporalWorkerDeployment",
+					Name:       "test-worker",
+					UID:        "worker-uid-123",
+				},
+			},
+		},
+	}
+
+	// Create fake client with the deployment that won't be deleted
+	client := testhelpers.SetupFakeClient(persistentDeployment)
+
+	reconciler := &TemporalWorkerDeploymentReconciler{
+		Client: client,
+		Scheme: testhelpers.SetupTestScheme(),
+	}
+
+	// Create a test logger using testlogr
+	logger := testlogr.New(t)
+
+	// Test with a very short timeout to simulate timeout condition
+	// This will use the actual waitForOwnedDeploymentsToBeDeleted method which has built-in timeout
+	err := reconciler.waitForOwnedDeploymentsToBeDeleted(ctx, logger, workerDeploy)
+
+	// Should timeout waiting for deployments to be deleted
+	if err == nil {
+		t.Error("Expected timeout error when deployments don't get deleted")
+	}
+
+	// Error message should indicate timeout
+	if err != nil && !contains(err.Error(), "timeout") {
+		t.Errorf("Expected timeout error, got: %v", err)
+	}
+}
+
+func TestPartialCleanupFailure(t *testing.T) {
+	ctx := context.Background()
+
+	// Create a TemporalWorkerDeployment using test helpers
+	workerDeploy := testhelpers.ModifyObj(testhelpers.MakeTWDWithName("test-worker", "default"), func(twd *temporaliov1alpha1.TemporalWorkerDeployment) *temporaliov1alpha1.TemporalWorkerDeployment {
+		twd.UID = "worker-uid-123"
+		return twd
+	})
+
+	// Create multiple deployments owned by the worker deployment
+	deployment1 := &appsv1.Deployment{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "deployment-1",
+			Namespace: "default",
+			OwnerReferences: []metav1.OwnerReference{
+				{
+					APIVersion: apiGVStr,
+					Kind:       "TemporalWorkerDeployment",
+					Name:       "test-worker",
+					UID:        "worker-uid-123",
+				},
+			},
+		},
+	}
+
+	deployment2 := &appsv1.Deployment{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "deployment-2",
+			Namespace: "default",
+			OwnerReferences: []metav1.OwnerReference{
+				{
+					APIVersion: apiGVStr,
+					Kind:       "TemporalWorkerDeployment",
+					Name:       "test-worker",
+					UID:        "worker-uid-123",
+				},
+			},
+		},
+	}
+
+	// Create fake client with multiple deployments
+	client := testhelpers.SetupFakeClient(deployment1, deployment2)
+
+	reconciler := &TemporalWorkerDeploymentReconciler{
+		Client: client,
+		Scheme: testhelpers.SetupTestScheme(),
+	}
+
+	// Create a test logger using testlogr
+	logger := testlogr.New(t)
+
+	// Delete one deployment manually to simulate partial cleanup
+	err := client.Delete(ctx, deployment1)
+	if err != nil {
+		t.Fatalf("Failed to delete deployment1: %v", err)
+	}
+
+	// Now test cleanup - it should handle the mixed state gracefully
+	// (one deployment already deleted, one still exists)
+	err = reconciler.cleanupManagedResources(ctx, logger, workerDeploy)
+
+	// This should eventually succeed as the cleanup logic should handle
+	// deployments that are already deleted gracefully
+	if err != nil && !contains(err.Error(), "timeout") {
+		t.Errorf("Cleanup should handle partial cleanup gracefully, got error: %v", err)
+	}
+}
+
+// Helper function to check if a string contains a substring
+func contains(s, substr string) bool {
+	return strings.Contains(s, substr)
+}

internal/controller/worker_controller.go

@@ -17,6 +17,7 @@ import (
 	appsv1 "k8s.io/api/apps/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/fields"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	ctrl "sigs.k8s.io/controller-runtime"
@@ -36,6 +37,10 @@ const (
 	buildIDLabel   = "temporal.io/build-id"
 	// TemporalWorkerDeploymentFinalizer is the finalizer used to ensure proper cleanup of resources
 	TemporalWorkerDeploymentFinalizer = "temporal.io/temporal-worker-deployment-finalizer"
+
+	// Cleanup timeout and polling constants
+	cleanupTimeout      = 2 * time.Minute
+	cleanupPollInterval = 5 * time.Second
 )
 
 // TemporalWorkerDeploymentReconciler reconciles a TemporalWorkerDeployment object
@@ -236,18 +241,29 @@ func (r *TemporalWorkerDeploymentReconciler) handleDeletion(ctx context.Context,
 func (r *TemporalWorkerDeploymentReconciler) cleanupManagedResources(ctx context.Context, l logr.Logger, workerDeploy *temporaliov1alpha1.TemporalWorkerDeployment) error {
 	l.Info("Cleaning up managed resources")
 
-	// List all deployments owned by this TemporalWorkerDeployment
-	deploymentList := &appsv1.DeploymentList{}
+	// Try to use field selector for efficient querying of owned deployments
+	// Fall back to listing all deployments if field selector is not available (e.g., in tests)
 	listOpts := &client.ListOptions{
-		Namespace: workerDeploy.Namespace,
+		Namespace:     workerDeploy.Namespace,
+		FieldSelector: fields.OneTermEqualSelector(deployOwnerKey, workerDeploy.Name),
 	}
 
-	if err := r.List(ctx, deploymentList, listOpts); err != nil {
-		return fmt.Errorf("failed to list deployments: %w", err)
+	deploymentList := &appsv1.DeploymentList{}
+	err := r.List(ctx, deploymentList, listOpts)
+	if err != nil {
+		// If field selector fails (common in tests), fall back to listing all deployments
+		l.Info("Field selector not available, falling back to listing all deployments", "error", err.Error())
+		listOpts = &client.ListOptions{
+			Namespace: workerDeploy.Namespace,
+		}
+		if err := r.List(ctx, deploymentList, listOpts); err != nil {
+			return fmt.Errorf("failed to list deployments: %w", err)
+		}
 	}
 
-	// Filter deployments owned by this TemporalWorkerDeployment and delete them
+	// Delete all owned deployments
 	for _, deployment := range deploymentList.Items {
+		// Check ownership for all deployments when not using field selector
 		if r.isOwnedByWorkerDeployment(&deployment, workerDeploy) {
 			l.Info("Deleting managed deployment", "deployment", deployment.Name)
 			if err := r.Delete(ctx, &deployment); err != nil && !apierrors.IsNotFound(err) {
@@ -256,29 +272,64 @@ func (r *TemporalWorkerDeploymentReconciler) cleanupManagedResources(ctx context
 		}
 	}
 
-	// Wait for all owned deployments to be deleted
-	for _, deployment := range deploymentList.Items {
-		if r.isOwnedByWorkerDeployment(&deployment, workerDeploy) {
-			// Check if deployment still exists
-			currentDeployment := &appsv1.Deployment{}
-			err := r.Get(ctx, types.NamespacedName{
-				Namespace: deployment.Namespace,
-				Name:      deployment.Name,
-			}, currentDeployment)
-
-			if err == nil {
-				// Deployment still exists, requeue to wait for deletion
-				l.Info("Waiting for deployment to be deleted", "deployment", deployment.Name)
-				return fmt.Errorf("still waiting for deployment %s to be deleted", deployment.Name)
-			} else if !apierrors.IsNotFound(err) {
-				return fmt.Errorf("failed to check deployment status %s: %w", deployment.Name, err)
+	// Wait for all owned deployments to be deleted with proper polling
+	return r.waitForOwnedDeploymentsToBeDeleted(ctx, l, workerDeploy)
+}
+
+// waitForOwnedDeploymentsToBeDeleted waits for all owned deployments to be deleted with proper polling and timeout
+func (r *TemporalWorkerDeploymentReconciler) waitForOwnedDeploymentsToBeDeleted(ctx context.Context, l logr.Logger, workerDeploy *temporaliov1alpha1.TemporalWorkerDeployment) error {
+	// Create a timeout context for cleanup operations
+	cleanupCtx, cancel := context.WithTimeout(ctx, cleanupTimeout)
+	defer cancel()
+
+	ticker := time.NewTicker(cleanupPollInterval)
+	defer ticker.Stop()
+
+	l.Info("Waiting for owned deployments to be deleted", "timeout", cleanupTimeout)
+
+	for {
+		select {
+		case <-cleanupCtx.Done():
+			if cleanupCtx.Err() == context.DeadlineExceeded {
+				return fmt.Errorf("timeout waiting for deployments to be deleted after %v", cleanupTimeout)
+			}
+			return fmt.Errorf("context cancelled while waiting for deployments to be deleted: %w", cleanupCtx.Err())
+
+		case <-ticker.C:
+			// Try to use field selector for efficient querying, with fallback
+			listOpts := &client.ListOptions{
+				Namespace:     workerDeploy.Namespace,
+				FieldSelector: fields.OneTermEqualSelector(deployOwnerKey, workerDeploy.Name),
+			}
+
+			deploymentList := &appsv1.DeploymentList{}
+			err := r.List(cleanupCtx, deploymentList, listOpts)
+			if err != nil {
+				// If field selector fails (common in tests), fall back to listing all deployments
+				listOpts = &client.ListOptions{
+					Namespace: workerDeploy.Namespace,
+				}
+				if err := r.List(cleanupCtx, deploymentList, listOpts); err != nil {
+					return fmt.Errorf("failed to list deployments during cleanup: %w", err)
+				}
+			}
+
+			// Check if any owned deployments still exist
+			hasOwnedDeployments := false
+			for _, deployment := range deploymentList.Items {
+				if r.isOwnedByWorkerDeployment(&deployment, workerDeploy) {
+					hasOwnedDeployments = true
+					l.Info("Still waiting for deployment to be deleted", "deployment", deployment.Name)
+					break
+				}
+			}
+
+			if !hasOwnedDeployments {
+				l.Info("All owned deployments have been deleted")
+				return nil
 			}
-			// IsNotFound error means deployment was successfully deleted
 		}
 	}
-
-	l.Info("All managed resources have been cleaned up")
-	return nil
 }
 
 // isOwnedByWorkerDeployment checks if a deployment is owned by the given TemporalWorkerDeployment

internal/tests/internal/integration_test.go

@@ -26,6 +26,24 @@ const (
 	testDrainageRefreshInterval       = time.Second
 )
 
+// waitForCondition polls a condition function until it returns true or timeout is reached
+func waitForCondition(condition func() bool, timeout, interval time.Duration) bool {
+	deadline := time.After(timeout)
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-deadline:
+			return false
+		case <-ticker.C:
+			if condition() {
+				return true
+			}
+		}
+	}
+}
+
 type testEnv struct {
 	k8sClient  client.Client
 	ts         *temporaltest.TestServer
@@ -349,19 +367,24 @@ func testDeploymentDeletionProtection(t *testing.T, k8sClient client.Client, ts
 		// Direct deletion failed as expected due to owner reference protection
 		t.Logf("Direct deletion failed as expected: %v", err)
 	} else {
-		// If deletion succeeded, verify the controller recreates it
-		time.Sleep(5 * time.Second) // Give controller time to recreate
+		// If deletion succeeded, verify the controller recreates it with proper polling
+		eventuallyRecreated := func() bool {
+			var recreatedDeployment appsv1.Deployment
+			err := k8sClient.Get(ctx, types.NamespacedName{
+				Name:      expectedDeploymentName,
+				Namespace: twd.Namespace,
+			}, &recreatedDeployment)
+
+			if err != nil {
+				return false // Deployment not found yet
+			}
 
-		var recreatedDeployment appsv1.Deployment
-		err = k8sClient.Get(ctx, types.NamespacedName{
-			Name:      expectedDeploymentName,
-			Namespace: twd.Namespace,
-		}, &recreatedDeployment)
+			// Check if it's a new deployment (different UID)
+			return recreatedDeployment.UID != originalUID
+		}
 
-		if err != nil {
-			assert.Fail(t, "Controller should have recreated the deployment after direct deletion", "Error: %v", err)
-		} else {
-			assert.NotEqual(t, originalUID, recreatedDeployment.UID, "Deployment should have been recreated with new UID")
+		if !waitForCondition(eventuallyRecreated, 30*time.Second, 1*time.Second) {
+			assert.Fail(t, "Controller should have recreated the deployment after direct deletion within 30 seconds")
 		}
 	}
 
@@ -371,22 +394,15 @@ func testDeploymentDeletionProtection(t *testing.T, k8sClient client.Client, ts
 	require.NoError(t, err, "failed to delete TemporalWorkerDeployment")
 
 	// Wait for the deployment to be cleaned up
-	deadline := time.Now().Add(30 * time.Second)
-	deploymentDeleted := false
-	for time.Now().Before(deadline) {
+	eventuallyDeleted := func() bool {
 		var checkDeployment appsv1.Deployment
-		err = k8sClient.Get(ctx, types.NamespacedName{
+		err := k8sClient.Get(ctx, types.NamespacedName{
 			Name:      expectedDeploymentName,
 			Namespace: twd.Namespace,
 		}, &checkDeployment)
-		if err != nil {
-			if client.IgnoreNotFound(err) == nil {
-				deploymentDeleted = true
-				break
-			}
-		}
-		time.Sleep(1 * time.Second)
+		return client.IgnoreNotFound(err) == nil // Returns true if deployment is not found (deleted)
 	}
 
+	deploymentDeleted := waitForCondition(eventuallyDeleted, 30*time.Second, 1*time.Second)
 	assert.True(t, deploymentDeleted, "Controller should have cleaned up the deployment when TWD was deleted")
 }