Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
272 changes: 163 additions & 109 deletions cmd/ateapi/internal/controlapi/syncer.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,25 +17,46 @@ package controlapi
import (
"context"
"errors"
"fmt"
"log/slog"
"maps"
"time"

"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
"github.com/agent-substrate/substrate/internal/resources"
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"
corev1 "k8s.io/api/core/v1"
"k8s.io/apimachinery/pkg/util/wait"
"k8s.io/client-go/tools/cache"
"k8s.io/client-go/util/workqueue"
)

// syncerWorkerCount is the number of goroutines draining the work queue. The
// queue never hands the same key to two workers concurrently, so per-key
// ordering is preserved.
const syncerWorkerCount = 2

// workerKey identifies a worker row in the store. pool is captured at enqueue
// time because once the pod is gone from the informer cache, the
// ate.dev/worker-pool label cannot be recovered from a namespace/name key.
type workerKey struct {
namespace string
pool string
name string
}

// WorkerPoolSyncer reconciles the state of worker pods from Kubernetes Informer
// into the store.
//
// Informer event handlers only enqueue keys; worker goroutines reconcile each
// key against the current informer cache state, requeuing with rate-limited
// backoff on transient failures such as store.ErrVersionConflict.
type WorkerPoolSyncer struct {
persistence store.Interface
workerInformer cache.SharedIndexInformer
workerPoolLister listersv1alpha1.WorkerPoolLister
queue workqueue.TypedRateLimitingInterface[workerKey]
}

// NewWorkerPoolSyncer creates a new WorkerPoolSyncer.
Expand All @@ -44,19 +65,28 @@ func NewWorkerPoolSyncer(persistence store.Interface, workerInformer cache.Share
persistence: persistence,
workerInformer: workerInformer,
workerPoolLister: workerPoolLister,
queue: workqueue.NewTypedRateLimitingQueue(workqueue.DefaultTypedControllerRateLimiter[workerKey]()),
}
}

// Start starts the background reconciliation loop.
// Start registers the event handlers and starts the background workers. The
// informer's initial list synthesizes Add events for every existing pod, so no
// explicit startup re-list is needed as long as Start is called before the
// informer factory is started.
func (s *WorkerPoolSyncer) Start(ctx context.Context) {
s.workerInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
AddFunc: func(obj interface{}) {
pod := obj.(*corev1.Pod)
s.syncWorkerToStore(ctx, pod)
s.enqueuePod(obj.(*corev1.Pod))
},
UpdateFunc: func(oldObj, newObj interface{}) {
pod := newObj.(*corev1.Pod)
s.syncWorkerToStore(ctx, pod)
oldPod := oldObj.(*corev1.Pod)
newPod := newObj.(*corev1.Pod)
// If the pool label changed, enqueue the old key too so its
// now-stale store row gets cleaned up.
if oldPod.Labels[workerPodLabel] != newPod.Labels[workerPodLabel] {
s.enqueuePod(oldPod)
}
s.enqueuePod(newPod)
},
DeleteFunc: func(obj interface{}) {
var pod *corev1.Pod
Expand All @@ -74,40 +104,76 @@ func (s *WorkerPoolSyncer) Start(ctx context.Context) {
slog.ErrorContext(ctx, "Unknown object type in delete handler", slog.Any("obj", obj))
return
}
slog.InfoContext(ctx, "Syncer: removing worker from store (pod deleted)", slog.String("worker", pod.Namespace+"/"+pod.Name))
// TODO: make this more robust. Informer event handlers cannot signal failure and
// the informer never retries them, so a cleanup that fails here is not retried
// until the next startup reconcile. The canonical fix is a rate-limited workqueue:
// enqueue the worker key and run this from a worker goroutine that requeues with
// backoff on error.
if err := s.reconcileDeadWorker(ctx, pod.Namespace, pod.Labels[workerPodLabel], pod.Name); err != nil {
slog.ErrorContext(ctx, "Failed to reconcile deleted worker", slog.String("worker", pod.Namespace+"/"+pod.Name), slog.Any("err", err))
}
s.enqueuePod(pod)
},
})

