Fix propogate issue to workers with worker pools changes

cmd/atecontroller/internal/workersync/informer.go

@@ -17,6 +17,7 @@ package workersync
 import (
 	"time"
 
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/informers"
 	"k8s.io/client-go/kubernetes"
@@ -33,5 +34,18 @@ func WorkerPodInformer(kc kubernetes.Interface) (informers.SharedInformerFactory
 			options.LabelSelector = workerPodLabel
 		}),
 	)
-	return factory, factory.Core().V1().Pods().Informer()
+	informer := factory.Core().V1().Pods().Informer()
+	if err := informer.AddIndexers(cache.Indexers{workerPoolIndex: workerPoolIndexFunc}); err != nil {
+		panic("adding worker pool pod index: " + err.Error())
+	}
+	return factory, informer
+}
+
+func workerPoolIndexFunc(obj interface{}) ([]string, error) {
+	pod := obj.(*corev1.Pod)
+	poolName := pod.Labels[workerPodLabel]
+	if poolName == "" {
+		return nil, nil
+	}
+	return []string{pod.Namespace + "/" + poolName}, nil
 }

cmd/atecontroller/internal/workersync/syncer.go

@@ -23,7 +23,7 @@ import (
 	"maps"
 	"time"
 
-	listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1"
+	atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1"
 	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -43,6 +43,10 @@ const syncerWorkerCount = 2
 // the pod informer is narrowed by.
 const workerPodLabel = "ate.dev/worker-pool"
 
+// workerPoolIndex maps a WorkerPool namespace/name to the worker Pods labeled
+// as members of that pool.
+const workerPoolIndex = "worker-pool"
+
 // workerKey identifies the pod incarnation a queued event concerns. namespace
 // and name locate the pod in the informer, which is indexed by namespace/name
 // rather than by UID.
@@ -90,19 +94,19 @@ func (k workerKey) logAttrs() []any {
 // key against the current informer cache state, requeuing with rate-limited
 // backoff on transient failures such as a lost version precondition.
 type WorkerPoolSyncer struct {
-	client           ateapipb.ControlClient
-	workerInformer   cache.SharedIndexInformer
-	workerPoolLister listersv1alpha1.WorkerPoolLister
-	queue            workqueue.TypedRateLimitingInterface[workerKey]
+	client             ateapipb.ControlClient
+	workerInformer     cache.SharedIndexInformer
+	workerPoolInformer cache.SharedIndexInformer
+	queue              workqueue.TypedRateLimitingInterface[workerKey]
 }
 
 // NewWorkerPoolSyncer creates a new WorkerPoolSyncer.
-func NewWorkerPoolSyncer(client ateapipb.ControlClient, workerInformer cache.SharedIndexInformer, workerPoolLister listersv1alpha1.WorkerPoolLister) *WorkerPoolSyncer {
+func NewWorkerPoolSyncer(client ateapipb.ControlClient, workerInformer, workerPoolInformer cache.SharedIndexInformer) *WorkerPoolSyncer {
 	return &WorkerPoolSyncer{
-		client:           client,
-		workerInformer:   workerInformer,
-		workerPoolLister: workerPoolLister,
-		queue:            workqueue.NewTypedRateLimitingQueue(workqueue.DefaultTypedControllerRateLimiter[workerKey]()),
+		client:             client,
+		workerInformer:     workerInformer,
+		workerPoolInformer: workerPoolInformer,
+		queue:              workqueue.NewTypedRateLimitingQueue(workqueue.DefaultTypedControllerRateLimiter[workerKey]()),
 	}
 }
 
@@ -146,6 +150,10 @@ func (s *WorkerPoolSyncer) Start(ctx context.Context) {
 			s.enqueuePod(pod)
 		},
 	})
+	s.workerPoolInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
+		AddFunc:    s.enqueueWorkerPool,
+		UpdateFunc: func(_, obj interface{}) { s.enqueueWorkerPool(obj) },
+	})
 
 	go func() {
 		defer s.queue.ShutDown()
@@ -169,6 +177,23 @@ func (s *WorkerPoolSyncer) Start(ctx context.Context) {
 	}()
 }
 
