pgalloc: one async page loader per pages file

This resolves the "future work" described in MemoryFile.LoadFrom().

PiperOrigin-RevId: 808400892

Author: nixprime

pkg/sentry/kernel/kernel_restore.go

@@ -16,7 +16,6 @@ package kernel
 
 import (
 	"bufio"
-	"errors"
 	"fmt"
 	"io"
 
@@ -262,9 +261,8 @@ type AsyncMFLoader struct {
 	metadataWg  sync.WaitGroup
 	metadataErr error
 
-	loadWg     sync.WaitGroup
-	loadErrsMu sync.Mutex
-	loadErrs   []error
+	loadWg  sync.WaitGroup
+	loadErr error
 }
 
 // NewAsyncMFLoader creates a new AsyncMFLoader. It takes ownership of
@@ -301,22 +299,16 @@ func (mfl *AsyncMFLoader) backgroundGoroutine(pagesMetadataFD, pagesFileFD *fd.F
 	// or compressio.NewSimpleReader().
 	pagesMetadata := bufio.NewReader(pagesMetadataFD)
 
+	mfl.loadWg.Add(1)
+	apfl := pgalloc.StartAsyncPagesFileLoad(int32(pagesFileFD.FD()), func(err error) {
+		defer mfl.loadWg.Done()
+		mfl.loadErr = err
+	}, timeline)
+	cu.Add(apfl.MemoryFilesDone)
+
 	opts := pgalloc.LoadOpts{
-		PagesFile: pagesFileFD,
-		OnAsyncPageLoadStart: func(mf *pgalloc.MemoryFile) {
-			mfl.loadWg.Add(1)
-			log.Infof("Starting async page load for %p", mf)
-		},
-		OnAsyncPageLoadDone: func(mf *pgalloc.MemoryFile, err error) {
-			defer mfl.loadWg.Done()
-			if err != nil {
-				log.Warningf("Async page load error for %p: %v", mf, err)
-				mfl.loadErrsMu.Lock()
-				mfl.loadErrs = append(mfl.loadErrs, fmt.Errorf("%p: async page load: %w", mf, err))
-				mfl.loadErrsMu.Unlock()
-			}
-		},
-		Timeline: timeline,
+		PagesFile: apfl,
+		Timeline:  timeline,
 	}
 
 	timeline.Reached("loading mainMF")
@@ -347,9 +339,9 @@ func (mfl *AsyncMFLoader) backgroundGoroutine(pagesMetadataFD, pagesFileFD *fd.F
 
 	// Wait for page loads to complete and report errors.
 	mfl.loadWg.Wait()
-	if loadErr := errors.Join(mfl.loadErrs...); loadErr != nil {
+	if mfl.loadErr != nil {
 		timeline.Invalidate("page load failed")
-		log.Warningf("Failed to load MemoryFile pages: %v", loadErr)
+		log.Warningf("Failed to load MemoryFile pages: %v", mfl.loadErr)
 		return
 	}
 	log.Infof("All MemoryFile pages have been loaded.")
@@ -381,5 +373,5 @@ func (mfl *AsyncMFLoader) Wait() error {
 		return err
 	}
 	mfl.loadWg.Wait()
-	return errors.Join(mfl.loadErrs...)
+	return mfl.loadErr
 }

pkg/sentry/pgalloc/BUILD

@@ -8,10 +8,17 @@ package(
 )
 
 declare_mutex(
-    name = "apl_shared_mutex",
-    out = "apl_shared_mutex.go",
+    name = "amfls_mutex",
+    out = "amfls_mutex.go",
     package = "pgalloc",
-    prefix = "aplShared",
+    prefix = "amfls",
+)
+
+declare_mutex(
+    name = "apfl_mutex",
+    out = "apfl_mutex.go",
+    package = "pgalloc",
+    prefix = "apfl",
 )
 
 declare_mutex(
@@ -21,6 +28,18 @@ declare_mutex(
     prefix = "memoryFile",
 )
 
+go_template_instance(
+    name = "amfl_list",
+    out = "amfl_list.go",
+    package = "pgalloc",
+    prefix = "asyncMemoryFileLoad",
+    template = "//pkg/ilist:generic_list",
+    types = {
+        "Element": "*asyncMemoryFileLoad",
+        "Linker": "*asyncMemoryFileLoad",
+    },
+)
+
 go_template_instance(
     name = "apl_unloaded_set",
     out = "apl_unloaded_set.go",
@@ -128,7 +147,9 @@ go_template_instance(
 go_library(
     name = "pgalloc",
     srcs = [
-        "apl_shared_mutex.go",
+        "amfl_list.go",
+        "amfls_mutex.go",
+        "apfl_mutex.go",
         "apl_unloaded_set.go",
         "context.go",
         "debug.go",
@@ -150,7 +171,6 @@ go_library(
         "//pkg/bitmap",
         "//pkg/context",
         "//pkg/errors/linuxerr",
-        "//pkg/fd",
         "//pkg/gohacks",
         "//pkg/goid",
         "//pkg/hostarch",

pkg/sentry/pgalloc/pgalloc.go

@@ -14,6 +14,12 @@
 
 // Package pgalloc contains the page allocator subsystem, which provides
 // allocatable memory that may be mapped into application address spaces.
+//
+// Lock order:
+//
+//	AsyncPagesFileLoad.amflsMu
+//	  MemoryFile.mu
+//	    AsyncPagesFileLoad.mu
 package pgalloc
 
 import (
@@ -190,7 +196,7 @@ type MemoryFile struct {
 
 	// If asyncPageLoad is non-nil, it tracks the state of in-progress or
 	// failed async page loading.
-	asyncPageLoad atomic.Pointer[aplShared]
+	asyncPageLoad atomic.Pointer[asyncMemoryFileLoad]
 
 	// file is the backing file. The file pointer is immutable.
 	file *os.File
@@ -1435,8 +1441,8 @@ func (f *MemoryFile) MapInternal(fr memmap.FileRange, at hostarch.AccessType) (s
 		return safemem.BlockSeq{}, linuxerr.EACCES
 	}
 
-	if apl := f.asyncPageLoad.Load(); apl != nil {
-		if err := apl.awaitLoad(f, fr); err != nil {
+	if amfl := f.asyncPageLoad.Load(); amfl != nil {
+		if err := amfl.awaitLoad(fr); err != nil {
 			return safemem.BlockSeq{}, err
 		}
 	}
@@ -1799,8 +1805,8 @@ func (f *MemoryFile) File() *os.File {
 
 // DataFD implements memmap.File.DataFD.
 func (f *MemoryFile) DataFD(fr memmap.FileRange) (int, error) {
-	if apl := f.asyncPageLoad.Load(); apl != nil {
-		if err := apl.awaitLoad(f, fr); err != nil {
+	if amfl := f.asyncPageLoad.Load(); amfl != nil {
+		if err := amfl.awaitLoad(fr); err != nil {
 			return -1, err
 		}
 	}