Handle restore failure more robustly.

- This change most notably adds a new loaderState "restoreFailed". As of
  writing, this is used to classify restore failures which occur after the last
  container's restore command is received. See future work section.
- `runsc wait --restore` now returns the restore error, if one occurred.
- Added more explicit checks for all values of loaderState and error out when
  an unexpected state is reached.
- `restorer` has been simplified to save a pointer to `cm` instead of using
  callback functions, which are harder to follow (and step-into).

Future work:
We still do not handle failures that occur when intermediate container's
restore command fails. Or if we receive an unexpected "start" command instead
of "restore". In such scenarios, it is unlikely that we will recover to have
a successful restore as it suggests some issue with the higher level tooling
issuing these commands.

PiperOrigin-RevId: 831757593

Author: ayushr2

runsc/boot/controller.go

@@ -584,12 +584,11 @@ func (cm *containerManager) Restore(o *RestoreOpts, _ *struct{}) error {
 	}
 
 	cm.restorer = &restorer{
-		readyToStart:  cm.onStart,
-		onRestoreDone: cm.onRestoreDone,
-		stateFile:     reader,
-		background:    o.Background,
-		timer:         timer,
-		mainMF:        mf,
+		cm:         cm,
+		stateFile:  reader,
+		background: o.Background,
+		timer:      timer,
+		mainMF:     mf,
 	}
 	cm.l.restoreDone = sync.NewCond(&cm.l.mu)
 	cm.l.state = restoringUnstarted
@@ -687,6 +686,15 @@ func getRestoreReadersForLocalCheckpointFiles(o *RestoreOpts) (io.ReadCloser, io
 		nil
 }
 
