Simplify state reconciliation and support changing Softnet settings (#364)

* Simplify state reconciliation and support changing Softnet settings

* Remove unused "updateFunc" parameter from syncOnDiskVMs()

* Don't take an address of a loop variable

* ensure → ensures

* updateVMState(): don't forget to update VMState

* Introduce TestSpecUpdateSoftnet integration test

* Update OpenAPI specification to include generation/observedGeneration

internal/worker/fsm.go

@@ -0,0 +1,58 @@
+package worker
+
+import (
+	"fmt"
+
+	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
+	"github.com/samber/mo"
+)
+
+type Action string
+
+const (
+	ActionIgnore         Action = "ignore"
+	ActionCreate         Action = "create"
+	ActionMonitorPending Action = "monitor-pending"
+	ActionReportRunning  Action = "report-running"
+	ActionMonitorRunning Action = "monitor-running"
+	ActionStop           Action = "stop"
+	ActionFail           Action = "fail"
+	ActionLostTrack      Action = "lost-track"
+	ActionImpossible     Action = "impossible"
+	ActionDelete         Action = "delete"
+)
+
+var transitions = map[mo.Option[v1.VMStatus]]map[mo.Option[v1.VMStatus]]Action{
+	mo.None[v1.VMStatus](): {
+		mo.None[v1.VMStatus]():      ActionIgnore,
+		mo.Some(v1.VMStatusPending): ActionDelete,
+		mo.Some(v1.VMStatusRunning): ActionDelete,
+		mo.Some(v1.VMStatusFailed):  ActionDelete,
+	},
+	mo.Some(v1.VMStatusPending): {
+		mo.None[v1.VMStatus]():      ActionCreate,
+		mo.Some(v1.VMStatusPending): ActionMonitorPending,
+		mo.Some(v1.VMStatusRunning): ActionReportRunning,
+		mo.Some(v1.VMStatusFailed):  ActionFail,
+	},
+	mo.Some(v1.VMStatusRunning): {
+		mo.None[v1.VMStatus]():      ActionLostTrack,
+		mo.Some(v1.VMStatusPending): ActionImpossible,
+		mo.Some(v1.VMStatusRunning): ActionMonitorRunning,
+		mo.Some(v1.VMStatusFailed):  ActionFail,
+	},
+	mo.Some(v1.VMStatusFailed): {
+		mo.None[v1.VMStatus]():      ActionIgnore,
+		mo.Some(v1.VMStatusPending): ActionStop,
+		mo.Some(v1.VMStatusRunning): ActionStop,
+		mo.Some(v1.VMStatusFailed):  ActionIgnore,
+	},
+}
+
+func optionToString[T any](option mo.Option[T]) string {
+	if option.IsNone() {
+		return "None"
+	}
+
+	return fmt.Sprintf("Some(%v)", option.MustGet())
+}

internal/worker/fsm_test.go

@@ -0,0 +1,27 @@
+package worker
+
+import (
+	"testing"
+
+	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
+	"github.com/samber/mo"
+	"github.com/stretchr/testify/require"
+)
+
+// TestExplicitStateTransitions ensures that all state transitions
+// yield a defined action (something other than ActionUndefined).
+func TestExplicitStateTransitions(t *testing.T) {
+	possibleStates := []mo.Option[v1.VMStatus]{
+		mo.None[v1.VMStatus](),
+		mo.Some(v1.VMStatusPending),
+		mo.Some(v1.VMStatusRunning),
+		mo.Some(v1.VMStatusFailed),
+	}
+
+	for _, remote := range possibleStates {
+		for _, local := range possibleStates {
+			require.Positivef(t, transitions[remote][local], "state transition %s -> %s is not defined",
+				optionToString(remote), optionToString(local))
+		}
+	}
+}

internal/worker/vmmanager/vm.go

@@ -46,8 +46,8 @@ type VM struct {
 	// Image FQN feature, see https://github.com/cirruslabs/orchard/issues/164
 	imageFQN atomic.Pointer[string]
 
-	status atomic.Pointer[string]
-	err    atomic.Pointer[error]
+	statusMessage atomic.Pointer[string]
+	err           atomic.Pointer[error]
 
 	ctx    context.Context
 	cancel context.CancelFunc
@@ -89,7 +89,7 @@ func NewVM(
 		defer vm.wg.Done()
 
 		if vmResource.ImagePullPolicy == v1.ImagePullPolicyAlways {
-			vm.setStatus("pulling VM image...")
+			vm.setStatusMessage("pulling VM image...")
 
 			pullStartedAt := time.Now()
 
@@ -123,38 +123,9 @@ func NewVM(
 		}
 
 		vm.cloned.Store(true)
-
-		// Launch the startup script goroutine as close as possible
-		// to the VM startup (below) to avoid "tart ip" timing out
-		if vm.Resource.StartupScript != nil {
-			vm.setStatus("VM started, running startup script...")
-
-			go vm.runScript(vm.Resource.StartupScript, eventStreamer)
-		} else {
-			vm.setStatus("VM started")
-		}
-
 		vm.started.Store(true)
 
-		if err := vm.run(vm.ctx); err != nil {
-			select {
-			case <-vm.ctx.Done():
-				// Do not return an error because it's the user's intent to cancel this VM
-			default:
-				vm.setErr(fmt.Errorf("%w: %v", ErrVMFailed, err))
-			}
-
-			return
-		}
-
-		select {
-		case <-vm.ctx.Done():
-			// Do not return an error because it's the user's intent to cancel this VM
-		default:
-			if !vm.stopping.Load() {
-				vm.setErr(fmt.Errorf("%w: VM exited unexpectedly", ErrVMFailed))
-			}
-		}
+		vm.run(vm.ctx, eventStreamer)
 	}()
 
 	return vm
@@ -176,8 +147,20 @@ func (vm *VM) id() string {
 	return vm.onDiskName.String()
 }
 
-func (vm *VM) Status() string {
-	status := vm.status.Load()
+func (vm *VM) Status() v1.VMStatus {
+	if vm.Err() != nil {
+		return v1.VMStatusFailed
+	}
+
+	if vm.Started() {
+		return v1.VMStatusRunning
+	}
+
+	return v1.VMStatusPending
+}
+
+func (vm *VM) StatusMessage() string {
+	status := vm.statusMessage.Load()
 
 	if status != nil {
 		return *status
@@ -186,9 +169,9 @@ func (vm *VM) Status() string {
 	return ""
 }
 
-func (vm *VM) setStatus(status string) {
+func (vm *VM) setStatusMessage(status string) {
 	vm.logger.Debugf(status)
-	vm.status.Store(&status)
+	vm.statusMessage.Store(&status)
 }
 
 func (vm *VM) Err() error {
@@ -206,7 +189,7 @@ func (vm *VM) setErr(err error) {
 }
 
 func (vm *VM) cloneAndConfigure(ctx context.Context) error {
-	vm.setStatus("cloning VM...")
+	vm.setStatusMessage("cloning VM...")
 
 	_, _, err := tart.Tart(ctx, vm.logger, "clone", vm.Resource.Image, vm.id())
 	if err != nil {
@@ -221,7 +204,7 @@ func (vm *VM) cloneAndConfigure(ctx context.Context) error {
 	}
 
 	// Set memory
-	vm.setStatus("configuring VM...")
+	vm.setStatusMessage("configuring VM...")
 
 	memory := vm.Resource.AssignedMemory
 
@@ -333,7 +316,17 @@ func (vm *VM) cloneAndConfigure(ctx context.Context) error {
 	return nil
 }
 
-func (vm *VM) run(ctx context.Context) error {
+func (vm *VM) run(ctx context.Context, eventStreamer *client.EventStreamer) {
+	// Launch the startup script goroutine as close as possible
+	// to the VM startup (below) to avoid "tart ip" timing out
+	if vm.Resource.StartupScript != nil {
+		vm.setStatusMessage("VM started, running startup script...")
+
+		go vm.runScript(vm.Resource.StartupScript, eventStreamer)
+	} else {
+		vm.setStatusMessage("VM started")
+	}
+
 	var runArgs = []string{"run"}
 
 	if vm.Resource.NetSoftnetDeprecated || vm.Resource.NetSoftnet {
@@ -364,10 +357,24 @@ func (vm *VM) run(ctx context.Context) error {
 	runArgs = append(runArgs, vm.id())
 	_, _, err := tart.Tart(ctx, vm.logger, runArgs...)
 	if err != nil {
-		return err
+		select {
+		case <-vm.ctx.Done():
+			// Do not return an error because it's the user's intent to cancel this VM
+		default:
+			vm.setErr(fmt.Errorf("%w: %v", ErrVMFailed, err))
+		}
+
+		return
 	}
 
-	return nil
+	select {
+	case <-vm.ctx.Done():
+		// Do not return an error because it's the user's intent to cancel this VM
+	default:
+		if !vm.stopping.Load() {
+			vm.setErr(fmt.Errorf("%w: VM exited unexpectedly", ErrVMFailed))
+		}
+	}
 }
 
 func (vm *VM) IP(ctx context.Context) (string, error) {
@@ -401,6 +408,7 @@ func (vm *VM) Stop() {
 	vm.logger.Debugf("stopping VM")
 
 	vm.stopping.Store(true)
+	defer vm.stopping.Store(false)
 
 	// Try to gracefully terminate the VM
 	_, _, _ = tart.Tart(context.Background(), zap.NewNop().Sugar(), "stop", "--timeout", "5", vm.id())
@@ -412,6 +420,19 @@ func (vm *VM) Stop() {
 	vm.logger.Debugf("VM stopped")
 }
 
+func (vm *VM) Reboot(eventStreamer *client.EventStreamer) {
+	vm.Stop()
+
+	vm.ctx, vm.cancel = context.WithCancel(context.Background())
+	vm.wg.Add(1)
+
+	go func() {
+		defer vm.wg.Done()
+
+		vm.run(vm.ctx, eventStreamer)
+	}()
+}
+
 func (vm *VM) Delete() error {
 	if !vm.cloned.Load() {
 		return nil

internal/worker/worker.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"slices"
 	"time"
 
 	"github.com/avast/retry-go/v4"
@@ -18,8 +19,11 @@ import (
 	"github.com/cirruslabs/orchard/pkg/client"
 	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
 	"github.com/cirruslabs/orchard/rpc"
+	mapset "github.com/deckarep/golang-set/v2"
 	"github.com/dustin/go-humanize"
 	"github.com/hashicorp/go-multierror"
+	"github.com/samber/lo"
+	"github.com/samber/mo"
 	"github.com/shirou/gopsutil/v4/cpu"
 	"github.com/shirou/gopsutil/v4/mem"
 	"go.opentelemetry.io/otel/metric"
@@ -191,6 +195,13 @@ func (worker *Worker) runNewSession(ctx context.Context) error {
 		return nil
 	}
 
+	// Backward compatibility with for older Orchard Controllers
+	updateFunc := worker.client.VMs().UpdateState
+
+	if !info.Capabilities.Has(v1.ControllerCapabilityVMStateEndpoint) {
+		updateFunc = worker.client.VMs().Update
+	}
+
 	for {
 		if err := worker.updateWorker(ctx); err != nil {
 			worker.logger.Errorf("failed to update worker resource: %v", err)
@@ -198,7 +209,7 @@ func (worker *Worker) runNewSession(ctx context.Context) error {
 			return nil
 		}
 
-		if err := worker.syncVMs(subCtx); err != nil {
+		if err := worker.syncVMs(subCtx, updateFunc); err != nil {
 			worker.logger.Warnf("failed to sync VMs: %v", err)
 
 			return nil
@@ -260,26 +271,144 @@ func (worker *Worker) updateWorker(ctx context.Context) error {
 }
 
 //nolint:nestif,gocognit // nested "if" and cognitive complexity is tolerable for now
-func (worker *Worker) syncVMs(ctx context.Context) error {
+func (worker *Worker) syncVMs(ctx context.Context, updateVM func(context.Context, v1.VM) (*v1.VM, error)) error {
+	allKeys := mapset.NewSet[ondiskname.OnDiskName]()
+
 	remoteVMs, err := worker.client.VMs().FindForWorker(ctx, worker.name)
 	if err != nil {
 		return err
 	}
-	remoteVMsIndex := map[ondiskname.OnDiskName]v1.VM{}
+	remoteVMsIndex := map[ondiskname.OnDiskName]*v1.VM{}
 	for _, remoteVM := range remoteVMs {
-		remoteVMsIndex[ondiskname.NewFromResource(remoteVM)] = remoteVM
+		onDiskName := ondiskname.NewFromResource(remoteVM)
+		allKeys.Add(onDiskName)
+		// Can't take an address of a loop variable
+		remoteVMCopy := remoteVM
+		remoteVMsIndex[onDiskName] = &remoteVMCopy
+	}
+
+	localVMsIndex := map[ondiskname.OnDiskName]*vmmanager.VM{}
+	for _, vm := range worker.vmm.List() {
+		onDiskName := vm.OnDiskName()
+		allKeys.Add(onDiskName)
+		localVMsIndex[onDiskName] = vm
 	}
 
 	worker.logger.Infof("syncing %d local VMs against %d remote VMs...",
-		worker.vmm.Len(), len(remoteVMsIndex))
+		len(localVMsIndex), len(remoteVMsIndex))
 
-	// It's important to check the remote VMs against local ones first
-	// to stop the failed VMs before we start the new VMs, otherwise we
-	// risk violating the resource constraints (e.g. a maximum of 2 VMs
-	// per host)
-	for _, vm := range worker.vmm.List() {
-		remoteVM, ok := remoteVMsIndex[vm.OnDiskName()]
-		if !ok {
+	var pairs []lo.Tuple3[ondiskname.OnDiskName, *v1.VM, *vmmanager.VM]
+
+	for onDiskName := range allKeys.Iter() {
+		vmResource := remoteVMsIndex[onDiskName]
+		vm := localVMsIndex[onDiskName]
+
+		pairs = append(pairs, lo.T3(onDiskName, vmResource, vm))
+	}
+
+	// It's important to process the remote VMs in failed state
+	// and local VMs that ceased to exist remotely first, otherwise
+	// we risk violating the scheduler resource assumptions
+	sortNonExistentAndFailedFirst(pairs)
+
+	for _, tuple := range pairs {
+		onDiskName, vmResource, vm := lo.Unpack3(tuple)
+
+		remoteState := mo.None[v1.VMStatus]()
+		if vmResource != nil {
+			remoteState = mo.Some(vmResource.Status)
+		}
+
+		localState := mo.None[v1.VMStatus]()
+		if vm != nil {
+			localState = mo.Some(vm.Status())
+		}
+
+		action := transitions[remoteState][localState]
+
+		worker.logger.Debugf("processing VM: %s, remote: %v, local: %v, action: %v\n", onDiskName,
+			optionToString(remoteState), optionToString(localState), action)
+
+		switch action {
+		case ActionCreate:
+			// Remote VM was created, but not the local VM
+			worker.createVM(onDiskName, *vmResource)
+		case ActionMonitorPending:
+			if vmResource.StatusMessage != vm.StatusMessage() {
+				vmResource.StatusMessage = vm.StatusMessage()
+
+				if _, err := updateVM(ctx, *vmResource); err != nil {
+					return err
+				}
+			}
+		case ActionReportRunning:
+			// Remote VM was created, and the local VM too,
+			// check if the local VM had already started
+			// and update the remote VM as accordingly
+
+			// Image FQN feature, see https://github.com/cirruslabs/orchard/issues/164
+			if imageFQN := vm.ImageFQN(); imageFQN != nil {
+				vmResource.ImageFQN = *imageFQN
+			}
+
+			// Mark the remote VM as started
+			vmResource.Status = v1.VMStatusRunning
+			vmResource.StatusMessage = vm.StatusMessage()
+
+			if _, err := updateVM(ctx, *vmResource); err != nil {
+				return err
+			}
+		case ActionMonitorRunning:
+			if vmResource.StatusMessage != vm.StatusMessage() {
+				vmResource.StatusMessage = vm.StatusMessage()
+
+				if _, err := updateVM(ctx, *vmResource); err != nil {
+					return err
+				}
+			}
+
+			if vmResource.Generation != vm.Resource.Generation {
+				// Something changed, reboot the VM for the changes to take effect
+				vm.Resource = *vmResource
+
+				eventStreamer := worker.client.VMs().StreamEvents(vmResource.Name)
+
+				vm.Reboot(eventStreamer)
+
+				vmResource.ObservedGeneration = vm.Resource.Generation
+
+				if _, err := updateVM(ctx, *vmResource); err != nil {
+					return err
+				}
+			}
+		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()
+		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()
+			}
+
+			var statusMessage string
+
+			switch action {
+			case ActionFail:
+				statusMessage = vm.Err().Error()
+			case ActionLostTrack:
+				statusMessage = "Worker lost track of VM"
+			case ActionImpossible:
+				statusMessage = "Encountered an impossible transition"
+			}
+
+			vmResource.Status = v1.VMStatusFailed
+			vmResource.StatusMessage = statusMessage
+			if _, err := updateVM(ctx, *vmResource); err != nil {
+				return err
+			}
+		case ActionDelete:
 			// Remote VM was deleted, delete local VM
 			//
 			// Note: this check needs to run for each VM
@@ -287,58 +416,6 @@ func (worker *Worker) syncVMs(ctx context.Context) error {
 			if err := worker.deleteVM(vm); err != nil {
 				return err
 			}
-		} else {
-			if remoteVM.Status == v1.VMStatusFailed {
-				// VM has failed on the remote side, stop it locally to prevent incorrect
-				// worker's resources calculation in the Controller's scheduler
-				vm.Stop()
-			} else if vm.Err() != nil {
-				// 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
-				vm.Stop()
-
-				// Report the VM as failed
-				remoteVM.Status = v1.VMStatusFailed
-				remoteVM.StatusMessage = vm.Err().Error()
-				if _, err := worker.client.VMs().Update(ctx, remoteVM); err != nil {
-					return err
-				}
-			} else if vm.Status() != remoteVM.StatusMessage {
-				// Report the new VM status message
-				remoteVM.StatusMessage = vm.Status()
-				if _, err := worker.client.VMs().Update(ctx, remoteVM); err != nil {
-					return err
-				}
-			}
-		}
-	}
-
-	for _, vmResource := range remoteVMsIndex {
-		odn := ondiskname.NewFromResource(vmResource)
-
-		if vmResource.Status != v1.VMStatusPending {
-			continue
-		}
-
-		if vm, ok := worker.vmm.Get(odn); ok {
-			// Remote VM was created, and the local VM too,
-			// check if the local VM had already started
-			// and update the remote VM as accordingly
-			if vm.Started() {
-				// Image FQN feature, see https://github.com/cirruslabs/orchard/issues/164
-				if imageFQN := vm.ImageFQN(); imageFQN != nil {
-					vmResource.ImageFQN = *imageFQN
-				}
-
-				// Mark the remote VM as started
-				vmResource.Status = v1.VMStatusRunning
-				if _, err := worker.client.VMs().Update(ctx, vmResource); err != nil {
-					return err
-				}
-			}
-		} else {
-			// Remote VM was created, but not the local VM
-			worker.createVM(odn, vmResource)
 		}
 	}
 
@@ -403,15 +480,6 @@ func (worker *Worker) syncOnDiskVMs(ctx context.Context) error {
 					worker.logger.Warnf("failed to stop")
 				}
 			}
-
-			if remoteVM.Status != v1.VMStatusFailed {
-				remoteVM.Status = v1.VMStatusFailed
-				remoteVM.StatusMessage = "Worker lost track of VM"
-				_, err := worker.client.VMs().Update(ctx, remoteVM)
-				if err != nil {
-					return err
-				}
-			}
 		}
 	}
 
@@ -454,3 +522,36 @@ func (worker *Worker) requestVMSyncing() {
 		worker.logger.Debugf("There's already a syncing request in the queue, skipping")
 	}
 }
+
+func sortNonExistentAndFailedFirst(input []lo.Tuple3[ondiskname.OnDiskName, *v1.VM, *vmmanager.VM]) {
+	slices.SortStableFunc(input, func(left, right lo.Tuple3[ondiskname.OnDiskName, *v1.VM, *vmmanager.VM]) int {
+		_, leftVM, _ := lo.Unpack3(left)
+		_, rightVM, _ := lo.Unpack3(right)
+
+		leftNonExistent := leftVM == nil
+		rightNonExistent := rightVM == nil
+
+		switch {
+		case leftNonExistent && rightNonExistent:
+			return 0
+		case leftNonExistent:
+			return -1
+		case rightNonExistent:
+			return 1
+		}
+
+		leftFailed := leftVM != nil && leftVM.Status == v1.VMStatusFailed
+		rightFailed := rightVM != nil && rightVM.Status == v1.VMStatusFailed
+
+		switch {
+		case leftFailed && rightFailed:
+			return 0
+		case leftFailed:
+			return -1
+		case rightFailed:
+			return 1
+		default:
+			return 0
+		}
+	})
+}

internal/worker/worker_test.go

@@ -0,0 +1,37 @@
+package worker
+
+import (
+	"testing"
+
+	"github.com/cirruslabs/orchard/internal/worker/ondiskname"
+	"github.com/cirruslabs/orchard/internal/worker/vmmanager"
+	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
+	"github.com/samber/lo"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSortNonExistentAndFailedFirst(t *testing.T) {
+	newVMTuple := func(name string, vmResource *v1.VM) lo.Tuple3[ondiskname.OnDiskName, *v1.VM, *vmmanager.VM] {
+		return lo.T3(ondiskname.New(name, name, 0), vmResource, &vmmanager.VM{})
+	}
+
+	target := []lo.Tuple3[ondiskname.OnDiskName, *v1.VM, *vmmanager.VM]{
+		newVMTuple("test1", &v1.VM{Status: v1.VMStatusFailed}),
+		newVMTuple("test2", &v1.VM{Status: v1.VMStatusPending}),
+		newVMTuple("test3", &v1.VM{Status: v1.VMStatusRunning}),
+		newVMTuple("test5", nil),
+		newVMTuple("test4", &v1.VM{Status: v1.VMStatusFailed}),
+	}
+
+	sortNonExistentAndFailedFirst(target)
+
+	expected := []lo.Tuple3[ondiskname.OnDiskName, *v1.VM, *vmmanager.VM]{
+		newVMTuple("test5", nil),
+		newVMTuple("test1", &v1.VM{Status: v1.VMStatusFailed}),
+		newVMTuple("test4", &v1.VM{Status: v1.VMStatusFailed}),
+		newVMTuple("test2", &v1.VM{Status: v1.VMStatusPending}),
+		newVMTuple("test3", &v1.VM{Status: v1.VMStatusRunning}),
+	}
+
+	require.Equal(t, expected, target)
+}