Wait for Tart and Vetu stop completion

internal/worker/vmmanager/base/base.go

@@ -103,10 +103,11 @@ func (vm *VM) ConditionsSet() mapset.Set[v1.ConditionType] {
 }
 
 func (vm *VM) Conditions() []v1.Condition {
-	// Only expose a minimum amount of conditions necessary
-	// for the Orchard Controller to make decisions
+	// The worker must observe transitions before applying a new specification.
 	return []v1.Condition{
 		vm.conditionTypeToCondition(v1.ConditionTypeRunning),
+		vm.conditionTypeToCondition(v1.ConditionTypeSuspending),
+		vm.conditionTypeToCondition(v1.ConditionTypeStopping),
 	}
 }
 

internal/worker/vmmanager/stop_test.go

@@ -0,0 +1,199 @@
+package vmmanager_test
+
+import (
+	"errors"
+	"os"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/cirruslabs/orchard/internal/worker/vmmanager"
+	"github.com/cirruslabs/orchard/internal/worker/vmmanager/tart"
+	"github.com/cirruslabs/orchard/internal/worker/vmmanager/vetu"
+	v1 "github.com/cirruslabs/orchard/pkg/resource/v1"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+)
+
+const stopTestTimeout = 15 * time.Second
+
+func TestStopWaitsForCommandAndRun(t *testing.T) {
+	runtimes := []struct {
+		name string
+		new  func(v1.VM) vmmanager.VM
+	}{
+		{name: "tart", new: func(resource v1.VM) vmmanager.VM {
+			return tart.NewVM(resource, nil, nil, nil, zap.NewNop().Sugar())
+		}},
+		{name: "vetu", new: func(resource v1.VM) vmmanager.VM {
+			return vetu.NewVM(resource, nil, nil, nil, zap.NewNop().Sugar())
+		}},
+	}
+
+	for _, runtime := range runtimes {
+		for _, commandFirst := range []bool{true, false} {
+			order := "run finishes first"
+			if commandFirst {
+				order = "stop command finishes first"
+			}
+			t.Run(runtime.name+"/"+order, func(t *testing.T) {
+				checkStopCompletion(t, runtime.name, runtime.new, commandFirst)
+			})
+		}
+	}
+}
+
+//nolint:exhaustruct_v5 // VM settings unrelated to lifecycle transitions use their defaults.
+func checkStopCompletion(t *testing.T, commandName string, newVM func(v1.VM) vmmanager.VM, commandFirst bool) {
+	t.Helper()
+
+	dir := installStopTestCommand(t, commandName)
+	runGate := filepath.Join(dir, "run.release")
+	stopGate := filepath.Join(dir, "stop.release")
+	release := func(path string) { require.NoError(t, os.WriteFile(path, nil, 0o600)) }
+	waitForFile := func(path string) {
+		require.Eventually(t, func() bool {
+			_, err := os.Stat(path)
+			return err == nil
+		}, stopTestTimeout, 5*time.Millisecond, "command did not reach %s", path)
+	}
+	waitForStop := func(done <-chan error) {
+		select {
+		case err := <-done:
+			require.NoError(t, err)
+		case <-time.After(stopTestTimeout):
+			t.Fatal("Stop did not complete after the commands finished")
+		}
+	}
+
+	resource := v1.VM{Name: "stop-test", UID: "test-uid", Image: "source-image"}
+	vm := newVM(resource)
+	var first <-chan error
+	t.Cleanup(func() {
+		if t.Failed() {
+			t.Logf("VM status before fixture cleanup: %s; error: %v", vm.StatusMessage(), vm.Err())
+		}
+		release(runGate)
+		release(stopGate)
+		if first != nil {
+			waitForStop(first)
+		}
+		waitForStop(vm.Stop())
+	})
+
+	// A restart must create a new Stop completion for the new run.
+	for run := range 2 {
+		if run != 0 {
+			for _, name := range []string{"run.release", "stop.release", "run.started", "stop.started", "stop.finished"} {
+				require.NoError(t, os.Remove(filepath.Join(dir, name)))
+			}
+			vm.Start(nil)
+		}
+		waitForFile(filepath.Join(dir, "run.started"))
+		first = vm.Stop()
+		waitForFile(filepath.Join(dir, "stop.started"))
+
+		if commandFirst {
+			release(stopGate)
+			waitForFile(filepath.Join(dir, "stop.finished"))
+			pidText, err := os.ReadFile(filepath.Join(dir, "run.started")) //nolint:gosec // Read a t.TempDir fixture.
+			require.NoError(t, err)
+			pid, err := strconv.Atoi(strings.TrimSpace(string(pidText)))
+			require.NoError(t, err)
+			// Cancellation kills the run process, but its child still holds
+			// stdout open, so Cmd.Run and the VM goroutine cannot finish yet.
+			require.Eventually(t, func() bool {
+				return errors.Is(syscall.Kill(pid, 0), syscall.ESRCH)
+			}, stopTestTimeout, 5*time.Millisecond, "Stop did not cancel the run process")
+		} else {
+			release(runGate)
+			require.Eventually(t, func() bool {
+				return v1.ConditionIsFalse(vm.Conditions(), v1.ConditionTypeRunning)
+			}, stopTestTimeout, 5*time.Millisecond, "the run did not finish")
+		}
+
+		callers := make(chan (<-chan error), 8)
+		for range cap(callers) {
+			go func() { callers <- vm.Stop() }()
+		}
+		stops := []<-chan error{first}
+		for range cap(callers) {
+			select {
+			case done := <-callers:
+				stops = append(stops, done)
+			case <-time.After(stopTestTimeout):
+				t.Fatal("concurrent Stop call did not return its completion channel")
+			}
+		}
+		for _, done := range stops {
+			select {
+			case <-done:
+				t.Fatal("Stop completed while shutdown was still in progress")
+			default:
+			}
+		}
+		// Reconciliation must still see Stopping when the run has ended
+		// but its stop command has not returned.
+		require.True(t, v1.ConditionIsTrue(vm.Conditions(), v1.ConditionTypeStopping))
+
+		release(runGate)
+		release(stopGate)
+		for _, done := range stops {
+			waitForStop(done)
+		}
+		waitForStop(vm.Stop())
+		require.False(t, v1.ConditionIsTrue(vm.Conditions(), v1.ConditionTypeRunning))
+		require.False(t, v1.ConditionIsTrue(vm.Conditions(), v1.ConditionTypeStopping))
+		require.NoError(t, vm.Err())
+	}
+	commands, err := os.ReadFile(filepath.Join(dir, "commands.log")) //nolint:gosec // Read a t.TempDir fixture.
+	require.NoError(t, err)
+	require.Equal(t, 2, strings.Count(string(commands), "stop --timeout 5 "),
+		"concurrent callers must share one stop command per run")
+}
+
+func installStopTestCommand(t *testing.T, name string) string {
+	t.Helper()
+
+	dir := t.TempDir()
+	const script = `#!/bin/sh
+set -eu
+wait_for_file() {
+  attempts=0
+  while [ ! -f "$1" ]; do
+    attempts=$((attempts + 1))
+    if [ "$attempts" -ge "$ORCHARD_STOP_TEST_WAIT_ATTEMPTS" ]; then exit 1; fi
+    /bin/sleep 0.01
+  done
+}
+printf '%s\n' "$*" >> "$ORCHARD_STOP_TEST_DIR/commands.log"
+case "$1" in
+  get)
+    printf '%s\n' '{"Running":false,"State":"stopped"}'
+    ;;
+  run)
+    # Keep the child's output pipes open after cancellation kills this shell.
+    wait_for_file "$ORCHARD_STOP_TEST_DIR/run.release" &
+    printf '%s\n' "$$" > "$ORCHARD_STOP_TEST_DIR/run.started"
+    wait
+    ;;
+  stop)
+    : > "$ORCHARD_STOP_TEST_DIR/stop.started"
+    wait_for_file "$ORCHARD_STOP_TEST_DIR/stop.release"
+    : > "$ORCHARD_STOP_TEST_DIR/stop.finished"
+    ;;
+esac
+`
+	commandPath := filepath.Join(dir, name)
+	require.NoError(t, os.WriteFile(commandPath, []byte(script), 0o600))
+	require.NoError(t, os.Chmod(commandPath, 0o700)) //nolint:gosec // Fake commands must be executable.
+	t.Setenv("PATH", dir+string(os.PathListSeparator)+os.Getenv("PATH"))
+	t.Setenv("ORCHARD_STOP_TEST_DIR", dir)
+	// A command can wait through several Go-side milestones before its gate opens.
+	// Keep it bounded, with enough time for every milestone under process load.
+	t.Setenv("ORCHARD_STOP_TEST_WAIT_ATTEMPTS", strconv.Itoa(int(8*stopTestTimeout/(10*time.Millisecond))))
+	return dir
+}

