Ensure that failed starts/restarts are cleaned up properly (#1421)

There were a number of places where a workload start could fail, and not set the status of the workload to error. This identifies and fixes those issues.

This change increased the cyclomatic complexity of the restart method beyond our limit, so I refactored the function to simplify it.

Author: dmjb

pkg/workloads/manager.go

@@ -275,6 +275,10 @@ func (d *defaultManager) RunWorkloadDetached(ctx context.Context, runConfig *run
 
 	// Start the detached process
 	if err := detachedCmd.Start(); err != nil {
+		// If the start failed, we need to set the status to error before returning.
+		if err := d.statuses.SetWorkloadStatus(ctx, runConfig.BaseName, rt.WorkloadStatusError, ""); err != nil {
+			logger.Warnf("Failed to set workload %s status to error: %v", runConfig.BaseName, err)
+		}
 		return fmt.Errorf("failed to start detached process: %v", err)
 	}
 
@@ -444,79 +448,130 @@ func (d *defaultManager) RestartWorkloads(ctx context.Context, names []string, f
 
 	for _, name := range names {
 		group.Go(func() error {
-			// Create a child context with a longer timeout
-			childCtx, cancel := context.WithTimeout(context.Background(), AsyncOperationTimeout)
-			defer cancel()
+			return d.restartSingleWorkload(ctx, name, foreground)
+		})
+	}
 
-			// NOTE: Once we have the status manager implemented, we can use it
-			// to ensure that the workload exists and is in the `stopped` state
-			// before restarting.
-			var containerBaseName string
-			var running bool
-			// Try to find the container.
-			container, err := d.runtime.GetWorkloadInfo(childCtx, name)
-			if err != nil {
-				if errors.Is(err, rt.ErrWorkloadNotFound) {
-					logger.Warnf("Warning: Failed to find container: %v", err)
-					logger.Warnf("Trying to find state with name %s directly...", name)
-
-					// Try to use the provided name as the base name
-					containerBaseName = name
-					running = false
-				} else {
-					return fmt.Errorf("failed to find workload %s: %v", name, err)
-				}
-			} else {
-				// Container found, check if it's running and get the base name,
-				running = container.IsRunning()
-				containerBaseName = labels.GetContainerBaseName(container.Labels)
-			}
+	return group, nil
+}
 
-			// Check if the proxy process is running
-			proxyRunning := proxy.IsRunning(containerBaseName)
+// restartSingleWorkload handles the restart logic for a single workload
+func (d *defaultManager) restartSingleWorkload(ctx context.Context, name string, foreground bool) error {
+	// Create a child context with a longer timeout
+	childCtx, cancel := context.WithTimeout(context.Background(), AsyncOperationTimeout)
+	defer cancel()
 
-			if running && proxyRunning {
-				logger.Infof("Container %s and proxy are already running", name)
-				return nil
-			}
+	// Get workload state information
+	workloadState, err := d.getWorkloadState(childCtx, name)
+	if err != nil {
+		return err
+	}
 
-			// Load the configuration from the state store
-			// This is done synchronously since it is relatively inexpensive operation
-			// and it allows for better error handling.
-			mcpRunner, err := d.loadRunnerFromState(childCtx, containerBaseName)
-			if err != nil {
-				// TODO: If the state file has gone missing, we should delete
-				// the workload - since there is no chance of recovery.
-				return fmt.Errorf("failed to load state for %s: %v", containerBaseName, err)
-			}
+	// Check if already running
+	if d.isWorkloadAlreadyRunning(name, workloadState) {
+		return nil
+	}
 
-			// At this point we're sure that the workload exists but is not running.
-			// Transition workload to `starting` state.
-			if err := d.statuses.SetWorkloadStatus(ctx, name, rt.WorkloadStatusStarting, ""); err != nil {
-				logger.Warnf("Failed to set workload %s status to starting: %v", name, err)
-			}
-			logger.Infof("Loaded configuration from state for %s", containerBaseName)
+	// Load runner configuration from state
+	mcpRunner, err := d.loadRunnerFromState(childCtx, workloadState.BaseName)
+	if err != nil {
+		return fmt.Errorf("failed to load state for %s: %v", workloadState.BaseName, err)
+	}
 
-			// Run the tooling server inside a detached process.
-			logger.Infof("Starting tooling server %s...", name)
+	// Set workload status to starting
+	if err := d.statuses.SetWorkloadStatus(ctx, name, rt.WorkloadStatusStarting, ""); err != nil {
+		logger.Warnf("Failed to set workload %s status to starting: %v", name, err)
+	}
+	logger.Infof("Loaded configuration from state for %s", workloadState.BaseName)
 
-			// If the container is running but the proxy is not, stop the container first
-			if running { // && !proxyRunning was previously here but is implied by previous if statement.
-				logger.Infof("Container %s is running but proxy is not. Stopping container...", name)
-				if err = d.runtime.StopWorkload(childCtx, name); err != nil {
-					return fmt.Errorf("failed to stop container %s: %v", name, err)
-				}
-				logger.Infof("Container %s stopped", name)
-			}
+	// Stop container if running but proxy is not
+	if err := d.stopContainerIfNeeded(childCtx, ctx, name, workloadState); err != nil {
+		return err
+	}
 
-			if foreground {
-				return d.RunWorkload(ctx, mcpRunner.Config)
-			}
-			return d.RunWorkloadDetached(ctx, mcpRunner.Config)
-		})
+	// Start the workload
+	return d.startWorkload(ctx, name, mcpRunner, foreground)
+}
+
+// workloadState holds the current state of a workload for restart operations
+type workloadState struct {
+	BaseName     string
+	Running      bool
+	ProxyRunning bool
+}
+
+// getWorkloadState retrieves the current state of a workload
+func (d *defaultManager) getWorkloadState(ctx context.Context, name string) (*workloadState, error) {
+	workloadSt := &workloadState{}
+
+	// Try to find the container
+	container, err := d.runtime.GetWorkloadInfo(ctx, name)
+	if err != nil {
+		if errors.Is(err, rt.ErrWorkloadNotFound) {
+			logger.Warnf("Warning: Failed to find container: %v", err)
+			logger.Warnf("Trying to find state with name %s directly...", name)
+			// Try to use the provided name as the base name
+			workloadSt.BaseName = name
+			workloadSt.Running = false
+		} else {
+			return nil, fmt.Errorf("failed to find workload %s: %v", name, err)
+		}
+	} else {
+		// Container found, check if it's running and get the base name
+		workloadSt.Running = container.IsRunning()
+		workloadSt.BaseName = labels.GetContainerBaseName(container.Labels)
 	}
 
-	return group, nil
+	// Check if the proxy process is running
+	workloadSt.ProxyRunning = proxy.IsRunning(workloadSt.BaseName)
+
+	return workloadSt, nil
+}
+
+// isWorkloadAlreadyRunning checks if the workload is already fully running
+func (*defaultManager) isWorkloadAlreadyRunning(name string, workloadSt *workloadState) bool {
+	if workloadSt.Running && workloadSt.ProxyRunning {
+		logger.Infof("Container %s and proxy are already running", name)
+		return true
+	}
+	return false
+}
+
+// stopContainerIfNeeded stops the container if it's running but proxy is not
+func (d *defaultManager) stopContainerIfNeeded(childCtx, ctx context.Context, name string, workloadSt *workloadState) error {
+	if !workloadSt.Running {
+		return nil
+	}
+
+	logger.Infof("Container %s is running but proxy is not. Stopping container...", name)
+	if err := d.runtime.StopWorkload(childCtx, name); err != nil {
+		if statusErr := d.statuses.SetWorkloadStatus(ctx, name, rt.WorkloadStatusError, ""); statusErr != nil {
+			logger.Warnf("Failed to set workload %s status to error: %v", name, statusErr)
+		}
+		return fmt.Errorf("failed to stop container %s: %v", name, err)
+	}
+	logger.Infof("Container %s stopped", name)
+	return nil
+}
+
+// startWorkload starts the workload in either foreground or background mode
+func (d *defaultManager) startWorkload(ctx context.Context, name string, mcpRunner *runner.Runner, foreground bool) error {
+	logger.Infof("Starting tooling server %s...", name)
+
+	var err error
+	if foreground {
+		err = d.RunWorkload(ctx, mcpRunner.Config)
+	} else {
+		err = d.RunWorkloadDetached(ctx, mcpRunner.Config)
+	}
+
+	if err != nil {
+		// If we could not start the workload, set the status to error before returning
+		if statusErr := d.statuses.SetWorkloadStatus(ctx, name, rt.WorkloadStatusError, ""); statusErr != nil {
+			logger.Warnf("Failed to set workload %s status to error: %v", name, statusErr)
+		}
+	}
+	return err
 }
 
 // TODO: Move to dedicated config management interface.