+// enqueueWorkerPool schedules every current pod in pool.
+func (s *WorkerPoolSyncer) enqueueWorkerPool(obj interface{}) {
+	pool, ok := obj.(*atev1alpha1.WorkerPool)
+	if !ok {
+		slog.Error("Syncer: unexpected WorkerPool informer object", slog.Any("obj", obj))
+		return
+	}
+	pods, err := s.workerInformer.GetIndexer().ByIndex(workerPoolIndex, pool.Namespace+"/"+pool.Name)
+	if err != nil {
+		slog.Error("Syncer: listing pods for WorkerPool update", "workerPool", pool.Namespace+"/"+pool.Name, "err", err)
+		return
+	}
+	for _, obj := range pods {
+		s.enqueuePod(obj.(*corev1.Pod))
+	}
+}
+
 func (s *WorkerPoolSyncer) enqueuePod(pod *corev1.Pod) {
 	s.queue.Add(workerKey{namespace: pod.Namespace, name: pod.Name, uid: string(pod.UID)})
 }
@@ -245,10 +270,17 @@ func (s *WorkerPoolSyncer) reconcile(ctx context.Context, key workerKey) error {
 
 func (s *WorkerPoolSyncer) createOrUpdateWorker(ctx context.Context, key workerKey, pod *corev1.Pod) error {
 	poolName := pod.Labels[workerPodLabel]
-	pool, err := s.workerPoolLister.WorkerPools(key.namespace).Get(poolName)
+	poolObject, exists, err := s.workerPoolInformer.GetIndexer().GetByKey(key.namespace + "/" + poolName)
 	if err != nil {
 		return fmt.Errorf("getting WorkerPool %s/%s: %w", key.namespace, poolName, err)
 	}
+	if !exists {
+		return fmt.Errorf("getting WorkerPool %s/%s: not found", key.namespace, poolName)
+	}
+	pool, ok := poolObject.(*atev1alpha1.WorkerPool)
+	if !ok {
+		return fmt.Errorf("getting WorkerPool %s/%s: unexpected object type %T", key.namespace, poolName, poolObject)
+	}
 
 	w, err := s.client.GetWorker(ctx, &ateapipb.GetWorkerRequest{Worker: key.workerRef()})
 	if status.Code(err) == codes.NotFound {

cmd/atecontroller/internal/workersync/syncer_test.go

@@ -25,7 +25,6 @@ import (
 	atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1"
 	atefake "github.com/agent-substrate/substrate/pkg/client/clientset/versioned/fake"
 	"github.com/agent-substrate/substrate/pkg/client/informers/externalversions"
-	listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1"
 	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -99,10 +98,10 @@ func registeredWorker(ns, poolName, podName, uid, ip string) *ateapipb.Worker {
 	}
 }
 
-// poolLister builds the WorkerPool lister the syncer reads, and returns the
-// indexer behind it so a test can seed and mutate pools synchronously rather
+// newWorkerPoolInformer builds the WorkerPool informer the syncer reads, and
+// returns its indexer so a test can seed and mutate pools synchronously rather
 // than starting a factory and waiting for a watch to deliver them.
-func poolLister(t *testing.T, initPools ...*atev1alpha1.WorkerPool) (listersv1alpha1.WorkerPoolLister, cache.Indexer) {
+func newWorkerPoolInformer(t *testing.T, initPools ...*atev1alpha1.WorkerPool) (cache.SharedIndexInformer, cache.Indexer) {
 	t.Helper()
 	//nolint:staticcheck // NewSimpleClientset is the only available fake clientset for versioned CRDs.
 	pools := externalversions.NewSharedInformerFactory(atefake.NewSimpleClientset(), 0).Api().V1alpha1().WorkerPools()
@@ -112,7 +111,7 @@ func poolLister(t *testing.T, initPools ...*atev1alpha1.WorkerPool) (listersv1al
 			t.Fatalf("seeding WorkerPool %s/%s: %v", pool.Namespace, pool.Name, err)
 		}
 	}
-	return pools.Lister(), indexer
+	return pools.Informer(), indexer
 }
 
 // setupSyncerTest wires a running syncer to a fake Control API and a fake
@@ -123,11 +122,11 @@ func setupSyncerTest(t *testing.T, ctx context.Context, api *fakeControl, initPo
 	//nolint:staticcheck // NewSimpleClientset is what the informer machinery takes.
 	fakeK8s := fake.NewSimpleClientset()
 	workerFactory, workerInformer := WorkerPodInformer(fakeK8s)
-	lister, _ := poolLister(t, initPools...)
+	workerPoolInformer, _ := newWorkerPoolInformer(t, initPools...)
 
 	// Start before the factory: the informer's initial list is what seeds the
 	// queue with the pods that already exist.
-	NewWorkerPoolSyncer(api, workerInformer, lister).Start(ctx)
+	NewWorkerPoolSyncer(api, workerInformer, workerPoolInformer).Start(ctx)
 	workerFactory.Start(ctx.Done())
 	workerFactory.WaitForCacheSync(ctx.Done())
 
@@ -142,9 +141,9 @@ func setupReconcileTest(t *testing.T, api *fakeControl, initPools ...*atev1alpha
 
 	//nolint:staticcheck // NewSimpleClientset is what the informer machinery takes.
 	_, workerInformer := WorkerPodInformer(fake.NewSimpleClientset())
-	lister, poolIndexer := poolLister(t, initPools...)
+	workerPoolInformer, poolIndexer := newWorkerPoolInformer(t, initPools...)
 
-	return NewWorkerPoolSyncer(api, workerInformer, lister), workerInformer.GetIndexer(), poolIndexer
+	return NewWorkerPoolSyncer(api, workerInformer, workerPoolInformer), workerInformer.GetIndexer(), poolIndexer
 }
 
 // seedPod puts a pod in the syncer's cache as though the informer had delivered
@@ -630,6 +629,36 @@ func TestSyncer_RequeueOnMissingWorkerPool(t *testing.T) {
 	}
 }
 
+// TestEnqueueWorkerPool verifies that a WorkerPool change requeues every pod
+// in that pool, so its Worker record receives updated labels promptly.
+func TestEnqueueWorkerPool(t *testing.T) {
+	api := newFakeControl()
+	s, pods, _ := setupReconcileTest(t, api)
+	pool := workerPool("ns-pool-update", "pool-a", "gvisor", nil)
+
+	first := seedPod(t, pods, workerPod(pool.Namespace, "worker-1", pool.Name, testPodUID, "10.0.0.1"))
+	second := seedPod(t, pods, workerPod(pool.Namespace, "worker-2", pool.Name, otherPodUID, "10.0.0.2"))
+	seedPod(t, pods, workerPod(pool.Namespace, "other-pool-worker", "pool-b", "33333333-3333-3333-3333-333333333333", "10.0.0.3"))
+
+	s.enqueueWorkerPool(pool)
+	if got := s.queue.Len(); got != 2 {
+		t.Fatalf("queued workers = %d, want 2", got)
+	}
+
+	got := map[workerKey]bool{}
+	for range 2 {
+		key, quit := s.queue.Get()
+		if quit {
+			t.Fatal("queue shut down while reading enqueued workers")
+		}
+		got[key] = true
+		s.queue.Done(key)
+	}
+	if !got[first] || !got[second] {
+		t.Errorf("queued workers = %v, want %v and %v", got, first, second)
+	}
+}
+
 // TestSyncer_SoftDelete_ViaInformer walks the whole transition through the
 // informer rather than seeding an already-registered worker: a pod is
 // registered ACTIVE, then enters graceful termination and flips to DRAINING

cmd/atecontroller/main.go

@@ -227,13 +227,13 @@ func main() {
 	// informer, so asking the shared cache for one would impose it on every other
 	// controller here too.
 	ateFactory := externalversions.NewSharedInformerFactory(ateClient, 0)
-	workerPoolLister := ateFactory.Api().V1alpha1().WorkerPools().Lister()
+	workerPoolInformer := ateFactory.Api().V1alpha1().WorkerPools()
 	workerPodInformerFactory, workerPodInformer := workersync.WorkerPodInformer(k8sClient)
 
 	// Start registers the informer event handlers, so it has to run before the
 	// factory does: the initial list then synthesizes an Add for every pod that
 	// already exists, and no explicit startup re-list is needed.
-	workersync.NewWorkerPoolSyncer(ateapiClient, workerPodInformer, workerPoolLister).Start(runCtx)
+	workersync.NewWorkerPoolSyncer(ateapiClient, workerPodInformer, workerPoolInformer.Informer()).Start(runCtx)
 
 	workerPodInformerFactory.Start(runCtx.Done())
 	ateFactory.Start(runCtx.Done())