internal/worker/vmmanager/tart/tart.go

@@ -32,6 +32,9 @@ type VM struct {
 
 	wg *sync.WaitGroup
 
+	stopMtx  sync.Mutex
+	stopDone chan error
+
 	dialer dialer.Dialer
 
 	*base.VM
@@ -267,7 +270,8 @@ func (vm *VM) cloneAndConfigure(ctx context.Context) error {
 }
 
 func (vm *VM) run(ctx context.Context, eventStreamer *client.EventStreamer) {
-	defer vm.ConditionsSet().RemoveAll(v1.ConditionTypeRunning, v1.ConditionTypeSuspending, v1.ConditionTypeStopping)
+	// Stop owns Stopping until both its command and this goroutine finish.
+	defer vm.ConditionsSet().RemoveAll(v1.ConditionTypeRunning, v1.ConditionTypeSuspending)
 
 	// Launch the startup script goroutine as close as possible
 	// to the VM startup (below) to avoid "tart ip" timing out
@@ -398,37 +402,44 @@ func (vm *VM) Suspend() <-chan error {
 }
 
 func (vm *VM) Stop() <-chan error {
-	errCh := make(chan error, 1)
-
-	select {
-	case <-vm.ctx.Done():
-		// VM is already suspended/stopped
-		errCh <- nil
-
-		return errCh
-	default:
-		// VM is still running
+	vm.stopMtx.Lock()
+	defer vm.stopMtx.Unlock()
+	if vm.stopDone != nil {
+		return vm.stopDone
 	}
 
+	done := make(chan error)
+	vm.stopDone = done
+	ctx, cancel, wg := vm.ctx, vm.cancel, vm.wg
 	vm.SetStatusMessage("Stopping VM")
 	vm.ConditionsSet().Add(v1.ConditionTypeStopping)
 
 	go func() {
-		// Try to gracefully terminate the VM
-		_, _, _ = Tart(context.Background(), zap.NewNop().Sugar(), "stop", "--timeout", "5", vm.id())
+		if ctx.Err() == nil {
+			// Try to gracefully terminate the VM.
+			_, _, _ = Tart(context.WithoutCancel(ctx), zap.NewNop().Sugar(), "stop", "--timeout", "5", vm.id())
+		}
 
-		// Terminate the VM goroutine ("tart pull", "tart clone", "tart run", etc.) via the context
-		vm.cancel()
-		vm.wg.Wait()
+		// Cancellation requests shutdown; it does not establish completion.
+		cancel()
+		wg.Wait()
 
-		// We don't return an error because we always terminate a VM
-		errCh <- nil
+		vm.stopMtx.Lock()
+		vm.ConditionsSet().Remove(v1.ConditionTypeStopping)
+		// Closing broadcasts successful completion to every Stop caller.
+		close(done)
+		vm.stopMtx.Unlock()
 	}()
 
-	return errCh
+	return done
 }
 
 func (vm *VM) Start(eventStreamer *client.EventStreamer) {
+	// The worker defers Start while Stopping is true.
+	vm.stopMtx.Lock()
+	defer vm.stopMtx.Unlock()
+	vm.stopDone = nil
+
 	vm.SetStatusMessage("Starting VM")
 	vm.ConditionsSet().Add(v1.ConditionTypeRunning)
 

internal/worker/vmmanager/vetu/vetu.go

@@ -32,6 +32,9 @@ type VM struct {
 
 	wg *sync.WaitGroup
 
+	stopMtx  sync.Mutex
+	stopDone chan error
+
 	dialer dialer.Dialer
 
 	*base.VM
@@ -197,7 +200,8 @@ func (vm *VM) cloneAndConfigure(ctx context.Context) error {
 }
 
 func (vm *VM) run(ctx context.Context, eventStreamer *client.EventStreamer) {
-	defer vm.ConditionsSet().RemoveAll(v1.ConditionTypeRunning, v1.ConditionTypeSuspending, v1.ConditionTypeStopping)
+	// Stop owns Stopping until both its command and this goroutine finish.
+	defer vm.ConditionsSet().RemoveAll(v1.ConditionTypeRunning, v1.ConditionTypeSuspending)
 
 	// Launch the startup script goroutine as close as possible
 	// to the VM startup (below) to avoid "vetu ip" timing out
@@ -257,37 +261,44 @@ func (vm *VM) Suspend() <-chan error {
 }
 
 func (vm *VM) Stop() <-chan error {
-	errCh := make(chan error, 1)
-
-	select {
-	case <-vm.ctx.Done():
-		// VM is already suspended/stopped
-		errCh <- nil
-
-		return errCh
-	default:
-		// VM is still running
+	vm.stopMtx.Lock()
+	defer vm.stopMtx.Unlock()
+	if vm.stopDone != nil {
+		return vm.stopDone
 	}
 
+	done := make(chan error)
+	vm.stopDone = done
+	ctx, cancel, wg := vm.ctx, vm.cancel, vm.wg
 	vm.SetStatusMessage("Stopping VM")
 	vm.ConditionsSet().Add(v1.ConditionTypeStopping)
 
 	go func() {
-		// Try to gracefully terminate the VM
-		_, _, _ = Vetu(context.Background(), zap.NewNop().Sugar(), "stop", "--timeout", "5", vm.id())
+		if ctx.Err() == nil {
+			// Try to gracefully terminate the VM.
+			_, _, _ = Vetu(context.WithoutCancel(ctx), zap.NewNop().Sugar(), "stop", "--timeout", "5", vm.id())
+		}
 
-		// Terminate the VM goroutine ("vetu pull", "vetu clone", "vetu run", etc.) via the context
-		vm.cancel()
-		vm.wg.Wait()
+		// Cancellation requests shutdown; it does not establish completion.
+		cancel()
+		wg.Wait()
 
-		// We don't return an error because we always terminate a VM
-		errCh <- nil
+		vm.stopMtx.Lock()
+		vm.ConditionsSet().Remove(v1.ConditionTypeStopping)
+		// Closing broadcasts successful completion to every Stop caller.
+		close(done)
+		vm.stopMtx.Unlock()
 	}()
 
-	return errCh
+	return done
 }
 
 func (vm *VM) Start(eventStreamer *client.EventStreamer) {
+	// The worker defers Start while Stopping is true.
+	vm.stopMtx.Lock()
+	defer vm.stopMtx.Unlock()
+	vm.stopDone = nil
+
 	vm.SetStatusMessage("Starting VM")
 	vm.ConditionsSet().Add(v1.ConditionTypeRunning)
 

internal/worker/worker.go

@@ -673,6 +673,7 @@ func (worker *Worker) reconcileRunningVM(vmResource *v1.VM, vm vmmanager.VM) {
 
 		if shouldStop {
 			vm.Stop()
+			return
 		} else {
 			vm.Suspend()
 		}