Add a runsc wait --checkpoint n command to wait for a checkpoint to complete.

This command waits for (n-1)th checkpoint to complete successfully. Then waits
for the next checkpoint attempt (which would increment checkpoint count to n)
and returns its status.

If sandbox checkpoint count has already reached n, it returns immediately.

PiperOrigin-RevId: 651884599

pkg/sentry/kernel/kernel.go

@@ -362,6 +362,10 @@ type Kernel struct {
 	// checkpointMu is used to protect the checkpointing related fields below.
 	checkpointMu sync.Mutex `state:"nosave"`
 
+	// checkpointCond is used to wait for a checkpoint to complete. It uses
+	// checkpointMu as its mutex.
+	checkpointCond sync.Cond `state:"nosave"`
+
 	// additionalCheckpointState stores additional state that needs
 	// to be checkpointed. It's protected by checkpointMu.
 	additionalCheckpointState map[any]any
@@ -370,9 +374,19 @@ type Kernel struct {
 	// asynchronous checkpointing. It's protected by checkpointMu.
 	saver Saver `state:"nosave"`
 
-	// checkpointCounter is the number of times the kernel has been checkpointed.
-	// It's protected by checkpointMu.
+	// checkpointCounter aims to track the number of times the kernel has been
+	// successfully checkpointed. It's updated via calls to OnCheckpointAttempt()
+	// and IncCheckpointCount(). Kernel checkpoint-ers must call these methods
+	// appropriately so the counter is accurate. It's protected by checkpointMu.
 	checkpointCounter uint32
+
+	// lastCheckpointStatus is the error value returned from the most recent
+	// checkpoint attempt. If this value is nil, then the `checkpointCounter`-th
+	// checkpoint attempt succeeded and no checkpoint attempt has completed since.
+	// If this value is non-nil, then the `checkpointCounter`-th checkpoint
+	// attempt succeeded, after which at least one more checkpoint attempt was
+	// made and failed with this error. It's protected by checkpointMu.
+	lastCheckpointStatus error `state:"nosave"`
 }
 
 // Saver is an interface for saving the kernel.
@@ -461,6 +475,7 @@ func (k *Kernel) Init(args InitKernelArgs) error {
 		k.rootNetworkNamespace = inet.NewRootNamespace(nil, nil, args.RootUserNamespace)
 	}
 	k.runningTasksCond.L = &k.runningTasksMu
+	k.checkpointCond.L = &k.checkpointMu
 	k.cpuClockTickerWakeCh = make(chan struct{}, 1)
 	k.cpuClockTickerStopCond.L = &k.runningTasksMu
 	k.applicationCores = args.ApplicationCores
@@ -745,6 +760,7 @@ func (k *Kernel) LoadFrom(ctx context.Context, r io.Reader, pagesMetadata, pages
 	}
 
 	k.runningTasksCond.L = &k.runningTasksMu
+	k.checkpointCond.L = &k.checkpointMu
 	k.cpuClockTickerWakeCh = make(chan struct{}, 1)
 	k.cpuClockTickerStopCond.L = &k.runningTasksMu
 
@@ -2139,3 +2155,48 @@ func (k *Kernel) CheckpointCount() uint32 {
 	defer k.checkpointMu.Unlock()
 	return k.checkpointCounter
 }
+
+// OnCheckpointAttempt is called when a checkpoint attempt is completed. err is
+// any checkpoint errors that may have occurred.
+func (k *Kernel) OnCheckpointAttempt(err error) {
+	k.checkpointMu.Lock()
+	defer k.checkpointMu.Unlock()
+	if err == nil {
+		k.checkpointCounter++
+	}
+	k.lastCheckpointStatus = err
+	k.checkpointCond.Broadcast()
+}
+
+// ResetCheckpointStatus resets the last checkpoint status, indicating a new
+// checkpoint is in progress. Caller must call OnCheckpointAttempt when the
+// checkpoint attempt is completed.
+func (k *Kernel) ResetCheckpointStatus() {
+	k.checkpointMu.Lock()
+	defer k.checkpointMu.Unlock()
+	k.lastCheckpointStatus = nil
+}
+
+// WaitCheckpoint waits for the Kernel to have been successfully checkpointed
+// n-1 times, then waits for either the n-th successful checkpoint (in which
+// case it returns nil) or any number of failed checkpoints (in which case it
+// returns an error returned by any such failure).
+func (k *Kernel) WaitCheckpoint(n uint32) error {
+	if n == 0 {
+		return nil
+	}
+	k.checkpointMu.Lock()
+	defer k.checkpointMu.Unlock()
+	if k.checkpointCounter >= n {
+		// n-th checkpoint already completed successfully.
+		return nil
+	}
+	for k.checkpointCounter < n {
+		if k.checkpointCounter == n-1 && k.lastCheckpointStatus != nil {
+			// n-th checkpoint was attempted but it had failed.
+			return k.lastCheckpointStatus
+		}
+		k.checkpointCond.Wait()
+	}
+	return nil
+}

runsc/boot/controller.go

@@ -92,6 +92,13 @@ const (
 	// return its ExitStatus.
 	ContMgrWaitPID = "containerManager.WaitPID"
 
+	// ContMgrWaitCheckpoint waits for the Kernel to have been successfully
+	// checkpointed n-1 times, then waits for either the n-th successful
+	// checkpoint (in which case it returns nil) or any number of failed
+	// checkpoints (in which case it returns an error returned by any such
+	// failure).
+	ContMgrWaitCheckpoint = "containerManager.WaitCheckpoint"
+
 	// ContMgrRootContainerStart starts a new sandbox with a root container.
 	ContMgrRootContainerStart = "containerManager.StartRoot"
 
@@ -692,6 +699,16 @@ func (cm *containerManager) WaitPID(args *WaitPIDArgs, waitStatus *uint32) error
 	return err
 }
 
+// WaitCheckpoint waits for the Kernel to have been successfully checkpointed
+// n-1 times, then waits for either the n-th successful checkpoint (in which
+// case it returns nil) or any number of failed checkpoints (in which case it
+// returns an error returned by any such failure).
+func (cm *containerManager) WaitCheckpoint(n *uint32, _ *struct{}) error {
+	err := cm.l.k.WaitCheckpoint(*n)
+	log.Debugf("containerManager.WaitCheckpoint, n = %d, err = %v", *n, err)
+	return err
+}
+
 // SignalDeliveryMode enumerates different signal delivery modes.
 type SignalDeliveryMode int
 

runsc/boot/restore.go

@@ -307,7 +307,13 @@ func (r *restorer) restore(l *Loader) error {
 	return nil
 }
 
-func (l *Loader) save(o *control.SaveOpts) error {
+func (l *Loader) save(o *control.SaveOpts) (err error) {
+	defer func() {
+		// This closure is required to capture the final value of err.
+		l.k.OnCheckpointAttempt(err)
+	}()
+	l.k.ResetCheckpointStatus()
+
 	// TODO(gvisor.dev/issues/6243): save/restore not supported w/ hostinet
 	if l.root.conf.Network == config.NetworkHost {
 		return errors.New("checkpoint not supported when using hostinet")
@@ -334,7 +340,6 @@ func (l *Loader) save(o *control.SaveOpts) error {
 		if err := postResumeImpl(l.k); err != nil {
 			return err
 		}
-		l.k.IncCheckpointCount()
 	}
 	return nil
 }

runsc/cmd/wait.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/google/subcommands"
 	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/log"
 	"gvisor.dev/gvisor/runsc/cmd/util"
 	"gvisor.dev/gvisor/runsc/config"
 	"gvisor.dev/gvisor/runsc/container"
@@ -33,8 +34,9 @@ const (
 
 // Wait implements subcommands.Command for the "wait" command.
 type Wait struct {
-	rootPID int
-	pid     int
+	rootPID    int
+	pid        int
+	checkpoint uint
 }
 
 // Name implements subcommands.Command.Name.
@@ -56,6 +58,7 @@ func (*Wait) Usage() string {
 func (wt *Wait) SetFlags(f *flag.FlagSet) {
 	f.IntVar(&wt.rootPID, "rootpid", unsetPID, "select a PID in the sandbox root PID namespace to wait on instead of the container's root process")
 	f.IntVar(&wt.pid, "pid", unsetPID, "select a PID in the container's PID namespace to wait on instead of the container's root process")
+	f.UintVar(&wt.checkpoint, "checkpoint", 0, "wait for (n-1)th checkpoint to complete successfully, then waits for the next checkpoint attempt and returns its status. When set to 0, it disables checkpoint waiting.")
 }
 
 // Execute implements subcommands.Command.Execute. It waits for a process in a
@@ -78,6 +81,16 @@ func (wt *Wait) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomm
 		util.Fatalf("loading container: %v", err)
 	}
 
+	if wt.checkpoint > 0 {
+		if wt.rootPID != unsetPID || wt.pid != unsetPID {
+			log.Warningf("waiting for checkpoint to complete, ignoring -pid and -rootpid")
+		}
+		if err := c.WaitCheckpoint(uint32(wt.checkpoint)); err != nil {
+			util.Fatalf("waiting for %d-th checkpoint to complete: %v", wt.checkpoint, err)
+		}
+		return subcommands.ExitSuccess
+	}
+
 	var waitStatus unix.WaitStatus
 	switch {
 	// Wait on the whole container.

runsc/container/container.go

@@ -653,6 +653,18 @@ func (c *Container) WaitPID(pid int32) (unix.WaitStatus, error) {
 	return c.Sandbox.WaitPID(c.ID, pid)
 }
 
+// WaitCheckpoint waits for the Kernel to have been successfully checkpointed
+// n-1 times, then waits for either the n-th successful checkpoint (in which
+// case it returns nil) or any number of failed checkpoints (in which case it
+// returns an error returned by any such failure).
+func (c *Container) WaitCheckpoint(n uint32) error {
+	log.Debugf("Wait on %d-th checkpoint to complete in container, cid: %s", n, c.ID)
+	if !c.IsSandboxRunning() {
+		return fmt.Errorf("sandbox is not running")
+	}
+	return c.Sandbox.WaitCheckpoint(n)
+}
+
 // SignalContainer sends the signal to the container. If all is true and signal
 // is SIGKILL, then waits for all processes to exit before returning.
 // SignalContainer returns an error if the container is already stopped.

runsc/container/multi_container_test.go

@@ -2760,11 +2760,26 @@ func testMultiContainerCheckpointRestore(t *testing.T, conf *config.Config, comp
 		t.Fatalf("Failed to wait for output file: %v", err)
 	}
 
+	checkpointWaiter := make(chan struct{}, 1)
+	go func() {
+		// WaitCheckpoint on the second container.
+		if err := conts[1].WaitCheckpoint(1); err != nil {
+			t.Errorf("error waiting for checkpoint to complete: %v", err)
+		}
+		checkpointWaiter <- struct{}{}
+	}()
+
 	// Checkpoint root container; save state into new file.
 	if err := conts[0].Checkpoint(dir, false /* direct */, statefile.Options{Compression: compression}, pgalloc.SaveOpts{}); err != nil {
 		t.Fatalf("error checkpointing container to empty file: %v", err)
 	}
-	defer os.RemoveAll(dir)
+
+	// WaitCheckpoint() should return after checkpoint is complete.
+	select {
+	case <-checkpointWaiter:
+	case <-time.After(10 * time.Second):
+		t.Errorf("timed out waiting for checkpoint to complete")
+	}
 
 	lastNum, err := readOutputNum(outputPath, -1)
 	if err != nil {

runsc/sandbox/sandbox.go

@@ -1273,6 +1273,15 @@ func (s *Sandbox) WaitPID(cid string, pid int32) (unix.WaitStatus, error) {
 	return ws, nil
 }
 
+// WaitCheckpoint waits for the Kernel to have been successfully checkpointed
+// n-1 times, then waits for either the n-th successful checkpoint (in which
+// case it returns nil) or any number of failed checkpoints (in which case it
+// returns an error returned by any such failure).
+func (s *Sandbox) WaitCheckpoint(n uint32) error {
+	log.Debugf("Waiting for %d-th checkpoint to complete in sandbox %q", n, s.ID)
+	return s.call(boot.ContMgrWaitCheckpoint, &n, nil)
+}
+
 // IsRootContainer returns true if the specified container ID belongs to the
 // root container.
 func (s *Sandbox) IsRootContainer(cid string) bool {