go func() {
defer s.queue.ShutDown()
if !cache.WaitForCacheSync(ctx.Done(), s.workerInformer.HasSynced) {
slog.ErrorContext(ctx, "Syncer: failed to sync informer cache")
return
}

slog.InfoContext(ctx, "Syncer: performing initial sync on startup")
objs := s.workerInformer.GetIndexer().List()
for _, obj := range objs {
pod := obj.(*corev1.Pod)
s.syncWorkerToStore(ctx, pod)
for range syncerWorkerCount {
go wait.UntilWithContext(ctx, s.runWorker, time.Second)
}

// Reconcile the other direction: clean up stored workers whose pods no
// longer exist. This recovers delete events missed while ate-api-server
// was down — neither the watch relist nor the resync period can replay a
// delete across a process restart, because the informer cache starts empty.
s.reconcileOrphanedWorkers(ctx)
// Reconcile the other direction: enqueue every stored worker so records
// whose pods no longer exist are cleaned up. This recovers delete events
// missed while ate-api-server was down — neither the watch relist nor
// the resync period can replay a delete across a process restart,
// because the informer cache starts empty. Runs after the cache sync so
// the indexer is an authoritative snapshot of live pods.
s.enqueueStoredWorkers(ctx)

<-ctx.Done()
}()
}

