[codex] Preserve running VMs across controller restarts

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
+}

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))
+		})
+	}
+}

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
+}