From 2da158908c848b007d4a7adf11eadc4b38bcb8e6 Mon Sep 17 00:00:00 2001 From: Yibo Zhuang Date: Tue, 18 Aug 2026 07:38:36 -0700 Subject: [PATCH] Wait for VM shutdown before freeing worker capacity (#465) --- internal/worker/worker.go | 17 +++- internal/worker/worker_stop_test.go | 135 ++++++++++++++++++++++++++++ 2 files changed, 150 insertions(+), 2 deletions(-) create mode 100644 internal/worker/worker_stop_test.go diff --git a/internal/worker/worker.go b/internal/worker/worker.go index 7cccc00..aa1cb79 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -475,12 +475,16 @@ func (worker *Worker) syncVMs(ctx context.Context, updateVM func(context.Context case ActionStop: // VM has failed on the remote side, stop it locally to prevent incorrect // worker's resources calculation in the Controller's scheduler - vm.Stop() + if err := waitForVMStop(ctx, vm); err != nil { + return fmt.Errorf("failed to stop VM: %w", err) + } case ActionFail, ActionLostTrack, ActionImpossible: // VM has failed on the local side, stop it before reporting as failed to prevent incorrect // worker's resources calculation in the Controller's scheduler if vm != nil { - vm.Stop() + if err := waitForVMStop(ctx, vm); err != nil { + return fmt.Errorf("failed to stop VM: %w", err) + } } var statusMessage string @@ -582,6 +586,15 @@ func (worker *Worker) syncOnDiskVMs(ctx context.Context) error { return nil } +func waitForVMStop(ctx context.Context, vm vmmanager.VM) error { + select { + case err := <-vm.Stop(): + return err + case <-ctx.Done(): + return ctx.Err() + } +} + func (worker *Worker) deleteVM(vm vmmanager.VM) error { <-vm.Stop() diff --git a/internal/worker/worker_stop_test.go b/internal/worker/worker_stop_test.go new file mode 100644 index 0000000..fdb51d6 --- /dev/null +++ b/internal/worker/worker_stop_test.go @@ -0,0 +1,135 @@ +package worker //nolint:testpackage // The regression test exercises unexported worker reconciliation. + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/cirruslabs/orchard/internal/worker/ondiskname" + "github.com/cirruslabs/orchard/internal/worker/vmmanager" + "github.com/cirruslabs/orchard/pkg/client" + v1 "github.com/cirruslabs/orchard/pkg/resource/v1" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +var errLocalVMFailed = errors.New("local VM failed") + +type delayedStopVM struct { + vmmanager.VM + + resource v1.VM + status v1.VMStatus + stopStarted chan struct{} + stopResult chan error +} + +func (vm *delayedStopVM) OnDiskName() ondiskname.OnDiskName { + return ondiskname.NewFromResource(vm.resource) +} + +func (vm *delayedStopVM) Status() v1.VMStatus { return vm.status } + +func (vm *delayedStopVM) Conditions() []v1.Condition { return nil } + +func (vm *delayedStopVM) Err() error { return errLocalVMFailed } + +func (vm *delayedStopVM) Stop() <-chan error { + close(vm.stopStarted) + return vm.stopResult +} + +func TestSyncVMsWaitsForVMShutdown(t *testing.T) { + tests := []struct { + name string + remoteStatus v1.VMStatus + localStatus v1.VMStatus + update bool + }{ + { + name: "remote failed VM stops before reconciliation continues", + remoteStatus: v1.VMStatusFailed, + localStatus: v1.VMStatusRunning, + }, + { + name: "local failed VM stops before reporting failure", + remoteStatus: v1.VMStatusRunning, + localStatus: v1.VMStatusFailed, + update: true, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + resource := v1.VM{ + Meta: v1.Meta{Name: "test-vm"}, + UID: "00112233-4455-6677-8899-aabbccddeeff", + Worker: "test-worker", + Status: test.remoteStatus, + } + + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + if request.Method != http.MethodGet || request.URL.Path != "/v1/vms" { + t.Errorf("unexpected request: %s %s", request.Method, request.URL.Path) + return + } + + response.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(response).Encode([]map[string]any{{ + "name": resource.Name, + "uid": resource.UID, + "worker": resource.Worker, + "status": resource.Status, + }}); err != nil { + t.Errorf("encode VM response: %v", err) + } + })) + defer server.Close() + + apiClient, err := client.New(client.WithAddress(server.URL)) + require.NoError(t, err) + vm := &delayedStopVM{ + resource: resource, + status: test.localStatus, + stopStarted: make(chan struct{}), + stopResult: make(chan error, 1), + } + manager := vmmanager.New() + manager.Put(vm.OnDiskName(), vm) + worker := &Worker{name: "test-worker", client: apiClient, vmm: manager, logger: zap.NewNop().Sugar()} + updated := make(chan v1.VM, 1) + finished := make(chan error, 1) + + go func() { + finished <- worker.syncVMs(context.Background(), func(_ context.Context, updatedVM v1.VM) error { + updated <- updatedVM + return nil + }) + }() + + select { + case <-vm.stopStarted: + case <-time.After(time.Second): + t.Fatal("VM shutdown was not requested") + } + + select { + case updatedVM := <-updated: + t.Fatalf("VM was reported %q before shutdown completed", updatedVM.Status) + case err := <-finished: + t.Fatalf("reconciliation continued before VM shutdown completed: %v", err) + case <-time.After(30 * time.Millisecond): + } + + vm.stopResult <- nil + require.NoError(t, <-finished) + if test.update { + require.Equal(t, v1.VMStatusFailed, (<-updated).Status) + } + }) + } +}