func (s *WorkerPoolSyncer) syncWorkerToStore(ctx context.Context, pod *corev1.Pod) {
func (s *WorkerPoolSyncer) enqueuePod(pod *corev1.Pod) {
s.queue.Add(workerKey{namespace: pod.Namespace, pool: pod.Labels[workerPodLabel], name: pod.Name})
}

func (s *WorkerPoolSyncer) runWorker(ctx context.Context) {
for s.processNextWorkItem(ctx) {
}
}

func (s *WorkerPoolSyncer) processNextWorkItem(ctx context.Context) bool {
key, quit := s.queue.Get()
if quit {
return false
}
defer s.queue.Done(key)

if err := s.reconcile(ctx, key); err != nil {
slog.ErrorContext(ctx, "Syncer: reconcile failed, requeueing",
slog.String("worker", key.namespace+"/"+key.name),
slog.String("pool", key.pool),
slog.Any("err", err))
s.queue.AddRateLimited(key)
return true
}
s.queue.Forget(key)
return true
}

// reconcile converges the store row for key with the current pod state in the
// informer cache. Returning an error requeues the key with backoff.
func (s *WorkerPoolSyncer) reconcile(ctx context.Context, key workerKey) error {
obj, exists, err := s.workerInformer.GetIndexer().GetByKey(key.namespace + "/" + key.name)
if err != nil {
return err
}
if !exists {
slog.InfoContext(ctx, "Syncer: removing worker from store (pod deleted)", slog.String("worker", key.namespace+"/"+key.name))
return s.reconcileDeadWorker(ctx, key.namespace, key.pool, key.name)
}
pod := obj.(*corev1.Pod)
if pod.Labels[workerPodLabel] != key.pool {
// The pod moved to a different pool; this key's store row is stale.
return s.reconcileDeadWorker(ctx, key.namespace, key.pool, key.name)
}
// Checked before eligibility: draining works off the stored record by name and
// never reads the pod IP, while a Terminating pod can legitimately report no
// IP once its sandbox is torn down. Gating on the IP first would drop the
Expand All @@ -118,89 +184,99 @@ func (s *WorkerPoolSyncer) syncWorkerToStore(ctx context.Context, pod *corev1.Po
// the bound actor here — inside the pod ateom has received SIGTERM and is
// gracefully shutting the actor down. Actor cleanup happens on the Pod
// Deleted event.
if err := s.markWorkerDraining(ctx, pod.Namespace, pod.Labels[workerPodLabel], pod.Name); err != nil {
slog.ErrorContext(ctx, "Failed to mark worker draining", slog.String("worker", pod.Namespace+"/"+pod.Name), slog.Any("err", err))
}
return
return s.markWorkerDraining(ctx, key.namespace, key.pool, key.name)
}

if !isWorkerEligible(pod) {
return
// No IP yet; a later update event re-enqueues the pod.
return nil
}
return s.createOrUpdateWorker(ctx, key, pod)
}

poolName := pod.Labels[workerPodLabel]
pool, err := s.workerPoolLister.WorkerPools(pod.Namespace).Get(poolName)
func (s *WorkerPoolSyncer) createOrUpdateWorker(ctx context.Context, key workerKey, pod *corev1.Pod) error {
pool, err := s.workerPoolLister.WorkerPools(key.namespace).Get(key.pool)
if err != nil {
slog.ErrorContext(ctx, "Failed to get WorkerPool for worker pod", slog.String("worker", pod.Namespace+"/"+pod.Name), slog.String("pool", poolName), slog.Any("err", err))
return
return fmt.Errorf("getting WorkerPool %s/%s: %w", key.namespace, key.pool, err)
}

w, err := s.persistence.GetWorker(ctx, pod.Namespace, poolName, pod.Name)
w, err := s.persistence.GetWorker(ctx, key.namespace, key.pool, key.name)
if err != nil {
if errors.Is(err, store.ErrNotFound) {
slog.InfoContext(ctx, "Syncer: creating worker in store", slog.String("worker", pod.Namespace+"/"+pod.Name))
worker := &ateapipb.Worker{
WorkerNamespace: pod.Namespace,
WorkerPool: poolName,
WorkerPod: pod.Name,
Ip: pod.Status.PodIP,
WorkerPodUid: string(pod.UID),
NodeName: pod.Spec.NodeName,
SandboxClass: string(pool.Spec.SandboxClass),
Labels: pool.GetLabels(),
State: ateapipb.Worker_STATE_ACTIVE,
}
// TODO(thockin): for now this is the only place Workers are
// created. If/when this becomes a regular API, validation should
// move there.
if errs := resources.ValidateWorker(worker, nil); len(errs) > 0 {
err := status.Error(codes.InvalidArgument, errs.ToAggregate().Error())
slog.ErrorContext(ctx, "Invalid worker", slog.Any("err", err))
return
}
err = s.persistence.CreateWorker(ctx, worker)
if err != nil && !errors.Is(err, store.ErrAlreadyExists) {
slog.ErrorContext(ctx, "Failed to create worker in store", slog.Any("err", err))
}
return
if !errors.Is(err, store.ErrNotFound) {
return fmt.Errorf("getting worker from store: %w", err)
}
slog.InfoContext(ctx, "Syncer: creating worker in store", slog.String("worker", key.namespace+"/"+key.name))
worker := &ateapipb.Worker{
WorkerNamespace: pod.Namespace,
WorkerPool: key.pool,
WorkerPod: pod.Name,
Ip: pod.Status.PodIP,
WorkerPodUid: string(pod.UID),
NodeName: pod.Spec.NodeName,
SandboxClass: string(pool.Spec.SandboxClass),
Labels: pool.GetLabels(),
State: ateapipb.Worker_STATE_ACTIVE,
}
// TODO(thockin): for now this is the only place Workers are
// created. If/when this becomes a regular API, validation should
// move there.
if errs := resources.ValidateWorker(worker, nil); len(errs) > 0 {
// Terminal: the inputs are deterministic, retrying cannot help. A
// future pod event re-enqueues the key.
slog.ErrorContext(ctx, "Invalid worker", slog.Any("err", errs.ToAggregate()))
return nil
}
// ErrAlreadyExists means we lost a create race; requeue and converge
// via the update path.
return s.persistence.CreateWorker(ctx, worker)
}

if w.WorkerPodUid != string(pod.UID) {
// The pod was deleted and recreated under the same name, and the queue
// coalesced the two events. The store row belongs to the dead pod:
// clean it up (releasing any actor bound to the old incarnation) and
// requeue to create a fresh row.
slog.InfoContext(ctx, "Syncer: worker in store belongs to a replaced pod, deleting", slog.String("worker", key.namespace+"/"+key.name))
if err := s.reconcileDeadWorker(ctx, key.namespace, key.pool, key.name); err != nil {
return err
}
slog.ErrorContext(ctx, "Failed to get worker from store", slog.Any("err", err))
return
return fmt.Errorf("worker %s/%s/%s belonged to replaced pod UID %s; requeueing to recreate", key.namespace, key.pool, key.name, w.WorkerPodUid)
}

changed := false
if w.Ip != pod.Status.PodIP {
// TODO: I don't think this is possible, but handling this case so we can log it just in case we can reproduce it.
slog.InfoContext(ctx, "Syncer: updating worker in store (IP changed)", slog.String("worker", pod.Namespace+"/"+pod.Name))
slog.InfoContext(ctx, "Syncer: updating worker in store (IP changed)", slog.String("worker", key.namespace+"/"+key.name))
w.Ip = pod.Status.PodIP
changed = true
}
if w.SandboxClass != string(pool.Spec.SandboxClass) {
slog.InfoContext(ctx, "Syncer: updating worker in store (SandboxClass changed)", slog.String("worker", pod.Namespace+"/"+pod.Name))
slog.InfoContext(ctx, "Syncer: updating worker in store (SandboxClass changed)", slog.String("worker", key.namespace+"/"+key.name))
w.SandboxClass = string(pool.Spec.SandboxClass)
changed = true
}
if !maps.Equal(w.Labels, pool.GetLabels()) {
slog.InfoContext(ctx, "Syncer: updating worker in store (labels changed)", slog.String("worker", pod.Namespace+"/"+pod.Name))
slog.InfoContext(ctx, "Syncer: updating worker in store (labels changed)", slog.String("worker", key.namespace+"/"+key.name))
w.Labels = pool.GetLabels()
changed = true
}

if changed {
if err = s.persistence.UpdateWorker(ctx, w, w.Version); err != nil {
slog.ErrorContext(ctx, "Failed to update worker in store", slog.Any("err", err))
}
if !changed {
return nil
}

// ErrVersionConflict requeues the key; the retry re-fetches the worker at
// its new version.
return s.persistence.UpdateWorker(ctx, w, w.Version)
}

func isWorkerEligible(pod *corev1.Pod) bool {
return pod.Status.PodIP != ""
}

// markWorkerDraining transitions a worker to STATE_DRAINING so the scheduler
// stops routing new actors to it while its pod is Terminating. Best-effort: if
// the worker is already gone or already draining, or a concurrent update wins,
// there is nothing more to do — the Pod Deleted event will clean up the record.
// stops routing new actors to it while its pod is Terminating. If the worker is
// already gone or already draining there is nothing more to do — the Pod
// Deleted event will clean up the record. A version conflict is returned so the
// caller requeues and retries against the updated record.
func (s *WorkerPoolSyncer) markWorkerDraining(ctx context.Context, namespace, pool, podName string) error {
worker, err := s.persistence.GetWorker(ctx, namespace, pool, podName)
if err != nil {
Expand Down Expand Up @@ -230,47 +306,25 @@ func (s *WorkerPoolSyncer) reconcileDeadWorker(ctx context.Context, namespace, p
return s.persistence.DeleteWorker(ctx, namespace, pool, podName)
}

// reconcileOrphanedWorkers cleans up stored worker records whose pods no longer
// exist. It runs once after the informer cache has synced, when the indexer is a
// fresh, authoritative snapshot of the live worker pods, so a worker missing
// from the indexer (or present under a different pod UID, i.e. name reuse) is
// stale.
func (s *WorkerPoolSyncer) reconcileOrphanedWorkers(ctx context.Context) {
var workers []*ateapipb.Worker
// enqueueStoredWorkers enqueues a key for every worker record in the store.
// Records whose pods are live and unchanged reconcile to a no-op; orphaned
// records (pod gone, or its name reused by a new pod UID) get cleaned up.
func (s *WorkerPoolSyncer) enqueueStoredWorkers(ctx context.Context) {
var pageToken string
for {
wPage, nextToken, err := s.persistence.ListWorkers(ctx, 1000, pageToken)
workers, nextToken, err := s.persistence.ListWorkers(ctx, 1000, pageToken)
if err != nil {
slog.ErrorContext(ctx, "Syncer: failed to list workers for orphan reconcile", slog.Any("err", err))
return
}
workers = append(workers, wPage...)
for _, w := range workers {
s.queue.Add(workerKey{namespace: w.GetWorkerNamespace(), pool: w.GetWorkerPool(), name: w.GetWorkerPod()})
}
if nextToken == "" {
break
return
}
pageToken = nextToken
}
indexer := s.workerInformer.GetIndexer()
for _, w := range workers {
key := w.GetWorkerNamespace() + "/" + w.GetWorkerPod()
obj, exists, err := indexer.GetByKey(key)
if err != nil {
slog.ErrorContext(ctx, "Syncer: indexer lookup failed during orphan reconcile", slog.String("worker", key), slog.Any("err", err))
continue
}
// The pod is still live only if it is present under the same UID the
// worker recorded; a different UID means the name was reused by a new pod
// and this record belongs to a dead incarnation.
if exists {
if pod, ok := obj.(*corev1.Pod); ok && string(pod.UID) == w.GetWorkerPodUid() {
continue
}
}
slog.InfoContext(ctx, "Syncer: reconciling orphaned worker (pod gone)", slog.String("worker", key))
if err := s.reconcileDeadWorker(ctx, w.GetWorkerNamespace(), w.GetWorkerPool(), w.GetWorkerPod()); err != nil {
slog.ErrorContext(ctx, "Syncer: failed to reconcile orphaned worker", slog.String("worker", key), slog.Any("err", err))
}
}
}

// releaseActorOnDeadWorker resets the actor bound to a vanishing worker pod. An
Expand Down
Loading