From 835c71ddee0154c0acc529e978c1f83ff1f7b80a Mon Sep 17 00:00:00 2001 From: Fedor Korotkov Date: Tue, 23 Jun 2026 15:20:42 +0100 Subject: [PATCH] Preserve running VMs across controller restarts --- internal/controller/scheduler/scheduler.go | 13 +- .../controller/scheduler/scheduler_test.go | 52 +++ internal/tests/controller_restart_test.go | 300 ++++++++++++++++++ 3 files changed, 364 insertions(+), 1 deletion(-) create mode 100644 internal/controller/scheduler/scheduler_test.go create mode 100644 internal/tests/controller_restart_test.go diff --git a/internal/controller/scheduler/scheduler.go b/internal/controller/scheduler/scheduler.go index e45c06a7..2d93796c 100644 --- a/internal/controller/scheduler/scheduler.go +++ b/internal/controller/scheduler/scheduler.go @@ -36,6 +36,7 @@ type Scheduler struct { store storepkg.Store notifier *notifier.Notifier workerOfflineTimeout time.Duration + startedAt time.Time logger *zap.SugaredLogger schedulingRequested chan bool @@ -87,6 +88,11 @@ func NewScheduler( } func (scheduler *Scheduler) Run() { + // Workers might have continued running VMs while the controller was + // unavailable. Give them a full offline timeout to reconnect before stale + // heartbeats can cause their VMs to be marked as failed. + scheduler.startedAt = time.Now() + for { // wait either the scheduling interval or a request to schedule select { @@ -550,7 +556,8 @@ func (scheduler *Scheduler) healthCheckVM(txn storepkg.Transaction, vm v1.VM) er return err } - if worker.Offline(scheduler.workerOfflineTimeout) && !vm.TerminalState() { + if scheduler.workerFailureChecksEnabled(time.Now()) && + worker.Offline(scheduler.workerOfflineTimeout) && !vm.TerminalState() { vm.Status = v1.VMStatusFailed vm.StatusMessage = "VM is assigned to a worker that " + "lost connection with the controller" @@ -588,3 +595,7 @@ func (scheduler *Scheduler) healthCheckVM(txn storepkg.Transaction, vm v1.VM) er return nil } + +func (scheduler *Scheduler) workerFailureChecksEnabled(now time.Time) bool { + return scheduler.startedAt.IsZero() || now.Sub(scheduler.startedAt) > scheduler.workerOfflineTimeout +} diff --git a/internal/controller/scheduler/scheduler_test.go b/internal/controller/scheduler/scheduler_test.go new file mode 100644 index 00000000..7269ab22 --- /dev/null +++ b/internal/controller/scheduler/scheduler_test.go @@ -0,0 +1,52 @@ +package scheduler + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestWorkerFailureChecksEnabled(t *testing.T) { + now := time.Unix(1_000_000, 0) + offlineTimeout := 3 * time.Minute + + testCases := []struct { + name string + startedAt time.Time + enabled bool + }{ + { + name: "scheduler-not-started", + enabled: true, + }, + { + name: "scheduler-just-started", + startedAt: now, + }, + { + name: "before-offline-timeout", + startedAt: now.Add(-offlineTimeout + time.Second), + }, + { + name: "at-offline-timeout", + startedAt: now.Add(-offlineTimeout), + }, + { + name: "after-offline-timeout", + startedAt: now.Add(-offlineTimeout - time.Second), + enabled: true, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + scheduler := Scheduler{ + workerOfflineTimeout: offlineTimeout, + startedAt: testCase.startedAt, + } + + require.Equal(t, testCase.enabled, scheduler.workerFailureChecksEnabled(now)) + }) + } +} diff --git a/internal/tests/controller_restart_test.go b/internal/tests/controller_restart_test.go new file mode 100644 index 00000000..16f14d25 --- /dev/null +++ b/internal/tests/controller_restart_test.go @@ -0,0 +1,300 @@ +package tests_test + +import ( + "bytes" + "context" + "errors" + "fmt" + "net" + "net/http" + "os" + "os/exec" + "os/signal" + "testing" + "time" + + "github.com/cirruslabs/orchard/internal/controller" + "github.com/cirruslabs/orchard/internal/tests/platformdependent" + "github.com/cirruslabs/orchard/internal/tests/wait" + "github.com/cirruslabs/orchard/internal/worker" + "github.com/cirruslabs/orchard/pkg/client" + v1 "github.com/cirruslabs/orchard/pkg/resource/v1" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +const ( + controllerRestartHelperEnv = "ORCHARD_CONTROLLER_RESTART_TEST_HELPER" + controllerRestartHelperDataDirEnv = "ORCHARD_CONTROLLER_RESTART_TEST_DATA_DIR" + controllerRestartHelperListenAddrEnv = "ORCHARD_CONTROLLER_RESTART_TEST_LISTEN_ADDR" + controllerRestartHelperOfflineTimeoutEnv = "ORCHARD_CONTROLLER_RESTART_TEST_OFFLINE_TIMEOUT" + controllerRestartWorkerOfflineTimeout = 20 * time.Second + controllerRestartWorkerDisconnectDuration = 21 * time.Second +) + +func TestControllerRestartDoesNotFailRunningVMs(t *testing.T) { + ctx := t.Context() + listenAddr := unusedTCPAddress(t) + dataDir := t.TempDir() + + devClient, err := client.New(client.WithAddress("http://" + listenAddr)) + require.NoError(t, err) + + controllerProcess := startControllerProcess(t, devClient, dataDir, listenAddr) + devWorker := startSyntheticWorker(t, devClient) + + var workerName string + require.True(t, wait.Wait(30*time.Second, func() bool { + workers, err := devClient.Workers().List(ctx) + if err != nil || len(workers) != 1 { + return false + } + + workerName = workers[0].Name + + return true + }), "failed to wait for the worker to register") + + const vmName = "controller-restart-vm" + + require.NoError(t, devClient.VMs().Create(ctx, platformdependent.VM(vmName))) + require.True(t, wait.Wait(30*time.Second, func() bool { + vm, err := devClient.VMs().Get(ctx, vmName) + if err != nil { + return false + } + + t.Logf("Waiting for the synthetic VM to start. Current status: %s", vm.Status) + + return vm.Status == v1.VMStatusRunning + }), "failed to wait for the synthetic VM to start") + + vmBeforeRestart, err := devClient.VMs().Get(ctx, vmName) + require.NoError(t, err) + require.Equal(t, workerName, vmBeforeRestart.Worker) + + // Stop the worker's controller session without stopping its running VMs, then + // keep the controller down until the persisted worker heartbeat is stale. + devWorker.stop(t) + require.NoError(t, controllerProcess.stop()) + time.Sleep(controllerRestartWorkerDisconnectDuration) + + restartedAt := time.Now() + startControllerProcess(t, devClient, dataDir, listenAddr) + + // Keep the worker disconnected beyond the first post-restart scheduler + // health check. The VM should remain running during the startup grace period. + observationDeadline := time.Now().Add(7 * time.Second) + for time.Now().Before(observationDeadline) { + vm, err := devClient.VMs().Get(ctx, vmName) + require.NoError(t, err) + require.Equal(t, v1.VMStatusRunning, vm.Status) + require.Equal(t, vmBeforeRestart.UID, vm.UID) + require.Equal(t, workerName, vm.Worker) + + time.Sleep(250 * time.Millisecond) + } + + devWorker.start(t) + require.True(t, wait.Wait(30*time.Second, func() bool { + workerResource, err := devClient.Workers().Get(ctx, workerName) + if err != nil || !workerResource.LastSeen.After(restartedAt) { + return false + } + + vm, err := devClient.VMs().Get(ctx, vmName) + + return err == nil && vm.Status == v1.VMStatusRunning + }), "worker did not reconnect with its VM still running after the controller restart") +} + +// TestControllerRestartHelperProcess runs an Orchard Controller in a separate +// process so that its Badger database is released and can be reopened during a +// realistic controller restart. +func TestControllerRestartHelperProcess(t *testing.T) { + if os.Getenv(controllerRestartHelperEnv) != "1" { + return + } + + dataDirPath := requireEnv(t, controllerRestartHelperDataDirEnv) + listenAddr := requireEnv(t, controllerRestartHelperListenAddrEnv) + offlineTimeout, err := time.ParseDuration(requireEnv(t, controllerRestartHelperOfflineTimeoutEnv)) + require.NoError(t, err) + + dataDir, err := controller.NewDataDir(dataDirPath) + require.NoError(t, err) + + devController, err := controller.New( + controller.WithDataDir(dataDir), + controller.WithListenAddr(listenAddr), + controller.WithInsecureAuthDisabled(), + controller.WithExperimentalRPCV2(), + controller.WithWorkerOfflineTimeout(offlineTimeout), + controller.WithLogger(zap.NewNop()), + ) + require.NoError(t, err) + + ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt) + defer cancel() + + err = devController.Run(ctx) + require.True(t, err == nil || errors.Is(err, context.Canceled) || errors.Is(err, http.ErrServerClosed), + "controller failed: %v", err) +} + +type controllerTestProcess struct { + cmd *exec.Cmd + output bytes.Buffer + stopped bool +} + +func startControllerProcess( + t *testing.T, + devClient *client.Client, + dataDir string, + listenAddr string, +) *controllerTestProcess { + t.Helper() + + process := &controllerTestProcess{} + process.cmd = exec.Command(os.Args[0], "-test.run=^TestControllerRestartHelperProcess$", "-test.v") + process.cmd.Env = append(os.Environ(), + controllerRestartHelperEnv+"=1", + controllerRestartHelperDataDirEnv+"="+dataDir, + controllerRestartHelperListenAddrEnv+"="+listenAddr, + controllerRestartHelperOfflineTimeoutEnv+"="+controllerRestartWorkerOfflineTimeout.String(), + ) + process.cmd.Stdout = &process.output + process.cmd.Stderr = &process.output + + require.NoError(t, process.cmd.Start()) + t.Cleanup(func() { + require.NoError(t, process.stop()) + }) + + if wait.Wait(10*time.Second, func() bool { + requestCtx, requestCancel := context.WithTimeout(t.Context(), 500*time.Millisecond) + defer requestCancel() + + _, err := devClient.Controller().Info(requestCtx) + + return err == nil + }) { + return process + } + + stopErr := process.stop() + t.Fatalf("controller failed to become ready: %v\n%s", stopErr, process.output.String()) + + return nil +} + +func (process *controllerTestProcess) stop() error { + if process.stopped { + return nil + } + process.stopped = true + + signalErr := process.cmd.Process.Signal(os.Interrupt) + waitCh := make(chan error, 1) + go func() { + waitCh <- process.cmd.Wait() + }() + + select { + case waitErr := <-waitCh: + if signalErr != nil && !errors.Is(signalErr, os.ErrProcessDone) { + return fmt.Errorf("failed to interrupt controller process: %w", signalErr) + } + if waitErr != nil { + return fmt.Errorf("controller process failed: %w\n%s", waitErr, process.output.String()) + } + + return nil + case <-time.After(10 * time.Second): + _ = process.cmd.Process.Kill() + <-waitCh + + return fmt.Errorf("timed out waiting for controller process to stop\n%s", process.output.String()) + } +} + +type syntheticTestWorker struct { + worker *worker.Worker + cancel context.CancelFunc + done chan error +} + +func startSyntheticWorker(t *testing.T, devClient *client.Client) *syntheticTestWorker { + t.Helper() + + devWorker, err := worker.New(devClient, worker.WithSynthetic(), worker.WithLogger(zap.NewNop())) + require.NoError(t, err) + + testWorker := &syntheticTestWorker{ + worker: devWorker, + } + testWorker.start(t) + t.Cleanup(func() { + testWorker.stop(t) + require.NoError(t, testWorker.worker.Close()) + }) + + return testWorker +} + +func (worker *syntheticTestWorker) start(t *testing.T) { + t.Helper() + require.Nil(t, worker.cancel, "synthetic test worker is already running") + + workerCtx, cancelWorker := context.WithCancel(context.Background()) + workerDone := make(chan error, 1) + worker.cancel = cancelWorker + worker.done = workerDone + + go func() { + workerDone <- worker.worker.Run(workerCtx) + }() +} + +func (worker *syntheticTestWorker) stop(t *testing.T) { + t.Helper() + if worker.cancel == nil { + return + } + + worker.cancel() + + select { + case err := <-worker.done: + if err != nil && !errors.Is(err, context.Canceled) { + t.Errorf("dev worker failed: %v", err) + } + case <-time.After(10 * time.Second): + t.Fatal("timed out waiting for dev worker to stop") + } + + worker.cancel = nil + worker.done = nil +} + +func unusedTCPAddress(t *testing.T) string { + t.Helper() + + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + listenAddr := listener.Addr().String() + require.NoError(t, listener.Close()) + + return listenAddr +} + +func requireEnv(t *testing.T, name string) string { + t.Helper() + + value := os.Getenv(name) + require.NotEmpty(t, value, "%s must be set", name) + + return value +}