From 43d6be837bd8b7492210ddd575ea226ec896be3a Mon Sep 17 00:00:00 2001 From: Lev Rado Date: Sun, 13 Sep 2026 17:53:42 +0300 Subject: [PATCH 1/2] Requeue workers when their WorkerPool changes --- .../internal/workersync/informer.go | 16 ++++++- .../internal/workersync/syncer.go | 46 +++++++++++++++---- .../internal/workersync/syncer_test.go | 42 ++++++++++++++--- cmd/atecontroller/main.go | 5 +- 4 files changed, 91 insertions(+), 18 deletions(-) diff --git a/cmd/atecontroller/internal/workersync/informer.go b/cmd/atecontroller/internal/workersync/informer.go index e587f123c8..dff246aa60 100644 --- a/cmd/atecontroller/internal/workersync/informer.go +++ b/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 } diff --git a/cmd/atecontroller/internal/workersync/syncer.go b/cmd/atecontroller/internal/workersync/syncer.go index 97fbd77d5e..a86b8e3696 100644 --- a/cmd/atecontroller/internal/workersync/syncer.go +++ b/cmd/atecontroller/internal/workersync/syncer.go @@ -23,6 +23,7 @@ import ( "maps" "time" + atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1" listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "google.golang.org/grpc/codes" @@ -43,6 +44,8 @@ const syncerWorkerCount = 2 // the pod informer is narrowed by. const workerPodLabel = "ate.dev/worker-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 +93,21 @@ 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 + workerPoolLister listersv1alpha1.WorkerPoolLister + 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 cache.SharedIndexInformer, workerPoolLister listersv1alpha1.WorkerPoolLister, workerPoolInformer cache.SharedIndexInformer) *WorkerPoolSyncer { return &WorkerPoolSyncer{ - client: client, - workerInformer: workerInformer, - workerPoolLister: workerPoolLister, - queue: workqueue.NewTypedRateLimitingQueue(workqueue.DefaultTypedControllerRateLimiter[workerKey]()), + client: client, + workerInformer: workerInformer, + workerPoolLister: workerPoolLister, + workerPoolInformer: workerPoolInformer, + queue: workqueue.NewTypedRateLimitingQueue(workqueue.DefaultTypedControllerRateLimiter[workerKey]()), } } @@ -146,6 +151,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 +178,25 @@ func (s *WorkerPoolSyncer) Start(ctx context.Context) { }() } +// enqueueWorkerPool schedules every current pod in pool. Worker records mirror +// their WorkerPool labels and sandbox class, neither of which causes a Pod +// event when changed. +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)}) } diff --git a/cmd/atecontroller/internal/workersync/syncer_test.go b/cmd/atecontroller/internal/workersync/syncer_test.go index 42c87cb26e..2e70c25c92 100644 --- a/cmd/atecontroller/internal/workersync/syncer_test.go +++ b/cmd/atecontroller/internal/workersync/syncer_test.go @@ -102,7 +102,7 @@ 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 // 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 poolLister(t *testing.T, initPools ...*atev1alpha1.WorkerPool) (listersv1alpha1.WorkerPoolLister, cache.Indexer, cache.SharedIndexInformer) { 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 +112,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.Lister(), indexer, pools.Informer() } // setupSyncerTest wires a running syncer to a fake Control API and a fake @@ -123,11 +123,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...) + lister, _, poolInformer := poolLister(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, lister, poolInformer).Start(ctx) workerFactory.Start(ctx.Done()) workerFactory.WaitForCacheSync(ctx.Done()) @@ -142,9 +142,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...) + lister, poolIndexer, poolInformer := poolLister(t, initPools...) - return NewWorkerPoolSyncer(api, workerInformer, lister), workerInformer.GetIndexer(), poolIndexer + return NewWorkerPoolSyncer(api, workerInformer, lister, poolInformer), workerInformer.GetIndexer(), poolIndexer } // seedPod puts a pod in the syncer's cache as though the informer had delivered @@ -630,6 +630,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 diff --git a/cmd/atecontroller/main.go b/cmd/atecontroller/main.go index 299cc5b853..afe111af86 100644 --- a/cmd/atecontroller/main.go +++ b/cmd/atecontroller/main.go @@ -227,13 +227,14 @@ 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() + workerPoolLister := workerPoolInformer.Lister() 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, workerPoolLister, workerPoolInformer.Informer()).Start(runCtx) workerPodInformerFactory.Start(runCtx.Done()) ateFactory.Start(runCtx.Done()) From 64d33b64e1e1b5ce534397b16e40a9e9b78f5eda Mon Sep 17 00:00:00 2001 From: Lev Rado Date: Sun, 13 Sep 2026 22:52:09 +0300 Subject: [PATCH 2/2] Read WorkerPool objects from informer indexer --- .../internal/workersync/syncer.go | 20 +++++++++++-------- .../internal/workersync/syncer_test.go | 17 ++++++++-------- cmd/atecontroller/main.go | 3 +-- 3 files changed, 21 insertions(+), 19 deletions(-) diff --git a/cmd/atecontroller/internal/workersync/syncer.go b/cmd/atecontroller/internal/workersync/syncer.go index a86b8e3696..78434b60c5 100644 --- a/cmd/atecontroller/internal/workersync/syncer.go +++ b/cmd/atecontroller/internal/workersync/syncer.go @@ -24,7 +24,6 @@ import ( "time" atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1" - 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" @@ -44,6 +43,8 @@ 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 @@ -95,17 +96,15 @@ func (k workerKey) logAttrs() []any { type WorkerPoolSyncer struct { client ateapipb.ControlClient workerInformer cache.SharedIndexInformer - workerPoolLister listersv1alpha1.WorkerPoolLister workerPoolInformer cache.SharedIndexInformer queue workqueue.TypedRateLimitingInterface[workerKey] } // NewWorkerPoolSyncer creates a new WorkerPoolSyncer. -func NewWorkerPoolSyncer(client ateapipb.ControlClient, workerInformer cache.SharedIndexInformer, workerPoolLister listersv1alpha1.WorkerPoolLister, workerPoolInformer cache.SharedIndexInformer) *WorkerPoolSyncer { +func NewWorkerPoolSyncer(client ateapipb.ControlClient, workerInformer, workerPoolInformer cache.SharedIndexInformer) *WorkerPoolSyncer { return &WorkerPoolSyncer{ client: client, workerInformer: workerInformer, - workerPoolLister: workerPoolLister, workerPoolInformer: workerPoolInformer, queue: workqueue.NewTypedRateLimitingQueue(workqueue.DefaultTypedControllerRateLimiter[workerKey]()), } @@ -178,9 +177,7 @@ func (s *WorkerPoolSyncer) Start(ctx context.Context) { }() } -// enqueueWorkerPool schedules every current pod in pool. Worker records mirror -// their WorkerPool labels and sandbox class, neither of which causes a Pod -// event when changed. +// enqueueWorkerPool schedules every current pod in pool. func (s *WorkerPoolSyncer) enqueueWorkerPool(obj interface{}) { pool, ok := obj.(*atev1alpha1.WorkerPool) if !ok { @@ -273,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 { diff --git a/cmd/atecontroller/internal/workersync/syncer_test.go b/cmd/atecontroller/internal/workersync/syncer_test.go index 2e70c25c92..a38f95c409 100644 --- a/cmd/atecontroller/internal/workersync/syncer_test.go +++ b/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, cache.SharedIndexInformer) { +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, pools.Informer() + 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, _, poolInformer := 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, poolInformer).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, poolInformer := poolLister(t, initPools...) + workerPoolInformer, poolIndexer := newWorkerPoolInformer(t, initPools...) - return NewWorkerPoolSyncer(api, workerInformer, lister, poolInformer), 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 diff --git a/cmd/atecontroller/main.go b/cmd/atecontroller/main.go index afe111af86..83b758928b 100644 --- a/cmd/atecontroller/main.go +++ b/cmd/atecontroller/main.go @@ -228,13 +228,12 @@ func main() { // controller here too. ateFactory := externalversions.NewSharedInformerFactory(ateClient, 0) workerPoolInformer := ateFactory.Api().V1alpha1().WorkerPools() - workerPoolLister := workerPoolInformer.Lister() 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, workerPoolInformer.Informer()).Start(runCtx) + workersync.NewWorkerPoolSyncer(ateapiClient, workerPodInformer, workerPoolInformer.Informer()).Start(runCtx) workerPodInformerFactory.Start(runCtx.Done()) ateFactory.Start(runCtx.Done())