Remove checkpoint_count from runsc wait --checkpoint

This is done because external callers are not able to know
the snapshot generation number from the outside.

PiperOrigin-RevId: 707979556

pkg/sentry/fsimpl/proc/BUILD

@@ -119,6 +119,7 @@ go_library(
         "//pkg/tcpip/header",
         "//pkg/tcpip/network/ipv4",
         "//pkg/usermem",
+        "//pkg/waiter",
     ],
 )
 

pkg/sentry/kernel/kernel.go

@@ -352,10 +352,6 @@ 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
@@ -364,19 +360,13 @@ type Kernel struct {
 	// asynchronous checkpointing. It's protected by checkpointMu.
 	saver Saver `state:"nosave"`
 
-	// 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
+	// CheckpointWait is used to wait for a checkpoint to complete.
+	CheckpointWait CheckpointWaitable
 
-	// 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"`
+	// checkpointGen aims to track the number of times the kernel has been
+	// successfully checkpointed. Callers of checkpoint must notify the kernel
+	// when checkpoint/restore are done. It's protected by checkpointMu.
+	checkpointGen CheckpointGeneration
 
 	// UnixSocketOpts stores configuration options for management of unix sockets.
 	UnixSocketOpts transport.UnixSocketOpts
@@ -466,7 +456,6 @@ 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
@@ -495,6 +484,7 @@ func (k *Kernel) Init(args InitKernelArgs) error {
 	}
 	k.MaxFDLimit.Store(args.MaxFDLimit)
 	k.containerNames = make(map[string]string)
+	k.CheckpointWait.k = k
 
 	ctx := k.SupervisorContext()
 	if err := k.vfs.Init(ctx); err != nil {
@@ -788,7 +778,6 @@ func (k *Kernel) LoadFrom(ctx context.Context, r, pagesMetadata io.Reader, pages
 	}
 
 	k.runningTasksCond.L = &k.runningTasksMu
-	k.checkpointCond.L = &k.checkpointMu
 	k.cpuClockTickerWakeCh = make(chan struct{}, 1)
 	k.cpuClockTickerStopCond.L = &k.runningTasksMu
 

pkg/sentry/kernel/kernel_restore.go

@@ -14,12 +14,29 @@
 
 package kernel
 
+import (
+	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/sync"
+)
+
 // Saver is an interface for saving the kernel.
 type Saver interface {
 	SaveAsync() error
 	SpecEnviron(containerName string) []string
 }
 
+// CheckpointGeneration stores information about the last checkpoint taken.
+//
+// +stateify savable
+type CheckpointGeneration struct {
+	// Count is incremented every time a checkpoint is triggered, even if the
+	// chekpoint failed.
+	Count uint32
+	// Restore indicates if the current instance resumed after the checkpoint or
+	// it was restored from a checkpoint.
+	Restore bool
+}
+
 // AddStateToCheckpoint adds a key-value pair to be additionally checkpointed.
 func (k *Kernel) AddStateToCheckpoint(key, v any) {
 	k.checkpointMu.Lock()
@@ -58,62 +75,120 @@ func (k *Kernel) Saver() Saver {
 	return k.saver
 }
 
-// IncCheckpointCount increments the checkpoint counter.
-func (k *Kernel) IncCheckpointCount() {
+// CheckpointGen returns the current checkpoint generation.
+func (k *Kernel) CheckpointGen() CheckpointGeneration {
 	k.checkpointMu.Lock()
 	defer k.checkpointMu.Unlock()
-	k.checkpointCounter++
+
+	return k.checkpointGen
 }
 
-// CheckpointCount returns the current checkpoint count. Note that the result
-// may be stale by the time the caller uses it.
-func (k *Kernel) CheckpointCount() uint32 {
+// OnRestoreDone is called to notify the kernel that a checkpoint restore has been
+// completed successfully.
+func (k *Kernel) OnRestoreDone() {
 	k.checkpointMu.Lock()
 	defer k.checkpointMu.Unlock()
-	return k.checkpointCounter
+
+	k.checkpointGen.Count++
+	k.checkpointGen.Restore = true
+
+	k.CheckpointWait.signal(k.checkpointGen, nil)
 }
 
 // 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++
+		log.Infof("Checkpoint completed successfully.")
+	} else {
+		log.Warningf("Checkpoint attempt failed with error: %v", err)
 	}
-	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
+
+	k.checkpointGen.Count++
+	k.checkpointGen.Restore = false
+
+	k.CheckpointWait.signal(k.checkpointGen, 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 (k *Kernel) WaitCheckpoint(n uint32) error {
-	if n == 0 {
-		return nil
+// WaitForCheckpoint waits for the Kernel to have been successfully checkpointed.
+func (k *Kernel) WaitForCheckpoint() error {
+	// Send checkpoint result to a channel and wait on it.
+	ch := make(chan error, 1)
+	callback := func(_ CheckpointGeneration, err error) { ch <- err }
+	key := k.CheckpointWait.Register(callback, k.CheckpointGen().Count+1)
+	defer k.CheckpointWait.Unregister(key)
+
+	return <-ch
+}
+
+type checkpointWaiter struct {
+	// count indicates the checkpoint generation that this waiter is interested in.
+	count uint32
+	// callback is the function that will be called when the checkpoint generation
+	// reaches the desired count. It is set to nil after the callback is called.
+	callback func(CheckpointGeneration, error)
+}
+
+// CheckpointWaitable is a waitable object that waits for a
+// checkpoint to complete.
+//
+// +stateify savable
+type CheckpointWaitable struct {
+	k *Kernel
+
+	mu sync.Mutex `state:"nosave"`
+
+	// Don't save the waiters, because they are repopulated after restore. It also
+	// allows for external entities to wait for the checkpoint.
+	waiters map[*checkpointWaiter]struct{} `state:"nosave"`
+}
+
+// Register registers a callback that is notified when the checkpoint generation count is higher
+// than the desired count.
+func (w *CheckpointWaitable) Register(cb func(CheckpointGeneration, error), count uint32) any {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	waiter := &checkpointWaiter{
+		count:    count,
+		callback: cb,
 	}
-	k.checkpointMu.Lock()
-	defer k.checkpointMu.Unlock()
-	if k.checkpointCounter >= n {
-		// n-th checkpoint already completed successfully.
-		return nil
+	if w.waiters == nil {
+		w.waiters = make(map[*checkpointWaiter]struct{})
 	}
-	for k.checkpointCounter < n {
-		if k.checkpointCounter == n-1 && k.lastCheckpointStatus != nil {
-			// n-th checkpoint was attempted but it had failed.
-			return k.lastCheckpointStatus
+	w.waiters[waiter] = struct{}{}
+
+	if gen := w.k.CheckpointGen(); count <= gen.Count {
+		// The checkpoint has already occurred. Signal immediately.
+		waiter.callback(gen, nil)
+		waiter.callback = nil
+	}
+	return waiter
+}
+
+// Unregister unregisters a waiter. It must be called even if the channel
+// was signalled.
+func (w *CheckpointWaitable) Unregister(key any) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	delete(w.waiters, key.(*checkpointWaiter))
+	if len(w.waiters) == 0 {
+		w.waiters = nil
+	}
+}
+
+func (w *CheckpointWaitable) signal(gen CheckpointGeneration, err error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	for waiter := range w.waiters {
+		if waiter.callback != nil && waiter.count <= gen.Count {
+			waiter.callback(gen, err)
+			waiter.callback = nil
 		}
-		k.checkpointCond.Wait()
 	}
-	return nil
 }

runsc/boot/controller.go

@@ -722,13 +722,11 @@ 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)
+// WaitCheckpoint waits for the Kernel to have been successfully checkpointed.
+func (cm *containerManager) WaitCheckpoint(*struct{}, *struct{}) error {
+	log.Debugf("containerManager.WaitCheckpoint")
+	err := cm.l.k.WaitForCheckpoint()
+	log.Debugf("containerManager.WaitCheckpoint done, err = %v", err)
 	return err
 }
 

runsc/boot/restore.go

@@ -735,7 +735,8 @@ func (r *restorer) restore(l *Loader) error {
 		// Restore was successful, so increment the checkpoint count manually. The
 		// count was saved while the previous kernel was being saved and checkpoint
 		// success was unknown at that time. Now we know the checkpoint succeeded.
-		l.k.IncCheckpointCount()
+		l.k.OnRestoreDone()
+
 		log.Infof("Restore successful")
 	}()
 	return nil
@@ -746,7 +747,6 @@ func (l *Loader) save(o *control.SaveOpts) (err error) {
 		// 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 {

runsc/cmd/wait.go

@@ -36,7 +36,7 @@ const (
 type Wait struct {
 	rootPID    int
 	pid        int
-	checkpoint uint
+	checkpoint bool
 }
 
 // Name implements subcommands.Command.Name.
@@ -58,7 +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.")
+	f.BoolVar(&wt.checkpoint, "checkpoint", false, "wait for the next checkpoint to complete")
 }
 
 // Execute implements subcommands.Command.Execute. It waits for a process in a
@@ -81,12 +81,12 @@ func (wt *Wait) Execute(_ context.Context, f *flag.FlagSet, args ...any) subcomm
 		util.Fatalf("loading container: %v", err)
 	}
 
-	if wt.checkpoint > 0 {
+	if wt.checkpoint {
 		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)
+		if err := c.WaitCheckpoint(); err != nil {
+			util.Fatalf("waiting for checkpoint to complete: %v", err)
 		}
 		return subcommands.ExitSuccess
 	}

runsc/container/container.go

@@ -633,16 +633,13 @@ 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)
+// WaitCheckpoint waits for the Kernel to have been successfully checkpointed.
+func (c *Container) WaitCheckpoint() error {
+	log.Debugf("Waiting for checkpoint to complete in container, cid: %s", c.ID)
 	if !c.IsSandboxRunning() {
 		return fmt.Errorf("sandbox is not running")
 	}
-	return c.Sandbox.WaitCheckpoint(n)
+	return c.Sandbox.WaitCheckpoint()
 }
 
 // SignalContainer sends the signal to the container. If all is true and signal

runsc/container/multi_container_test.go

@@ -2763,7 +2763,7 @@ func testMultiContainerCheckpointRestore(t *testing.T, conf *config.Config, comp
 	checkpointWaiter := make(chan struct{}, 1)
 	go func() {
 		// WaitCheckpoint on the second container.
-		if err := conts[1].WaitCheckpoint(1); err != nil {
+		if err := conts[1].WaitCheckpoint(); err != nil {
 			t.Errorf("error waiting for checkpoint to complete: %v", err)
 		}
 		checkpointWaiter <- struct{}{}

runsc/sandbox/sandbox.go

@@ -1329,13 +1329,10 @@ 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)
+// WaitCheckpoint waits for the Kernel to have been successfully checkpointed.
+func (s *Sandbox) WaitCheckpoint() error {
+	log.Debugf("Waiting for checkpoint to complete in sandbox %q", s.ID)
+	return s.call(boot.ContMgrWaitCheckpoint, nil, nil)
 }
 
 // IsRootContainer returns true if the specified container ID belongs to the