Add runsc wait --fscheckpoint/fsrestore.

Note that runsc wait --fsrestore only waits for filesystem restore operations
to complete for the given container, rather than the whole sandbox (like wait
--fscheckpoint, --checkpoint, and --restore do).

PiperOrigin-RevId: 886982088

Author: nixprime

pkg/sentry/kernel/BUILD

@@ -101,6 +101,13 @@ declare_mutex(
     prefix = "fsContext",
 )
 
+declare_mutex(
+    name = "fs_save_mutex",
+    out = "fs_save_mutex.go",
+    package = "kernel",
+    prefix = "fsSave",
+)
+
 declare_mutex(
     name = "running_tasks_mutex",
     out = "running_tasks_mutex.go",
@@ -254,6 +261,7 @@ go_library(
         "fs_context.go",
         "fs_context_mutex.go",
         "fs_context_refs.go",
+        "fs_save_mutex.go",
         "fscheckpoint.go",
         "ipc_namespace.go",
         "kcov.go",

pkg/sentry/kernel/fscheckpoint.go

@@ -77,7 +77,7 @@ func (k *Kernel) FSSave(ctx context.Context, opts *FSSaveOpts) error {
 	if opts.ExitAfterSaving {
 		defer k.Kill(linux.WaitStatusExit(0)) // consistent with sentry/state.SaveOpts.Save
 	}
-	return k.quiescePausedAnd(ctx, func() error {
+	err := k.quiescePausedAnd(ctx, func() error {
 		var (
 			asyncPageSaveWg  sync.WaitGroup
 			asyncPageSaveErr error
@@ -226,6 +226,13 @@ func (k *Kernel) FSSave(ctx context.Context, opts *FSSaveOpts) error {
 		}
 		return nil
 	})
+	k.fsSaveMu.Lock()
+	defer k.fsSaveMu.Unlock()
+	for _, c := range k.fsSaveWaiters {
+		c <- err
+	}
+	k.fsSaveWaiters = nil
+	return err
 }
 
 type countingWriter struct {
@@ -239,3 +246,16 @@ func (cw *countingWriter) Write(src []byte) (int, error) {
 	cw.count += uint64(n)
 	return n, err
 }
+
+// WaitForFSSave waits for a call to k.FSSave() to complete, then returns the
+// error returned by that call.
+//
+// This API is difficult to use without races, but is consistent with
+// k.WaitForCheckpoint().
+func (k *Kernel) WaitForFSSave() error {
+	c := make(chan error, 1)
+	k.fsSaveMu.Lock()
+	k.fsSaveWaiters = append(k.fsSaveWaiters, c)
+	k.fsSaveMu.Unlock()
+	return <-c
+}

pkg/sentry/kernel/kernel.go

@@ -401,6 +401,11 @@ type Kernel struct {
 
 	// MaxKeySetSize is the maximum number of keys in a key set.
 	MaxKeySetSize atomicbitops.Int32
+
+	// fsSaveWaiters holds waiters for Kernel.WaitForFSSave. fsSaveWaiters is
+	// protected by fsSaveMu.
+	fsSaveMu      fsSaveMutex  `state:"nosave"`
+	fsSaveWaiters []chan error `state:"nosave"`
 }
 
 // InitKernelArgs holds arguments to Init.

pkg/sentry/pgalloc/save_restore.go

@@ -905,9 +905,15 @@ type LoadOpts struct {
 	// returns, PagesFileOffset will be updated to the offset of the first byte
 	// in PagesFile after this MemoryFile's contents.
 	//
+	// Reading from PagesFile may continue after LoadFrom returns. If
+	// DoneCallback is not nil, it will be called when reading for this
+	// MemoryFile completes. DoneCallback will be called whether or not
+	// LoadFrom returns a non-nil error.
+	//
 	// Invariant: PagesFileOffset must be page-aligned.
 	PagesFile       *AsyncPagesFileLoad
 	PagesFileOffset uint64
+	DoneCallback    func(error)
 
 	// Optional timeline for the restore process.
 	// If async page loading is enabled, a forked timeline will be created, so
@@ -917,10 +923,16 @@ type LoadOpts struct {
 }
 
 // LoadFrom loads MemoryFile state from the given stream.
-func (f *MemoryFile) LoadFrom(ctx context.Context, r io.Reader, opts *LoadOpts) error {
+func (f *MemoryFile) LoadFrom(ctx context.Context, r io.Reader, opts *LoadOpts) (err error) {
 	mfTimeline := opts.Timeline.Fork(fmt.Sprintf("mf:%p", f)).Lease()
 	defer mfTimeline.End()
 
+	defer func() {
+		if opts.DoneCallback != nil {
+			opts.DoneCallback(err)
+		}
+	}()
+
 	// Load metadata.
 	timeMetadataStart := gohacks.Nanotime()
 	var mfs memoryFileSaved
@@ -995,10 +1007,11 @@ func (f *MemoryFile) LoadFrom(ctx context.Context, r io.Reader, opts *LoadOpts)
 			}
 		}
 		amfl = &asyncMemoryFileLoad{
-			f:        f,
-			pf:       opts.PagesFile,
-			df:       df,
-			timeline: mfTimeline.Transfer(),
+			f:            f,
+			pf:           opts.PagesFile,
+			df:           df,
+			doneCallback: opts.DoneCallback,
+			timeline:     mfTimeline.Transfer(),
 		}
 		amfl.pf.amflsMu.Lock()
 		if err := amfl.pf.err(); err != nil {
@@ -1008,6 +1021,7 @@ func (f *MemoryFile) LoadFrom(ctx context.Context, r io.Reader, opts *LoadOpts)
 		amfl.pf.amfls.PushBack(amfl)
 		amfl.pf.amflsMu.Unlock()
 		f.asyncPageLoad.Store(amfl)
+		opts.DoneCallback = nil
 		defer func() {
 			amfl.pf.amflsMu.Lock()
 			defer amfl.pf.amflsMu.Unlock()
@@ -1025,6 +1039,10 @@ func (f *MemoryFile) LoadFrom(ctx context.Context, r io.Reader, opts *LoadOpts)
 				amfl.pf.amfls.Remove(amfl)
 				amfl.f.asyncPageLoad.Store(nil)
 				amfl.timeline.End()
+				if amfl.doneCallback != nil {
+					amfl.doneCallback(nil)
+					amfl.doneCallback = nil
+				}
 			}
 		}()
 	}
@@ -1205,10 +1223,11 @@ func (apfl *AsyncPagesFileLoad) err() error {
 // asyncMemoryFileLoad holds async page loading state for a single MemoryFile.
 type asyncMemoryFileLoad struct {
 	// Immutable fields:
-	f        *MemoryFile
-	pf       *AsyncPagesFileLoad
-	df       stateio.DestinationFile
-	timeline *timing.Timeline
+	f            *MemoryFile
+	pf           *AsyncPagesFileLoad
+	df           stateio.DestinationFile
+	doneCallback func(error)
+	timeline     *timing.Timeline
 
 	// minUnloaded is the MemoryFile offset of the first unloaded byte.
 	minUnloaded atomicbitops.Uint64
@@ -1597,6 +1616,10 @@ func (apfl *AsyncPagesFileLoad) main() {
 				ul.started = false
 				ul.waiters = nil
 			}
+			if amfl.doneCallback != nil {
+				amfl.doneCallback(apfl.err())
+				amfl.doneCallback = nil
+			}
 		}
 		apfl.mu.Unlock()
 		apfl.amflsMu.Unlock()
@@ -1876,6 +1899,10 @@ func (apfl *AsyncPagesFileLoad) main() {
 					apfl.amfls.Remove(amfl)
 					amfl.f.asyncPageLoad.Store(nil)
 					amfl.timeline.End()
+					if amfl.doneCallback != nil {
+						amfl.doneCallback(nil)
+						amfl.doneCallback = nil
+					}
 				} else {
 					amfl.minUnloaded.Store(amfl.unloaded.FirstSegment().Start())
 				}

runsc/boot/controller.go

@@ -110,6 +110,14 @@ const (
 	// ContMgrWaitRestore waits for the Kernel restore to complete.
 	ContMgrWaitRestore = "containerManager.WaitRestore"
 
+	// ContMgrWaitFSCheckpoint waits for the next filesystem checkpoint save to
+	// complete.
+	ContMgrWaitFSCheckpoint = "containerManager.WaitFSCheckpoint"
+
+	// ContMgrWaitFSRestore waits for filesystem checkpoint restore to complete
+	// for all current containers.
+	ContMgrWaitFSRestore = "containerManager.WaitFSRestore"
+
 	// ContMgrRootContainerStart starts a new sandbox with a root container.
 	ContMgrRootContainerStart = "containerManager.StartRoot"
 
@@ -870,6 +878,26 @@ func (cm *containerManager) WaitRestore(*struct{}, *struct{}) error {
 	return err
 }
 
+func (cm *containerManager) WaitFSCheckpoint(*struct{}, *struct{}) error {
+	log.Debugf("containerManager.WaitFSCheckpoint")
+	err := cm.l.k.WaitForFSSave()
+	log.Debugf("containerManager.WaitFSCheckpoint done, err = %v", err)
+	return err
+}
+
+// WaitFSRestoreArgs holds arguments to containerManager.WaitFSRestore.
+type WaitFSRestoreArgs struct {
+	// CID is the container ID.
+	CID string
+}
+
+func (cm *containerManager) WaitFSRestore(args *WaitFSRestoreArgs, _ *struct{}) error {
+	log.Debugf("containerManager.WaitFSRestore")
+	err := cm.l.fsRestore.wait(args.CID)
+	log.Debugf("containerManager.WaitFSRestore done, err = %v", err)
+	return err
+}
+
 // SignalDeliveryMode enumerates different signal delivery modes.
 type SignalDeliveryMode int
 

runsc/boot/fscheckpoint.go

@@ -84,6 +84,16 @@ type fsRestore struct {
 	apfl        *pgalloc.AsyncPagesFileLoad
 	mfs         map[checkpoint.ResourceID]*fscheckpoint.MemoryFile
 	tmpfs       map[checkpoint.ResourceID]*fscheckpoint.Tmpfs
+
+	waitMu  sync.Mutex
+	waitMap map[string]*fsRestoreContainer // key is container ID
+}
+
+type fsRestoreContainer struct {
+	// These fields are protected by fsRestore.waitMu.
+	err        error
+	asyncLoads int // number of MemoryFiles currently in async page loading
+	cond       sync.Cond
 }
 
 // fsRestoreOpts holds options to startFSRestore.
@@ -111,8 +121,9 @@ func makeFSRestoreOptsForLocalCheckpoint(args *Args) (fsRestoreOpts, error) {
 // startFSRestore takes ownership of resources in opts.
 func startFSRestore(opts *fsRestoreOpts) (*fsRestore, error) {
 	fsr := &fsRestore{
-		mfs:   make(map[checkpoint.ResourceID]*fscheckpoint.MemoryFile),
-		tmpfs: make(map[checkpoint.ResourceID]*fscheckpoint.Tmpfs),
+		mfs:     make(map[checkpoint.ResourceID]*fscheckpoint.MemoryFile),
+		tmpfs:   make(map[checkpoint.ResourceID]*fscheckpoint.Tmpfs),
+		waitMap: make(map[string]*fsRestoreContainer),
 	}
 
 	// TODO: NOLINT - Currently we read the whole pages metadata file into a
@@ -225,32 +236,71 @@ func startFSRestore(opts *fsRestoreOpts) (*fsRestore, error) {
 	return fsr, nil
 }
 
-func (fsr *fsRestore) memoryFileLoadArgs(id checkpoint.ResourceID) (io.Reader, uint64, error) {
+// +checklocks:fsr.waitMu
+func (fsr *fsRestore) ensureContainer(cid string) *fsRestoreContainer {
+	c := fsr.waitMap[cid]
+	if c == nil {
+		c = &fsRestoreContainer{}
+		c.cond.L = &fsr.waitMu
+		fsr.waitMap[cid] = c
+	}
+	return c
+}
+
+func (c *fsRestoreContainer) setError(err error) error {
+	if c.err == nil && err != nil {
+		c.err = err
+		c.cond.Broadcast()
+	}
+	return err
+}
+
+func (fsr *fsRestore) memoryFileLoadArgs(id checkpoint.ResourceID, cid string) (io.Reader, uint64, func(error), error) {
 	if fsr == nil {
-		return nil, 0, nil
+		return nil, 0, func(error) {}, nil
 	}
+
 	fsr.wg.Wait()
 	if fsr.manifestErr != nil {
-		return nil, 0, fsr.manifestErr
+		return nil, 0, nil, fsr.manifestErr
 	}
 	mmf := fsr.mfs[id]
 	if mmf == nil {
-		return nil, 0, nil
+		return nil, 0, func(error) {}, nil
 	}
 	pagesMetadata, err := fsr.getPagesMetadata()
-	if mmf.PagesMetadataEnd <= uint64(len(pagesMetadata)) {
-		return bytes.NewReader(pagesMetadata[mmf.PagesMetadataStart:mmf.PagesMetadataEnd]), mmf.PagesStart, nil
-	}
-	if err != nil {
-		return nil, 0, fmt.Errorf("failed to read pages metadata: %w", err)
+
+	fsr.waitMu.Lock()
+	defer fsr.waitMu.Unlock()
+	c := fsr.ensureContainer(cid)
+	if mmf.PagesMetadataEnd > uint64(len(pagesMetadata)) {
+		if err != nil {
+			return nil, 0, nil, c.setError(fmt.Errorf("failed to read pages metadata: %w", err))
+		}
+		return nil, 0, nil, c.setError(fmt.Errorf("MemoryFile %q has pages metadata range [%d, %d) beyond pages metadata file size %d", mmf.ResourceID, mmf.PagesMetadataStart, mmf.PagesMetadataEnd, len(pagesMetadata)))
 	}
-	return nil, 0, fmt.Errorf("MemoryFile %q has pages metadata range [%d, %d) beyond pages metadata file size %d", mmf.ResourceID, mmf.PagesMetadataStart, mmf.PagesMetadataEnd, len(pagesMetadata))
+	c.asyncLoads++
+	return bytes.NewReader(pagesMetadata[mmf.PagesMetadataStart:mmf.PagesMetadataEnd]), mmf.PagesStart, func(err error) {
+		fsr.waitMu.Lock()
+		defer fsr.waitMu.Unlock()
+		c.asyncLoads--
+		switch {
+		case err != nil:
+			if c.err == nil {
+				c.err = err
+			}
+			fallthrough
+		case c.asyncLoads == 0:
+			c.cond.Broadcast()
+		}
+	}, nil
 }
 
-func (fsr *fsRestore) tmpfsSourceTar(id checkpoint.ResourceID) (io.ReadCloser, error) {
+func (fsr *fsRestore) tmpfsSourceTar(id checkpoint.ResourceID, cid string) (io.ReadCloser, error) {
 	if fsr == nil {
 		return nil, nil
 	}
+
 	fsr.wg.Wait()
 	if fsr.manifestErr != nil {
 		return nil, fsr.manifestErr
@@ -260,11 +310,43 @@ func (fsr *fsRestore) tmpfsSourceTar(id checkpoint.ResourceID) (io.ReadCloser, e
 		return nil, nil
 	}
 	multiTar, err := fsr.getMultiTar()
+
+	fsr.waitMu.Lock()
+	defer fsr.waitMu.Unlock()
 	if mt.TarEnd <= uint64(len(multiTar)) {
 		return io.NopCloser(bytes.NewReader(multiTar[mt.TarStart:mt.TarEnd])), nil
 	}
+	c := fsr.ensureContainer(cid)
 	if err != nil {
-		return nil, fmt.Errorf("failed to read tar archive: %w", err)
+		return nil, c.setError(fmt.Errorf("failed to read tar archive: %w", err))
+	}
+	return nil, c.setError(fmt.Errorf("tmpfs %q has tar range [%d, %d) beyond multi-tar file size %d", mt.ResourceID, mt.TarStart, mt.TarEnd, len(multiTar)))
+}
+
+// wait blocks until either all filesystems have been restored for the
+// container with the given ID, or an error occurs while restoring filesystems
+// for that container.
+func (fsr *fsRestore) wait(cid string) error {
+	if fsr == nil {
+		return fmt.Errorf("filesystem restore is not enabled")
+	}
+	fsr.wg.Wait()
+	if fsr.manifestErr != nil {
+		return fsr.manifestErr
+	}
+	fsr.waitMu.Lock()
+	defer fsr.waitMu.Unlock()
+	c := fsr.waitMap[cid]
+	if c == nil {
+		return fmt.Errorf("no filesystems restored for container %s", cid)
+	}
+	for {
+		if c.err != nil {
+			return c.err
+		}
+		if c.asyncLoads == 0 {
+			return nil
+		}
+		c.cond.Wait()
 	}
-	return nil, fmt.Errorf("tmpfs %q has tar range [%d, %d) beyond multi-tar file size %d", mt.ResourceID, mt.TarStart, mt.TarEnd, len(multiTar))
 }