Make sandbox_class in workers immutable

cmd/ateapi/internal/controlapi/worker_test.go

@@ -289,7 +289,6 @@ func TestUpdateWorker(t *testing.T) {
 
 	got, err := svc.UpdateWorker(ctx, &ateapipb.UpdateWorkerRequest{
 		Worker: updateFrom(seeded, func(w *ateapipb.Worker) {
-			w.SandboxClass = "microvm"
 			w.Labels = map[string]string{"tier": "batch"}
 		}),
 	})
@@ -298,7 +297,6 @@ func TestUpdateWorker(t *testing.T) {
 	}
 
 	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 != "" {
@@ -321,16 +319,12 @@ func TestUpdateWorker_OmittedMutableFieldIsCleared(t *testing.T) {
 
 	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())
 	}
@@ -347,7 +341,7 @@ func TestUpdateWorker_LeavesStatusAlone(t *testing.T) {
 
 	got, err := svc.UpdateWorker(ctx, &ateapipb.UpdateWorkerRequest{
 		Worker: updateFrom(assigned, func(w *ateapipb.Worker) {
-			w.SandboxClass = "microvm"
+			w.Labels = map[string]string{"tier": "batch"}
 			// 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}
@@ -369,7 +363,7 @@ func TestUpdateWorker_Preconditions(t *testing.T) {
 	update := func(bend func(*ateapipb.ResourceMetadata)) error {
 		_, err := svc.UpdateWorker(ctx, &ateapipb.UpdateWorkerRequest{
 			Worker: updateFrom(seeded, func(w *ateapipb.Worker) {
-				w.SandboxClass = "microvm"
+				w.Labels = map[string]string{"tier": "batch"}
 				bend(w.Metadata)
 			}),
 		})
@@ -436,10 +430,12 @@ func TestUpdateWorker_Errors(t *testing.T) {
 		{"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},
+		{"sandbox_class changed", func(w *ateapipb.Worker) { w.SandboxClass = "microvm" }, 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},
+		{"sandbox_class omitted", func(w *ateapipb.Worker) { w.SandboxClass = "" }, codes.InvalidArgument},
 	}
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
@@ -455,9 +451,9 @@ func TestUpdateWorker_Errors(t *testing.T) {
 	}
 }
 
-// A draining worker can still have everything else about it updated; only its
-// status is frozen.
-func TestUpdateWorker_DrainingWorkerKeepsOtherFieldsMutable(t *testing.T) {
+// A draining worker can still have its labels updated; only its status is
+// frozen.
+func TestUpdateWorker_DrainingWorkerKeepsLabelsMutable(t *testing.T) {
 	ctx := context.Background()
 	svc, persistence := newWorkerAPIService(t)
 	seedAPIWorker(t, ctx, persistence, newAPIWorker(apiWorkerName))
@@ -467,13 +463,13 @@ func TestUpdateWorker_DrainingWorkerKeepsOtherFieldsMutable(t *testing.T) {
 	}
 
 	got, err := svc.UpdateWorker(ctx, &ateapipb.UpdateWorkerRequest{
-		Worker: updateFrom(drained, func(w *ateapipb.Worker) { w.SandboxClass = "microvm" }),
+		Worker: updateFrom(drained, func(w *ateapipb.Worker) { w.Labels = map[string]string{"tier": "batch"} }),
 	})
 	if err != nil {
 		t.Fatalf("UpdateWorker() failed: %v", err)
 	}
-	if got.GetSandboxClass() != "microvm" {
-		t.Errorf("sandbox_class = %q, want microvm", got.GetSandboxClass())
+	if got.GetLabels()["tier"] != "batch" {
+		t.Errorf("labels = %v, want tier=batch", got.GetLabels())
 	}
 	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)

cmd/ateapi/internal/store/store.go

@@ -300,6 +300,7 @@ func CheckWorkerMutation(stored, mutated *ateapipb.Worker) error {
 		{"worker_pod_uid", stored.GetWorkerPodUid(), mutated.GetWorkerPodUid()},
 		{"node_name", stored.GetNodeName(), mutated.GetNodeName()},
 		{"ip", stored.GetIp(), mutated.GetIp()},
+		{"sandbox_class", stored.GetSandboxClass(), mutated.GetSandboxClass()},
 	} {
 		if f.stored != f.mutated {
 			return fmt.Errorf("%w: %s changed from %q to %q", ErrImmutableField, f.name, f.stored, f.mutated)

cmd/ateapi/internal/store/storecontract/contract.go

@@ -1257,7 +1257,7 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa
 
 		_, 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"
+			toUpdate.Labels = map[string]string{"tier": "edited-anyway"}
 			return nil
 		})
 		if !errors.Is(err, store.ErrUIDConflict) {
@@ -1318,7 +1318,7 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa
 
 		sentinel := errors.New("nothing to do")
 		_, err = s.UpdateWorker(ctx, testWorkerName, store.PreconditionFrom(created), func(toUpdate *ateapipb.Worker) error {
-			toUpdate.SandboxClass = "edited-anyway"
+			toUpdate.Labels = map[string]string{"tier": "edited-anyway"}
 			return sentinel
 		})
 		if !errors.Is(err, sentinel) {
@@ -1329,8 +1329,8 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa
 		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 len(got.GetLabels()) != 0 {
+			t.Errorf("aborted mutation was written: labels are %v", got.GetLabels())
 		}
 		if got.GetMetadata().GetVersion() != 1 {
 			t.Errorf("aborted mutation bumped the version to %d, want 1", got.GetMetadata().GetVersion())
@@ -1362,6 +1362,11 @@ func runWorkerContractTests(t *testing.T, setup func(t *testing.T) store.Interfa
 			{"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" }},
+			// A pool's sandboxClass drives its pods' shape, so editing it
+			// replaces every pod rather than reclassifying any. Accepting the
+			// change here would advertise the old pod's shape wrongly to the
+			// scheduler for as long as it survives the rollout.
+			{"sandbox_class", "sandbox_class", func(w *ateapipb.Worker) { w.SandboxClass = "microvm" }},
 			{"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.

cmd/atecontroller/internal/workersync/fakecontrol_test.go

@@ -195,12 +195,10 @@ func (f *fakeControl) UpdateWorker(_ context.Context, in *ateapipb.UpdateWorkerR
 	updated.Metadata = proto.Clone(stored.GetMetadata()).(*ateapipb.ResourceMetadata)
 	updated.Status = proto.Clone(stored.GetStatus()).(*ateapipb.WorkerStatus)
 
-	// sandbox_class and labels are the only fields an update may change, so
-	// pinning those two to what is stored leaves any remaining difference on a
-	// field that is immutable after create — including one the request cleared
-	// by omitting it.
+	// labels is the only field an update may change, so pinning it to what is
+	// stored leaves any remaining difference on a field that is immutable after
+	// create — including one the request cleared by omitting it.
 	probe := proto.Clone(updated).(*ateapipb.Worker)
-	probe.SandboxClass = stored.GetSandboxClass()
 	probe.Labels = stored.GetLabels()
 	if !proto.Equal(probe, stored) {
 		return nil, status.Error(codes.InvalidArgument, "update changed a field that is immutable after create")

cmd/atecontroller/internal/workersync/syncer.go

@@ -279,23 +279,28 @@ func (s *WorkerPoolSyncer) createOrUpdateWorker(ctx context.Context, key workerK
 		return fmt.Errorf("getting worker: %w", err)
 	}
 
-	// UpdateWorker replaces the whole resource, so the two mutable fields are
+	// UpdateWorker replaces the whole resource, so the one mutable field is
 	// edited onto the Worker as it was read and the rest is sent back unchanged
 	// — anything else altered here, including a field cleared by omission, is
 	// rejected as INVALID_ARGUMENT. Everything else on a Worker is immutable
 	// after create, so drift there cannot be repaired by an update; it takes a
 	// new pod, which arrives under a new key.
 	var changed bool
-	if w.GetSandboxClass() != string(pool.Spec.SandboxClass) {
-		slog.InfoContext(ctx, "Syncer: updating worker (SandboxClass changed)", key.logAttrs()...)
-		w.SandboxClass = string(pool.Spec.SandboxClass)
-		changed = true
-	}
 	if !maps.Equal(w.GetLabels(), pool.GetLabels()) {
 		slog.InfoContext(ctx, "Syncer: updating worker (labels changed)", key.logAttrs()...)
 		w.Labels = pool.GetLabels()
 		changed = true
 	}
+	if w.GetSandboxClass() != string(pool.Spec.SandboxClass) {
+		// Expected mid-rollout: sandboxClass drives the worker pod's shape, so
+		// editing it on the pool replaces every pod rather than reclassifying
+		// any. This pod predates that edit and is on its way out; its successor
+		// registers under a new key with the new class. Writing the pool's value
+		// back would be rejected, and would misreport this pod's shape to the
+		// scheduler if it were not.
+		slog.DebugContext(ctx, "Syncer: registered worker sandbox class predates its pool",
+			append(key.logAttrs(), slog.String("registered", w.GetSandboxClass()), slog.String("pool", string(pool.Spec.SandboxClass)))...)
+	}
 	if w.GetIp() != pod.Status.PodIP {
 		// TODO: I don't think this is possible, but handling this case so we can
 		// log it just in case we can reproduce it. It is logged rather than

cmd/atecontroller/internal/workersync/syncer_test.go

@@ -526,12 +526,12 @@ func TestSyncer_UpdateWorker_RetryOnVersionConflict(t *testing.T) {
 
 	// Change a mutable worker field on the pool, so the next reconcile has an
 	// update to make.
-	if err := poolIndexer.Update(workerPool(ns, poolName, "microvm", map[string]string{"foo": "bar"})); err != nil {
+	if err := poolIndexer.Update(workerPool(ns, poolName, "gvisor", map[string]string{"foo": "baz"})); err != nil {
 		t.Fatalf("updating pool: %v", err)
 	}
 
 	// Land a concurrent version bump the moment the syncer calls UpdateWorker.
-	// The injected change is a drain: it is mutable, and unlike sandbox_class the
+	// The injected change is a drain: it is mutable, and unlike the labels the
 	// syncer's update path does not write it, so it survives the retry only if
 	// the retry re-read.
 	conflicted := false
@@ -556,14 +556,45 @@ func TestSyncer_UpdateWorker_RetryOnVersionConflict(t *testing.T) {
 	// The retry re-reads, so both changes end up on the record.
 	mustReconcile(t, ctx, s, key)
 	got := api.get(testPodUID)
-	if got.GetSandboxClass() != "microvm" {
-		t.Errorf("worker sandbox class = %q, want microvm", got.GetSandboxClass())
+	if got.GetLabels()["foo"] != "baz" {
+		t.Errorf("worker labels = %v, want foo=baz", got.GetLabels())
 	}
 	if got.GetStatus().GetState() != ateapipb.WorkerState_WORKER_STATE_DRAINING {
 		t.Errorf("worker state = %v, want the concurrently injected DRAINING to survive the retry", got.GetStatus().GetState())
 	}
 }
 
+// Editing a pool's sandboxClass rolls its pods rather than reclassifying them,
+// so a pod that outlives the edit keeps the class it was built with. Writing the
+// pool's new value back would be rejected as an immutable-field change, and
+// would misreport this pod's shape to the scheduler if it were not.
+func TestSyncer_SandboxClassDriftLeavesWorkerAlone(t *testing.T) {
+	ctx := context.Background()
+
+	ns, podName, poolName := "ns-syncer-class", "worker-unit-class", "pool-class"
+
+	api := newFakeControl()
+	s, pods, poolIndexer := setupReconcileTest(t, api, workerPool(ns, poolName, "gvisor", nil))
+	key := seedPod(t, pods, workerPod(ns, podName, poolName, testPodUID, "10.0.0.1"))
+
+	mustReconcile(t, ctx, s, key)
+	registered := api.get(testPodUID)
+
+	if err := poolIndexer.Update(workerPool(ns, poolName, "microvm", nil)); err != nil {
+		t.Fatalf("updating pool: %v", err)
+	}
+	mustReconcile(t, ctx, s, key)
+
+	got := api.get(testPodUID)
+	if got.GetSandboxClass() != "gvisor" {
+		t.Errorf("worker sandbox class = %q, want it left at gvisor", got.GetSandboxClass())
+	}
+	if got.GetMetadata().GetVersion() != registered.GetMetadata().GetVersion() {
+		t.Errorf("worker version = %d, want %d: the drift must not provoke a write",
+			got.GetMetadata().GetVersion(), registered.GetMetadata().GetVersion())
+	}
+}
+
 // TestSyncer_RequeueOnMissingWorkerPool verifies that a pod whose WorkerPool is
 // not yet in the lister is requeued rather than dropped, and converges once the
 // pool appears.

pkg/proto/ateapipb/ateapi.proto

@@ -1040,12 +1040,12 @@ message UpdateWorkerRequest {
   // 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.
+  // labels is the only field 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;
 }
 
@@ -1094,10 +1094,10 @@ message ListActorsResponse {
 // 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.
+// labels is the only mutable field; 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 {
   // Output-only: name, uid, version and timestamps are all server-assigned.
   // uid and version are echoed back on UpdateWorker as preconditions.
@@ -1111,9 +1111,11 @@ message Worker {
   string worker_pod_uid = 5;
   string node_name = 6;
   string ip = 7;
-
-  // Mutable.
   string sandbox_class = 8;
+
+  // The owning WorkerPool's labels, which the scheduler matches actor and
+  // template worker selectors against. Mutable: a pool's labels can be edited
+  // without disturbing its pods.
   map<string, string> labels = 9;
 
   // The compute capacity this worker can give an actor sandbox. Immutable, set