Wait for main MF loading to start before writing to it during restore.

When using async page loading (the checkpoint image has compression=none), the
kernel and memory files are loaded in parallel. A step in kernel loading tries
writing to the main MemoryFile. Ensure that step happens after the main
MemoryFile has started being loaded (i.e. mainMF.LoadFrom() method has returned
successfully). MemoryFile.asyncPageLoad field needs to be initialized before
it can be written to.

Fixes #11842

PiperOrigin-RevId: 779995419

Author: ayushr2

pkg/sentry/kernel/kernel.go

@@ -758,7 +758,7 @@ func (k *Kernel) invalidateUnsavableMappings(ctx context.Context) error {
 }
 
 // LoadFrom returns a new Kernel loaded from args.
-func (k *Kernel) LoadFrom(ctx context.Context, r io.Reader, loadMFs bool, timeReady chan struct{}, net inet.Stack, clocks sentrytime.Clocks, vfsOpts *vfs.CompleteRestoreOptions, saveRestoreNet bool) error {
+func (k *Kernel) LoadFrom(ctx context.Context, r io.Reader, asyncMFLoader *AsyncMFLoader, timeReady chan struct{}, net inet.Stack, clocks sentrytime.Clocks, vfsOpts *vfs.CompleteRestoreOptions, saveRestoreNet bool) error {
 	loadStart := time.Now()
 
 	k.runningTasksCond.L = &k.runningTasksMu
@@ -794,14 +794,6 @@ func (k *Kernel) LoadFrom(ctx context.Context, r io.Reader, loadMFs bool, timeRe
 	log.Infof("Kernel load stats: %s", stats.String())
 	log.Infof("Kernel load took [%s].", time.Since(kernelStart))
 
-	if loadMFs {
-		mfStart := time.Now()
-		if err := k.loadMemoryFiles(ctx, r); err != nil {
-			return fmt.Errorf("failed to load memory files: %w", err)
-		}
-		log.Infof("Memory files load took [%s].", time.Since(mfStart))
-	}
-
 	if !saveRestoreNet {
 		// rootNetworkNamespace and stack should be populated after
 		// loading the state file. Reset the stack before restoring the
@@ -810,6 +802,21 @@ func (k *Kernel) LoadFrom(ctx context.Context, r io.Reader, loadMFs bool, timeRe
 		k.rootNetworkNamespace.RestoreRootStack(net)
 	}
 
+	if asyncMFLoader == nil {
+		mfStart := time.Now()
+		if err := k.loadMemoryFiles(ctx, r); err != nil {
+			return fmt.Errorf("failed to load memory files: %w", err)
+		}
+		log.Infof("Memory files load took [%s].", time.Since(mfStart))
+	} else {
+		// Timekeeper restore below tries writing to the main MF. If async page
+		// loading is being used, we need to make sure that the main MF loading has
+		// started before we try to write to it.
+		if err := asyncMFLoader.WaitMainMFStart(); err != nil {
+			return fmt.Errorf("main MF start failed: %w", err)
+		}
+	}
+
 	k.Timekeeper().SetClocks(clocks, k.vdsoParams)
 
 	if timeReady != nil {

pkg/sentry/kernel/kernel_restore.go

@@ -256,6 +256,9 @@ type AsyncMFLoader struct {
 	// MemoryFiles, once they are known. This channel is written to exactly once.
 	privateMFsChan chan map[string]*pgalloc.MemoryFile
 
+	mainMFStartWg   sync.WaitGroup
+	mainMetadataErr error
+
 	metadataWg  sync.WaitGroup
 	metadataErr error
 
@@ -275,6 +278,7 @@ func NewAsyncMFLoader(pagesMetadata, pagesFile *fd.FD, mainMF *pgalloc.MemoryFil
 	mfl := &AsyncMFLoader{
 		privateMFsChan: make(chan map[string]*pgalloc.MemoryFile, 1),
 	}
+	mfl.mainMFStartWg.Add(1)
 	mfl.metadataWg.Add(1)
 	mfl.loadWg.Add(1)
 	go mfl.backgroundGoroutine(pagesMetadata, pagesFile, mainMF, timeline)
@@ -318,9 +322,12 @@ func (mfl *AsyncMFLoader) backgroundGoroutine(pagesMetadataFD, pagesFileFD *fd.F
 	timeline.Reached("loading mainMF")
 	log.Infof("Loading metadata for main MemoryFile: %p", mainMF)
 	ctx := context.Background()
-	if err := mainMF.LoadFrom(ctx, pagesMetadata, &opts); err != nil {
+	err := mainMF.LoadFrom(ctx, pagesMetadata, &opts)
+	mfl.metadataErr = err
+	mfl.mainMetadataErr = err
+	mfl.mainMFStartWg.Done()
+	if err != nil {
 		log.Warningf("Failed to load main MemoryFile %p: %v", mainMF, err)
-		mfl.metadataErr = err
 		return
 	}
 	timeline.Reached("waiting for privateMF info")
@@ -353,6 +360,13 @@ func (mfl *AsyncMFLoader) KickoffPrivate(mfmap map[string]*pgalloc.MemoryFile) {
 	mfl.privateMFsChan <- mfmap
 }
 
+// WaitMainMFStart waits for the background goroutine to successfully start
+// asynchronously loading the main MemoryFile.
+func (mfl *AsyncMFLoader) WaitMainMFStart() error {
+	mfl.mainMFStartWg.Wait()
+	return mfl.mainMetadataErr
+}
+
 // WaitMetadata waits for the background goroutine to successfully complete
 // reading all MemoryFile metadata and report any errors.
 func (mfl *AsyncMFLoader) WaitMetadata() error {

runsc/boot/restore.go

@@ -276,7 +276,7 @@ func (r *restorer) restore(l *Loader) error {
 
 	// Load the state.
 	r.timer.Reached("loading kernel")
-	if err := l.k.LoadFrom(ctx, r.stateFile, r.asyncMFLoader == nil, nil, oldInetStack, time.NewCalibratedClocks(), &vfs.CompleteRestoreOptions{}, l.saveRestoreNet); err != nil {
+	if err := l.k.LoadFrom(ctx, r.stateFile, r.asyncMFLoader, nil, oldInetStack, time.NewCalibratedClocks(), &vfs.CompleteRestoreOptions{}, l.saveRestoreNet); err != nil {
 		return fmt.Errorf("failed to load kernel: %w", err)
 	}
 	r.timer.Reached("kernel loaded")