Worker: introduce --experimental-vm-stop-timeout (#505)

internal/command/worker/run.go

@@ -48,6 +48,7 @@ var defaultMemory uint64
 var username string
 var addressPprof string
 var debug bool
+var experimentalVMStopTimeout uint16
 
 // Hidden flags
 var synthetic bool
@@ -88,6 +89,9 @@ func newRunCommand() *cobra.Command {
 	cmd.Flags().StringVar(&addressPprof, "listen-pprof", "",
 		"start pprof HTTP server on localhost:6060 for diagnostic purposes (e.g. \"localhost:6060\")")
 	cmd.Flags().BoolVar(&debug, "debug", false, "enable debug logging")
+	cmd.Flags().Uint16Var(&experimentalVMStopTimeout, "experimental-vm-stop-timeout",
+		worker.DefaultVMStopTimeoutSeconds,
+		"seconds to wait for graceful shutdown of Tart and Vetu VMs before forceful termination")
 
 	// Hidden flags
 	cmd.Flags().BoolVar(&synthetic, "synthetic", false,
@@ -107,6 +111,7 @@ func runWorker(cmd *cobra.Command, args []string) (err error) {
 		worker.WithName(name),
 		worker.WithLabels(labels),
 		worker.WithDefaultCPUAndMemory(defaultCPU, defaultMemory),
+		worker.WithVMStopTimeout(experimentalVMStopTimeout),
 	}
 
 	// Run the macOS "Local Network" permission helper

internal/worker/option.go

@@ -59,6 +59,12 @@ func WithSoftnetPolicyUpdates(enabled bool) Option {
 	}
 }
 
+func WithVMStopTimeout(seconds uint16) Option {
+	return func(worker *Worker) {
+		worker.vmStopTimeoutSeconds = seconds
+	}
+}
+
 func WithLogger(logger *zap.Logger) Option {
 	return func(worker *Worker) {
 		worker.logger = logger.Sugar()

internal/worker/runtime/tart.go

@@ -12,10 +12,12 @@ import (
 	"go.uber.org/zap"
 )
 
-type Tart struct{}
+type Tart struct {
+	stopTimeoutSeconds uint16
+}
 
-func NewTart() *Tart {
+func NewTart(stopTimeoutSeconds uint16) *Tart {
-	return &Tart{}
+	return &Tart{stopTimeoutSeconds: stopTimeoutSeconds}
 }
 
 func (tart *Tart) ID() v1.Runtime {
@@ -34,7 +36,8 @@ func (tart *Tart) NewVM(
 	softnetPolicyUpdates bool,
 	logger *zap.SugaredLogger,
 ) vmmanager.VM {
-	return tartpkg.NewVM(vmResource, eventStreamer, vmPullTimeHistogram, dialer, softnetPolicyUpdates, logger)
+	return tartpkg.NewVM(vmResource, eventStreamer, vmPullTimeHistogram, dialer,
+		softnetPolicyUpdates, logger, tart.stopTimeoutSeconds)
 }
 
 func (tart *Tart) ListVMs(ctx context.Context, logger *zap.SugaredLogger) ([]vmmanager.VMInfo, error) {

internal/worker/runtime/vetu.go

@@ -12,10 +12,12 @@ import (
 	"go.uber.org/zap"
 )
 
-type Vetu struct{}
+type Vetu struct {
+	stopTimeoutSeconds uint16
+}
 
-func NewVetu() *Vetu {
+func NewVetu(stopTimeoutSeconds uint16) *Vetu {
-	return &Vetu{}
+	return &Vetu{stopTimeoutSeconds: stopTimeoutSeconds}
 }
 
 func (vetu *Vetu) ID() v1.Runtime {
@@ -34,7 +36,7 @@ func (vetu *Vetu) NewVM(
 	_ bool,
 	logger *zap.SugaredLogger,
 ) vmmanager.VM {
-	return vetupkg.NewVM(vmResource, eventStreamer, vmPullTimeHistogram, dialer, logger)
+	return vetupkg.NewVM(vmResource, eventStreamer, vmPullTimeHistogram, dialer, logger, vetu.stopTimeoutSeconds)
 }
 
 func (vetu *Vetu) ListVMs(ctx context.Context, logger *zap.SugaredLogger) ([]vmmanager.VMInfo, error) {

internal/worker/vmmanager/stop_test.go

@@ -26,10 +26,10 @@ func TestStopWaitsForCommandAndRun(t *testing.T) {
 		new  func(v1.VM) vmmanager.VM
 	}{
 		{name: "tart", new: func(resource v1.VM) vmmanager.VM {
-			return tart.NewVM(resource, nil, nil, nil, false, zap.NewNop().Sugar())
+			return tart.NewVM(resource, nil, nil, nil, false, zap.NewNop().Sugar(), 5)
 		}},
 		{name: "vetu", new: func(resource v1.VM) vmmanager.VM {
-			return vetu.NewVM(resource, nil, nil, nil, zap.NewNop().Sugar())
+			return vetu.NewVM(resource, nil, nil, nil, zap.NewNop().Sugar(), 5)
 		}},
 	}
 

internal/worker/vmmanager/tart/tart.go

@@ -38,8 +38,9 @@ type VM struct {
 
 	wg *sync.WaitGroup
 
-	stopMtx  sync.Mutex
+	stopMtx            sync.Mutex
-	stopDone chan error
+	stopDone           chan error
+	stopTimeoutSeconds uint16
 
 	dialer dialer.Dialer
 
@@ -57,6 +58,7 @@ func NewVM(
 	dialer dialer.Dialer,
 	softnetPolicyUpdates bool,
 	logger *zap.SugaredLogger,
+	stopTimeoutSeconds uint16,
 ) *VM {
 	vmContext, vmContextCancel := context.WithCancel(context.Background())
 
@@ -75,7 +77,8 @@ func NewVM(
 		ctx:    vmContext,
 		cancel: vmContextCancel,
 
-		wg: &sync.WaitGroup{},
+		wg:                 &sync.WaitGroup{},
+		stopTimeoutSeconds: stopTimeoutSeconds,
 
 		dialer:               dialer,
 		softnetPolicyUpdates: softnetPolicyUpdates,
@@ -490,7 +493,8 @@ func (vm *VM) Stop() <-chan error {
 	go func() {
 		if ctx.Err() == nil {
 			// Try to gracefully terminate the VM.
-			_, _, _ = Tart(context.WithoutCancel(ctx), zap.NewNop().Sugar(), "stop", "--timeout", "5", vm.id())
+			_, _, _ = Tart(context.WithoutCancel(ctx), zap.NewNop().Sugar(), "stop", "--timeout",
+				strconv.FormatUint(uint64(vm.stopTimeoutSeconds), 10), vm.id())
 		}
 
 		// Cancellation requests shutdown; it does not establish completion.

internal/worker/vmmanager/vetu/vetu.go

@@ -33,8 +33,9 @@ type VM struct {
 
 	wg *sync.WaitGroup
 
-	stopMtx  sync.Mutex
+	stopMtx            sync.Mutex
-	stopDone chan error
+	stopDone           chan error
+	stopTimeoutSeconds uint16
 
 	dialer dialer.Dialer
 
@@ -47,6 +48,7 @@ func NewVM(
 	vmPullTimeHistogram metric.Float64Histogram,
 	dialer dialer.Dialer,
 	logger *zap.SugaredLogger,
+	stopTimeoutSeconds uint16,
 ) *VM {
 	vmContext, vmContextCancel := context.WithCancel(context.Background())
 
@@ -62,7 +64,8 @@ func NewVM(
 		ctx:    vmContext,
 		cancel: vmContextCancel,
 
-		wg: &sync.WaitGroup{},
+		wg:                 &sync.WaitGroup{},
+		stopTimeoutSeconds: stopTimeoutSeconds,
 
 		dialer: dialer,
 
@@ -281,7 +284,8 @@ func (vm *VM) Stop() <-chan error {
 	go func() {
 		if ctx.Err() == nil {
 			// Try to gracefully terminate the VM.
-			_, _, _ = Vetu(context.WithoutCancel(ctx), zap.NewNop().Sugar(), "stop", "--timeout", "5", vm.id())
+			_, _, _ = Vetu(context.WithoutCancel(ctx), zap.NewNop().Sugar(), "stop", "--timeout",
+				strconv.FormatUint(uint64(vm.stopTimeoutSeconds), 10), vm.id())
 		}
 
 		// Cancellation requests shutdown; it does not establish completion.

internal/worker/worker.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"os"
 	"slices"
+	"strconv"
 	"time"
 
 	goruntime "runtime"
@@ -48,6 +49,8 @@ const (
 	onDiskVMSyncTimeout          = 30 * time.Second
 
 	tartVersionSoftnetPolicyUpdates = "2.34.0"
+
+	DefaultVMStopTimeoutSeconds = 5
 )
 
 var (
@@ -72,6 +75,7 @@ type Worker struct {
 	runtime runtime.Runtime
 
 	softnetPolicyUpdates mo.Option[bool]
+	vmStopTimeoutSeconds uint16
 
 	vmPullTimeHistogram metric.Float64Histogram
 
@@ -82,11 +86,12 @@ type Worker struct {
 
 func New(client *client.Client, opts ...Option) (*Worker, error) {
 	worker := &Worker{
-		client:        client,
+		client:               client,
-		pollTicker:    time.NewTicker(pollInterval),
+		pollTicker:           time.NewTicker(pollInterval),
-		recoveredVMs:  make(map[ondiskname.OnDiskName]time.Time),
+		recoveredVMs:         make(map[ondiskname.OnDiskName]time.Time),
-		vmm:           vmmanager.New(),
+		vmm:                  vmmanager.New(),
-		syncRequested: make(chan bool, 1),
+		syncRequested:        make(chan bool, 1),
+		vmStopTimeoutSeconds: DefaultVMStopTimeoutSeconds,
 	}
 
 	// Apply options
@@ -110,9 +115,9 @@ func New(client *client.Client, opts ...Option) (*Worker, error) {
 
 	if worker.runtime == nil {
 		if goruntime.GOOS == "linux" {
-			worker.runtime = runtime.NewVetu()
+			worker.runtime = runtime.NewVetu(worker.vmStopTimeoutSeconds)
 		} else {
-			worker.runtime = runtime.NewTart()
+			worker.runtime = runtime.NewTart(worker.vmStopTimeoutSeconds)
 		}
 	}
 
@@ -920,7 +925,8 @@ func (worker *Worker) syncOnDiskVMsWithInventory(ctx context.Context, vmInfos []
 			// On-disk VM doesn't exist on the controller nor in the Worker's VM manager,
 			// stop it (if applicable) and delete it
 			if vmInfo.Running {
-				_, _, err := worker.runtime.Cmd(ctx, worker.logger, "stop", vmInfo.Name)
+				_, _, err := worker.runtime.Cmd(ctx, worker.logger, "stop", "--timeout",
+					strconv.FormatUint(uint64(worker.vmStopTimeoutSeconds), 10), vmInfo.Name)
 				if err != nil {
 					worker.logger.Warnf("failed to stop")
 				}
@@ -935,7 +941,8 @@ func (worker *Worker) syncOnDiskVMsWithInventory(ctx context.Context, vmInfos []
 			// but we've lost track of it, so shut it down (if applicable)
 			// and report the error (if not failed yet)
 			if vmInfo.Running {
-				_, _, err := worker.runtime.Cmd(ctx, worker.logger, "stop", vmInfo.Name)
+				_, _, err := worker.runtime.Cmd(ctx, worker.logger, "stop", "--timeout",
+					strconv.FormatUint(uint64(worker.vmStopTimeoutSeconds), 10), vmInfo.Name)
 				if err != nil {
 					worker.logger.Warnf("failed to stop")
 				}

internal/worker/worker_test.go

@@ -232,7 +232,7 @@ func newWorkerWithFakeTart(
 		vmm:           vmmanager.New(),
 		pollTicker:    pollTicker,
 		syncRequested: make(chan bool, 1),
-		runtime:       runtime.NewTart(),
+		runtime:       runtime.NewTart(DefaultVMStopTimeoutSeconds),
 		logger:        zap.NewNop().Sugar(),
 	}
 }