chore: Patch in race condition for eviction queue fixes to v0.32.x (#1035)

Author: jonathan-innis

pkg/controllers/node/termination/machine_test.go

@@ -568,7 +568,7 @@ var _ = Describe("Machine/Termination", func() {
 			ExpectReconcileSucceeded(ctx, terminationController, client.ObjectKeyFromObject(node))
 
 			// Expect that the old pod's key still exists in the queue
-			Expect(queue.Has(terminator.NewQueueKey(pod)))
+			Expect(queue.Has(pod)).To(BeTrue())
 
 			// Re-create the pod and node, it should now have the same name, but a different UUID
 			node = test.Node(test.NodeOptions{

pkg/controllers/node/termination/nodeclaim_test.go

@@ -570,7 +570,7 @@ var _ = Describe("NodeClaim/Termination", func() {
 			ExpectReconcileSucceeded(ctx, terminationController, client.ObjectKeyFromObject(node))
 
 			// Expect that the old pod's key still exists in the queue
-			Expect(queue.Has(terminator.NewQueueKey(pod)))
+			Expect(queue.Has(pod)).To(BeTrue())
 
 			// Re-create the pod and node, it should now have the same name, but a different UUID
 			node = test.Node(test.NodeOptions{

pkg/controllers/node/termination/suite_test.go

@@ -20,6 +20,9 @@ import (
 	"testing"
 	"time"
 
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
 	"github.com/samber/lo"
 	clock "k8s.io/utils/clock/testing"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -34,8 +37,6 @@ import (
 	"github.com/aws/karpenter-core/pkg/operator/scheme"
 	"github.com/aws/karpenter-core/pkg/test"
 
-	. "github.com/onsi/ginkgo/v2"
-	. "github.com/onsi/gomega"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	. "knative.dev/pkg/logging/testing"
@@ -75,7 +76,7 @@ var _ = AfterSuite(func() {
 func ExpectNotEnqueuedForEviction(e *terminator.Queue, pods ...*v1.Pod) {
 	GinkgoHelper()
 	for _, pod := range pods {
-		Expect(e.Has(terminator.NewQueueKey(pod))).To(BeFalse())
+		Expect(e.Has(pod)).To(BeFalse())
 	}
 }
 

pkg/controllers/node/termination/terminator/eviction.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"sync"
 	"time"
 
 	"github.com/samber/lo"
@@ -74,7 +75,9 @@ func NewQueueKey(pod *v1.Pod) QueueKey {
 
 type Queue struct {
 	workqueue.RateLimitingInterface
-	sets.Set[QueueKey]
+
+	mu  sync.Mutex
+	set sets.Set[QueueKey]
 
 	kubeClient client.Client
 	recorder   events.Recorder
@@ -83,7 +86,7 @@ type Queue struct {
 func NewQueue(kubeClient client.Client, recorder events.Recorder) *Queue {
 	queue := &Queue{
 		RateLimitingInterface: workqueue.NewRateLimitingQueue(workqueue.NewItemExponentialFailureRateLimiter(evictionQueueBaseDelay, evictionQueueMaxDelay)),
-		Set:                   sets.New[QueueKey](),
+		set:                   sets.New[QueueKey](),
 		kubeClient:            kubeClient,
 		recorder:              recorder,
 	}
@@ -100,15 +103,25 @@ func (q *Queue) Builder(_ context.Context, m manager.Manager) controller.Builder
 
 // Add adds pods to the Queue
 func (q *Queue) Add(pods ...*v1.Pod) {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
 	for _, pod := range pods {
 		qk := NewQueueKey(pod)
-		if !q.Set.Has(qk) {
-			q.Set.Insert(qk)
+		if !q.set.Has(qk) {
+			q.set.Insert(qk)
 			q.RateLimitingInterface.Add(qk)
 		}
 	}
 }
 
+func (q *Queue) Has(pod *v1.Pod) bool {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
+	return q.set.Has(NewQueueKey(pod))
+}
+
 func (q *Queue) Reconcile(ctx context.Context, _ reconcile.Request) (reconcile.Result, error) {
 	// Check if the queue is empty. client-go recommends not using this function to gate the subsequent
 	// get call, but since we're popping items off the queue synchronously, there should be no synchonization
@@ -126,7 +139,9 @@ func (q *Queue) Reconcile(ctx context.Context, _ reconcile.Request) (reconcile.R
 	// Evict pod
 	if q.Evict(ctx, qk) {
 		q.RateLimitingInterface.Forget(qk)
-		q.Set.Delete(qk)
+		q.mu.Lock()
+		q.set.Delete(qk)
+		q.mu.Unlock()
 		return reconcile.Result{RequeueAfter: controller.Immediately}, nil
 	}
 	// Requeue pod if eviction failed
@@ -170,6 +185,9 @@ func (q *Queue) Evict(ctx context.Context, key QueueKey) bool {
 }
 
 func (q *Queue) Reset() {
+	q.mu.Lock()
+	defer q.mu.Unlock()
+
 	q.RateLimitingInterface = workqueue.NewRateLimitingQueue(workqueue.NewItemExponentialFailureRateLimiter(evictionQueueBaseDelay, evictionQueueMaxDelay))
-	q.Set = sets.New[QueueKey]()
+	q.set = sets.New[QueueKey]()
 }

pkg/controllers/node/termination/terminator/suite_test.go

@@ -16,20 +16,20 @@ package terminator_test
 
 import (
 	"context"
+	"sync"
 	"testing"
 
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
 	"github.com/samber/lo"
+	v1 "k8s.io/api/core/v1"
 	policyv1 "k8s.io/api/policy/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
 	"k8s.io/apimachinery/pkg/util/uuid"
 	. "knative.dev/pkg/logging/testing"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
-	v1 "k8s.io/api/core/v1"
-
 	"github.com/aws/karpenter-core/pkg/controllers/node/termination/terminator"
 
 	"github.com/aws/karpenter-core/pkg/apis"
@@ -125,5 +125,31 @@ var _ = Describe("Eviction/Queue", func() {
 			ExpectApplied(ctx, env.Client, pdb, pdb2, pod)
 			Expect(queue.Evict(ctx, terminator.NewQueueKey(pod))).To(BeFalse())
 		})
+		It("should ensure that calling Evict() is valid while making Add() calls", func() {
+			cancelCtx, cancel := context.WithCancel(ctx)
+			wg := sync.WaitGroup{}
+			DeferCleanup(func() {
+				cancel()
+				wg.Wait() // Ensure that we wait for reconcile loop to finish so that we don't get a RACE
+			})
+
+			// Keep calling Reconcile() for the entirety of this test
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+
+				for {
+					ExpectReconcileSucceeded(ctx, queue, client.ObjectKey{})
+					if cancelCtx.Err() != nil {
+						return
+					}
+				}
+			}()
+
+			// Ensure that we add enough pods to the queue while we are pulling items off of the queue (enough to trigger a DATA RACE)
+			for i := 0; i < 10000; i++ {
+				queue.Add(test.Pod())
+			}
+		})
 	})
 })