+func (cm *containerManager) onRestoreFailed(err error) {
+	cm.l.mu.Lock()
+	cm.l.state = restoreFailed
+	cm.l.restoreErr = err
+	cm.l.mu.Unlock()
+	cm.l.restoreDone.Broadcast()
+	cm.restorer = nil
+}
+
 func (cm *containerManager) onRestoreDone() {
 	cm.l.mu.Lock()
 	cm.l.state = restored

runsc/boot/loader.go

@@ -167,14 +167,17 @@ type loaderState int
 const (
 	// created indicates that the Loader has been created, but not started yet.
 	created loaderState = iota
-	// started indicates that the Loader has been started.
+	// started indicates that the Loader has been started. This means that the
+	// root container is running. Subsequent containers may still be unstarted.
 	started
 	// restoringUnstarted indicates that the Loader has been created and is
 	// restoring containers, but not started yet.
 	restoringUnstarted
-	// restoringStarted indicates that the Loader has been created and started,
-	// while restore continues in the background.
+	// restoringStarted indicates that the Loader has been created and started
+	// along with all containers, while restore continues in the background.
 	restoringStarted
+	// restoreFailed indicates that the Loader has failed to restore.
+	restoreFailed
 	// restored indicates that the Loader has been fully restored.
 	restored
 )
@@ -190,6 +193,8 @@ func (s loaderState) String() string {
 		return "restoringUnstarted"
 	case restoringStarted:
 		return "restoringStarted"
+	case restoreFailed:
+		return "restoreFailed"
 	case restored:
 		return "restored"
 	default:
@@ -283,6 +288,11 @@ type Loader struct {
 	// during restore.
 	saveRestoreNet bool
 
+	// restoreErr is the error that occurred during restore.
+	//
+	// +checklocks:mu
+	restoreErr error
+
 	LoaderExtra
 }
 
@@ -1007,8 +1017,8 @@ func (l *Loader) run() error {
 		return fmt.Errorf("trying to start deleted container %q", l.sandboxID)
 	}
 
-	// If we are restoring, we do not want to create a process.
-	if l.state != restoringUnstarted {
+	switch l.state {
+	case created:
 		if l.root.conf.ProfileEnable {
 			pprof.Initialize()
 		}
@@ -1049,6 +1059,10 @@ func (l *Loader) run() error {
 				return c.ContainerStart(context.Background(), fields, &evt)
 			})
 		}
+	case restoringUnstarted:
+		// If we are restoring, we do not want to create a process.
+	default:
+		return fmt.Errorf("Loader.Run() called in unexpected state=%s", l.state)
 	}
 
 	ep.tg = l.k.GlobalInit()
@@ -1083,10 +1097,13 @@ func (l *Loader) run() error {
 	if err := l.k.Start(); err != nil {
 		return err
 	}
-	if l.state == restoringUnstarted {
-		l.state = restoringStarted
-	} else {
+	switch l.state {
+	case created:
 		l.state = started
+	case restoringUnstarted:
+		l.state = restoringStarted
+	default:
+		panic(fmt.Sprintf("state=%s in Loader.run() should be impossible", l.state))
 	}
 	return nil
 }
@@ -1478,29 +1495,43 @@ func (l *Loader) executeAsync(args *control.ExecArgs) (kernel.ThreadID, error) {
 
 // waitContainer waits for the init process of a container to exit.
 func (l *Loader) waitContainer(cid string, waitStatus *uint32) error {
-	// Don't defer unlock, as doing so would make it impossible for
-	// multiple clients to wait on the same container.
-	key := execID{cid: cid}
-	tg, err := l.threadGroupFromID(key)
-	if err != nil {
-		l.mu.Lock()
-		// Extra handling is needed if the restoring container has not started yet.
-		if l.state != restoringUnstarted {
-			l.mu.Unlock()
-			return err
-		}
-		// Container could be restoring, first check if container exists.
-		if _, err := l.findProcessLocked(key); err != nil {
-			l.mu.Unlock()
-			return err
-		}
+	l.mu.Lock()
+	state := l.state
+	if state == restoringUnstarted {
 		log.Infof("Waiting for the container to restore, CID: %q", cid)
 		l.restoreDone.Wait()
 		l.mu.Unlock()
-
 		log.Infof("Restore is completed, trying to wait for container %q again.", cid)
 		return l.waitContainer(cid, waitStatus)
 	}
+	tg, err := l.tryThreadGroupFromIDLocked(execID{cid: cid})
+	l.mu.Unlock()
+	if err != nil {
+		// The container does not exist.
+		return err
+	}
+	if tg == nil {
+		// The container has not been started.
+		switch state {
+		case created, started:
+			// Note that state=started means the root container has been started,
+			// but other containers may not have started yet.
+			return fmt.Errorf("container %q not started", cid)
+		case restoringStarted, restored:
+			// The container has restored, we *should* have found the init process...
+			return fmt.Errorf("could not find init process of restored container %q in state %q", cid, state)
+		case restoreFailed:
+			// If restore failed, we should return the a non-zero exit status here to
+			// indicate that the container failed and transition to "stopped" state.
+			log.Warningf("Restore failed, returning from waitContainer with non-zero exit status")
+			*waitStatus = 1
+			return nil
+		case restoringUnstarted:
+			panic("impossible")
+		default:
+			panic(fmt.Sprintf("Invalid state: %s", state))
+		}
+	}
 
 	// If the thread either has already exited or exits during waiting,
 	// consider the container exited.
@@ -1520,15 +1551,15 @@ func (l *Loader) waitContainer(cid string, waitStatus *uint32) error {
 func (l *Loader) waitRestore() error {
 	l.mu.Lock()
 	defer l.mu.Unlock()
-	if l.state == restored {
-		return nil
+	if l.state == restored || l.state == restoreFailed {
+		return l.restoreErr
 	}
 	if l.state != restoringUnstarted && l.state != restoringStarted {
 		return fmt.Errorf("sandbox is not being restored, cannot wait for restore: state=%s", l.state)
 	}
 	log.Infof("Waiting for the sandbox to restore")
 	l.restoreDone.Wait()
-	return nil
+	return l.restoreErr
 }
 
 func (l *Loader) waitPID(tgid kernel.ThreadID, cid string, waitStatus *uint32) error {
@@ -2097,7 +2128,10 @@ func (l *Loader) containerRuntimeState(cid string) ContainerRuntimeState {
 		return RuntimeStateStopped
 	}
 	if exec.tg == nil {
-		// Container has no thread group assigned, so it has started yet.
+		if l.state == restoreFailed {
+			return RuntimeStateStopped
+		}
+		// Container has no thread group assigned, so it has not started yet.
 		return RuntimeStateCreating
 	}
 	if exec.tg.Leader().ExitState() == kernel.TaskExitNone {

runsc/boot/restore.go

@@ -98,11 +98,8 @@ type restorer struct {
 	// deviceFile is the required to start the platform.
 	deviceFile *fd.FD
 
-	// readyToStart is a callback triggered when the sandbox is ready to start.
-	readyToStart func() error
-
-	// onRestoreDone is a callback triggered when the restore is done.
-	onRestoreDone func()
+	// cm is the container manager that is used to restore the sandbox.
+	cm *containerManager
 
 	// checkpointedSpecs contains the map of container specs used during
 	// checkpoint.
@@ -160,13 +157,11 @@ func (r *restorer) restoreContainerInfo(l *Loader, info *containerInfo, containe
 	// Non-container-specific restore work:
 
 	if len(r.containers) == r.totalContainers {
-		if err := specutils.RestoreValidateSpec(r.checkpointedSpecs, l.GetContainerSpecs(), l.root.conf); err != nil {
-			return fmt.Errorf("failed to handle restore spec validation: %w", err)
-		}
-		r.timer.Reached("specs validated")
-
 		// Trigger the restore if this is the last container.
-		return r.restore(l)
+		if err := r.restore(l); err != nil {
+			r.cm.onRestoreFailed(err)
+			return err
+		}
 	}
 	return nil
 }
@@ -183,6 +178,13 @@ func createNetworkStackForRestore(l *Loader) (*stack.Stack, inet.Stack) {
 func (r *restorer) restore(l *Loader) error {
 	log.Infof("Starting to restore %d containers", len(r.containers))
 
+	// Validate the container specs to ensure they did not meaningfully change
+	// between checkpoint and restore.
+	if err := specutils.RestoreValidateSpec(r.checkpointedSpecs, l.GetContainerSpecs(), l.root.conf); err != nil {
+		return fmt.Errorf("failed to handle restore spec validation: %w", err)
+	}
+	r.timer.Reached("specs validated")
+
 	// Create a new root network namespace with the network stack of the
 	// old kernel to preserve the existing network configuration.
 	oldStack, oldInetStack := createNetworkStackForRestore(l)
@@ -353,7 +355,7 @@ func (r *restorer) restore(l *Loader) error {
 	cu.Clean()
 
 	r.timer.Reached("Starting sandbox")
-	if err := r.readyToStart(); err != nil {
+	if err := r.cm.onStart(); err != nil {
 		return fmt.Errorf("restorer.readyToStart callback failed: %w", err)
 	}
 
@@ -386,7 +388,7 @@ func (r *restorer) restore(l *Loader) error {
 			}
 		}
 
-		r.onRestoreDone()
+		r.cm.onRestoreDone()
 		postRestoreThread.Reached("kernel notified")
 
 		log.Infof("Restore successful")