diff --git a/cmd/ateapi/internal/controlapi/workflow_pause.go b/cmd/ateapi/internal/controlapi/workflow_pause.go index 0eae9fd6e..2b632a6c4 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause.go +++ b/cmd/ateapi/internal/controlapi/workflow_pause.go @@ -203,12 +203,21 @@ func (s *FinalizePausedStep) Execute(ctx context.Context, input *PauseInput, sta } // TODO(dberkov) - what if InProgressSnapshot is empty? That shouldn't be possible. if latestActor.InProgressSnapshot != "" { + // If the worker record was already gone we don't know which node + // holds the checkpoint. Record no locality rather than an empty + // node name: findFreeWorker treats any non-empty restriction list + // as a hard filter, and no worker has an empty node name, so [""] + // would make the actor permanently unschedulable. + var nodesWithSnapshot []string + if nodeName != "" { + nodesWithSnapshot = []string{nodeName} + } latestActor.LatestSnapshotInfo = &ateapipb.SnapshotInfo{ Type: ateapipb.SnapshotType_SNAPSHOT_TYPE_LOCAL, Data: &ateapipb.SnapshotInfo_Local{ Local: &ateapipb.LocalSnapshotInfo{ SnapshotPrefix: latestActor.InProgressSnapshot, - NodeVmsWithLocalSnapshots: []string{nodeName}, + NodeVmsWithLocalSnapshots: nodesWithSnapshot, }, }, } diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index 2d42b710f..105e53e4e 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume.go @@ -170,13 +170,13 @@ func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, stat // If not, find a free one using randomized shuffling if assignedWorker == nil { - pickedWorker := s.findFreeWorker(workers, eligible, state.Actor.GetLatestSnapshotInfo().GetLocal().GetNodeVmsWithLocalSnapshots()) - if pickedWorker == nil { - return status.Errorf(codes.FailedPrecondition, "no free workers available") + scan := s.findFreeWorker(workers, eligible, state.Actor.GetLatestSnapshotInfo().GetLocal().GetNodeVmsWithLocalSnapshots()) + if scan.picked == nil { + return status.Errorf(codes.FailedPrecondition, "%s", scan.noWorkerMessage()) } - assignedWorker = pickedWorker - slog.InfoContext(ctx, "Picked worker", slog.Any("worker", pickedWorker.String())) + assignedWorker = scan.picked + slog.InfoContext(ctx, "Picked worker", slog.Any("worker", assignedWorker.String())) } // Workers() returns pointers directly from the cache so we need to clone before @@ -219,16 +219,58 @@ func (s *AssignWorkerStep) RetryBackoff() *wait.Backoff { } } -func (s *AssignWorkerStep) findFreeWorker(workers []*ateapipb.Worker, eligible map[types.NamespacedName]struct{}, nodesRestrictions []string) *ateapipb.Worker { +// workerScan is the outcome of a findFreeWorker pass. When picked is nil, +// restrictedNodes and sawWorkerOnRestrictedNode (gathered during the same +// pass) let the error distinguish a snapshot-locality conflict from plain +// capacity exhaustion. +type workerScan struct { + picked *ateapipb.Worker + // restrictedNodes is the snapshot-locality restriction in effect, + // with empty entries dropped. + restrictedNodes []string + // sawWorkerOnRestrictedNode reports whether any eligible worker exists on + // a restricted node; when picked is nil, all such workers are busy. + sawWorkerOnRestrictedNode bool +} + +// noWorkerMessage explains a failed scan. A bare "no free workers available" +// points operators at cluster capacity, which is the wrong place to look when +// the real conflict is snapshot locality — or worse, a removed node that held +// the actor's only checkpoint. +func (sc workerScan) noWorkerMessage() string { + if len(sc.restrictedNodes) == 0 { + return "no free workers available" + } + if !sc.sawWorkerOnRestrictedNode { + return fmt.Sprintf( + "actor's local snapshot is on node(s) %v but no eligible workers exist on those nodes (the node(s) may have been removed)", + sc.restrictedNodes) + } + return fmt.Sprintf( + "actor's local snapshot is on node(s) %v but all eligible workers on those nodes are busy", + sc.restrictedNodes) +} + +func (s *AssignWorkerStep) findFreeWorker(workers []*ateapipb.Worker, eligible map[types.NamespacedName]struct{}, nodesRestrictions []string) workerScan { + // Drop empty node names from the restriction list. Actor records written + // before FinalizePausedStep guarded against an unknown node contain [""], + // which no worker's NodeName can ever match; treating that as "no + // restriction" keeps those actors schedulable. + scan := workerScan{ + restrictedNodes: slices.DeleteFunc(slices.Clone(nodesRestrictions), func(n string) bool { return n == "" }), + } + var freeWorkers []*ateapipb.Worker for _, worker := range workers { - if worker.Assignment != nil { + if _, ok := eligible[types.NamespacedName{Namespace: worker.GetWorkerNamespace(), Name: worker.GetWorkerPool()}]; !ok { continue } - if _, ok := eligible[types.NamespacedName{Namespace: worker.GetWorkerNamespace(), Name: worker.GetWorkerPool()}]; !ok { + if slices.Contains(scan.restrictedNodes, worker.GetNodeName()) { + scan.sawWorkerOnRestrictedNode = true + } else if len(scan.restrictedNodes) > 0 { continue } - if len(nodesRestrictions) == 0 || slices.Contains(nodesRestrictions, worker.GetNodeName()) { + if worker.Assignment == nil { freeWorkers = append(freeWorkers, worker) } } @@ -237,9 +279,9 @@ func (s *AssignWorkerStep) findFreeWorker(workers []*ateapipb.Worker, eligible m rand.Shuffle(len(freeWorkers), func(i, j int) { freeWorkers[i], freeWorkers[j] = freeWorkers[j], freeWorkers[i] }) - return freeWorkers[0] + scan.picked = freeWorkers[0] } - return nil + return scan } type CallAteletRestoreStep struct { diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index 733c99674..5981654ea 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume_test.go @@ -185,3 +185,150 @@ func TestEligibleWorkerPools(t *testing.T) { }) } } + +func freeWorker(pool, pod, node string) *ateapipb.Worker { + return &ateapipb.Worker{ + WorkerNamespace: "ns", + WorkerPool: pool, + WorkerPod: pod, + NodeName: node, + } +} + +func TestFindFreeWorker_NodeRestrictions(t *testing.T) { + eligible := map[types.NamespacedName]struct{}{ + {Namespace: "ns", Name: "pool1"}: {}, + } + workers := []*ateapipb.Worker{ + freeWorker("pool1", "w1", "node1"), + freeWorker("pool1", "w2", "node2"), + } + + tests := []struct { + name string + nodesRestrictions []string + wantNodes []string // acceptable NodeNames for the picked worker + wantNil bool + }{ + { + name: "no restrictions picks any worker", + nodesRestrictions: nil, + wantNodes: []string{"node1", "node2"}, + }, + { + name: "restriction filters to the matching node", + nodesRestrictions: []string{"node2"}, + wantNodes: []string{"node2"}, + }, + { + name: "restriction on an absent node matches nothing", + nodesRestrictions: []string{"node3"}, + wantNil: true, + }, + { + // Actor records written before FinalizePausedStep guarded against + // an unknown node contain [""]; it must not exclude every worker. + name: "empty node name is ignored, not treated as a restriction", + nodesRestrictions: []string{""}, + wantNodes: []string{"node1", "node2"}, + }, + { + name: "empty node name alongside a real one keeps the real restriction", + nodesRestrictions: []string{"", "node1"}, + wantNodes: []string{"node1"}, + }, + } + + s := &AssignWorkerStep{} + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := s.findFreeWorker(workers, eligible, tt.nodesRestrictions).picked + if tt.wantNil { + if got != nil { + t.Fatalf("expected no worker, got %v", got) + } + return + } + if got == nil { + t.Fatalf("expected a worker, got nil") + } + found := false + for _, n := range tt.wantNodes { + if got.GetNodeName() == n { + found = true + } + } + if !found { + t.Errorf("picked worker on node %q, want one of %v", got.GetNodeName(), tt.wantNodes) + } + }) + } +} + +func busyWorker(pool, pod, node string) *ateapipb.Worker { + w := freeWorker(pool, pod, node) + w.Assignment = &ateapipb.Assignment{ + Actor: &ateapipb.ActorRef{Atespace: "ate", Name: "other"}, + } + return w +} + +func TestNoWorkerMessage(t *testing.T) { + eligible := map[types.NamespacedName]struct{}{ + {Namespace: "ns", Name: "pool1"}: {}, + } + + tests := []struct { + name string + workers []*ateapipb.Worker + nodesRestrictions []string + wantMessage string + }{ + { + name: "no restriction keeps the generic message", + workers: []*ateapipb.Worker{busyWorker("pool1", "w1", "node1")}, + nodesRestrictions: nil, + wantMessage: "no free workers available", + }, + { + name: "snapshot node busy", + workers: []*ateapipb.Worker{ + busyWorker("pool1", "w1", "node1"), + freeWorker("pool1", "w2", "node2"), + }, + nodesRestrictions: []string{"node1"}, + wantMessage: "actor's local snapshot is on node(s) [node1] but all eligible workers on those nodes are busy", + }, + { + name: "snapshot node gone", + workers: []*ateapipb.Worker{ + freeWorker("pool1", "w2", "node2"), + }, + nodesRestrictions: []string{"node1"}, + wantMessage: "actor's local snapshot is on node(s) [node1] but no eligible workers exist on those nodes (the node(s) may have been removed)", + }, + { + // A worker on the snapshot's node in an ineligible pool must not + // flip the message from "gone" to "busy". + name: "ineligible pools don't count as workers on the node", + workers: []*ateapipb.Worker{ + busyWorker("pool2", "w1", "node1"), + }, + nodesRestrictions: []string{"node1"}, + wantMessage: "actor's local snapshot is on node(s) [node1] but no eligible workers exist on those nodes (the node(s) may have been removed)", + }, + } + + s := &AssignWorkerStep{} + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + scan := s.findFreeWorker(tt.workers, eligible, tt.nodesRestrictions) + if scan.picked != nil { + t.Fatalf("expected no worker to be picked, got %v", scan.picked) + } + if got := scan.noWorkerMessage(); got != tt.wantMessage { + t.Errorf("noWorkerMessage() = %q, want %q", got, tt.wantMessage) + } + }) + } +}