Merge pull request #11360 from Snowflake-Labs:xuzhoyin-vmas-fix

PiperOrigin-RevId: 720834745

Author: gvisor-bot

pkg/sentry/mm/mm_test.go

@@ -29,18 +29,23 @@ import (
 	"gvisor.dev/gvisor/pkg/usermem"
 )
 
-func testMemoryManager(ctx context.Context) *MemoryManager {
+func testMemoryManagerWithMmapDirection(ctx context.Context, mmapDirection arch.MmapDirection) *MemoryManager {
 	p := platform.FromContext(ctx)
 	mm := NewMemoryManager(p, pgalloc.MemoryFileFromContext(ctx), false)
 	mm.layout = arch.MmapLayout{
-		MinAddr:      p.MinUserAddress(),
-		MaxAddr:      p.MaxUserAddress(),
-		BottomUpBase: p.MinUserAddress(),
-		TopDownBase:  p.MaxUserAddress(),
+		MinAddr:          p.MinUserAddress(),
+		MaxAddr:          p.MaxUserAddress(),
+		BottomUpBase:     p.MinUserAddress(),
+		TopDownBase:      p.MaxUserAddress(),
+		DefaultDirection: mmapDirection,
 	}
 	return mm
 }
 
+func testMemoryManager(ctx context.Context) *MemoryManager {
+	return testMemoryManagerWithMmapDirection(ctx, arch.MmapBottomUp)
+}
+
 func (mm *MemoryManager) realUsageAS() uint64 {
 	return uint64(mm.vmas.Span())
 }
@@ -272,3 +277,66 @@ func TestAIOLookupAfterDestroy(t *testing.T) {
 		t.Errorf("AIOContext found even after AIOContext manager is destroyed")
 	}
 }
+
+func TestGetAllocationDirection(t *testing.T) {
+	testCases := []struct {
+		name          string
+		mmapDirection arch.MmapDirection
+		ar            hostarch.AddrRange
+		vma           *vma
+		expected      pgalloc.Direction
+	}{
+		{
+			"No last fault in vma with mmap direction BottomUp",
+			arch.MmapBottomUp,
+			hostarch.AddrRange{123, 456},
+			&vma{lastFault: 0},
+			pgalloc.BottomUp,
+		},
+		{
+			"No last fault in vma with mmap direction TopDown",
+			arch.MmapTopDown,
+			hostarch.AddrRange{123, 456},
+			&vma{lastFault: 0},
+			pgalloc.TopDown,
+		},
+		{
+			"Last fault in vma equals to addr range, with mmap direction BottomUp",
+			arch.MmapBottomUp,
+			hostarch.AddrRange{123, 456},
+			&vma{lastFault: 123},
+			pgalloc.BottomUp,
+		},
+		{
+			"Last fault in vma equals to addr range, with mmap direction TopDown",
+			arch.MmapTopDown,
+			hostarch.AddrRange{123, 456},
+			&vma{lastFault: 123},
+			pgalloc.TopDown,
+		},
+		{
+			"Last fault in vma greater than addr range",
+			arch.MmapTopDown,
+			hostarch.AddrRange{123, 456},
+			&vma{lastFault: 456},
+			pgalloc.TopDown,
+		},
+		{
+			"Last fault in vma smaller than addr range",
+			arch.MmapTopDown,
+			hostarch.AddrRange{123, 456},
+			&vma{lastFault: 100},
+			pgalloc.BottomUp,
+		},
+	}
+	for _, test := range testCases {
+		t.Run(test.name, func(t *testing.T) {
+			ctx := contexttest.Context(t)
+			mm := testMemoryManagerWithMmapDirection(ctx, test.mmapDirection)
+			actual := mm.getAllocationDirection(test.ar, test.vma)
+			if actual != test.expected {
+				t.Errorf("Unexpected allocation direction. Expected: %s, Actual: %s", test.expected, actual)
+			}
+		})
+	}
+}

pkg/sentry/mm/pma.go

@@ -24,6 +24,7 @@ import (
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/safecopy"
 	"gvisor.dev/gvisor/pkg/safemem"
+	"gvisor.dev/gvisor/pkg/sentry/arch"
 	"gvisor.dev/gvisor/pkg/sentry/memmap"
 	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
 	"gvisor.dev/gvisor/pkg/sentry/usage"
@@ -185,6 +186,33 @@ func (mm *MemoryManager) getVecPMAsLocked(ctx context.Context, ars hostarch.Addr
 	return ars, nil
 }
 
+// Gets the default memory file offset allocation direction aligned with
+// default address space allocation direction.
+func (mm *MemoryManager) getDefaultAllocationDirection() pgalloc.Direction {
+	if mm.layout.DefaultDirection == arch.MmapTopDown {
+		return pgalloc.TopDown
+	}
+	return pgalloc.BottomUp
+}
+
+// Gets the memory file offset allocation direction based on address space allocation direction
+// and memory access order.
+func (mm *MemoryManager) getAllocationDirection(ar hostarch.AddrRange, vma *vma) pgalloc.Direction {
+	lastFault := atomic.LoadUintptr(&vma.lastFault)
+	arStart := uintptr(ar.Start)
+	// If this VMA does not have last page fault or last page fault equals to arStart,
+	// use the default allocation direction.
+	if lastFault == 0 || lastFault == arStart {
+		return mm.getDefaultAllocationDirection()
+	}
+	// Detect cases where memory is accessed downwards and change memory file
+	// allocation order to increase the chances that pages are coalesced.
+	if arStart < lastFault {
+		return pgalloc.TopDown
+	}
+	return pgalloc.BottomUp
+}
+
 // getPMAsInternalLocked is equivalent to getPMAsLocked, with the following
 // exceptions:
 //
@@ -220,12 +248,7 @@ func (mm *MemoryManager) getPMAsInternalLocked(ctx context.Context, vseg vmaIter
 
 	vma := vseg.ValuePtr()
 	memCgID := pgalloc.MemoryCgroupIDFromContext(ctx)
-	allocDir := pgalloc.BottomUp
-	if uintptr(ar.Start) < atomic.LoadUintptr(&vma.lastFault) {
-		// Detect cases where memory is accessed downwards and change memory file
-		// allocation order to increase the chances that pages are coalesced.
-		allocDir = pgalloc.TopDown
-	}
+	allocDir := mm.getAllocationDirection(ar, vma)
 	atomic.StoreUintptr(&vma.lastFault, uintptr(ar.Start))
 
 	// Limit the range we allocate to ar, aligned to hugepage boundaries.

pkg/sentry/mm/vma.go

@@ -484,6 +484,13 @@ func (vmaSetFunctions) Merge(ar1 hostarch.AddrRange, vma1 vma, ar2 hostarch.Addr
 		// need to worry about whether we're in a mm.mappingMu critical section.
 		vma2.id.DecRef(context.Background())
 	}
+
+	// If the existing vma (vma2) has non-zero lastFault address,
+	// we should preserve it to the resulting merged-VMA
+	if vma1.lastFault == 0 {
+		vma1.lastFault = vma2.lastFault
+	}
+
 	return vma1, true
 }