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..fbbbba173 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume.go @@ -220,6 +220,12 @@ func (s *AssignWorkerStep) RetryBackoff() *wait.Backoff { } func (s *AssignWorkerStep) findFreeWorker(workers []*ateapipb.Worker, eligible map[types.NamespacedName]struct{}, nodesRestrictions []string) *ateapipb.Worker { + // 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. + nodesRestrictions = slices.DeleteFunc(slices.Clone(nodesRestrictions), func(n string) bool { return n == "" }) + var freeWorkers []*ateapipb.Worker for _, worker := range workers { if worker.Assignment != nil { diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index 733c99674..834a4f625 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume_test.go @@ -185,3 +185,82 @@ 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) + 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) + } + }) + } +}