From a15559c56cc00c1abc6581937147014fbdb3710c Mon Sep 17 00:00:00 2001
From: Jamie Liu <jamieliu@google.com>
Date: Mon, 30 Sep 2024 13:34:48 -0700
Subject: [PATCH] mm: limit AddressSpace overmapping during async page loading

On platforms that do not create page table entries in
platform.AddressSpace.MapFile(precommit=false), i.e. platforms for which
platform.Platform.MapUnit() == 0, platform.AddressSpace.MapFile() is generally
implemented as some form of host mmap(), which only synchronously creates host
kernel VMAs (virtual memory areas) and creates page table entries lazily in
response to application faults. On such platforms, MM.mapAsLocked() creates the
largest possible host VMAs since doing so reduces future sentry-handled page
faults and has effectively no additional cost. However, when async page loading
is active, this must wait for all mapped pages to be loaded, which may result
in the faulting application blocking for significantly longer than expected (in
experiments, a single page fault could result in waiting for up to 64GB of data
to be loaded). In such cases, additionally constrain mapped sizes to limit wait
times.

PiperOrigin-RevId: 680699946
---
 pkg/sentry/mm/address_space.go     | 22 ++++++++++++++++------
 pkg/sentry/pgalloc/save_restore.go |  6 ++++++
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/pkg/sentry/mm/address_space.go b/pkg/sentry/mm/address_space.go
index 496bd4ca28..df30235f5f 100644
--- a/pkg/sentry/mm/address_space.go
+++ b/pkg/sentry/mm/address_space.go
@@ -20,6 +20,7 @@ import (
 	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/sentry/memmap"
+	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
 	"gvisor.dev/gvisor/pkg/sentry/platform"
 )
 
@@ -175,6 +176,14 @@ func (mm *MemoryManager) mapASLocked(pseg pmaIterator, ar hostarch.AddrRange, pl
 	// By default, map entire pmas at a time, under the assumption that there
 	// is no cost to mapping more of a pma than necessary.
 	mapAR := hostarch.AddrRange{0, ^hostarch.Addr(hostarch.PageSize - 1)}
+	setMapUnit := func(mapUnit uint64) {
+		mapMask := hostarch.Addr(mapUnit - 1)
+		mapAR.Start = ar.Start &^ mapMask
+		// If rounding ar.End up overflows, just keep the existing mapAR.End.
+		if end := (ar.End + mapMask) &^ mapMask; end >= ar.End {
+			mapAR.End = end
+		}
+	}
 	if platformEffect != memmap.PlatformEffectDefault {
 		// When explicitly committing, only map ar, since overmapping may incur
 		// unexpected resource usage. When explicitly populating, do the same
@@ -183,12 +192,13 @@ func (mm *MemoryManager) mapASLocked(pseg pmaIterator, ar hostarch.AddrRange, pl
 		mapAR = ar
 	} else if mapUnit := mm.p.MapUnit(); mapUnit != 0 {
 		// Limit the range we map to ar, aligned to mapUnit.
-		mapMask := hostarch.Addr(mapUnit - 1)
-		mapAR.Start = ar.Start &^ mapMask
-		// If rounding ar.End up overflows, just keep the existing mapAR.End.
-		if end := (ar.End + mapMask) &^ mapMask; end >= ar.End {
-			mapAR.End = end
-		}
+		setMapUnit(mapUnit)
+	} else if mf, ok := pseg.ValuePtr().file.(*pgalloc.MemoryFile); ok && mf.IsAsyncLoading() {
+		// Impose an arbitrary mapUnit in order to avoid calling
+		// platform.AddressSpace.MapFile() => mf.DataFD() or mf.MapInternal()
+		// with unnecessarily large ranges, resulting in unnecessarily long
+		// waits.
+		setMapUnit(32 << 20)
 	}
 	if checkInvariants {
 		if !mapAR.IsSupersetOf(ar) {
diff --git a/pkg/sentry/pgalloc/save_restore.go b/pkg/sentry/pgalloc/save_restore.go
index 0495aa602c..2a12db715f 100644
--- a/pkg/sentry/pgalloc/save_restore.go
+++ b/pkg/sentry/pgalloc/save_restore.go
@@ -494,6 +494,12 @@ var aplWaiterPool = sync.Pool{
 	},
 }
 
+// IsAsyncLoading returns true if async page loading is in progress or has
+// failed permanently.
+func (f *MemoryFile) IsAsyncLoading() bool {
+	return f.asyncPageLoad.Load() != nil
+}
+
 // AwaitLoadAll blocks until async page loading has completed. If async page
 // loading is not in progress, AwaitLoadAll returns immediately.
 func (f *MemoryFile) AwaitLoadAll() error {