From f09e35f3e7e7f2c5cd1462f98297b180f966bccc Mon Sep 17 00:00:00 2001 From: Scott Funkenhauser Date: Thu, 20 Aug 2026 13:14:29 -0400 Subject: [PATCH] Implementation for the worker API --- .../internal/actoridentity/actoridentity.go | 6 +- .../actoridentity/actoridentity_test.go | 23 +- cmd/ateapi/internal/controlapi/atespace.go | 20 +- cmd/ateapi/internal/controlapi/crash.go | 8 +- cmd/ateapi/internal/controlapi/crash_test.go | 10 +- .../controlapi/functionaltest/actor_test.go | 6 +- cmd/ateapi/internal/controlapi/service.go | 4 + cmd/ateapi/internal/controlapi/syncer.go | 41 +- cmd/ateapi/internal/controlapi/syncer_test.go | 90 ++- cmd/ateapi/internal/controlapi/worker.go | 268 ++++++- cmd/ateapi/internal/controlapi/worker_test.go | 706 +++++++++++++++++- cmd/ateapi/internal/controlapi/workflow.go | 2 +- .../internal/controlapi/workflow_pause.go | 7 +- .../controlapi/workflow_pause_test.go | 2 +- .../internal/controlapi/workflow_resume.go | 33 +- .../controlapi/workflow_resume_test.go | 27 +- .../controlapi/workflow_suspend_test.go | 2 +- cmd/ateapi/internal/store/atepg/atepg.go | 133 ++-- cmd/ateapi/internal/store/atepg/outbox.go | 30 +- .../internal/store/atepg/outbox_test.go | 22 +- cmd/ateapi/internal/store/store.go | 97 ++- .../internal/store/storecontract/contract.go | 335 ++++++++- internal/resources/validate.go | 176 +---- internal/resources/validate_test.go | 527 ++----------- pkg/proto/ateapipb/ateapi.pb.go | 45 +- pkg/proto/ateapipb/ateapi.proto | 37 +- pkg/proto/ateapipb/ateapi_grpc.pb.go | 2 + 27 files changed, 1791 insertions(+), 868 deletions(-) diff --git a/cmd/ateapi/internal/actoridentity/actoridentity.go b/cmd/ateapi/internal/actoridentity/actoridentity.go index 9251e1aa55..0210cdf79b 100644 --- a/cmd/ateapi/internal/actoridentity/actoridentity.go +++ b/cmd/ateapi/internal/actoridentity/actoridentity.go @@ -307,11 +307,7 @@ func authenticateAtelet(ctx context.Context) (*ateletCaller, error) { // validateWorkerRef checks the reference to the Worker the certificate is // minted for. Workers are global-scoped, so the reference carries no atespace. func validateWorkerRef(worker *ateapipb.ObjectRef) error { - fldPath := field.NewPath("worker") - if worker == nil { - return field.Required(fldPath, "") - } - return resources.ValidateGlobalObjectRef(worker, fldPath).ToAggregate() + return resources.ValidateGlobalObjectRef(worker, field.NewPath("worker")).ToAggregate() } // authorizeActor resolves the actor from the authenticated worker and verifies diff --git a/cmd/ateapi/internal/actoridentity/actoridentity_test.go b/cmd/ateapi/internal/actoridentity/actoridentity_test.go index 688c6e7d03..35ec4e38de 100644 --- a/cmd/ateapi/internal/actoridentity/actoridentity_test.go +++ b/cmd/ateapi/internal/actoridentity/actoridentity_test.go @@ -212,14 +212,17 @@ func TestMintCertReadsThroughStaleWorkerCache(t *testing.T) { if err != nil { t.Fatalf("read seeded worker: %v", err) } - if worker.Status == nil { - worker.Status = &ateapipb.WorkerStatus{} - } - worker.Status.Assignment = &ateapipb.ActorAssignment{ - Actor: (resources.ActorRef{Atespace: testAtespace, Name: testActorName}).ToObjectRef(), - ActorUid: actor.GetMetadata().GetUid(), - } - if err := st.UpdateWorker(ctx, worker, worker.GetMetadata().GetVersion()); err != nil { + _, err = st.UpdateWorker(ctx, testWorkerName, store.PreconditionFrom(worker), func(toUpdate *ateapipb.Worker) error { + if toUpdate.Status == nil { + toUpdate.Status = &ateapipb.WorkerStatus{} + } + toUpdate.Status.Assignment = &ateapipb.ActorAssignment{ + Actor: (resources.ActorRef{Atespace: testAtespace, Name: testActorName}).ToObjectRef(), + ActorUid: actor.GetMetadata().GetUid(), + } + return nil + }) + if err != nil { t.Fatalf("assign worker in store: %v", err) } } @@ -273,7 +276,7 @@ func TestMintCertReadsThroughWorkerCacheMiss(t *testing.T) { if workerInStore { // Phase 2: register and assign the worker in the store only, // after the cache stopped listening. - if err := st.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := st.CreateWorker(ctx, &ateapipb.Worker{ Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, WorkerNamespace: testPodNS, WorkerPool: testPool, @@ -461,7 +464,7 @@ func seedActor(t *testing.T, ctx context.Context, st store.Interface, f actorFix if f.unassigned { worker.Status.Assignment = nil } - if err := st.CreateWorker(ctx, worker); err != nil { + if _, err := st.CreateWorker(ctx, worker); err != nil { t.Fatalf("seed worker: %v", err) } } diff --git a/cmd/ateapi/internal/controlapi/atespace.go b/cmd/ateapi/internal/controlapi/atespace.go index 0c509886fc..a889380a84 100644 --- a/cmd/ateapi/internal/controlapi/atespace.go +++ b/cmd/ateapi/internal/controlapi/atespace.go @@ -92,15 +92,7 @@ func (s *RPCService) GetAtespace(ctx context.Context, req *ateapipb.GetAtespaceR func validateGetAtespaceRequest(req *ateapipb.GetAtespaceRequest) field.ErrorList { var fldPath *field.Path - var errs field.ErrorList - - if val, fldPath := req.Atespace, fldPath.Child("atespace"); val == nil { - errs = append(errs, field.Required(fldPath, "")) - } else { - errs = append(errs, resources.ValidateGlobalObjectRef(val, fldPath)...) - } - - return errs + return resources.ValidateGlobalObjectRef(req.GetAtespace(), fldPath.Child("atespace")) } func (s *RPCService) ListAtespaces(ctx context.Context, req *ateapipb.ListAtespacesRequest) (*ateapipb.ListAtespacesResponse, error) { @@ -159,13 +151,5 @@ func (s *RPCService) DeleteAtespace(ctx context.Context, req *ateapipb.DeleteAte func validateDeleteAtespaceRequest(req *ateapipb.DeleteAtespaceRequest) field.ErrorList { var fldPath *field.Path - var errs field.ErrorList - - if val, fldPath := req.Atespace, fldPath.Child("atespace"); val == nil { - errs = append(errs, field.Required(fldPath, "")) - } else { - errs = append(errs, resources.ValidateGlobalObjectRef(val, fldPath)...) - } - - return errs + return resources.ValidateGlobalObjectRef(req.GetAtespace(), fldPath.Child("atespace")) } diff --git a/cmd/ateapi/internal/controlapi/crash.go b/cmd/ateapi/internal/controlapi/crash.go index 0b257bfdb1..b4bffdbec1 100644 --- a/cmd/ateapi/internal/controlapi/crash.go +++ b/cmd/ateapi/internal/controlapi/crash.go @@ -109,7 +109,7 @@ type crashActorStore interface { GetActor(ctx context.Context, actorRef resources.ActorRef) (*ateapipb.Actor, error) UpdateActor(ctx context.Context, actorRef resources.ActorRef, precondition store.Precondition, mutate func(toUpdate *ateapipb.Actor) error) (*ateapipb.Actor, error) GetWorker(ctx context.Context, name string) (*ateapipb.Worker, error) - UpdateWorker(ctx context.Context, worker *ateapipb.Worker, expectedVersion int64) error + UpdateWorker(ctx context.Context, name string, precondition store.Precondition, mutate func(toUpdate *ateapipb.Worker) error) (*ateapipb.Worker, error) } // releaseWorker clears the worker's assignment if it still points at the given @@ -145,8 +145,10 @@ func releaseWorker(ctx context.Context, st crashActorStore, actor *ateapipb.Acto return sandboxClass, nil } - worker.Status.Assignment = nil - if err := st.UpdateWorker(ctx, worker, worker.GetMetadata().GetVersion()); err != nil { + if _, err := st.UpdateWorker(ctx, workerName, store.PreconditionFrom(worker), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.Assignment = nil + return nil + }); err != nil { return sandboxClass, fmt.Errorf("while releasing worker: %w", err) } return sandboxClass, nil diff --git a/cmd/ateapi/internal/controlapi/crash_test.go b/cmd/ateapi/internal/controlapi/crash_test.go index 58081f644c..1e107c726d 100644 --- a/cmd/ateapi/internal/controlapi/crash_test.go +++ b/cmd/ateapi/internal/controlapi/crash_test.go @@ -80,7 +80,7 @@ func seedWorker(t *testing.T, ctx context.Context, st store.Interface, actorRef } } } - if err := st.CreateWorker(ctx, worker); err != nil { + if _, err := st.CreateWorker(ctx, worker); err != nil { t.Fatalf("seed worker: %v", err) } } @@ -200,7 +200,7 @@ func TestCrashActor(t *testing.T) { }, }, } - if err := st.CreateWorker(ctx, worker); err != nil { + if _, err := st.CreateWorker(ctx, worker); err != nil { t.Fatalf("CreateWorker: %v", err) } }, @@ -432,7 +432,7 @@ func TestCrashActor_Metrics(t *testing.T) { }, }, } - if err := st.CreateWorker(ctx, worker); err != nil { + if _, err := st.CreateWorker(ctx, worker); err != nil { t.Fatalf("CreateWorker: %v", err) } @@ -513,8 +513,8 @@ type failingUpdateWorkerStore struct { err error } -func (f failingUpdateWorkerStore) UpdateWorker(context.Context, *ateapipb.Worker, int64) error { - return f.err +func (f failingUpdateWorkerStore) UpdateWorker(context.Context, string, store.Precondition, func(*ateapipb.Worker) error) (*ateapipb.Worker, error) { + return nil, f.err } // A transient failure releasing the worker must not move the actor to the diff --git a/cmd/ateapi/internal/controlapi/functionaltest/actor_test.go b/cmd/ateapi/internal/controlapi/functionaltest/actor_test.go index 840a764048..2311fb26a2 100644 --- a/cmd/ateapi/internal/controlapi/functionaltest/actor_test.go +++ b/cmd/ateapi/internal/controlapi/functionaltest/actor_test.go @@ -2454,8 +2454,10 @@ func TestResumeActor_CrashesIfAssignedWorkerIsDraining(t *testing.T) { if err != nil { t.Fatalf("GetWorker(%s) failed: %v", assignedPod, err) } - assigned.Status.State = ateapipb.WorkerState_WORKER_STATE_DRAINING - if err := tc.persistence.UpdateWorker(context.Background(), assigned, assigned.GetMetadata().GetVersion()); err != nil { + if _, err := tc.persistence.UpdateWorker(context.Background(), assigned.GetMetadata().GetName(), store.PreconditionFrom(assigned), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.State = ateapipb.WorkerState_WORKER_STATE_DRAINING + return nil + }); err != nil { t.Fatalf("marking worker %s draining failed: %v", assignedPod, err) } diff --git a/cmd/ateapi/internal/controlapi/service.go b/cmd/ateapi/internal/controlapi/service.go index cf9a6cb390..c2f0b94d71 100644 --- a/cmd/ateapi/internal/controlapi/service.go +++ b/cmd/ateapi/internal/controlapi/service.go @@ -103,6 +103,10 @@ type serviceStore interface { ListActorTemplates(ctx context.Context, atespace string, opts store.ListOptions) (store.ListResponse[*ateapipb.ActorTemplate], error) DeleteActorTemplate(ctx context.Context, templateRef resources.ActorTemplateRef) (*ateapipb.ActorTemplate, error) ListWorkers(ctx context.Context, opts store.ListOptions) (store.ListResponse[*ateapipb.Worker], error) + GetWorker(ctx context.Context, name string) (*ateapipb.Worker, error) + CreateWorker(ctx context.Context, worker *ateapipb.Worker) (*ateapipb.Worker, error) + UpdateWorker(ctx context.Context, name string, precondition store.Precondition, mutate func(toUpdate *ateapipb.Worker) error) (*ateapipb.Worker, error) + DeleteWorker(ctx context.Context, name string, pre store.DeletePreconditions) (*ateapipb.Worker, error) AcquireLock(ctx context.Context, key string) (*store.Lock, error) } diff --git a/cmd/ateapi/internal/controlapi/syncer.go b/cmd/ateapi/internal/controlapi/syncer.go index 0e119aa2a6..be87fe75f6 100644 --- a/cmd/ateapi/internal/controlapi/syncer.go +++ b/cmd/ateapi/internal/controlapi/syncer.go @@ -91,9 +91,9 @@ type workerPoolSyncerStore interface { GetActor(ctx context.Context, actorRef resources.ActorRef) (*ateapipb.Actor, error) UpdateActor(ctx context.Context, actorRef resources.ActorRef, precondition store.Precondition, mutate func(toUpdate *ateapipb.Actor) error) (*ateapipb.Actor, error) GetWorker(ctx context.Context, name string) (*ateapipb.Worker, error) - CreateWorker(ctx context.Context, worker *ateapipb.Worker) error - UpdateWorker(ctx context.Context, worker *ateapipb.Worker, expectedVersion int64) error - DeleteWorker(ctx context.Context, name string) error + CreateWorker(ctx context.Context, worker *ateapipb.Worker) (*ateapipb.Worker, error) + UpdateWorker(ctx context.Context, name string, precondition store.Precondition, mutate func(toUpdate *ateapipb.Worker) error) (*ateapipb.Worker, error) + DeleteWorker(ctx context.Context, name string, pre store.DeletePreconditions) (*ateapipb.Worker, error) ListWorkers(ctx context.Context, opts store.ListOptions) (store.ListResponse[*ateapipb.Worker], error) } @@ -263,10 +263,10 @@ func (s *WorkerPoolSyncer) createOrUpdateWorker(ctx context.Context, key workerK State: ateapipb.WorkerState_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 { + // TODO: validateWorker now lives next to CreateWorker, which applies it + // too. Once this path calls the RPC instead of the store, the check + // here goes away and the errors below arrive as INVALID_ARGUMENT. + if errs := 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", append(key.logAttrs(), slog.Any("err", errs.ToAggregate()))...) @@ -274,7 +274,8 @@ func (s *WorkerPoolSyncer) createOrUpdateWorker(ctx context.Context, key workerK } // ErrAlreadyExists means we lost a create race; requeue and converge // via the update path. - return s.persistence.CreateWorker(ctx, worker) + _, err := s.persistence.CreateWorker(ctx, worker) + return err } changed := false @@ -300,7 +301,13 @@ func (s *WorkerPoolSyncer) createOrUpdateWorker(ctx context.Context, key workerK // ErrVersionConflict requeues the key; the retry re-fetches the worker at // its new version. - return s.persistence.UpdateWorker(ctx, w, w.GetMetadata().GetVersion()) + _, err = s.persistence.UpdateWorker(ctx, key.workerName(), store.PreconditionFrom(w), func(toUpdate *ateapipb.Worker) error { + toUpdate.Ip = w.GetIp() + toUpdate.SandboxClass = w.GetSandboxClass() + toUpdate.Labels = w.GetLabels() + return nil + }) + return err } func isWorkerEligible(pod *corev1.Pod) bool { @@ -356,8 +363,11 @@ func (s *WorkerPoolSyncer) markWorkerDraining(ctx context.Context, key workerKey return nil } slog.InfoContext(ctx, "Syncer: marking worker draining (pod deleting)", key.logAttrs()...) - worker.Status.State = ateapipb.WorkerState_WORKER_STATE_DRAINING - return s.persistence.UpdateWorker(ctx, worker, worker.GetMetadata().GetVersion()) + _, err = s.persistence.UpdateWorker(ctx, key.workerName(), store.PreconditionFrom(worker), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.State = ateapipb.WorkerState_WORKER_STATE_DRAINING + return nil + }) + return err } // reconcileDeadWorker cleans up a worker whose pod is gone. It releases the @@ -370,7 +380,14 @@ func (s *WorkerPoolSyncer) reconcileDeadWorker(ctx context.Context, name string) if err := s.releaseActorOnDeadWorker(ctx, name); err != nil { return err } - return s.persistence.DeleteWorker(ctx, name) + // The delete now reports absence rather than succeeding silently, but a + // worker already gone is exactly the state this is driving towards. + // Idempotency lives here, at the caller, so re-driving a reconcile is safe. + _, err := s.persistence.DeleteWorker(ctx, name, store.DeletePreconditions{}) + if errors.Is(err, store.ErrNotFound) { + return nil + } + return err } // storedWorkerListBackoff and storedWorkerListCap are the exponential backoff diff --git a/cmd/ateapi/internal/controlapi/syncer_test.go b/cmd/ateapi/internal/controlapi/syncer_test.go index 7cb6d1b0e5..67eeb95280 100644 --- a/cmd/ateapi/internal/controlapi/syncer_test.go +++ b/cmd/ateapi/internal/controlapi/syncer_test.go @@ -277,16 +277,21 @@ func TestSyncer_DeleteBoundWorker_ClearsActor(t *testing.T) { LatestSnapshot: &ateapipb.ObjectRef{Atespace: "team-orphan", Name: "last"}, }, }) - w, _ := persistence.GetWorker(ctx, testPodUID) - w.Status.Assignment = &ateapipb.ActorAssignment{ - ActorTemplate: &ateapipb.KubeNamespacedObjectRef{ - Namespace: ns, - Name: "tmpl", - }, - Actor: &ateapipb.ObjectRef{Atespace: createdActor.GetMetadata().GetAtespace(), Name: createdActor.GetMetadata().GetName()}, - ActorUid: createdActor.GetMetadata().GetUid(), + observedWorker, err := persistence.GetWorker(ctx, testPodUID) + if err != nil { + t.Fatalf("get worker: %v", err) } - if err := persistence.UpdateWorker(ctx, w, w.GetMetadata().GetVersion()); err != nil { + if _, err := persistence.UpdateWorker(ctx, testPodUID, store.PreconditionFrom(observedWorker), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.Assignment = &ateapipb.ActorAssignment{ + ActorTemplate: &ateapipb.KubeNamespacedObjectRef{ + Namespace: ns, + Name: "tmpl", + }, + Actor: &ateapipb.ObjectRef{Atespace: createdActor.GetMetadata().GetAtespace(), Name: createdActor.GetMetadata().GetName()}, + ActorUid: createdActor.GetMetadata().GetUid(), + } + return nil + }); err != nil { t.Fatalf("update worker: %v", err) } @@ -425,7 +430,7 @@ func TestSyncer_SoftDelete_MarksDraining(t *testing.T) { s := setupReconcileTest(t, persistence) ns, pool, pod, ip := "ns-drain", "pool1", "worker-drain", "10.0.0.2" - if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := persistence.CreateWorker(ctx, &ateapipb.Worker{ Metadata: &ateapipb.ResourceMetadata{Name: testPodUID}, WorkerNamespace: ns, WorkerPool: pool, WorkerPod: pod, Ip: ip, WorkerPodUid: testPodUID, NodeName: "node1", @@ -472,7 +477,7 @@ func TestSyncer_SoftDelete_NoPodIP(t *testing.T) { s := setupReconcileTest(t, persistence) ns, pool, pod := "ns-drain-noip", "pool1", "worker-drain-noip" - if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := persistence.CreateWorker(ctx, &ateapipb.Worker{ Metadata: &ateapipb.ResourceMetadata{Name: testPodUID}, WorkerNamespace: ns, WorkerPool: pool, WorkerPod: pod, Ip: "10.0.0.3", WorkerPodUid: testPodUID, NodeName: "node1", @@ -540,7 +545,7 @@ func TestMarkWorkerDraining(t *testing.T) { t.Run("already draining returns nil", func(t *testing.T) { persistence, cleanup := storetest.SetupTestStore(t) defer cleanup() - if err := persistence.CreateWorker(ctx, newWorker(ateapipb.WorkerState_WORKER_STATE_DRAINING)); err != nil { + if _, err := persistence.CreateWorker(ctx, newWorker(ateapipb.WorkerState_WORKER_STATE_DRAINING)); err != nil { t.Fatalf("create worker: %v", err) } s := &WorkerPoolSyncer{persistence: persistence} @@ -552,7 +557,7 @@ func TestMarkWorkerDraining(t *testing.T) { t.Run("active worker marked draining", func(t *testing.T) { persistence, cleanup := storetest.SetupTestStore(t) defer cleanup() - if err := persistence.CreateWorker(ctx, newWorker(ateapipb.WorkerState_WORKER_STATE_ACTIVE)); err != nil { + if _, err := persistence.CreateWorker(ctx, newWorker(ateapipb.WorkerState_WORKER_STATE_ACTIVE)); err != nil { t.Fatalf("create worker: %v", err) } s := &WorkerPoolSyncer{persistence: persistence} @@ -589,7 +594,7 @@ func TestReconcileDeadWorker(t *testing.T) { }, }, }) - if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := persistence.CreateWorker(ctx, &ateapipb.Worker{ Metadata: &ateapipb.ResourceMetadata{Name: testPodUID}, WorkerNamespace: ns, WorkerPool: pool, WorkerPod: pod, Ip: "10.0.0.5", WorkerPodUid: testPodUID, NodeName: "node1", @@ -641,7 +646,7 @@ func TestReconcileDeadWorker_IgnoresStaleIncarnationAssignment(t *testing.T) { }, }, }) - if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := persistence.CreateWorker(ctx, &ateapipb.Worker{ Metadata: &ateapipb.ResourceMetadata{Name: uid}, WorkerNamespace: ns, WorkerPool: pool, WorkerPod: pod, WorkerPodUid: uid, @@ -709,7 +714,7 @@ func TestSyncer_ReconcileOrphanedWorkers(t *testing.T) { } // A worker whose pod is live must be preserved. - if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := persistence.CreateWorker(ctx, &ateapipb.Worker{ Metadata: &ateapipb.ResourceMetadata{Name: liveUID}, WorkerNamespace: ns, WorkerPool: pool, WorkerPod: "worker-live", Ip: "10.0.0.9", WorkerPodUid: liveUID, NodeName: "node1", Status: &ateapipb.WorkerStatus{State: ateapipb.WorkerState_WORKER_STATE_ACTIVE}, @@ -730,7 +735,7 @@ func TestSyncer_ReconcileOrphanedWorkers(t *testing.T) { }, }, }) - if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := persistence.CreateWorker(ctx, &ateapipb.Worker{ Metadata: &ateapipb.ResourceMetadata{Name: orphanUID}, WorkerNamespace: ns, WorkerPool: pool, WorkerPod: "worker-orphan", Ip: "10.0.0.10", WorkerPodUid: orphanUID, NodeName: "node1", @@ -803,7 +808,7 @@ func TestSyncer_EnqueueStoredWorkers_RetriesTransientListError(t *testing.T) { defer cleanup() ns, pool := "ns-enq-retry", "pool1" - if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := persistence.CreateWorker(ctx, &ateapipb.Worker{ WorkerNamespace: ns, WorkerPool: pool, WorkerPod: "worker-1", Ip: "10.0.0.10", WorkerPodUid: "22222222-2222-2222-2222-222222222222", NodeName: "node1", Status: &ateapipb.WorkerStatus{ @@ -957,7 +962,7 @@ func TestReleaseActorOnDeadWorker_StateTransitions(t *testing.T) { }, }, }) - if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := persistence.CreateWorker(ctx, &ateapipb.Worker{ Metadata: &ateapipb.ResourceMetadata{Name: testPodUID}, WorkerNamespace: ns, WorkerPool: pool, WorkerPod: pod, Ip: ip, WorkerPodUid: testPodUID, NodeName: "node1", @@ -1000,15 +1005,14 @@ func TestReleaseActorOnDeadWorker_StateTransitions(t *testing.T) { type conflictStore struct { store.Interface conflictTriggered atomic.Bool - shouldInject func(worker *ateapipb.Worker) bool - onUpdate func(ctx context.Context, worker *ateapipb.Worker) + onUpdate func(ctx context.Context, name string) } -func (c *conflictStore) UpdateWorker(ctx context.Context, worker *ateapipb.Worker, expectedVersion int64) error { - if c.shouldInject != nil && c.shouldInject(worker) && c.conflictTriggered.CompareAndSwap(false, true) { - c.onUpdate(ctx, worker) +func (c *conflictStore) UpdateWorker(ctx context.Context, name string, precondition store.Precondition, mutate func(toUpdate *ateapipb.Worker) error) (*ateapipb.Worker, error) { + if c.onUpdate != nil && c.conflictTriggered.CompareAndSwap(false, true) { + c.onUpdate(ctx, name) } - return c.Interface.UpdateWorker(ctx, worker, expectedVersion) + return c.Interface.UpdateWorker(ctx, name, precondition, mutate) } func TestSyncer_UpdateWorker_RetryOnVersionConflict(t *testing.T) { @@ -1031,20 +1035,9 @@ func TestSyncer_UpdateWorker_RetryOnVersionConflict(t *testing.T) { var cs *conflictStore persistence, fakeK8s, fakeAte, syncer, cleanup := setupSyncerTestWithStore(t, ctx, func(s store.Interface) store.Interface { - // Configure the injector before the syncer starts. It only fires for the - // update under test, so the initial worker creation cannot consume it. - cs = &conflictStore{ - Interface: s, - shouldInject: func(w *ateapipb.Worker) bool { - return w.GetSandboxClass() == "microvm" - }, - onUpdate: func(c context.Context, w *ateapipb.Worker) { - if cw, err := s.GetWorker(c, testPodUID); err == nil { - cw.NodeName = "node2" - _ = s.UpdateWorker(c, cw, cw.GetMetadata().GetVersion()) - } - }, - } + // The injector is wired in below, once the worker exists and the pool + // change is staged, so the initial worker creation cannot consume it. + cs = &conflictStore{Interface: s} return cs }, pool) defer func() { @@ -1114,6 +1107,23 @@ func TestSyncer_UpdateWorker_RetryOnVersionConflict(t *testing.T) { t.Fatalf("pool informer cache failed to update: %v", err) } + // Configure conflictStore to inject a concurrent version bump in the store + // when the syncer calls UpdateWorker. The injected change lands on status.state: + // it is mutable, and unlike sandbox_class or labels the syncer's update path + // does not write it, so it survives the retry only if the retry re-read + // rather than replayed its stale copy. + cs.onUpdate = func(c context.Context, name string) { + observed, err := cs.Interface.GetWorker(c, name) + if err != nil { + t.Errorf("injecting conflict: get worker: %v", err) + return + } + _, _ = cs.Interface.UpdateWorker(c, name, store.PreconditionFrom(observed), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.State = ateapipb.WorkerState_WORKER_STATE_DRAINING + return nil + }) + } + // Touch the pod ONCE in K8s so the syncer reconciles it. The first reconcile's // UpdateWorker hits ErrVersionConflict (injected by conflictStore), which requeues // the key with backoff; the retry re-fetches the latest version from the store. @@ -1138,7 +1148,7 @@ func TestSyncer_UpdateWorker_RetryOnVersionConflict(t *testing.T) { } return false, err } - return w.SandboxClass == "microvm" && w.NodeName == "node2", nil + return w.SandboxClass == "microvm" && w.GetStatus().GetState() == ateapipb.WorkerState_WORKER_STATE_DRAINING, nil }) if err != nil { t.Fatalf("Worker failed to update SandboxClass after version conflict: %v", err) diff --git a/cmd/ateapi/internal/controlapi/worker.go b/cmd/ateapi/internal/controlapi/worker.go index b34a50ddf8..f1454baf53 100644 --- a/cmd/ateapi/internal/controlapi/worker.go +++ b/cmd/ateapi/internal/controlapi/worker.go @@ -16,12 +16,16 @@ package controlapi import ( "context" + "errors" "fmt" "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" + "github.com/agent-substrate/substrate/internal/resources" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + "k8s.io/apimachinery/pkg/api/validate/content" "k8s.io/apimachinery/pkg/util/validation/field" ) @@ -52,21 +56,275 @@ func validateListWorkersRequest(req *ateapipb.ListWorkersRequest) field.ErrorLis } func (s *RPCService) GetWorker(ctx context.Context, req *ateapipb.GetWorkerRequest) (*ateapipb.Worker, error) { - return nil, status.Error(codes.Unimplemented, "GetWorker is not implemented yet") + if errs := validateGetWorkerRequest(req); len(errs) > 0 { + return nil, toGRPCStatusError(errs) + } + name := req.GetWorker().GetName() + + worker, err := s.persistence.GetWorker(ctx, name) + if errors.Is(err, store.ErrNotFound) { + return nil, status.Errorf(codes.NotFound, "Worker %s not found", name) + } + if err != nil { + return nil, fmt.Errorf("while getting worker: %w", err) + } + return worker, nil +} + +func validateGetWorkerRequest(req *ateapipb.GetWorkerRequest) field.ErrorList { + var fldPath *field.Path + return resources.ValidateGlobalObjectRef(req.GetWorker(), fldPath.Child("worker")) } func (s *RPCService) CreateWorker(ctx context.Context, req *ateapipb.CreateWorkerRequest) (*ateapipb.Worker, error) { - return nil, status.Error(codes.Unimplemented, "CreateWorker is not implemented yet") + if errs := validateCreateWorkerRequest(req); len(errs) > 0 { + return nil, toGRPCStatusError(errs) + } + name := req.GetWorker().GetMetadata().GetName() + + // status is output-only, so whatever the request carried there is replaced + // rather than rejected. + worker := proto.Clone(req.GetWorker()).(*ateapipb.Worker) + worker.Status = &ateapipb.WorkerStatus{State: ateapipb.WorkerState_WORKER_STATE_ACTIVE} + + created, err := s.persistence.CreateWorker(ctx, worker) + if errors.Is(err, store.ErrAlreadyExists) { + return nil, status.Errorf(codes.AlreadyExists, "Worker %s already exists", name) + } + if err != nil { + return nil, fmt.Errorf("while creating worker: %w", err) + } + return created, nil +} + +func validateCreateWorkerRequest(req *ateapipb.CreateWorkerRequest) field.ErrorList { + var fldPath *field.Path + + worker, workerPath := req.GetWorker(), fldPath.Child("worker") + if worker == nil { + return field.ErrorList{field.Required(workerPath, "")} + } + return validateWorker(worker, workerPath) } +// UpdateWorker replaces the stored Worker with the one the request carries. +// Only sandbox_class and labels are the caller's to change; a request that +// alters an immutable field — including by leaving it unset, which would clear +// it — is rejected. The store enforces that, since only it holds the stored +// worker to compare against. func (s *RPCService) UpdateWorker(ctx context.Context, req *ateapipb.UpdateWorkerRequest) (*ateapipb.Worker, error) { - return nil, status.Error(codes.Unimplemented, "UpdateWorker is not implemented yet") + if errs := validateUpdateWorkerRequest(req); len(errs) > 0 { + return nil, toGRPCStatusError(errs) + } + in := req.GetWorker() + + return s.mutateWorker(ctx, in.GetMetadata().GetName(), store.PreconditionFrom(in), func(toUpdate *ateapipb.Worker) error { + // Status and metadata are server-owned fields. + status, metadata := toUpdate.GetStatus(), toUpdate.GetMetadata() + // Reset + merge from the input worker. + proto.Reset(toUpdate) + proto.Merge(toUpdate, in) + // Restore status and metadata from the server. + toUpdate.Status = status + toUpdate.Metadata = metadata + return nil + }) +} + +func validateUpdateWorkerRequest(req *ateapipb.UpdateWorkerRequest) field.ErrorList { + var fldPath *field.Path + + worker, workerPath := req.GetWorker(), fldPath.Child("worker") + if worker == nil { + return field.ErrorList{field.Required(workerPath, "")} + } + + // Only the metadata guards are checked here. The rest of the worker is + // pinned to what create stored — validateWorker already passed on it, and + // an update that changed any of it does not get written. + return resources.ValidateGlobalUpdateMetadataRef(worker.GetMetadata(), workerPath.Child("metadata")) } func (s *RPCService) DeleteWorker(ctx context.Context, req *ateapipb.DeleteWorkerRequest) (*ateapipb.Worker, error) { - return nil, status.Error(codes.Unimplemented, "DeleteWorker is not implemented yet") + if errs := validateDeleteWorkerRequest(req); len(errs) > 0 { + return nil, toGRPCStatusError(errs) + } + name := req.GetWorker().GetName() + + worker, err := s.persistence.DeleteWorker(ctx, name, store.DeletePreconditions{ + UID: req.GetOptions().GetUid(), + Version: req.GetOptions().GetVersion(), + }) + if err != nil { + switch { + case errors.Is(err, store.ErrNotFound): + return nil, status.Errorf(codes.NotFound, "Worker %s not found", name) + case errors.Is(err, store.ErrUIDConflict): + return nil, status.Errorf(codes.Aborted, "Worker %s does not have uid %s", name, req.GetOptions().GetUid()) + case errors.Is(err, store.ErrVersionConflict): + return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry") + } + return nil, fmt.Errorf("while deleting worker: %w", err) + } + return worker, nil +} + +func validateDeleteWorkerRequest(req *ateapipb.DeleteWorkerRequest) field.ErrorList { + var fldPath *field.Path + + errs := resources.ValidateGlobalObjectRef(req.GetWorker(), fldPath.Child("worker")) + + // Delete carries its preconditions in options, and each is optional: a zero + // value waives that guard. Absent options waive both, so nil needs no + // special case. + opts, optsPath := req.GetOptions(), fldPath.Child("options") + if val, p := opts.GetUid(), optsPath.Child("uid"); val != "" { + errs = append(errs, resources.ValidateUUID(val, p)...) + } + if val, p := opts.GetVersion(), optsPath.Child("version"); val < 0 { + errs = append(errs, field.Invalid(p, val, "must not be negative")) + } + + return errs } func (s *RPCService) DrainWorker(ctx context.Context, req *ateapipb.DrainWorkerRequest) (*ateapipb.Worker, error) { - return nil, status.Error(codes.Unimplemented, "DrainWorker is not implemented yet") + if errs := validateDrainWorkerRequest(req); len(errs) > 0 { + return nil, toGRPCStatusError(errs) + } + name := req.GetWorker().GetName() + + // A DrainWorkerRequest names a worker and carries no guards, so the ones + // the store requires come from a read here rather than from the client. A + // write that lands in between is reported as a conflict for the caller to + // retry, the same as any other guarded update. + observed, err := s.persistence.GetWorker(ctx, name) + if errors.Is(err, store.ErrNotFound) { + return nil, status.Errorf(codes.NotFound, "Worker %s not found", name) + } + if err != nil { + return nil, fmt.Errorf("while getting worker to drain: %w", err) + } + + return s.mutateWorker(ctx, name, store.PreconditionFrom(observed), func(toUpdate *ateapipb.Worker) error { + if toUpdate.GetStatus().GetState() == ateapipb.WorkerState_WORKER_STATE_DRAINING { + // already draining, do nothing + return &workerUnchanged{worker: proto.Clone(toUpdate).(*ateapipb.Worker)} + } + toUpdate.Status.State = ateapipb.WorkerState_WORKER_STATE_DRAINING + // status.assignment is deliberately left alone: a draining Worker keeps + // hosting the Actor bound to it until something releases it. Draining + // only stops the scheduler routing new Actors here. + return nil + }) +} + +func validateDrainWorkerRequest(req *ateapipb.DrainWorkerRequest) field.ErrorList { + var fldPath *field.Path + return resources.ValidateGlobalObjectRef(req.GetWorker(), fldPath.Child("worker")) +} + +// mutateWorker runs mutate against the named Worker and translates what comes +// back into the RPC's result. A mutation that found nothing to do reports the +// Worker it saw; anything else is a store error. +func (s *RPCService) mutateWorker(ctx context.Context, name string, precondition store.Precondition, mutate func(toUpdate *ateapipb.Worker) error) (*ateapipb.Worker, error) { + worker, err := s.persistence.UpdateWorker(ctx, name, precondition, mutate) + if err == nil { + return worker, nil + } + + var unchanged *workerUnchanged + if errors.As(err, &unchanged) { + return unchanged.worker, nil + } + switch { + case errors.Is(err, store.ErrNotFound): + return nil, status.Errorf(codes.NotFound, "Worker %s not found", name) + case errors.Is(err, store.ErrUIDConflict): + return nil, status.Errorf(codes.Aborted, "Worker %s is not the one the request describes", name) + case errors.Is(err, store.ErrVersionConflict): + return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry") + case errors.Is(err, store.ErrImmutableField): + return nil, status.Errorf(codes.InvalidArgument, "while updating worker %s: %v", name, err) + case errors.Is(err, store.ErrPreconditionRequired): + return nil, status.Errorf(codes.InvalidArgument, "while updating worker %s: %v", name, err) + } + return nil, fmt.Errorf("while updating worker: %w", err) +} + +// workerUnchanged ends an UpdateWorker mutation that found its work already +// done. The store hands a mutation's error straight back and leaves the Worker — +// and its version — untouched, which is what lets DrainWorker be idempotent: a +// call with nothing left to do costs no version bump. worker is a copy, because +// the store is free to reuse or discard the message once mutate returns. +type workerUnchanged struct { + worker *ateapipb.Worker +} + +func (u *workerUnchanged) Error() string { return "worker is already in the requested state" } + +// validateWorker checks that the caller-controlled fields of a Worker are +// well-formed. It is the create-time check: every field it covers is immutable +// afterwards, so no update path re-runs it. +func validateWorker(worker *ateapipb.Worker, fldPath *field.Path) field.ErrorList { + var errs field.ErrorList + + // Worker is global-scoped: metadata.atespace must be empty, name required + + // valid. uid and version are server-assigned, so a create ignores whatever + // the request carried in them. + metaPath := fldPath.Child("metadata") + if val, p := worker.GetMetadata().GetAtespace(), metaPath.Child("atespace"); val != "" { + errs = append(errs, field.Invalid(p, val, "must be empty for a global-scoped resource")) + } + if val, p := worker.GetMetadata().GetName(), metaPath.Child("name"); val == "" { + errs = append(errs, field.Required(p, "")) + } else { + errs = append(errs, resources.ValidateResourceName(val, p)...) + } + + if val, fldPath := worker.GetWorkerNamespace(), fldPath.Child("worker_namespace"); val == "" { + errs = append(errs, field.Required(fldPath, "")) + } else { + for _, msg := range content.IsDNS1123Label(val) { + errs = append(errs, field.Invalid(fldPath, val, msg)) + } + } + + if val, fldPath := worker.GetWorkerPool(), fldPath.Child("worker_pool"); val == "" { + errs = append(errs, field.Required(fldPath, "")) + } else { + for _, msg := range content.IsDNS1123Subdomain(val) { + errs = append(errs, field.Invalid(fldPath, val, msg)) + } + } + + if val, fldPath := worker.GetWorkerPod(), fldPath.Child("worker_pod"); val == "" { + errs = append(errs, field.Required(fldPath, "")) + } else { + for _, msg := range content.IsDNS1123Subdomain(val) { + errs = append(errs, field.Invalid(fldPath, val, msg)) + } + } + + if val, fldPath := worker.GetIp(), fldPath.Child("ip"); val == "" { + errs = append(errs, field.Required(fldPath, "")) + } else { + errs = append(errs, resources.ValidateIP(val, fldPath)...) + } + + if val, fldPath := worker.GetWorkerPodUid(), fldPath.Child("worker_pod_uid"); val == "" { + errs = append(errs, field.Required(fldPath, "")) + } else { + errs = append(errs, resources.ValidateUUID(val, fldPath)...) + } + + if val, fldPath := worker.GetNodeName(), fldPath.Child("node_name"); val == "" { + errs = append(errs, field.Required(fldPath, "")) + } else { + for _, msg := range content.IsDNS1123Subdomain(val) { + errs = append(errs, field.Invalid(fldPath, val, msg)) + } + } + + return errs } diff --git a/cmd/ateapi/internal/controlapi/worker_test.go b/cmd/ateapi/internal/controlapi/worker_test.go index 3809837f7a..ed9bcbee18 100644 --- a/cmd/ateapi/internal/controlapi/worker_test.go +++ b/cmd/ateapi/internal/controlapi/worker_test.go @@ -16,14 +16,110 @@ package controlapi import ( "context" + "strings" "testing" + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store/storetest" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/google/go-cmp/cmp" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/testing/protocmp" "k8s.io/apimachinery/pkg/util/validation/field" ) +// Worker names are pod UIDs, which are opaque to everything above the syncer. +const ( + apiWorkerName = "5f2c1a90-7b34-4e6d-8a11-0c3e9d5b7f42" + apiOtherWorkerName = "1a7e4c83-6d20-4f95-b3c8-9e0a2f6d4b17" +) + +// newAPIWorker returns a Worker in the shape CreateWorker accepts: named, with +// its pod coordinates filled in and no status — status is output-only. +func newAPIWorker(name string) *ateapipb.Worker { + return &ateapipb.Worker{ + Metadata: &ateapipb.ResourceMetadata{Name: name}, + WorkerNamespace: "ate-system", + WorkerPool: "pool-1", + WorkerPod: "worker-pod-1", + WorkerPodUid: name, + NodeName: "node-1", + Ip: "10.1.2.3", + SandboxClass: "gvisor", + Capacity: &ateapipb.WorkerCapacity{CpuMilli: 2000, MemoryBytes: 4 << 30}, + } +} + +func newAPIAssignment(actorUID string) *ateapipb.ActorAssignment { + return &ateapipb.ActorAssignment{ + ActorTemplate: &ateapipb.KubeNamespacedObjectRef{Namespace: "ate-system", Name: "tmpl"}, + Actor: &ateapipb.ObjectRef{Atespace: "team-a", Name: "actor-1"}, + ActorUid: actorUID, + } +} + +// newWorkerAPIService returns a service backed by a real store, which is what +// makes the compare-and-set assertions below meaningful — a fake would decide +// the outcome the test is trying to observe. +func newWorkerAPIService(t *testing.T) (*RPCService, store.Interface) { + t.Helper() + persistence, cleanup := storetest.SetupTestStore(t) + t.Cleanup(cleanup) + return &RPCService{persistence: persistence}, persistence +} + +// seedAPIWorker registers a worker directly through the store and returns it as +// stored, so tests start from a known uid and version. +func seedAPIWorker(t *testing.T, ctx context.Context, persistence store.Interface, worker *ateapipb.Worker) *ateapipb.Worker { + t.Helper() + worker = proto.Clone(worker).(*ateapipb.Worker) + if worker.GetStatus() == nil { + worker.Status = &ateapipb.WorkerStatus{State: ateapipb.WorkerState_WORKER_STATE_ACTIVE} + } + created, err := persistence.CreateWorker(ctx, worker) + if err != nil { + t.Fatalf("seeding worker %s: %v", worker.GetMetadata().GetName(), err) + } + return created +} + +// assignAPIWorker binds an Actor to a worker the way the resume workflow does: +// in-process, through the store. There is no AssignWorker RPC to go through. +func assignAPIWorker(t *testing.T, ctx context.Context, persistence store.Interface, name, actorUID string) *ateapipb.Worker { + t.Helper() + observed, err := persistence.GetWorker(ctx, name) + if err != nil { + t.Fatalf("getting worker %s to assign: %v", name, err) + } + assigned, err := persistence.UpdateWorker(ctx, name, store.PreconditionFrom(observed), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.Assignment = newAPIAssignment(actorUID) + return nil + }) + if err != nil { + t.Fatalf("assigning worker %s: %v", name, err) + } + return assigned +} + +func workerRef(name string) *ateapipb.ObjectRef { + return &ateapipb.ObjectRef{Name: name} +} + +// updateFrom builds the body of an UpdateWorker request the way a client does: +// read the worker, change what it means to change, send the whole thing back. +// The metadata comes along as the uid and version guards every update requires, +// and so does everything else — an update replaces the stored worker, so an +// immutable field the request drops reads as a request to clear it. +func updateFrom(observed *ateapipb.Worker, mutate func(*ateapipb.Worker)) *ateapipb.Worker { + worker := proto.Clone(observed).(*ateapipb.Worker) + if mutate != nil { + mutate(worker) + } + return worker +} + func TestValidateListWorkersRequest(t *testing.T) { tests := []struct { name string @@ -49,46 +145,596 @@ func TestValidateListWorkersRequest(t *testing.T) { } } -// The Worker CRUD methods are declared but not implemented. This pins that they -// report UNIMPLEMENTED rather than panicking on their nil dependencies, and -// will fail loudly as each one is filled in — at which point the corresponding -// case moves to a real test. -func TestWorkerAPIUnimplemented(t *testing.T) { - s := &RPCService{} +func TestGetWorker(t *testing.T) { ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + want := seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + + got, err := svc.GetWorker(ctx, &ateapipb.GetWorkerRequest{Worker: workerRef(apiWorkerName)}) + if err != nil { + t.Fatalf("GetWorker() failed: %v", err) + } + if diff := cmp.Diff(want, got, protocmp.Transform()); diff != "" { + t.Errorf("GetWorker() mismatch (-want +got):\n%s", diff) + } +} + +func TestGetWorker_Errors(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) tests := []struct { name string - call func() error + req *ateapipb.GetWorkerRequest + want codes.Code }{ - {"GetWorker", func() error { - _, err := s.GetWorker(ctx, &ateapipb.GetWorkerRequest{}) - return err - }}, - {"CreateWorker", func() error { - _, err := s.CreateWorker(ctx, &ateapipb.CreateWorkerRequest{}) - return err - }}, - {"UpdateWorker", func() error { - _, err := s.UpdateWorker(ctx, &ateapipb.UpdateWorkerRequest{}) - return err - }}, - {"DeleteWorker", func() error { - _, err := s.DeleteWorker(ctx, &ateapipb.DeleteWorkerRequest{}) - return err - }}, - {"DrainWorker", func() error { - _, err := s.DrainWorker(ctx, &ateapipb.DrainWorkerRequest{}) - return err - }}, + {"absent", &ateapipb.GetWorkerRequest{Worker: workerRef(apiOtherWorkerName)}, codes.NotFound}, + {"no ref", &ateapipb.GetWorkerRequest{}, codes.InvalidArgument}, + {"no name", &ateapipb.GetWorkerRequest{Worker: &ateapipb.ObjectRef{}}, codes.InvalidArgument}, + // Workers are global-scoped, so naming an atespace is a client bug + // rather than a lookup that happens to miss. + {"atespace set", &ateapipb.GetWorkerRequest{Worker: &ateapipb.ObjectRef{Atespace: "team-a", Name: apiWorkerName}}, codes.InvalidArgument}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, err := svc.GetWorker(ctx, tc.req) + if got := status.Code(err); got != tc.want { + t.Errorf("GetWorker() code = %v (err %v), want %v", got, err, tc.want) + } + }) + } +} + +func TestCreateWorker(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + + got, err := svc.CreateWorker(ctx, &ateapipb.CreateWorkerRequest{Worker: newAPIWorker(apiWorkerName)}) + if err != nil { + t.Fatalf("CreateWorker() failed: %v", err) + } + if got.GetMetadata().GetVersion() != 1 { + t.Errorf("created worker version = %d, want 1", got.GetMetadata().GetVersion()) + } + if got.GetMetadata().GetUid() == "" { + t.Error("created worker has no uid; the store is meant to assign one") + } + // A Worker is registered only once its pod is Ready and has an IP, which + // makes ACTIVE the only state it can be born in. + if got.GetStatus().GetState() != ateapipb.WorkerState_WORKER_STATE_ACTIVE { + t.Errorf("created worker state = %v, want %v", got.GetStatus().GetState(), ateapipb.WorkerState_WORKER_STATE_ACTIVE) + } + + stored, err := persistence.GetWorker(ctx, apiWorkerName) + if err != nil { + t.Fatalf("GetWorker() failed: %v", err) + } + if diff := cmp.Diff(stored, got, protocmp.Transform()); diff != "" { + t.Errorf("CreateWorker() returned something other than what it stored (-stored +returned):\n%s", diff) + } +} + +// status is output-only, so a request that carries one has it replaced rather +// than rejected. +func TestCreateWorker_IgnoresRequestStatus(t *testing.T) { + ctx := context.Background() + svc, _ := newWorkerAPIService(t) + + in := newAPIWorker(apiWorkerName) + in.Status = &ateapipb.WorkerStatus{ + State: ateapipb.WorkerState_WORKER_STATE_DRAINING, + Assignment: newAPIAssignment("actor-uid-1"), + } + + got, err := svc.CreateWorker(ctx, &ateapipb.CreateWorkerRequest{Worker: in}) + if err != nil { + t.Fatalf("CreateWorker() failed: %v", err) + } + want := &ateapipb.WorkerStatus{State: ateapipb.WorkerState_WORKER_STATE_ACTIVE} + if diff := cmp.Diff(want, got.GetStatus(), protocmp.Transform()); diff != "" { + t.Errorf("created worker status mismatch (-want +got):\n%s", diff) } +} +func TestCreateWorker_AlreadyExists(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + + _, err := svc.CreateWorker(ctx, &ateapipb.CreateWorkerRequest{Worker: newAPIWorker(apiWorkerName)}) + if got := status.Code(err); got != codes.AlreadyExists { + t.Errorf("CreateWorker() code = %v (err %v), want %v", got, err, codes.AlreadyExists) + } +} + +func TestCreateWorker_InvalidArgument(t *testing.T) { + ctx := context.Background() + svc, _ := newWorkerAPIService(t) + + tests := []struct { + name string + mutate func(*ateapipb.Worker) // nil sends no worker at all + }{ + {name: "no worker"}, + {name: "no name", mutate: func(w *ateapipb.Worker) { w.Metadata = &ateapipb.ResourceMetadata{} }}, + {name: "atespace set", mutate: func(w *ateapipb.Worker) { w.Metadata.Atespace = "team-a" }}, + {name: "no ip", mutate: func(w *ateapipb.Worker) { w.Ip = "" }}, + {name: "bad ip", mutate: func(w *ateapipb.Worker) { w.Ip = "not-an-ip" }}, + {name: "no node", mutate: func(w *ateapipb.Worker) { w.NodeName = "" }}, + {name: "no pool", mutate: func(w *ateapipb.Worker) { w.WorkerPool = "" }}, + {name: "no pod", mutate: func(w *ateapipb.Worker) { w.WorkerPod = "" }}, + {name: "pod uid not a uuid", mutate: func(w *ateapipb.Worker) { w.WorkerPodUid = "not-a-uuid" }}, + } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - err := tc.call() - if got := status.Code(err); got != codes.Unimplemented { - t.Errorf("%s: got code %v (err %v), want %v", tc.name, got, err, codes.Unimplemented) + req := &ateapipb.CreateWorkerRequest{} + if tc.mutate != nil { + worker := newAPIWorker(apiWorkerName) + tc.mutate(worker) + req.Worker = worker + } + _, err := svc.CreateWorker(ctx, req) + if got := status.Code(err); got != codes.InvalidArgument { + t.Errorf("CreateWorker() code = %v (err %v), want %v", got, err, codes.InvalidArgument) + } + }) + } +} + +func TestUpdateWorker(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seeded := seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + + got, err := svc.UpdateWorker(ctx, &ateapipb.UpdateWorkerRequest{ + Worker: updateFrom(seeded, func(w *ateapipb.Worker) { + w.SandboxClass = "microvm" + w.Labels = map[string]string{"tier": "batch"} + }), + }) + if err != nil { + t.Fatalf("UpdateWorker() failed: %v", err) + } + + want := proto.Clone(seeded).(*ateapipb.Worker) + want.SandboxClass = "microvm" + want.Labels = map[string]string{"tier": "batch"} + want.Metadata = got.GetMetadata() + if diff := cmp.Diff(want, got, protocmp.Transform()); diff != "" { + t.Errorf("UpdateWorker() mismatch (-want +got):\n%s", diff) + } + if got.GetMetadata().GetVersion() != 2 { + t.Errorf("updated worker version = %d, want 2", got.GetMetadata().GetVersion()) + } +} + +// Update replaces rather than patches, so a mutable field the request leaves +// unset is cleared. Immutable fields are the exception: dropping one of those +// is an error rather than a clear, which TestUpdateWorker_Errors covers. +func TestUpdateWorker_OmittedMutableFieldIsCleared(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + labelled := newAPIWorker(apiWorkerName) + labelled.Labels = map[string]string{"tier": "batch"} + seeded := seedAPIWorker(t, ctx, persistence, labelled) + + got, err := svc.UpdateWorker(ctx, &ateapipb.UpdateWorkerRequest{ + Worker: updateFrom(seeded, func(w *ateapipb.Worker) { + w.SandboxClass = "microvm" + w.Labels = nil + }), + }) + if err != nil { + t.Fatalf("UpdateWorker() failed: %v", err) + } + if got.GetSandboxClass() != "microvm" { + t.Errorf("sandbox_class = %q, want microvm", got.GetSandboxClass()) + } + if len(got.GetLabels()) != 0 { + t.Errorf("labels = %v, want them cleared: the request carried none", got.GetLabels()) + } +} + +// status is output-only, so the server keeps its own no matter what the request +// carries. That is what protects the in-process Actor binding, which lives +// under status and is written by the actor workflows rather than over the API. +func TestUpdateWorker_LeavesStatusAlone(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + assigned := assignAPIWorker(t, ctx, persistence, apiWorkerName, "actor-uid-1") + + got, err := svc.UpdateWorker(ctx, &ateapipb.UpdateWorkerRequest{ + Worker: updateFrom(assigned, func(w *ateapipb.Worker) { + w.SandboxClass = "microvm" + // A forged status: drained, and with the Actor released out from + // under the workflow that bound it. Neither may land. + w.Status = &ateapipb.WorkerStatus{State: ateapipb.WorkerState_WORKER_STATE_DRAINING} + }), + }) + if err != nil { + t.Fatalf("UpdateWorker() failed: %v", err) + } + if diff := cmp.Diff(assigned.GetStatus(), got.GetStatus(), protocmp.Transform()); diff != "" { + t.Errorf("UpdateWorker() disturbed status (-want +got):\n%s", diff) + } +} + +func TestUpdateWorker_Preconditions(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seeded := seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + + update := func(bend func(*ateapipb.ResourceMetadata)) error { + _, err := svc.UpdateWorker(ctx, &ateapipb.UpdateWorkerRequest{ + Worker: updateFrom(seeded, func(w *ateapipb.Worker) { + w.SandboxClass = "microvm" + bend(w.Metadata) + }), + }) + return err + } + + t.Run("stale version", func(t *testing.T) { + err := update(func(md *ateapipb.ResourceMetadata) { md.Version += 7 }) + if got := status.Code(err); got != codes.Aborted { + t.Errorf("UpdateWorker() code = %v, want %v", got, codes.Aborted) + } + }) + + t.Run("foreign uid", func(t *testing.T) { + err := update(func(md *ateapipb.ResourceMetadata) { md.Uid = apiOtherWorkerName }) + if got := status.Code(err); got != codes.Aborted { + t.Errorf("UpdateWorker() code = %v, want %v", got, codes.Aborted) + } + }) + + // Both guards are required: an update that pins neither is a blind write, + // which is rejected before it reaches the store. + t.Run("missing uid", func(t *testing.T) { + err := update(func(md *ateapipb.ResourceMetadata) { md.Uid = "" }) + if got := status.Code(err); got != codes.InvalidArgument { + t.Errorf("UpdateWorker() code = %v, want %v", got, codes.InvalidArgument) + } + }) + + t.Run("missing version", func(t *testing.T) { + err := update(func(md *ateapipb.ResourceMetadata) { md.Version = 0 }) + if got := status.Code(err); got != codes.InvalidArgument { + t.Errorf("UpdateWorker() code = %v, want %v", got, codes.InvalidArgument) + } + }) + + t.Run("matching", func(t *testing.T) { + if err := update(func(*ateapipb.ResourceMetadata) {}); err != nil { + t.Errorf("UpdateWorker() with matching preconditions failed: %v", err) + } + }) +} + +func TestUpdateWorker_Errors(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seeded := seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + + // Every case below carries the guards an update requires and the worker as + // stored, so the rule it is named for is the one that rejects it. + tests := []struct { + name string + mutate func(*ateapipb.Worker) // nil sends no worker at all + want codes.Code + }{ + {"no worker", nil, codes.InvalidArgument}, + {"atespace set", func(w *ateapipb.Worker) { w.Metadata.Atespace = "team-a" }, codes.InvalidArgument}, + {"absent", func(w *ateapipb.Worker) { + w.Metadata.Name = "9d1f7b06-3c58-4a2e-8b40-5f7c1e9a2d63" + }, codes.NotFound}, + // Immutable fields, changed. A replacement update carries the whole + // worker, so these are the cases where it carries a different one. + {"ip changed", func(w *ateapipb.Worker) { w.Ip = "10.9.9.9" }, codes.InvalidArgument}, + {"worker_pod changed", func(w *ateapipb.Worker) { w.WorkerPod = "worker-pod-2" }, codes.InvalidArgument}, + {"node_name changed", func(w *ateapipb.Worker) { w.NodeName = "node-2" }, codes.InvalidArgument}, + {"capacity changed", func(w *ateapipb.Worker) { w.Capacity.CpuMilli = 4000 }, codes.InvalidArgument}, + // And immutable fields dropped, which a replacement update reads as a + // request to clear them. Rejected rather than silently applied. + {"ip omitted", func(w *ateapipb.Worker) { w.Ip = "" }, codes.InvalidArgument}, + {"capacity omitted", func(w *ateapipb.Worker) { w.Capacity = nil }, codes.InvalidArgument}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + req := &ateapipb.UpdateWorkerRequest{} + if tc.mutate != nil { + req.Worker = updateFrom(seeded, tc.mutate) + } + _, err := svc.UpdateWorker(ctx, req) + if got := status.Code(err); got != tc.want { + t.Errorf("UpdateWorker() code = %v (err %v), want %v", got, err, tc.want) + } + }) + } +} + +// A draining worker can still have everything else about it updated; only its +// status is frozen. +func TestUpdateWorker_DrainingWorkerKeepsOtherFieldsMutable(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + drained, err := svc.DrainWorker(ctx, &ateapipb.DrainWorkerRequest{Worker: workerRef(apiWorkerName)}) + if err != nil { + t.Fatalf("DrainWorker() failed: %v", err) + } + + got, err := svc.UpdateWorker(ctx, &ateapipb.UpdateWorkerRequest{ + Worker: updateFrom(drained, func(w *ateapipb.Worker) { w.SandboxClass = "microvm" }), + }) + if err != nil { + t.Fatalf("UpdateWorker() failed: %v", err) + } + if got.GetSandboxClass() != "microvm" { + t.Errorf("sandbox_class = %q, want microvm", got.GetSandboxClass()) + } + if got.GetStatus().GetState() != ateapipb.WorkerState_WORKER_STATE_DRAINING { + t.Errorf("state = %v, want it still %v", got.GetStatus().GetState(), ateapipb.WorkerState_WORKER_STATE_DRAINING) + } +} + +func TestDeleteWorker(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seeded := seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + + got, err := svc.DeleteWorker(ctx, &ateapipb.DeleteWorkerRequest{Worker: workerRef(apiWorkerName)}) + if err != nil { + t.Fatalf("DeleteWorker() failed: %v", err) + } + if diff := cmp.Diff(seeded, got, protocmp.Transform()); diff != "" { + t.Errorf("DeleteWorker() returned something other than the worker it removed (-want +got):\n%s", diff) + } + if _, err := persistence.GetWorker(ctx, apiWorkerName); err == nil { + t.Error("worker still readable after DeleteWorker") + } +} + +// Delete reports absence rather than succeeding silently. Callers that want +// idempotence, like reconcileDeadWorker, opt into it by treating NOT_FOUND as +// success. +func TestDeleteWorker_Absent(t *testing.T) { + ctx := context.Background() + svc, _ := newWorkerAPIService(t) + + _, err := svc.DeleteWorker(ctx, &ateapipb.DeleteWorkerRequest{Worker: workerRef(apiWorkerName)}) + if got := status.Code(err); got != codes.NotFound { + t.Errorf("DeleteWorker() code = %v (err %v), want %v", got, err, codes.NotFound) + } +} + +// An assigned worker deletes like any other: the delete does not cascade, and +// an Actor pointing at a Worker that is gone is an expected steady state. +func TestDeleteWorker_AssignedWorkerDeletesAnyway(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + assignAPIWorker(t, ctx, persistence, apiWorkerName, "actor-uid-1") + + got, err := svc.DeleteWorker(ctx, &ateapipb.DeleteWorkerRequest{Worker: workerRef(apiWorkerName)}) + if err != nil { + t.Fatalf("DeleteWorker() failed: %v", err) + } + if got.GetStatus().GetAssignment().GetActorUid() != "actor-uid-1" { + t.Errorf("deleted worker assignment = %v, want the one it was holding", got.GetStatus().GetAssignment()) + } +} + +func TestDeleteWorker_Preconditions(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seeded := seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + + t.Run("stale version", func(t *testing.T) { + _, err := svc.DeleteWorker(ctx, &ateapipb.DeleteWorkerRequest{ + Worker: workerRef(apiWorkerName), + Options: &ateapipb.DeleteOptions{Version: seeded.GetMetadata().GetVersion() + 7}, + }) + if got := status.Code(err); got != codes.Aborted { + t.Errorf("DeleteWorker() code = %v (err %v), want %v", got, err, codes.Aborted) + } + }) + + t.Run("foreign uid", func(t *testing.T) { + _, err := svc.DeleteWorker(ctx, &ateapipb.DeleteWorkerRequest{ + Worker: workerRef(apiWorkerName), + Options: &ateapipb.DeleteOptions{Uid: apiOtherWorkerName}, + }) + if got := status.Code(err); got != codes.Aborted { + t.Errorf("DeleteWorker() code = %v (err %v), want %v", got, err, codes.Aborted) + } + }) + + // A refused delete must leave the worker where it was. + if _, err := persistence.GetWorker(ctx, apiWorkerName); err != nil { + t.Fatalf("worker gone after two refused deletes: %v", err) + } + + t.Run("matching", func(t *testing.T) { + if _, err := svc.DeleteWorker(ctx, &ateapipb.DeleteWorkerRequest{ + Worker: workerRef(apiWorkerName), + Options: &ateapipb.DeleteOptions{ + Uid: seeded.GetMetadata().GetUid(), + Version: seeded.GetMetadata().GetVersion(), + }, + }); err != nil { + t.Errorf("DeleteWorker() with matching preconditions failed: %v", err) + } + }) +} + +func TestDrainWorker(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + + got, err := svc.DrainWorker(ctx, &ateapipb.DrainWorkerRequest{Worker: workerRef(apiWorkerName)}) + if err != nil { + t.Fatalf("DrainWorker() failed: %v", err) + } + if got.GetStatus().GetState() != ateapipb.WorkerState_WORKER_STATE_DRAINING { + t.Errorf("state = %v, want %v", got.GetStatus().GetState(), ateapipb.WorkerState_WORKER_STATE_DRAINING) + } + if got.GetMetadata().GetVersion() != 2 { + t.Errorf("version = %d, want 2", got.GetMetadata().GetVersion()) + } + + // Draining again is a no-op, and specifically must not bump the version: + // callers re-drive drain on every pod event. + again, err := svc.DrainWorker(ctx, &ateapipb.DrainWorkerRequest{Worker: workerRef(apiWorkerName)}) + if err != nil { + t.Fatalf("second DrainWorker() failed: %v", err) + } + if diff := cmp.Diff(got, again, protocmp.Transform()); diff != "" { + t.Errorf("second DrainWorker() changed the worker (-first +second):\n%s", diff) + } +} + +// Drain deliberately leaves the bound Actor alone: it stops the scheduler +// routing new Actors here, it does not evict the one already running. +func TestDrainWorker_KeepsAssignment(t *testing.T) { + ctx := context.Background() + svc, persistence := newWorkerAPIService(t) + seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName)) + assignAPIWorker(t, ctx, persistence, apiWorkerName, "actor-uid-1") + + got, err := svc.DrainWorker(ctx, &ateapipb.DrainWorkerRequest{Worker: workerRef(apiWorkerName)}) + if err != nil { + t.Fatalf("DrainWorker() failed: %v", err) + } + if got.GetStatus().GetAssignment().GetActorUid() != "actor-uid-1" { + t.Errorf("assignment = %v, want it left in place", got.GetStatus().GetAssignment()) + } +} + +func TestDrainWorker_Errors(t *testing.T) { + ctx := context.Background() + svc, _ := newWorkerAPIService(t) + + tests := []struct { + name string + req *ateapipb.DrainWorkerRequest + want codes.Code + }{ + {"absent", &ateapipb.DrainWorkerRequest{Worker: workerRef(apiWorkerName)}, codes.NotFound}, + {"no ref", &ateapipb.DrainWorkerRequest{}, codes.InvalidArgument}, + {"atespace set", &ateapipb.DrainWorkerRequest{Worker: &ateapipb.ObjectRef{Atespace: "team-a", Name: apiWorkerName}}, codes.InvalidArgument}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, err := svc.DrainWorker(ctx, tc.req) + if got := status.Code(err); got != tc.want { + t.Errorf("DrainWorker() code = %v (err %v), want %v", got, err, tc.want) + } + }) + } +} + +// TestValidateWorker pins the field paths validateWorker reports. +// TestCreateWorker_InvalidArgument drives the same rules through the RPC, but +// only observes the status code. +func TestValidateWorker(t *testing.T) { + tests := []struct { + name string + mutate func(*ateapipb.Worker) // nil leaves the worker valid + wantMsg string // empty means valid + }{{ + name: "valid unassigned worker", + }, { + // status is output-only and every caller sets it itself, so it is not + // validated at all: a thoroughly malformed one still passes. + name: "status is not validated", + mutate: func(w *ateapipb.Worker) { + w.Status = &ateapipb.WorkerStatus{ + State: ateapipb.WorkerState(99), + Assignment: &ateapipb.ActorAssignment{Actor: &ateapipb.ObjectRef{Name: "actor"}}, + } + }, + }, { + name: "missing worker_namespace", + mutate: func(w *ateapipb.Worker) { w.WorkerNamespace = "" }, + wantMsg: "worker.worker_namespace: Required value", + }, { + name: "invalid worker_namespace", + mutate: func(w *ateapipb.Worker) { w.WorkerNamespace = "NS-1" }, + wantMsg: "worker.worker_namespace: Invalid value", + }, { + name: "missing worker_pool", + mutate: func(w *ateapipb.Worker) { w.WorkerPool = "" }, + wantMsg: "worker.worker_pool: Required value", + }, { + name: "missing worker_pod", + mutate: func(w *ateapipb.Worker) { w.WorkerPod = "" }, + wantMsg: "worker.worker_pod: Required value", + }, { + name: "missing ip", + mutate: func(w *ateapipb.Worker) { w.Ip = "" }, + wantMsg: "worker.ip: Required value", + }, { + name: "invalid ip", + mutate: func(w *ateapipb.Worker) { w.Ip = "not-an-ip" }, + wantMsg: "worker.ip: Invalid value", + }, { + name: "missing worker_pod_uid", + mutate: func(w *ateapipb.Worker) { w.WorkerPodUid = "" }, + wantMsg: "worker.worker_pod_uid: Required value", + }, { + name: "invalid worker_pod_uid", + mutate: func(w *ateapipb.Worker) { w.WorkerPodUid = "INVALID-UUID" }, + wantMsg: "worker.worker_pod_uid: Invalid value", + }, { + name: "missing node_name", + mutate: func(w *ateapipb.Worker) { w.NodeName = "" }, + wantMsg: "worker.node_name: Required value", + }, { + name: "invalid node_name", + mutate: func(w *ateapipb.Worker) { w.NodeName = "NODE_NAME" }, + wantMsg: "worker.node_name: Invalid value", + }, { + name: "missing metadata", + mutate: func(w *ateapipb.Worker) { w.Metadata = nil }, + wantMsg: "worker.metadata.name: Required value", + }, { + name: "missing metadata.name", + mutate: func(w *ateapipb.Worker) { w.Metadata = &ateapipb.ResourceMetadata{} }, + wantMsg: "worker.metadata.name: Required value", + }, { + name: "invalid metadata.name", + mutate: func(w *ateapipb.Worker) { w.Metadata.Name = "Not A Name" }, + wantMsg: "worker.metadata.name: Invalid value", + }, { + name: "metadata.atespace set on a global-scoped Worker", + mutate: func(w *ateapipb.Worker) { w.Metadata.Atespace = "team-a" }, + wantMsg: "worker.metadata.atespace: Invalid value", + }} + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + worker := newAPIWorker(apiWorkerName) + if tc.mutate != nil { + tc.mutate(worker) + } + errs := validateWorker(worker, field.NewPath("worker")) + if tc.wantMsg == "" { + if len(errs) > 0 { + t.Fatalf("validateWorker() = %v, want no errors", errs) + } + return + } + // Any error may match: a case can trip more than one rule, so the + // wanted error is not always the first one reported. + for _, err := range errs { + if strings.Contains(err.Error(), tc.wantMsg) { + return + } } + t.Errorf("validateWorker() = %v, want an error containing %q", errs, tc.wantMsg) }) } } diff --git a/cmd/ateapi/internal/controlapi/workflow.go b/cmd/ateapi/internal/controlapi/workflow.go index 7954e7ef51..4ce571cecc 100644 --- a/cmd/ateapi/internal/controlapi/workflow.go +++ b/cmd/ateapi/internal/controlapi/workflow.go @@ -116,7 +116,7 @@ type actorWorkflowStore interface { UpdateActor(ctx context.Context, actorRef resources.ActorRef, precondition store.Precondition, mutate func(toUpdate *ateapipb.Actor) error) (*ateapipb.Actor, error) DeleteActor(ctx context.Context, actorRef resources.ActorRef) (*ateapipb.Actor, error) GetWorker(ctx context.Context, name string) (*ateapipb.Worker, error) - UpdateWorker(ctx context.Context, worker *ateapipb.Worker, expectedVersion int64) error + UpdateWorker(ctx context.Context, name string, precondition store.Precondition, mutate func(toUpdate *ateapipb.Worker) error) (*ateapipb.Worker, error) GetActorSnapshot(ctx context.Context, atespace, name string) (*ateapipb.ActorSnapshot, error) CreateActorSnapshot(ctx context.Context, snapshot *ateapipb.ActorSnapshot) (*ateapipb.ActorSnapshot, error) AcquireLock(ctx context.Context, key string) (*store.Lock, error) diff --git a/cmd/ateapi/internal/controlapi/workflow_pause.go b/cmd/ateapi/internal/controlapi/workflow_pause.go index 3d2fd5910a..87b7f0cbba 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause.go +++ b/cmd/ateapi/internal/controlapi/workflow_pause.go @@ -240,8 +240,11 @@ func (w *ActorWorkflow) ensurePausedFinalized(ctx context.Context, actorRef reso if wass := worker.GetStatus().GetAssignment(); wass != nil { if wass.GetActorUid() == latestActor.GetMetadata().GetUid() { - worker.Status.Assignment = nil - if err := w.store.UpdateWorker(ctx, worker, worker.GetMetadata().GetVersion()); err != nil { + _, err := w.store.UpdateWorker(ctx, worker.GetMetadata().GetName(), store.PreconditionFrom(worker), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.Assignment = nil + return nil + }) + if err != nil { if errors.Is(err, store.ErrVersionConflict) { return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry") } diff --git a/cmd/ateapi/internal/controlapi/workflow_pause_test.go b/cmd/ateapi/internal/controlapi/workflow_pause_test.go index f8860fe1fc..15b9828530 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_pause_test.go @@ -121,7 +121,7 @@ func TestEnsurePausedFinalized_RecordsContentScope(t *testing.T) { InProgressLocalSnapshotName: "snap-prefix", }, }) - if err := st.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := st.CreateWorker(ctx, &ateapipb.Worker{ WorkerNamespace: "default", WorkerPool: "pool1", WorkerPod: "worker-pod-1", diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index ce2b4efd87..7a99c1b88b 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume.go @@ -402,13 +402,14 @@ func (w *ActorWorkflow) validateAssignedWorker(ctx context.Context, actorRef res } if !w.scheduler.Applies(worker, constraints) { slog.ErrorContext(ctx, "crashing actor because previously assigned worker is not eligible anymore") - release := proto.Clone(worker).(*ateapipb.Worker) - release.Status.Assignment = nil // If that worker's pool is no longer eligible (e.g. the actor's // worker_selector was updated after the failed attempt), release it back // to the free pool instead of leaving it claimed forever — nothing else // reclaims a healthy worker whose actor moved on to a different pool. - if err := w.store.UpdateWorker(ctx, release, release.GetMetadata().GetVersion()); err != nil { + if _, err := w.store.UpdateWorker(ctx, worker.GetMetadata().GetName(), store.PreconditionFrom(worker), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.Assignment = nil + return nil + }); err != nil { return nil, fmt.Errorf("while releasing stale worker assignment: %w", err) } if cerr := crashActor(ctx, w.store, actorRef, ateattr.OperationResume, ateattr.ReasonCorruptedAssignment); cerr != nil { @@ -473,10 +474,10 @@ func (w *ActorWorkflow) assignWorkerAttempt(ctx context.Context, actorRef resour assignedWorker = worker break } - // Workers() returns pointers directly from the cache so we need to clone before - // mutating so that the cache is not corrupted if UpdateWorker fails. + // Workers() returns pointers directly from the cache, so clone before + // handing the worker to the goroutine: the mutation runs against the + // store's own copy, but the precondition and the log below read this one. releaseWorker := proto.Clone(worker).(*ateapipb.Worker) - releaseWorker.Status.Assignment = nil // The claimed worker is no longer eligible (e.g. the actor's // worker_selector changed after the failed attempt); release it back // to the free pool — nothing else reclaims a healthy worker whose @@ -484,7 +485,10 @@ func (w *ActorWorkflow) assignWorkerAttempt(ctx context.Context, actorRef resour go func(release *ateapipb.Worker) { bgCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() - if err := w.store.UpdateWorker(bgCtx, release, release.GetMetadata().GetVersion()); err != nil { + if _, err := w.store.UpdateWorker(bgCtx, release.GetMetadata().GetName(), store.PreconditionFrom(release), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.Assignment = nil + return nil + }); err != nil { slog.ErrorContext(bgCtx, "Failed to release stale worker assignment", slog.String("worker", release.GetWorkerNamespace()+"/"+release.GetWorkerPod()), slog.Any("err", err)) @@ -505,10 +509,7 @@ func (w *ActorWorkflow) assignWorkerAttempt(ctx context.Context, actorRef resour slog.InfoContext(ctx, "Picked worker", slog.Any("worker", pickedWorker.String())) } - // Workers() returns pointers directly from the cache so we need to clone before - // mutating so that the cache is not corrupted if UpdateWorker fails. - assignedWorker = proto.Clone(assignedWorker).(*ateapipb.Worker) - assignedWorker.Status.Assignment = &ateapipb.ActorAssignment{ + assignment := &ateapipb.ActorAssignment{ ActorTemplate: &ateapipb.KubeNamespacedObjectRef{ Namespace: actor.GetActorTemplateNamespace(), Name: actor.GetActorTemplateName(), @@ -520,13 +521,21 @@ func (w *ActorWorkflow) assignWorkerAttempt(ctx context.Context, actorRef resour ActorUid: actor.GetMetadata().GetUid(), } - if err := w.store.UpdateWorker(ctx, assignedWorker, assignedWorker.GetMetadata().GetVersion()); err != nil { + // Workers() returns pointers directly from the cache, so the claim is written + // by mutating the store's own copy; the cached one is only read, for the + // version this claim is conditioned on. + stored, err := w.store.UpdateWorker(ctx, assignedWorker.GetMetadata().GetName(), store.PreconditionFrom(assignedWorker), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.Assignment = assignment + return nil + }) + if err != nil { if errors.Is(err, store.ErrNotFound) { w.workerCache.Forget(assignedWorker.GetMetadata().GetName()) return nil, nil, fmt.Errorf("selected worker disappeared before claim: %w", store.ErrVersionConflict) } return nil, nil, err } + assignedWorker = stored newAssignment := workerAssignmentFrom(assignedWorker) storedActor, err := w.store.UpdateActor(ctx, actorRef, store.PreconditionFrom(actor), func(toUpdate *ateapipb.Actor) error { diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index ebaa308bf2..a84be1fbe3 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume_test.go @@ -102,8 +102,8 @@ type updateWorkerErrorStore struct { err error } -func (s *updateWorkerErrorStore) UpdateWorker(context.Context, *ateapipb.Worker, int64) error { - return s.err +func (s *updateWorkerErrorStore) UpdateWorker(context.Context, string, store.Precondition, func(*ateapipb.Worker) error) (*ateapipb.Worker, error) { + return nil, s.err } func TestAssignWorkerAttempt_MissingSelectedWorkerIsRetried(t *testing.T) { @@ -162,7 +162,7 @@ func TestAssignWorkerAttempt_SkipsWorkerAssignedInOtherAtespace(t *testing.T) { }, }, } - if err := persistence.CreateWorker(ctx, worker); err != nil { + if _, err := persistence.CreateWorker(ctx, worker); err != nil { t.Fatalf("CreateWorker: %v", err) } @@ -239,7 +239,7 @@ func TestAssignWorkerAttempt_ReleasesIneligibleStaleWorkerInBackground(t *testin }, } for _, w := range []*ateapipb.Worker{stale, free} { - if err := persistence.CreateWorker(ctx, w); err != nil { + if _, err := persistence.CreateWorker(ctx, w); err != nil { t.Fatalf("CreateWorker(%s): %v", w.GetWorkerPod(), err) } } @@ -314,7 +314,7 @@ func TestAssignWorkerAttempt_RetryAfterConflictPicksFreshWorker(t *testing.T) { }, } for _, w := range []*ateapipb.Worker{contested, fallback} { - if err := persistence.CreateWorker(ctx, w); err != nil { + if _, err := persistence.CreateWorker(ctx, w); err != nil { t.Fatalf("CreateWorker(%s): %v", w.GetWorkerPod(), err) } } @@ -327,12 +327,13 @@ func TestAssignWorkerAttempt_RetryAfterConflictPicksFreshWorker(t *testing.T) { // A concurrent resume of another actor wins the contested worker, bumping // its stored version past the failed attempt's snapshot. - claimed := proto.Clone(beforeClaim).(*ateapipb.Worker) - claimed.Status.Assignment = &ateapipb.ActorAssignment{ - Actor: &ateapipb.ObjectRef{Atespace: "team-a", Name: "other"}, - ActorUid: "other-actor-uid", - } - if err := persistence.UpdateWorker(ctx, claimed, claimed.GetMetadata().GetVersion()); err != nil { + if _, err := persistence.UpdateWorker(ctx, beforeClaim.GetMetadata().GetName(), store.PreconditionFrom(beforeClaim), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.Assignment = &ateapipb.ActorAssignment{ + Actor: &ateapipb.ObjectRef{Atespace: "team-a", Name: "other"}, + ActorUid: "other-actor-uid", + } + return nil + }); err != nil { t.Fatalf("UpdateWorker (concurrent claim): %v", err) } @@ -410,7 +411,7 @@ func (c *conflictInjectingStore) UpdateActorSnapshotTag(ctx context.Context, ate // returns the actor plus a started worker cache. func seedAssignFixture(t *testing.T, ctx context.Context, persistence store.Interface) (*ateapipb.Actor, *workercache.Cache) { t.Helper() - if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := persistence.CreateWorker(ctx, &ateapipb.Worker{ Metadata: &ateapipb.ResourceMetadata{Name: testWorkerUID("pod-1")}, WorkerNamespace: "worker-ns", WorkerPool: "pool", @@ -785,7 +786,7 @@ func TestValidateAssignedWorker_WorkerOwnership(t *testing.T) { ctx := context.Background() persistence := newTestPersistence(t) - if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := persistence.CreateWorker(ctx, &ateapipb.Worker{ Metadata: &ateapipb.ResourceMetadata{Name: testWorkerUID("pod-1")}, WorkerNamespace: "worker-ns", WorkerPool: "pool", diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go index a37470826d..e386aeb441 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go @@ -409,7 +409,7 @@ func TestEnsureSuspendedFinalized_ReleasesOnlyOwnWorker(t *testing.T) { }, }, } - if err := persistence.CreateWorker(ctx, worker); err != nil { + if _, err := persistence.CreateWorker(ctx, worker); err != nil { t.Fatalf("CreateWorker: %v", err) } diff --git a/cmd/ateapi/internal/store/atepg/atepg.go b/cmd/ateapi/internal/store/atepg/atepg.go index a0d3e523ac..d917c407e2 100644 --- a/cmd/ateapi/internal/store/atepg/atepg.go +++ b/cmd/ateapi/internal/store/atepg/atepg.go @@ -1176,33 +1176,33 @@ func (p *Persistence) DeleteActorSnapshotTag(ctx context.Context, atespace, name // --- Workers --- -func (p *Persistence) CreateWorker(ctx context.Context, worker *ateapipb.Worker) error { +func (p *Persistence) CreateWorker(ctx context.Context, worker *ateapipb.Worker) (*ateapipb.Worker, error) { dbWorker := proto.Clone(worker).(*ateapipb.Worker) // Workers are global-scoped, so the atespace is always empty. dbWorker.Metadata = newCreateMetadata("", worker.GetMetadata().GetName()) protoBytes, err := proto.Marshal(dbWorker) if err != nil { - return fmt.Errorf("marshaling worker: %w", err) + return nil, fmt.Errorf("marshaling worker: %w", err) } - err = p.writeAndAppendEvent(ctx, store.WorkerEventCreated, dbWorker, func(ctx context.Context, tx pgx.Tx) (bool, error) { + created, err := p.writeAndAppendEvent(ctx, store.WorkerEventCreated, func(ctx context.Context, tx pgx.Tx) (*ateapipb.Worker, error) { _, err := tx.Exec(ctx, ` INSERT INTO workers (name, uid, version, proto) VALUES ($1, $2, $3, $4)`, dbWorker.GetMetadata().GetName(), dbWorker.GetMetadata().GetUid(), dbWorker.GetMetadata().GetVersion(), protoBytes) if err != nil { - return false, err + return nil, err } - return true, nil + return dbWorker, nil }) if err != nil { if isUniqueViolation(err) { - return store.ErrAlreadyExists + return nil, store.ErrAlreadyExists } - return fmt.Errorf("creating worker: %w", err) + return nil, fmt.Errorf("creating worker: %w", err) } - return nil + return created, nil } func getWorkerRow(ctx context.Context, q querier, name string) (*ateapipb.Worker, error) { @@ -1225,61 +1225,94 @@ func (p *Persistence) GetWorker(ctx context.Context, name string) (*ateapipb.Wor return getWorkerRow(ctx, p.pool, name) } -func (p *Persistence) UpdateWorker(ctx context.Context, worker *ateapipb.Worker, expectedVersion int64) error { - name := worker.GetMetadata().GetName() - - dbWorker := proto.Clone(worker).(*ateapipb.Worker) - dbWorker.Metadata = newUpdateMetadata(worker.GetMetadata()) - dbWorker.Metadata.Version = expectedVersion + 1 +// getWorkerRowForUpdate reads the worker and holds its row lock for the rest of +// tx, so nothing else can write the row between this read and the write that +// follows it. +func getWorkerRowForUpdate(ctx context.Context, tx pgx.Tx, name string) (*ateapipb.Worker, error) { + var protoBytes []byte + if err := tx.QueryRow(ctx, `SELECT proto FROM workers WHERE name = $1 FOR UPDATE`, name).Scan(&protoBytes); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return nil, store.ErrNotFound + } + return nil, fmt.Errorf("locking worker %s for update: %w", name, err) + } + out := &ateapipb.Worker{} + if err := proto.Unmarshal(protoBytes, out); err != nil { + return nil, fmt.Errorf("unmarshaling worker: %w", err) + } + return out, nil +} - protoBytes, err := proto.Marshal(dbWorker) - if err != nil { - return fmt.Errorf("marshaling worker: %w", err) +// UpdateWorker runs mutate against the worker read FOR UPDATE inside the write +// transaction, so a concurrent writer blocks on the row lock rather than +// interleaving. That is what makes an occupancy test inside mutate a +// compare-and-set. The predicate cannot be pushed into SQL: the row stores an +// opaque marshaled proto, so assignment is not addressable in a WHERE clause. +func (p *Persistence) UpdateWorker(ctx context.Context, name string, precondition store.Precondition, mutate func(*ateapipb.Worker) error) (*ateapipb.Worker, error) { + if err := precondition.Validate(); err != nil { + return nil, err } + return p.writeAndAppendEvent(ctx, store.WorkerEventUpdated, func(ctx context.Context, tx pgx.Tx) (*ateapipb.Worker, error) { + dbWorker, err := getWorkerRowForUpdate(ctx, tx, name) + if err != nil { + return nil, err + } + if err := precondition.Check(dbWorker.GetMetadata()); err != nil { + return nil, err + } - return p.writeAndAppendEvent(ctx, store.WorkerEventUpdated, dbWorker, func(ctx context.Context, tx pgx.Tx) (bool, error) { - var returned []byte - err := tx.QueryRow(ctx, ` - UPDATE workers - SET version = $1, proto = $2 - WHERE name = $3 AND version = $4 - RETURNING proto`, - dbWorker.GetMetadata().GetVersion(), protoBytes, name, expectedVersion, - ).Scan(&returned) - if err == nil { - return true, nil + // Snapshot the stored state before handing the worker to mutate. + // mutate is free to edit anything it is given. + workerBeforeMutation := proto.Clone(dbWorker).(*ateapipb.Worker) + if err := mutate(dbWorker); err != nil { + return nil, err + } + if err := store.CheckWorkerMutation(workerBeforeMutation, dbWorker); err != nil { + return nil, err } - if !errors.Is(err, pgx.ErrNoRows) { - return false, fmt.Errorf("updating worker %s: %w", name, err) + // Stored metadata is authoritative; discard any metadata edits made by + // the closure and derive the next revision from the row we locked. + dbWorker.Metadata = newUpdateMetadata(workerBeforeMutation.GetMetadata()) + + protoBytes, err := proto.Marshal(dbWorker) + if err != nil { + return nil, fmt.Errorf("marshaling worker: %w", err) } - current, getErr := getWorkerRow(ctx, tx, name) - if getErr != nil { - return false, getErr + commandTag, err := tx.Exec(ctx, ` + UPDATE workers + SET version = $1, proto = $2 + WHERE name = $3`, + dbWorker.GetMetadata().GetVersion(), protoBytes, name) + if err != nil { + return nil, fmt.Errorf("updating worker %s: %w", name, err) } - if current.GetMetadata().GetVersion() != expectedVersion { - return false, store.ErrVersionConflict + if commandTag.RowsAffected() != 1 { + return nil, fmt.Errorf("updating worker %s affected %d rows, want 1", name, commandTag.RowsAffected()) } - return false, fmt.Errorf("update worker %s: no row matched but current state is otherwise consistent", name) + return dbWorker, nil }) } -func (p *Persistence) DeleteWorker(ctx context.Context, name string) error { - deletedEvent := &ateapipb.Worker{Metadata: &ateapipb.ResourceMetadata{Name: name}} - return p.writeAndAppendEvent(ctx, store.WorkerEventDeleted, deletedEvent, func(ctx context.Context, tx pgx.Tx) (bool, error) { - var protoBytes []byte - err := tx.QueryRow(ctx, ` - DELETE FROM workers - WHERE name = $1 - RETURNING proto`, name).Scan(&protoBytes) +func (p *Persistence) DeleteWorker(ctx context.Context, name string, pre store.DeletePreconditions) (*ateapipb.Worker, error) { + return p.writeAndAppendEvent(ctx, store.WorkerEventDeleted, func(ctx context.Context, tx pgx.Tx) (*ateapipb.Worker, error) { + // Locked rather than plainly read so the incarnation pre was evaluated + // against is the one the DELETE removes. + deleted, err := getWorkerRowForUpdate(ctx, tx, name) if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - // Idempotent: nothing existed, so no event to publish either. - return false, nil - } - return false, fmt.Errorf("deleting worker %s: %w", name, err) + return nil, err + } + if err := pre.Check(deleted.GetMetadata()); err != nil { + return nil, err + } + commandTag, err := tx.Exec(ctx, `DELETE FROM workers WHERE name = $1`, name) + if err != nil { + return nil, fmt.Errorf("deleting worker %s: %w", name, err) + } + if commandTag.RowsAffected() != 1 { + return nil, fmt.Errorf("deleting worker %s affected %d rows, want 1", name, commandTag.RowsAffected()) } - return true, nil + return deleted, nil }) } diff --git a/cmd/ateapi/internal/store/atepg/outbox.go b/cmd/ateapi/internal/store/atepg/outbox.go index 89d2d27f2c..aa9b981272 100644 --- a/cmd/ateapi/internal/store/atepg/outbox.go +++ b/cmd/ateapi/internal/store/atepg/outbox.go @@ -68,36 +68,40 @@ func unmarshalWorkerEvent(payload []byte) (store.WorkerEvent, error) { return store.WorkerEvent{Type: eventType, Worker: worker}, nil } -// writeAndAppendEvent runs fn inside a transaction, then--only if fn -// reports a change worth publishing--appends the event to the worker_outbox -// table in the same transaction, so watchers see it if and only if the -// transaction commits. -func (p *Persistence) writeAndAppendEvent(ctx context.Context, eventType store.WorkerEventType, worker *ateapipb.Worker, fn func(ctx context.Context, tx pgx.Tx) (changed bool, err error)) error { +// writeAndAppendEvent runs fn inside a transaction, then--only if fn reports a +// worker worth publishing--appends the event to the worker_outbox table in the +// same transaction, so watchers see it if and only if the transaction commits. +// fn returns the worker the event carries, or nil to skip the event; it is +// returned from writeAndAppendEvent so callers get back what actually +// committed. The worker comes from fn rather than from the caller because an +// update only knows what it wrote once its mutation has run inside the +// transaction. +func (p *Persistence) writeAndAppendEvent(ctx context.Context, eventType store.WorkerEventType, fn func(ctx context.Context, tx pgx.Tx) (*ateapipb.Worker, error)) (*ateapipb.Worker, error) { tx, err := p.pool.Begin(ctx) if err != nil { - return fmt.Errorf("beginning transaction: %w", err) + return nil, fmt.Errorf("beginning transaction: %w", err) } defer tx.Rollback(ctx) //nolint:errcheck // no-op once committed - changed, err := fn(ctx, tx) + worker, err := fn(ctx, tx) if err != nil { - return err + return nil, err } - if changed { + if worker != nil { payload, err := marshalWorkerEvent(eventType, worker) if err != nil { - return fmt.Errorf("marshaling worker event: %w", err) + return nil, fmt.Errorf("marshaling worker event: %w", err) } if _, err := tx.Exec(ctx, `INSERT INTO worker_outbox (payload) VALUES ($1)`, payload); err != nil { - return fmt.Errorf("appending worker outbox: %w", err) + return nil, fmt.Errorf("appending worker outbox: %w", err) } } if err := tx.Commit(ctx); err != nil { - return fmt.Errorf("committing transaction: %w", err) + return nil, fmt.Errorf("committing transaction: %w", err) } - return nil + return worker, nil } const ( diff --git a/cmd/ateapi/internal/store/atepg/outbox_test.go b/cmd/ateapi/internal/store/atepg/outbox_test.go index 8c778a162a..89865d9570 100644 --- a/cmd/ateapi/internal/store/atepg/outbox_test.go +++ b/cmd/ateapi/internal/store/atepg/outbox_test.go @@ -73,7 +73,7 @@ func TestConnect_DedicatedWatchPool(t *testing.T) { WorkerPool: "pool", WorkerPod: "watchpool-pod", } - if err := p.CreateWorker(ctx, worker); err != nil { + if _, err := p.CreateWorker(ctx, worker); err != nil { t.Fatalf("CreateWorker failed: %v", err) } select { @@ -141,7 +141,7 @@ func TestWorkerEvent_OnlyAfterCommit(t *testing.T) { } // The equivalent committed write must produce an event. - if err := s.CreateWorker(ctx, worker); err != nil { + if _, err := s.CreateWorker(ctx, worker); err != nil { t.Fatalf("CreateWorker failed: %v", err) } select { @@ -359,7 +359,7 @@ func TestWorkerEvents_OneRowPerTransaction(t *testing.T) { WorkerPool: "pool", WorkerPod: "pod", } - if err := s.CreateWorker(ctx, worker); err != nil { + if _, err := s.CreateWorker(ctx, worker); err != nil { t.Fatalf("CreateWorker failed: %v", err) } for i := 0; i < 10; i++ { @@ -367,11 +367,13 @@ func TestWorkerEvents_OneRowPerTransaction(t *testing.T) { if err != nil { t.Fatalf("GetWorker failed: %v", err) } - if err := s.UpdateWorker(ctx, stored, stored.GetMetadata().GetVersion()); err != nil { + if _, err := s.UpdateWorker(ctx, "one-row-worker", store.PreconditionFrom(stored), func(*ateapipb.Worker) error { + return nil + }); err != nil { t.Fatalf("UpdateWorker %d failed: %v", i, err) } } - if err := s.DeleteWorker(ctx, "one-row-worker"); err != nil { + if _, err := s.DeleteWorker(ctx, "one-row-worker", store.DeletePreconditions{}); err != nil { t.Fatalf("DeleteWorker failed: %v", err) } @@ -413,7 +415,7 @@ func TestWatchWorkers_DeliveryFencedByOldestTransaction(t *testing.T) { WorkerPool: "pool", WorkerPod: "fenced", } - if err := s.CreateWorker(ctx, worker); err != nil { + if _, err := s.CreateWorker(ctx, worker); err != nil { t.Fatalf("CreateWorker failed: %v", err) } @@ -619,7 +621,7 @@ func TestWatchWorkers_ClosesWhenTrimmedPastCursor(t *testing.T) { WorkerPool: "pool", WorkerPod: "pod", } - if err := s.CreateWorker(ctx, worker); err != nil { + if _, err := s.CreateWorker(ctx, worker); err != nil { t.Fatalf("CreateWorker failed: %v", err) } select { @@ -704,7 +706,7 @@ func TestPartitionCreation_UnwedgesFromStrayedDefault(t *testing.T) { WorkerPool: "pool", WorkerPod: "wedge-pod", } - if err := s.CreateWorker(ctx, worker); err != nil { + if _, err := s.CreateWorker(ctx, worker); err != nil { t.Fatalf("CreateWorker during stall: %v", err) } var strays bool @@ -787,7 +789,7 @@ func TestOutboxMaintenance_NoDeadlockWithConcurrentWriters(t *testing.T) { Metadata: &ateapipb.ResourceMetadata{Name: fmt.Sprintf("ddl-worker-%d-%d", g, i)}, WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: fmt.Sprintf("ddl-pod-%d-%d", g, i), } - if err := s.CreateWorker(writerCtx, w); err != nil && writerCtx.Err() == nil { + if _, err := s.CreateWorker(writerCtx, w); err != nil && writerCtx.Err() == nil { select { case writerErr <- fmt.Errorf("writer %d iteration %d: %w", g, i, err): default: @@ -915,7 +917,7 @@ func TestWatchWorkers_BaselineDoesNotMaskOwedTrims(t *testing.T) { // C: a LATER xid that commits BEFORE the subscribe — the baseline poison: // a visible outbox row whose xid exceeds W's. - if err := s.CreateWorker(ctx, &ateapipb.Worker{ + if _, err := s.CreateWorker(ctx, &ateapipb.Worker{ Metadata: &ateapipb.ResourceMetadata{Name: "baseline-poison"}, WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: "baseline-pod", }); err != nil { diff --git a/cmd/ateapi/internal/store/store.go b/cmd/ateapi/internal/store/store.go index 4ca1296e48..5c119c058e 100644 --- a/cmd/ateapi/internal/store/store.go +++ b/cmd/ateapi/internal/store/store.go @@ -23,6 +23,7 @@ import ( "github.com/agent-substrate/substrate/internal/resources" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/protobuf/proto" ) var ( @@ -190,8 +191,10 @@ type Interface interface { // ErrNotFound if missing. DeleteActorTemplate(ctx context.Context, templateRef resources.ActorTemplateRef) (*ateapipb.ActorTemplate, error) - // Registers a new idle worker. Returns ErrAlreadyExists if already registered. - CreateWorker(ctx context.Context, worker *ateapipb.Worker) error + // Registers a new idle worker and returns the stored resource with + // server-assigned metadata (uid, version, timestamps). The input is not + // mutated. Returns ErrAlreadyExists if already registered. + CreateWorker(ctx context.Context, worker *ateapipb.Worker) (*ateapipb.Worker, error) // Fetches worker state by name. Returns ErrNotFound if missing. GetWorker(ctx context.Context, name string) (*ateapipb.Worker, error) @@ -199,13 +202,23 @@ type Interface interface { // Lists workers. ListWorkers(ctx context.Context, opts ListOptions) (ListResponse[*ateapipb.Worker], error) - // Updates worker state with optimistic concurrency check, keyed by - // worker.metadata.name. Returns ErrNotFound if missing, or - // ErrVersionConflict on version mismatch. - UpdateWorker(ctx context.Context, worker *ateapipb.Worker, expectedVersion int64) error + // UpdateWorker performs a transactional read-modify-write and returns the + // stored worker with advanced metadata (version, update_time). + // + // precondition guards the write against landing on unexpected state: it is + // checked against the stored worker before mutate runs. Both the uid and + // version guards are required. + // + // Returns ErrPreconditionRequired if the precondition omits either guard, + // ErrNotFound if missing, ErrUIDConflict or ErrVersionConflict if the + // precondition no longer holds, ErrVersionConflict if the retry budget is + // exhausted, or the mutate's error verbatim otherwise. + UpdateWorker(ctx context.Context, name string, precondition Precondition, mutate func(toUpdate *ateapipb.Worker) error) (*ateapipb.Worker, error) - // Removes a worker by name. Idempotent: does nothing if worker is not found. - DeleteWorker(ctx context.Context, name string) error + // Removes a worker by name and returns the deleted resource. Returns + // ErrNotFound if missing, or ErrUIDConflict/ErrVersionConflict if pre does + // not describe the worker the caller observed. + DeleteWorker(ctx context.Context, name string, pre DeletePreconditions) (*ateapipb.Worker, error) // WatchWorkers returns an active subscription to track worker state changes. // The watch's Events channel is closed when the caller calls Close, the @@ -233,6 +246,74 @@ type Precondition struct { Version int64 } +// DeletePreconditions pins the object incarnation a delete may act on. Unlike +// Precondition, whose guards an update requires, each guard here is +// independently waivable: the zero value pins nothing, which is what an +// unguarded delete wants. +type DeletePreconditions struct { + // UID accepts only the object carrying it; empty accepts whichever object + // holds the name at delete time. + UID string + // Version accepts only that revision; zero accepts whatever revision the + // store is at. + Version int64 +} + +// Check reports whether md still describes the object the caller observed. +// A waived guard is not checked. The uid is reported first: a new incarnation +// makes the version meaningless. +// +// Returns ErrUIDConflict or ErrVersionConflict, which the delete surfaces +// verbatim. +func (p DeletePreconditions) Check(md *ateapipb.ResourceMetadata) error { + if p.UID != "" && p.UID != md.GetUid() { + return ErrUIDConflict + } + if p.Version != 0 && p.Version != md.GetVersion() { + return ErrVersionConflict + } + return nil +} + +// CheckWorkerMutation reports whether an UpdateWorker mutation left the +// worker's immutable identity fields alone. A backend calls it between running +// the mutation and writing the result. It lives here, above any one backend, +// so the rule is stated once and a second backend inherits it rather than +// restating it. +// +// metadata is not checked: a backend re-stamps it from the object it read, so +// whatever the mutation made of it is discarded either way. +// +// capacity is checked along with the rest because UpdateWorker replaces the +// worker rather than patching it: a request that omits capacity is asking to +// clear it, and silently losing a worker's compute capacity is worse than +// rejecting the write. A future pod resize has to relax this rule first. +// +// A rejection wraps ErrImmutableField, so a backend can return it as-is and +// callers still get the sentinel they map to INVALID_ARGUMENT. +func CheckWorkerMutation(stored, mutated *ateapipb.Worker) error { + for _, f := range []struct { + name string + stored string + mutated string + }{ + {"worker_namespace", stored.GetWorkerNamespace(), mutated.GetWorkerNamespace()}, + {"worker_pool", stored.GetWorkerPool(), mutated.GetWorkerPool()}, + {"worker_pod", stored.GetWorkerPod(), mutated.GetWorkerPod()}, + {"worker_pod_uid", stored.GetWorkerPodUid(), mutated.GetWorkerPodUid()}, + {"node_name", stored.GetNodeName(), mutated.GetNodeName()}, + {"ip", stored.GetIp(), mutated.GetIp()}, + } { + if f.stored != f.mutated { + return fmt.Errorf("%w: %s changed from %q to %q", ErrImmutableField, f.name, f.stored, f.mutated) + } + } + if !proto.Equal(stored.GetCapacity(), mutated.GetCapacity()) { + return fmt.Errorf("%w: capacity changed from %v to %v", ErrImmutableField, stored.GetCapacity(), mutated.GetCapacity()) + } + return nil +} + // hasResourceMetadata is an object the store addresses by atespace and name, // and whose identity a caller can guard with a Precondition. type hasResourceMetadata interface { diff --git a/cmd/ateapi/internal/store/storecontract/contract.go b/cmd/ateapi/internal/store/storecontract/contract.go index d7ceea9380..056dd05c0d 100644 --- a/cmd/ateapi/internal/store/storecontract/contract.go +++ b/cmd/ateapi/internal/store/storecontract/contract.go @@ -20,6 +20,8 @@ import ( "context" "errors" "fmt" + "strings" + "sync" "testing" "time" @@ -88,6 +90,7 @@ func newTestWorker(name, pod string) *ateapipb.Worker { WorkerPool: "pool-1", WorkerPod: pod, WorkerPodUid: testWorkerPodUID, + Capacity: &ateapipb.WorkerCapacity{CpuMilli: 2000, MemoryBytes: 4 << 30}, Status: &ateapipb.WorkerStatus{}, } } @@ -1088,7 +1091,8 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa defer watch.Close() worker := newTestWorker(testWorkerName, "pod-1") - if err := s.CreateWorker(ctx, worker); err != nil { + created, err := s.CreateWorker(ctx, worker) + if err != nil { t.Fatalf("CreateWorker failed: %v", err) } @@ -1096,6 +1100,9 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa if err != nil { t.Fatalf("GetWorker failed: %v", err) } + if diff := cmp.Diff(got, created, protocmp.Transform()); diff != "" { + t.Errorf("CreateWorker returned a different worker than it stored (-stored +returned):\n%s", diff) + } if got.GetMetadata().GetUid() == "" { t.Errorf("CreateWorker stored an empty uid; want server-assigned uid") } @@ -1129,10 +1136,10 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa ctx := context.Background() worker := newTestWorker(testWorkerName, "pod-1") - if err := s.CreateWorker(ctx, worker); err != nil { + if _, err := s.CreateWorker(ctx, worker); err != nil { t.Fatalf("CreateWorker failed: %v", err) } - if err := s.CreateWorker(ctx, worker); !errors.Is(err, store.ErrAlreadyExists) { + if _, err := s.CreateWorker(ctx, worker); !errors.Is(err, store.ErrAlreadyExists) { t.Errorf("expected ErrAlreadyExists, got %v", err) } }) @@ -1142,7 +1149,8 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa ctx := context.Background() worker := newTestWorker(testWorkerName, "pod-1") - if err := s.CreateWorker(ctx, worker); err != nil { + created, err := s.CreateWorker(ctx, worker) + if err != nil { t.Fatalf("CreateWorker failed: %v", err) } @@ -1153,11 +1161,15 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa } defer watch.Close() - worker.Status.Assignment = &ateapipb.ActorAssignment{ + assignment := &ateapipb.ActorAssignment{ ActorTemplate: &ateapipb.KubeNamespacedObjectRef{Namespace: "default", Name: "test-template"}, Actor: &ateapipb.ObjectRef{Name: "session-1"}, } - if err := s.UpdateWorker(ctx, worker, 1); err != nil { + updated, err := s.UpdateWorker(ctx, testWorkerName, store.PreconditionFrom(created), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.Assignment = assignment + return nil + }) + if err != nil { t.Fatalf("UpdateWorker failed: %v", err) } @@ -1165,11 +1177,15 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa if err != nil { t.Fatalf("GetWorker failed: %v", err) } + if diff := cmp.Diff(got, updated, protocmp.Transform()); diff != "" { + t.Errorf("UpdateWorker returned a different worker than it stored (-stored +returned):\n%s", diff) + } if got.GetMetadata().GetVersion() != 2 { t.Errorf("expected version 2, got %d", got.GetMetadata().GetVersion()) } want := proto.Clone(worker).(*ateapipb.Worker) + want.Status.Assignment = assignment want.Metadata.Version = 2 if diff := cmp.Diff(want, got, protocmp.Transform(), ignoreUID, ignoreTimestamps); diff != "" { t.Errorf("UpdateWorker yielded unexpected state in DB (-want +got):\n%s", diff) @@ -1184,40 +1200,279 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa } }) - t.Run("UpdateWorker_Conflict", func(t *testing.T) { + t.Run("UpdateWorker_NotFound", func(t *testing.T) { s := setup(t) ctx := context.Background() - if err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod-1")); err != nil { + // A well-formed precondition, so it is the missing worker rather than the + // guard that decides the error. + pre := store.Precondition{UID: otherTestWorkerName, Version: 1} + _, err := s.UpdateWorker(ctx, testWorkerName, pre, func(*ateapipb.Worker) error { + t.Error("mutate ran for a worker that does not exist") + return nil + }) + if !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound, got %v", err) + } + }) + + t.Run("UpdateWorker_MissingPrecondition", func(t *testing.T) { + s := setup(t) + ctx := context.Background() + + created, err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod-1")) + if err != nil { t.Fatalf("CreateWorker failed: %v", err) } - worker1, err := s.GetWorker(ctx, testWorkerName) + for _, tt := range []struct { + name string + precondition store.Precondition + }{ + {"no precondition", store.Precondition{}}, + {"guarding on only a uid", store.Precondition{UID: created.GetMetadata().GetUid()}}, + {"guarding on only a version", store.Precondition{Version: created.GetMetadata().GetVersion()}}, + } { + t.Run(tt.name, func(t *testing.T) { + _, err := s.UpdateWorker(ctx, testWorkerName, tt.precondition, func(*ateapipb.Worker) error { + t.Error("mutate ran for a blind write") + return nil + }) + if !errors.Is(err, store.ErrPreconditionRequired) { + t.Errorf("UpdateWorker error = %v, want one matching store.ErrPreconditionRequired", err) + } + }) + } + }) + + // A worker name is a pod UID, so a name is only reused when the same pod is + // re-registered. The store still hands out a fresh uid, and a guard naming + // the old one must not reach the new incarnation. + t.Run("UpdateWorker_UIDConflict", func(t *testing.T) { + s := setup(t) + ctx := context.Background() + + original, err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod-1")) if err != nil { - t.Fatalf("GetWorker failed: %v", err) + t.Fatalf("CreateWorker failed: %v", err) + } + if _, err := s.DeleteWorker(ctx, testWorkerName, store.DeletePreconditions{}); err != nil { + t.Fatalf("DeleteWorker failed: %v", err) } - worker2, err := s.GetWorker(ctx, testWorkerName) + recreated, err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod-1")) + if err != nil { + t.Fatalf("recreate CreateWorker failed: %v", err) + } + if recreated.GetMetadata().GetUid() == original.GetMetadata().GetUid() { + t.Fatalf("recreated worker reused uid %s, want a fresh one", recreated.GetMetadata().GetUid()) + } + + _, err = s.UpdateWorker(ctx, testWorkerName, store.PreconditionFrom(original), func(toUpdate *ateapipb.Worker) error { + t.Error("mutate ran past its precondition once the guarded incarnation was gone") + toUpdate.SandboxClass = "edited-anyway" + return nil + }) + if !errors.Is(err, store.ErrUIDConflict) { + t.Errorf("UpdateWorker error = %v, want one matching store.ErrUIDConflict", err) + } + + got, err := s.GetWorker(ctx, testWorkerName) if err != nil { t.Fatalf("GetWorker failed: %v", err) } + if diff := cmp.Diff(recreated, got, protocmp.Transform()); diff != "" { + t.Errorf("rejected update changed the stored worker (-recreated +got):\n%s", diff) + } + }) + + t.Run("UpdateWorker_Conflict", func(t *testing.T) { + s := setup(t) + ctx := context.Background() - worker1.Status.Assignment = &ateapipb.ActorAssignment{Actor: &ateapipb.ObjectRef{Name: "session-1"}} - if err := s.UpdateWorker(ctx, worker1, worker1.GetMetadata().GetVersion()); err != nil { + if _, err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod-1")); err != nil { + t.Fatalf("CreateWorker failed: %v", err) + } + + // Both readers observe version 1; the first update moves the worker + // past it, so the second one's precondition can no longer hold. + observed, err := s.GetWorker(ctx, testWorkerName) + if err != nil { + t.Fatalf("GetWorker failed: %v", err) + } + if _, err := s.UpdateWorker(ctx, testWorkerName, store.PreconditionFrom(observed), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.Assignment = &ateapipb.ActorAssignment{Actor: &ateapipb.ObjectRef{Name: "session-1"}} + return nil + }); err != nil { t.Fatalf("UpdateWorker failed: %v", err) } - worker2.Status.Assignment = &ateapipb.ActorAssignment{Actor: &ateapipb.ObjectRef{Name: "session-2"}} - err = s.UpdateWorker(ctx, worker2, worker2.GetMetadata().GetVersion()) + _, err = s.UpdateWorker(ctx, testWorkerName, store.PreconditionFrom(observed), func(toUpdate *ateapipb.Worker) error { + toUpdate.Status.Assignment = &ateapipb.ActorAssignment{Actor: &ateapipb.ObjectRef{Name: "session-2"}} + return nil + }) if !errors.Is(err, store.ErrVersionConflict) { t.Errorf("expected ErrVersionConflict, got %v", err) } }) + // A mutation that reports an error leaves the worker exactly as it was, at + // the version it was already at. Callers depend on this to report "already + // in the desired state" without a write: DrainWorker on a worker already + // DRAINING, and the in-process release of a worker that is not assigned. + t.Run("UpdateWorker_MutateError", func(t *testing.T) { + s := setup(t) + ctx := context.Background() + + created, err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod-1")) + if err != nil { + t.Fatalf("CreateWorker failed: %v", err) + } + + sentinel := errors.New("nothing to do") + _, err = s.UpdateWorker(ctx, testWorkerName, store.PreconditionFrom(created), func(toUpdate *ateapipb.Worker) error { + toUpdate.SandboxClass = "edited-anyway" + return sentinel + }) + if !errors.Is(err, sentinel) { + t.Errorf("expected the mutate's error verbatim, got %v", err) + } + + got, err := s.GetWorker(ctx, testWorkerName) + if err != nil { + t.Fatalf("GetWorker failed: %v", err) + } + if got.GetSandboxClass() != "" { + t.Errorf("aborted mutation was written: sandbox_class is %q", got.GetSandboxClass()) + } + if got.GetMetadata().GetVersion() != 1 { + t.Errorf("aborted mutation bumped the version to %d, want 1", got.GetMetadata().GetVersion()) + } + }) + + // Every backend must reject a mutation that touches an immutable field, and + // must name the field it rejected on. This is the case that holds any new + // backend to that. + t.Run("UpdateWorker_ImmutableFields", func(t *testing.T) { + s := setup(t) + ctx := context.Background() + + // Every case below is rejected, so nothing writes and this stays the + // current incarnation for all of them. + created, err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod-1")) + if err != nil { + t.Fatalf("CreateWorker failed: %v", err) + } + + for _, tc := range []struct { + name string + field string + mutate func(*ateapipb.Worker) + }{ + {"worker_namespace", "worker_namespace", func(w *ateapipb.Worker) { w.WorkerNamespace = "other-ns" }}, + {"worker_pool", "worker_pool", func(w *ateapipb.Worker) { w.WorkerPool = "other-pool" }}, + {"worker_pod", "worker_pod", func(w *ateapipb.Worker) { w.WorkerPod = "other-pod" }}, + {"worker_pod_uid", "worker_pod_uid", func(w *ateapipb.Worker) { w.WorkerPodUid = otherTestWorkerName }}, + {"node_name", "node_name", func(w *ateapipb.Worker) { w.NodeName = "other-node" }}, + {"ip", "ip", func(w *ateapipb.Worker) { w.Ip = "10.0.0.9" }}, + {"capacity_changed", "capacity", func(w *ateapipb.Worker) { w.Capacity.CpuMilli = 4000 }}, + // An update replaces the worker, so a caller that leaves capacity + // out is asking to clear it. That is a change like any other. + {"capacity_cleared", "capacity", func(w *ateapipb.Worker) { w.Capacity = nil }}, + } { + t.Run(tc.name, func(t *testing.T) { + _, err := s.UpdateWorker(ctx, testWorkerName, store.PreconditionFrom(created), func(toUpdate *ateapipb.Worker) error { + tc.mutate(toUpdate) + return nil + }) + if !errors.Is(err, store.ErrImmutableField) { + t.Fatalf("changing %s returned %v, want ErrImmutableField", tc.field, err) + } + if !strings.Contains(err.Error(), tc.field) { + t.Errorf("error %v does not name the offending field %s", err, tc.field) + } + got, err := s.GetWorker(ctx, testWorkerName) + if err != nil { + t.Fatalf("GetWorker failed: %v", err) + } + if got.GetMetadata().GetVersion() != 1 { + t.Errorf("rejected mutation bumped the version to %d, want 1", got.GetMetadata().GetVersion()) + } + }) + } + }) + + // Claimants that all observed the same free worker must not all win. Two + // things keep that true and this exercises both: the precondition rejects + // every claimant whose read the winner has since invalidated, and the + // occupancy test inside mutate runs against the state the write lands on + // rather than the state the claimant read. + t.Run("UpdateWorker_ConcurrentAssign", func(t *testing.T) { + s := setup(t) + ctx := context.Background() + + created, err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod-1")) + if err != nil { + t.Fatalf("CreateWorker failed: %v", err) + } + + const claimants = 8 + errTaken := errors.New("already assigned") + var wg sync.WaitGroup + won := make([]bool, claimants) + for i := range claimants { + wg.Add(1) + go func() { + defer wg.Done() + _, err := s.UpdateWorker(ctx, testWorkerName, store.PreconditionFrom(created), func(toUpdate *ateapipb.Worker) error { + if toUpdate.GetStatus().GetAssignment() != nil { + return errTaken + } + toUpdate.Status.Assignment = &ateapipb.ActorAssignment{ + Actor: &ateapipb.ObjectRef{Atespace: "team-a", Name: fmt.Sprintf("actor-%d", i)}, + ActorUid: fmt.Sprintf("uid-%d", i), + } + return nil + }) + switch { + case err == nil: + won[i] = true + case errors.Is(err, errTaken), errors.Is(err, store.ErrVersionConflict): + default: + t.Errorf("claimant %d: unexpected error %v", i, err) + } + }() + } + wg.Wait() + + winners := 0 + for _, w := range won { + if w { + winners++ + } + } + if winners != 1 { + t.Fatalf("%d of %d claimants won the assignment, want exactly 1", winners, claimants) + } + + got, err := s.GetWorker(ctx, testWorkerName) + if err != nil { + t.Fatalf("GetWorker failed: %v", err) + } + if uid := got.GetStatus().GetAssignment().GetActorUid(); !strings.HasPrefix(uid, "uid-") { + t.Errorf("stored assignment names %q, want one of the claimants", uid) + } + // One winning write on top of the create, and no partial ones. + if got.GetMetadata().GetVersion() != 2 { + t.Errorf("worker is at version %d, want 2 (create plus the single winning assign)", got.GetMetadata().GetVersion()) + } + }) + t.Run("DeleteWorker", func(t *testing.T) { s := setup(t) ctx := context.Background() - if err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod-1")); err != nil { + created, err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod-1")) + if err != nil { t.Fatalf("CreateWorker failed: %v", err) } @@ -1227,9 +1482,13 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa } defer watch.Close() - if err := s.DeleteWorker(ctx, testWorkerName); err != nil { + deleted, err := s.DeleteWorker(ctx, testWorkerName, store.DeletePreconditions{}) + if err != nil { t.Fatalf("DeleteWorker failed: %v", err) } + if diff := cmp.Diff(created, deleted, protocmp.Transform()); diff != "" { + t.Errorf("DeleteWorker returned something other than what it removed (-want +got):\n%s", diff) + } if _, err := s.GetWorker(ctx, testWorkerName); !errors.Is(err, store.ErrNotFound) { t.Errorf("expected ErrNotFound after delete, got %v", err) } @@ -1243,12 +1502,40 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa } }) - t.Run("DeleteWorker_Idempotent", func(t *testing.T) { + // Absence is reported, not swallowed. Deletes of Workers used to succeed + // silently, unlike every other Delete on the interface; callers that want + // re-drivable cleanup treat ErrNotFound as success themselves. + t.Run("DeleteWorker_NotFound", func(t *testing.T) { + s := setup(t) + ctx := context.Background() + + if _, err := s.DeleteWorker(ctx, testWorkerName, store.DeletePreconditions{}); !errors.Is(err, store.ErrNotFound) { + t.Errorf("expected ErrNotFound deleting a missing worker, got %v", err) + } + }) + + t.Run("DeleteWorker_Preconditions", func(t *testing.T) { s := setup(t) ctx := context.Background() - if err := s.DeleteWorker(ctx, testWorkerName); err != nil { - t.Errorf("DeleteWorker of a missing worker should be a no-op, got %v", err) + created, err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod-1")) + if err != nil { + t.Fatalf("CreateWorker failed: %v", err) + } + uid, version := created.GetMetadata().GetUid(), created.GetMetadata().GetVersion() + + if _, err := s.DeleteWorker(ctx, testWorkerName, store.DeletePreconditions{Version: version + 1}); !errors.Is(err, store.ErrVersionConflict) { + t.Errorf("expected ErrVersionConflict for a stale version, got %v", err) + } + if _, err := s.DeleteWorker(ctx, testWorkerName, store.DeletePreconditions{UID: otherTestWorkerName}); !errors.Is(err, store.ErrUIDConflict) { + t.Errorf("expected ErrUIDConflict for a foreign uid, got %v", err) + } + if _, err := s.GetWorker(ctx, testWorkerName); err != nil { + t.Fatalf("a rejected delete removed the worker anyway: %v", err) + } + + if _, err := s.DeleteWorker(ctx, testWorkerName, store.DeletePreconditions{UID: uid, Version: version}); err != nil { + t.Errorf("DeleteWorker with matching preconditions failed: %v", err) } }) @@ -1276,10 +1563,10 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa s := setup(t) ctx := context.Background() - if err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod1")); err != nil { + if _, err := s.CreateWorker(ctx, newTestWorker(testWorkerName, "pod1")); err != nil { t.Fatalf("failed to create worker1: %v", err) } - if err := s.CreateWorker(ctx, newTestWorker(otherTestWorkerName, "pod2")); err != nil { + if _, err := s.CreateWorker(ctx, newTestWorker(otherTestWorkerName, "pod2")); err != nil { t.Fatalf("failed to create worker2: %v", err) } @@ -1325,7 +1612,7 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa for i := 0; i < 5; i++ { worker := newTestWorker(fmt.Sprintf("bb2e6a1c-0000-4000-8000-00000000000%d", i), fmt.Sprintf("pod%d", i)) - if err := s.CreateWorker(ctx, worker); err != nil { + if _, err := s.CreateWorker(ctx, worker); err != nil { t.Fatalf("failed to create worker %d: %v", i, err) } } @@ -1678,7 +1965,7 @@ func runDebugContractTests(t *testing.T, setup func(t *testing.T) store.Interfac }); err != nil { t.Fatalf("CreateActor failed: %v", err) } - if err := s.CreateWorker(ctx, &ateapipb.Worker{WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: "pod"}); err != nil { + if _, err := s.CreateWorker(ctx, &ateapipb.Worker{WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: "pod"}); err != nil { t.Fatalf("CreateWorker failed: %v", err) } lock, err := s.AcquireLock(ctx, "lock-1") diff --git a/internal/resources/validate.go b/internal/resources/validate.go index 5895cb33ad..a7ec1d4e2f 100644 --- a/internal/resources/validate.go +++ b/internal/resources/validate.go @@ -107,13 +107,53 @@ func ValidateUpdateMetadataRef(meta *ateapipb.ResourceMetadata, fldPath *field.P return errs } +// ValidateGlobalUpdateMetadataRef is the global-scoped counterpart of +// ValidateUpdateMetadataRef, and enforces the same rules: name identifies the +// resource, uid and version guard the incarnation and revision the update was +// written against, and all three are required. It differs only in atespace, +// which must be empty because a global resource belongs to none. It does not +// check the server-managed timestamps, which clients may not set. Unlike +// ValidateObjectRef, nil metadata is an error rather than a no-op: a request +// that names no resource cannot be served. +func ValidateGlobalUpdateMetadataRef(meta *ateapipb.ResourceMetadata, fldPath *field.Path) field.ErrorList { + var errs field.ErrorList + + if val, fldPath := meta.GetAtespace(), fldPath.Child("atespace"); val != "" { + errs = append(errs, field.Invalid(fldPath, val, "must be empty for a global-scoped resource")) + } + + if val, fldPath := meta.GetName(), fldPath.Child("name"); val == "" { + errs = append(errs, field.Required(fldPath, "")) + } else { + errs = append(errs, ValidateResourceName(val, fldPath)...) + } + + if val, fldPath := meta.GetUid(), fldPath.Child("uid"); val == "" { + errs = append(errs, field.Required(fldPath, "")) + } else { + errs = append(errs, ValidateUUID(val, fldPath)...) + } + + if val, fldPath := meta.GetVersion(), fldPath.Child("version"); val == 0 { + errs = append(errs, field.Required(fldPath, "")) + } else if val < 0 { + errs = append(errs, field.Invalid(fldPath, val, "must not be negative")) + } + + return errs +} + // ValidateGlobalObjectRef checks that a reference to a global-scoped resource is // well-formed: its atespace must be empty (global resources do not belong to an // atespace) and its name must be a valid resource name. It does not check that // the referenced resource actually exists. +// +// Unlike ValidateObjectRef, and like ValidateUpdateMetadataRef, nil is an +// error rather than a no-op: every global ref in the API names the resource a +// request acts on, and a request that names nothing cannot be served. func ValidateGlobalObjectRef(ref *ateapipb.ObjectRef, fldPath *field.Path) field.ErrorList { if ref == nil { - return nil + return field.ErrorList{field.Required(fldPath, "")} } var errs field.ErrorList @@ -214,140 +254,6 @@ func ValidateSnapshotLocation(location string) error { return nil } -// ValidateWorker checks that the worker message is well-formed. -func ValidateWorker(worker *ateapipb.Worker, fldPath *field.Path) field.ErrorList { - var errs field.ErrorList - - errs = append(errs, validateWorkerMetadata(worker, fldPath.Child("metadata"))...) - - if val, fldPath := worker.WorkerNamespace, fldPath.Child("worker_namespace"); val == "" { - errs = append(errs, field.Required(fldPath, "")) - } else { - for _, msg := range content.IsDNS1123Label(val) { - errs = append(errs, field.Invalid(fldPath, val, msg)) - } - } - - if val, fldPath := worker.WorkerPool, fldPath.Child("worker_pool"); val == "" { - errs = append(errs, field.Required(fldPath, "")) - } else { - for _, msg := range content.IsDNS1123Subdomain(val) { - errs = append(errs, field.Invalid(fldPath, val, msg)) - } - } - - if val, fldPath := worker.WorkerPod, fldPath.Child("worker_pod"); val == "" { - errs = append(errs, field.Required(fldPath, "")) - } else { - for _, msg := range content.IsDNS1123Subdomain(val) { - errs = append(errs, field.Invalid(fldPath, val, msg)) - } - } - - if val := worker.GetStatus().GetAssignment(); val != nil { - errs = append(errs, ValidateAssignment(val, fldPath.Child("status", "assignment"))...) - } - - if val, fldPath := worker.Ip, fldPath.Child("ip"); val == "" { - errs = append(errs, field.Required(fldPath, "")) - } else { - errs = append(errs, ValidateIP(val, fldPath)...) - } - - if val, fldPath := worker.WorkerPodUid, fldPath.Child("worker_pod_uid"); val == "" { - errs = append(errs, field.Required(fldPath, "")) - } else { - errs = append(errs, ValidateUUID(val, fldPath)...) - } - - if val, fldPath := worker.NodeName, fldPath.Child("node_name"); val == "" { - errs = append(errs, field.Required(fldPath, "")) - } else { - for _, msg := range content.IsDNS1123Subdomain(val) { - errs = append(errs, field.Invalid(fldPath, val, msg)) - } - } - - // state is server-managed; accept any defined enum value (the unset/zero - // WORKER_STATE_UNSPECIFIED is tolerated for backward compatibility), reject - // unknowns. - if val, fldPath := worker.GetStatus().GetState(), fldPath.Child("status", "state"); ateapipb.WorkerState_name[int32(val)] == "" { - errs = append(errs, field.NotSupported(fldPath, val, []string{ - ateapipb.WorkerState_WORKER_STATE_ACTIVE.String(), - ateapipb.WorkerState_WORKER_STATE_DRAINING.String(), - })) - } - - return errs -} - -// validateWorkerMetadata checks the Worker's identity. Workers are -// global-scoped, so atespace must be empty. -func validateWorkerMetadata(worker *ateapipb.Worker, fldPath *field.Path) field.ErrorList { - var errs field.ErrorList - - meta := worker.GetMetadata() - if meta == nil { - return append(errs, field.Required(fldPath, "")) - } - - if val, fldPath := meta.GetAtespace(), fldPath.Child("atespace"); val != "" { - errs = append(errs, field.Invalid(fldPath, val, "must be empty; Workers are global-scoped")) - } - - if val, fldPath := meta.GetName(), fldPath.Child("name"); val == "" { - errs = append(errs, field.Required(fldPath, "")) - } else { - errs = append(errs, ValidateResourceName(val, fldPath)...) - } - - if val, fldPath := meta.GetUid(), fldPath.Child("uid"); val != "" { - errs = append(errs, ValidateUUID(val, fldPath)...) - } - - if val, fldPath := meta.GetVersion(), fldPath.Child("version"); val < 0 { - errs = append(errs, field.Invalid(fldPath, val, "must not be negative")) - } - - return errs -} - -func ValidateAssignment(assignment *ateapipb.ActorAssignment, fldPath *field.Path) field.ErrorList { - var errs field.ErrorList - - if val, fldPath := assignment.ActorTemplate, fldPath.Child("actor_template"); val == nil { - errs = append(errs, field.Required(fldPath, "")) - } else { - if val, fldPath := assignment.ActorTemplate.Namespace, fldPath.Child("namespace"); val == "" { - errs = append(errs, field.Required(fldPath, "")) - } else { - for _, msg := range content.IsDNS1123Label(val) { - errs = append(errs, field.Invalid(fldPath, val, msg)) - } - } - - if val, fldPath := assignment.ActorTemplate.Name, fldPath.Child("name"); val == "" { - errs = append(errs, field.Required(fldPath, "")) - } else { - for _, msg := range content.IsDNS1123Subdomain(val) { - errs = append(errs, field.Invalid(fldPath, val, msg)) - } - } - } - - if val, fldPath := assignment.Actor, fldPath.Child("actor"); val == nil { - errs = append(errs, field.Required(fldPath, "")) - } else { - errs = append(errs, ValidateObjectRef(val, fldPath)...) - } - - if val, fldPath := assignment.ActorUid, fldPath.Child("actor_uid"); val == "" { - errs = append(errs, field.Required(fldPath, "")) - } - - return errs -} - // ValidateIP checks that the given string is a valid IP address, is not an // IPv4-mapped IPv6 address, and is in canonical form. func ValidateIP(ip string, fldPath *field.Path) field.ErrorList { diff --git a/internal/resources/validate_test.go b/internal/resources/validate_test.go index d310c19c57..49d735ca4b 100644 --- a/internal/resources/validate_test.go +++ b/internal/resources/validate_test.go @@ -169,6 +169,82 @@ func TestValidateUpdateMetadataRef(t *testing.T) { } } +func TestValidateGlobalUpdateMetadataRef(t *testing.T) { + const uid = "8bf5b1a2-3c4d-4e5f-8a9b-0c1d2e3f4a5b" + tests := []struct { + name string + input *ateapipb.ResourceMetadata + wantError field.ErrorList + }{ + { + name: "valid", + input: &ateapipb.ResourceMetadata{Name: "id1", Uid: uid, Version: 7}, + wantError: nil, + }, + { + name: "nil metadata", + input: nil, + wantError: field.ErrorList{ + field.Required(field.NewPath("path", "name"), ""), + field.Required(field.NewPath("path", "uid"), ""), + field.Required(field.NewPath("path", "version"), ""), + }, + }, + { + name: "no preconditions", + input: &ateapipb.ResourceMetadata{Name: "id1"}, + wantError: field.ErrorList{ + field.Required(field.NewPath("path", "uid"), ""), + field.Required(field.NewPath("path", "version"), ""), + }, + }, + { + // The global counterpart of ValidateUpdateMetadataRef's "missing + // atespace": a global resource belongs to none, so naming one is the + // error. + name: "atespace set", + input: &ateapipb.ResourceMetadata{Atespace: "ns1", Name: "id1", Uid: uid, Version: 7}, + wantError: field.ErrorList{field.Invalid(field.NewPath("path", "atespace"), "ns1", "")}, + }, + { + name: "missing name", + input: &ateapipb.ResourceMetadata{Uid: uid, Version: 7}, + wantError: field.ErrorList{field.Required(field.NewPath("path", "name"), "")}, + }, + { + name: "invalid name", + input: &ateapipb.ResourceMetadata{Name: "ID1", Uid: uid, Version: 7}, + wantError: field.ErrorList{field.Invalid(field.NewPath("path", "name"), "ID1", "")}, + }, + { + name: "missing uid", + input: &ateapipb.ResourceMetadata{Name: "id1", Version: 7}, + wantError: field.ErrorList{field.Required(field.NewPath("path", "uid"), "")}, + }, + { + name: "invalid uid", + input: &ateapipb.ResourceMetadata{Name: "id1", Uid: "not-a-uuid", Version: 7}, + wantError: field.ErrorList{field.Invalid(field.NewPath("path", "uid"), "not-a-uuid", "")}, + }, + { + name: "missing version", + input: &ateapipb.ResourceMetadata{Name: "id1", Uid: uid}, + wantError: field.ErrorList{field.Required(field.NewPath("path", "version"), "")}, + }, + { + name: "negative version", + input: &ateapipb.ResourceMetadata{Name: "id1", Uid: uid, Version: -1}, + wantError: field.ErrorList{field.Invalid(field.NewPath("path", "version"), int64(-1), "")}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := ValidateGlobalUpdateMetadataRef(tt.input, field.NewPath("path")) + field.ErrorMatcher{}.ByType().ByField().ByValue().Test(t, tt.wantError, got) + }) + } +} + func TestValidateGlobalObjectRef(t *testing.T) { tests := []struct { name string @@ -178,6 +254,12 @@ func TestValidateGlobalObjectRef(t *testing.T) { "valid global ref", &ateapipb.ObjectRef{Name: "team-a"}, "", + }, { + // Unlike ValidateObjectRef, a nil global ref is an error: it names the + // resource the request acts on. + "missing ref", + nil, + "path: Required value", }, { "atespace must be empty", &ateapipb.ObjectRef{Atespace: "ns1", Name: "team-a"}, @@ -317,449 +399,6 @@ func TestValidateSnapshotLocation(t *testing.T) { } } -// testWorkerPodUID is the pod UID the worker cases share. -const testWorkerPodUID = "123e4567-e89b-12d3-a456-426614174000" - -// testWorkerName is the resource name the worker cases share. Worker names are -// server-assigned and opaque; nothing here may assume a relationship to -// testWorkerPodUID. -const testWorkerName = "7c1b8d4a-92e6-4f30-a5c7-2b8f61d0e934" - -func TestValidateWorker(t *testing.T) { - tests := []struct { - name string - worker *ateapipb.Worker - wantMsg string // empty means valid - }{{ - name: "valid unassigned worker", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - }, - wantMsg: "", - }, { - name: "valid assigned worker", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - Assignment: &ateapipb.ActorAssignment{ - ActorTemplate: &ateapipb.KubeNamespacedObjectRef{ - Namespace: "actor-ns", - Name: "actor-template", - }, - Actor: &ateapipb.ObjectRef{ - Atespace: "actor-ns", - Name: "actor", - }, - ActorUid: "actor-uid", - }, - }, - }, - wantMsg: "", - }, { - name: "partially assigned worker, missing actor_template", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - Assignment: &ateapipb.ActorAssignment{ - Actor: &ateapipb.ObjectRef{ - Atespace: "actor-ns", - Name: "actor", - }, - ActorUid: "actor-uid", - }, - }, - }, - wantMsg: "worker.status.assignment.actor_template: Required value", - }, { - name: "partially assigned worker, missing actor_template.namespace", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - Assignment: &ateapipb.ActorAssignment{ - ActorTemplate: &ateapipb.KubeNamespacedObjectRef{ - Name: "actor-template", - }, - Actor: &ateapipb.ObjectRef{ - Atespace: "actor-ns", - Name: "actor", - }, - ActorUid: "actor-uid", - }, - }, - }, - wantMsg: "worker.status.assignment.actor_template.namespace: Required value", - }, { - name: "partially assigned worker, missing actor_template.name", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - Assignment: &ateapipb.ActorAssignment{ - ActorTemplate: &ateapipb.KubeNamespacedObjectRef{ - Namespace: "actor-ns", - }, - Actor: &ateapipb.ObjectRef{ - Atespace: "actor-ns", - Name: "actor", - }, - ActorUid: "actor-uid", - }, - }, - }, - wantMsg: "worker.status.assignment.actor_template.name: Required value", - }, { - name: "partially assigned worker, missing actor", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - Assignment: &ateapipb.ActorAssignment{ - ActorTemplate: &ateapipb.KubeNamespacedObjectRef{ - Name: "actor-template", - Namespace: "actor-ns", - }, - ActorUid: "actor-uid", - }, - }, - }, - wantMsg: "worker.status.assignment.actor: Required value", - }, { - name: "partially assigned worker, missing actor.name", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - Assignment: &ateapipb.ActorAssignment{ - ActorTemplate: &ateapipb.KubeNamespacedObjectRef{ - Name: "actor-template", - Namespace: "actor-ns", - }, - Actor: &ateapipb.ObjectRef{ - Atespace: "actor-ns", - }, - ActorUid: "actor-uid", - }, - }, - }, - wantMsg: "worker.status.assignment.actor.name: Required value", - }, { - name: "partially assigned worker, missing actor.atespace", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - Assignment: &ateapipb.ActorAssignment{ - ActorTemplate: &ateapipb.KubeNamespacedObjectRef{ - Name: "actor-template", - Namespace: "actor-ns", - }, - Actor: &ateapipb.ObjectRef{ - Name: "actor", - }, - ActorUid: "actor-uid", - }, - }, - }, - wantMsg: "worker.status.assignment.actor.atespace: Required value", - }, { - name: "partially assigned worker, missing actor_uid", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - Assignment: &ateapipb.ActorAssignment{ - ActorTemplate: &ateapipb.KubeNamespacedObjectRef{ - Name: "actor-template", - Namespace: "actor-ns", - }, - Actor: &ateapipb.ObjectRef{ - Atespace: "actor-ns", - Name: "actor", - }, - }, - }, - }, - wantMsg: "worker.status.assignment.actor_uid: Required value", - }, { - name: "missing worker_namespace", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1", - }, - wantMsg: "worker_namespace: Required value", - }, { - name: "invalid worker_namespace", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "NS-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1", - }, - wantMsg: "worker_namespace: Invalid value", - }, { - name: "missing ip", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1", - }, - wantMsg: "ip: Required value", - }, { - name: "invalid ip", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "not-an-ip", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1", - }, - wantMsg: "ip: Invalid value", - }, { - name: "missing worker_pod_uid", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - NodeName: "node-1", - }, - wantMsg: "worker_pod_uid: Required value", - }, { - name: "invalid worker_pod_uid", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: "INVALID-UUID", - NodeName: "node-1", - }, - wantMsg: "worker_pod_uid: Invalid value", - }, { - name: "missing node_name", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - }, - wantMsg: "node_name: Required value", - }, { - name: "invalid node_name", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "NODE_NAME", - }, - wantMsg: "node_name: Invalid value", - }, { - name: "valid active worker state", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - State: ateapipb.WorkerState_WORKER_STATE_ACTIVE, - }, - }, - wantMsg: "", - }, { - name: "valid draining worker state", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - State: ateapipb.WorkerState_WORKER_STATE_DRAINING, - }, - }, - wantMsg: "", - }, { - // The zero value (STATE_UNSPECIFIED) is tolerated for backward - // compatibility with worker records written before the state field existed. - name: "unset worker state is tolerated", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - State: ateapipb.WorkerState_WORKER_STATE_UNSPECIFIED, - }, - }, - wantMsg: "", - }, { - name: "invalid worker state", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - Status: &ateapipb.WorkerStatus{ - State: ateapipb.WorkerState(99), - }, - }, - wantMsg: "state: Unsupported value", - }, { - name: "missing metadata", - worker: &ateapipb.Worker{ - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - }, - wantMsg: "worker.metadata: Required value", - }, { - name: "missing metadata.name", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - }, - wantMsg: "worker.metadata.name: Required value", - }, { - name: "invalid metadata.name", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: "Not A Name"}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - }, - wantMsg: "worker.metadata.name: Invalid value", - }, { - name: "metadata.atespace set on a global-scoped Worker", - worker: &ateapipb.Worker{ - Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName, Atespace: "team-a"}, - WorkerNamespace: "ns-1", - WorkerPool: "pool-1", - WorkerPod: "pod-1", - Ip: "10.0.0.1", - WorkerPodUid: testWorkerPodUID, - NodeName: "node-1.example.com", - }, - wantMsg: "worker.metadata.atespace: Invalid value", - }} - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - errs := ValidateWorker(tt.worker, field.NewPath("worker")) - if tt.wantMsg == "" { - if len(errs) > 0 { - t.Fatalf("expected 0 errors, got %v", errs) - } - } else { - if len(errs) == 0 { - t.Fatalf("expected error matching %q, got 0", tt.wantMsg) - } - // Any error may match: a case can trip more than one rule, so - // the wanted error is not always the first one reported. - matched := false - for _, err := range errs { - ok, matchErr := regexp.MatchString(tt.wantMsg, err.Error()) - if matchErr != nil { - t.Fatalf("failed to compile regex %q: %v", tt.wantMsg, matchErr) - } - matched = matched || ok - } - if !matched { - t.Errorf("expected message matching %q, got %v", tt.wantMsg, errs) - } - } - }) - } -} - func TestValidateIP(t *testing.T) { tests := []struct { name string @@ -801,7 +440,7 @@ func TestValidateUUID(t *testing.T) { uuid string wantMsg string // empty means valid }{ - {"valid", testWorkerPodUID, ""}, + {"valid", "123e4567-e89b-12d3-a456-426614174000", ""}, {"too short", "123e4567", "must be a lowercase UUID"}, {"too long", "123e4567-e89b-12d3-a456-4266141740001", "must be a lowercase UUID"}, {"missing dashes", "123e4567e89b12d3a456426614174000", "must be a lowercase UUID"}, diff --git a/pkg/proto/ateapipb/ateapi.pb.go b/pkg/proto/ateapipb/ateapi.pb.go index 6d9a098ed5..08dbd35074 100644 --- a/pkg/proto/ateapipb/ateapi.pb.go +++ b/pkg/proto/ateapipb/ateapi.pb.go @@ -4217,9 +4217,16 @@ func (x *CreateWorkerRequest) GetWorker() *Worker { type UpdateWorkerRequest struct { state protoimpl.MessageState `protogen:"open.v1"` // The Worker to update. - // worker.metadata.name identifies which resource to update. - // worker.metadata.version and worker.metadata.uid are optional preconditions - // and zero values skip the check. + // worker.metadata.name identifies which resource to update. atespace is + // always empty; Workers are global-scoped. + // worker.metadata.version and worker.metadata.uid are required preconditions. + // + // sandbox_class and labels are the only fields an update may change. Every + // other field is replaced with what the request carries, and a field left + // unset is cleared — so read the Worker, change what you mean to change, and + // send the whole thing back. A request that alters an immutable field, by + // changing it or by omitting it, returns INVALID_ARGUMENT naming the field. + // status is output-only and whatever it carries is ignored. Worker *Worker `protobuf:"bytes,1,opt,name=worker,proto3" json:"worker,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache @@ -4488,22 +4495,32 @@ func (x *ListActorsResponse) GetNextPageToken() string { // Global-scoped: metadata.atespace is always empty. metadata.name is assigned // by the control plane and is opaque to clients — never parse it or derive it // from anything else; read pod identity from the named fields below. +// +// sandbox_class and labels are the only mutable fields; every other field is +// either immutable after creation or output-only. UpdateWorker replaces the +// whole resource, so an immutable field that a request changes — including by +// omitting it, which would clear it — is rejected with INVALID_ARGUMENT. type Worker struct { state protoimpl.MessageState `protogen:"open.v1"` - // Common resource metadata: name, uid, version, timestamps. + // Output-only: name, uid, version and timestamps are all server-assigned. + // uid and version are echoed back on UpdateWorker as preconditions. Metadata *ResourceMetadata `protobuf:"bytes,1,opt,name=metadata,proto3" json:"metadata,omitempty"` // Kubernetes coordinates. Immutable, set at creation. - WorkerNamespace string `protobuf:"bytes,2,opt,name=worker_namespace,json=workerNamespace,proto3" json:"worker_namespace,omitempty"` - WorkerPool string `protobuf:"bytes,3,opt,name=worker_pool,json=workerPool,proto3" json:"worker_pool,omitempty"` - WorkerPod string `protobuf:"bytes,4,opt,name=worker_pod,json=workerPod,proto3" json:"worker_pod,omitempty"` - WorkerPodUid string `protobuf:"bytes,5,opt,name=worker_pod_uid,json=workerPodUid,proto3" json:"worker_pod_uid,omitempty"` - NodeName string `protobuf:"bytes,6,opt,name=node_name,json=nodeName,proto3" json:"node_name,omitempty"` - Ip string `protobuf:"bytes,7,opt,name=ip,proto3" json:"ip,omitempty"` - SandboxClass string `protobuf:"bytes,8,opt,name=sandbox_class,json=sandboxClass,proto3" json:"sandbox_class,omitempty"` - Labels map[string]string `protobuf:"bytes,9,rep,name=labels,proto3" json:"labels,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` - // The compute capacity this worker can give an actor sandbox. + WorkerNamespace string `protobuf:"bytes,2,opt,name=worker_namespace,json=workerNamespace,proto3" json:"worker_namespace,omitempty"` + WorkerPool string `protobuf:"bytes,3,opt,name=worker_pool,json=workerPool,proto3" json:"worker_pool,omitempty"` + WorkerPod string `protobuf:"bytes,4,opt,name=worker_pod,json=workerPod,proto3" json:"worker_pod,omitempty"` + WorkerPodUid string `protobuf:"bytes,5,opt,name=worker_pod_uid,json=workerPodUid,proto3" json:"worker_pod_uid,omitempty"` + NodeName string `protobuf:"bytes,6,opt,name=node_name,json=nodeName,proto3" json:"node_name,omitempty"` + Ip string `protobuf:"bytes,7,opt,name=ip,proto3" json:"ip,omitempty"` + // Mutable. + SandboxClass string `protobuf:"bytes,8,opt,name=sandbox_class,json=sandboxClass,proto3" json:"sandbox_class,omitempty"` + Labels map[string]string `protobuf:"bytes,9,rep,name=labels,proto3" json:"labels,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + // The compute capacity this worker can give an actor sandbox. Immutable, set + // at creation: a worker pod's limits are fixed for its lifetime. Capacity *WorkerCapacity `protobuf:"bytes,10,opt,name=capacity,proto3" json:"capacity,omitempty"` - // Server-managed state. Absent from Create/Update request payloads. + // Output-only server-managed state. Absent from Create/Update request + // payloads; whatever a request carries here is ignored. DrainWorker is the + // only way a client moves state, and assignment is the scheduler's. Status *WorkerStatus `protobuf:"bytes,11,opt,name=status,proto3" json:"status,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache diff --git a/pkg/proto/ateapipb/ateapi.proto b/pkg/proto/ateapipb/ateapi.proto index b2c93bb8ec..e4437c520d 100644 --- a/pkg/proto/ateapipb/ateapi.proto +++ b/pkg/proto/ateapipb/ateapi.proto @@ -82,6 +82,7 @@ service Control { // Mark a Worker as terminating so the scheduler stops routing new Actors to // it. Idempotent; one-way. Deliberately leaves any bound Actor alone. + // Returns ABORTED if another write lands on the Worker first; retry. rpc DrainWorker(DrainWorkerRequest) returns (Worker) {} // List Actors. @@ -789,10 +790,9 @@ message ListWorkersResponse { string next_page_token = 2; } -// TODO: every method below returns UNIMPLEMENTED. Workers are still created, -// updated, and deleted by writing directly to the store from the -// WorkerPoolSyncer and the actor workflows; migrating those callers onto these -// RPCs lands in a follow-up change. +// TODO: Workers are still created, updated, and deleted by writing directly to +// the store from the WorkerPoolSyncer and the actor workflows; migrating those +// callers onto the RPCs below lands in a follow-up change. message GetWorkerRequest { // The Worker to fetch. atespace is always empty; Workers are global-scoped. @@ -806,9 +806,16 @@ message CreateWorkerRequest { message UpdateWorkerRequest { // The Worker to update. - // worker.metadata.name identifies which resource to update. - // worker.metadata.version and worker.metadata.uid are optional preconditions - // and zero values skip the check. + // worker.metadata.name identifies which resource to update. atespace is + // always empty; Workers are global-scoped. + // worker.metadata.version and worker.metadata.uid are required preconditions. + // + // sandbox_class and labels are the only fields an update may change. Every + // other field is replaced with what the request carries, and a field left + // unset is cleared — so read the Worker, change what you mean to change, and + // send the whole thing back. A request that alters an immutable field, by + // changing it or by omitting it, returns INVALID_ARGUMENT naming the field. + // status is output-only and whatever it carries is ignored. Worker worker = 1; } @@ -854,8 +861,14 @@ message ListActorsResponse { // Global-scoped: metadata.atespace is always empty. metadata.name is assigned // by the control plane and is opaque to clients — never parse it or derive it // from anything else; read pod identity from the named fields below. +// +// sandbox_class and labels are the only mutable fields; every other field is +// either immutable after creation or output-only. UpdateWorker replaces the +// whole resource, so an immutable field that a request changes — including by +// omitting it, which would clear it — is rejected with INVALID_ARGUMENT. message Worker { - // Common resource metadata: name, uid, version, timestamps. + // Output-only: name, uid, version and timestamps are all server-assigned. + // uid and version are echoed back on UpdateWorker as preconditions. ResourceMetadata metadata = 1; // Kubernetes coordinates. Immutable, set at creation. @@ -866,13 +879,17 @@ message Worker { string node_name = 6; string ip = 7; + // Mutable. string sandbox_class = 8; map labels = 9; - // The compute capacity this worker can give an actor sandbox. + // The compute capacity this worker can give an actor sandbox. Immutable, set + // at creation: a worker pod's limits are fixed for its lifetime. WorkerCapacity capacity = 10; - // Server-managed state. Absent from Create/Update request payloads. + // Output-only server-managed state. Absent from Create/Update request + // payloads; whatever a request carries here is ignored. DrainWorker is the + // only way a client moves state, and assignment is the scheduler's. WorkerStatus status = 11; } diff --git a/pkg/proto/ateapipb/ateapi_grpc.pb.go b/pkg/proto/ateapipb/ateapi_grpc.pb.go index 2af8dc21f3..5157d1930b 100644 --- a/pkg/proto/ateapipb/ateapi_grpc.pb.go +++ b/pkg/proto/ateapipb/ateapi_grpc.pb.go @@ -111,6 +111,7 @@ type ControlClient interface { DeleteWorker(ctx context.Context, in *DeleteWorkerRequest, opts ...grpc.CallOption) (*Worker, error) // Mark a Worker as terminating so the scheduler stops routing new Actors to // it. Idempotent; one-way. Deliberately leaves any bound Actor alone. + // Returns ABORTED if another write lands on the Worker first; retry. DrainWorker(ctx context.Context, in *DrainWorkerRequest, opts ...grpc.CallOption) (*Worker, error) // List Actors. ListActors(ctx context.Context, in *ListActorsRequest, opts ...grpc.CallOption) (*ListActorsResponse, error) @@ -467,6 +468,7 @@ type ControlServer interface { DeleteWorker(context.Context, *DeleteWorkerRequest) (*Worker, error) // Mark a Worker as terminating so the scheduler stops routing new Actors to // it. Idempotent; one-way. Deliberately leaves any bound Actor alone. + // Returns ABORTED if another write lands on the Worker first; retry. DrainWorker(context.Context, *DrainWorkerRequest) (*Worker, error) // List Actors. ListActors(context.Context, *ListActorsRequest) (*ListActorsResponse, error)