Handle guest OOM during syscall copies

pkg/sentry/kernel/task_run.go

@@ -410,6 +410,16 @@ func (t *Task) Yield() {
 	runtime.Gosched()
 }
 
+// HandleMemoryOOM handles an exhausted guest allocation budget after MM locks
+// have been released. The caller must be running on the task goroutine.
+func (t *Task) HandleMemoryOOM() bool {
+	if t.k.mf.MemoryLimit() == 0 || t.killed() || !t.killMemoryOOMVictim() {
+		return false
+	}
+	time.Sleep(time.Millisecond)
+	return !t.killed()
+}
+
 // killMemoryOOMVictim selects the largest resident process, adjusted by the
 // Linux oom_score_adj, instead of killing a small process that happens to fault.
 // Call only after releasing the faulting MemoryManager's locks.

pkg/sentry/mm/io.go

@@ -108,6 +108,20 @@ func translateIOError(ctx context.Context, err error) error {
 	return linuxerr.EFAULT
 }
 
+// memoryOOMHandler is supplied only by contexts running on the task goroutine.
+// Invoke it after releasing MM locks, before converting allocation errors to EFAULT.
+type memoryOOMHandler interface {
+	HandleMemoryOOM() bool
+}
+
+func (mm *MemoryManager) memoryOOMHandler(ctx context.Context, err error) memoryOOMHandler {
+	if err != linuxerr.ENOMEM || mm.mf.MemoryLimit() == 0 {
+		return nil
+	}
+	handler, _ := ctx.(memoryOOMHandler)
+	return handler
+}
+
 // CopyOut implements usermem.IO.CopyOut.
 func (mm *MemoryManager) CopyOut(ctx context.Context, addr hostarch.Addr, src []byte, opts usermem.IOOpts) (int, error) {
 	ar, ok := mm.CheckIORange(addr, int64(len(src)))
@@ -520,6 +534,7 @@ func (mm *MemoryManager) LoadUint32(ctx context.Context, addr hostarch.Addr, opt
 //   - ioar.Length() != 0.
 //   - ioar.Contains(addr).
 func (mm *MemoryManager) handleASIOFault(ctx context.Context, addr hostarch.Addr, ioar hostarch.AddrRange, at hostarch.AccessType) error {
+retry:
 	// Try to map all remaining pages in the I/O operation. This RoundUp can't
 	// overflow because otherwise it would have been caught by CheckIORange.
 	end, _ := ioar.End.RoundUp()
@@ -545,6 +560,13 @@ func (mm *MemoryManager) handleASIOFault(ctx context.Context, addr hostarch.Addr
 	mm.activeMu.Lock()
 	pseg, pend, err := mm.getPMAsLocked(ctx, vseg, ar, at, true /* callerIndirectCommit */)
 	mm.mappingMu.RUnlock()
+	if handler := mm.memoryOOMHandler(ctx, err); handler != nil {
+		mm.activeMu.Unlock()
+		if handler.HandleMemoryOOM() {
+			goto retry
+		}
+		return linuxerr.EFAULT
+	}
 	if pendaddr := pend.Start(); pendaddr < ar.End {
 		if pendaddr <= ar.Start {
 			mm.activeMu.Unlock()
@@ -573,6 +595,7 @@ func (mm *MemoryManager) handleASIOFault(ctx context.Context, addr hostarch.Addr
 //
 // Preconditions: 0 < ar.Length() <= math.MaxInt64.
 func (mm *MemoryManager) withInternalMappings(ctx context.Context, ar hostarch.AddrRange, at hostarch.AccessType, ignorePermissions bool, f func(safemem.BlockSeq) (uint64, error)) (int64, error) {
+retry:
 	// If pmas are already available, we can do IO without touching mm.vmas or
 	// mm.mappingMu.
 	mm.activeMu.RLock()
@@ -599,6 +622,13 @@ func (mm *MemoryManager) withInternalMappings(ctx context.Context, ar hostarch.A
 	mm.activeMu.Lock()
 	pseg, pend, perr := mm.getPMAsLocked(ctx, vseg, ar, at, true /* callerIndirectCommit */)
 	mm.mappingMu.RUnlock()
+	if handler := mm.memoryOOMHandler(ctx, perr); handler != nil {
+		mm.activeMu.Unlock()
+		if handler.HandleMemoryOOM() {
+			goto retry
+		}
+		return 0, linuxerr.EFAULT
+	}
 	if pendaddr := pend.Start(); pendaddr < ar.End {
 		if pendaddr <= ar.Start {
 			mm.activeMu.Unlock()
@@ -650,6 +680,7 @@ func (mm *MemoryManager) withVecInternalMappings(ctx context.Context, ars hostar
 		return mm.withInternalMappings(ctx, ars.Head(), at, ignorePermissions, f)
 	}
 
+retry:
 	// If pmas are already available, we can do IO without touching mm.vmas or
 	// mm.mappingMu.
 	mm.activeMu.RLock()
@@ -673,6 +704,13 @@ func (mm *MemoryManager) withVecInternalMappings(ctx context.Context, ars hostar
 	mm.activeMu.Lock()
 	pars, perr := mm.getVecPMAsLocked(ctx, vars, at, true /* callerIndirectCommit */)
 	mm.mappingMu.RUnlock()
+	if handler := mm.memoryOOMHandler(ctx, perr); handler != nil {
+		mm.activeMu.Unlock()
+		if handler.HandleMemoryOOM() {
+			goto retry
+		}
+		return 0, linuxerr.EFAULT
+	}
 	if pars.NumBytes() == 0 {
 		mm.activeMu.Unlock()
 		return 0, translateIOError(ctx, perr)

pkg/sentry/mm/mm_test.go

@@ -22,6 +22,7 @@ import (
 	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/errors/linuxerr"
 	"gvisor.dev/gvisor/pkg/hostarch"
+	"gvisor.dev/gvisor/pkg/safemem"
 	"gvisor.dev/gvisor/pkg/sentry/arch"
 	"gvisor.dev/gvisor/pkg/sentry/contexttest"
 	"gvisor.dev/gvisor/pkg/sentry/limits"
@@ -110,6 +111,64 @@ func TestMemoryBudgetFallsBackFromHugePagePrefetch(t *testing.T) {
 	}
 }
 
+type memoryOOMContext struct {
+	context.Context
+	handle func() bool
+}
+
+func (ctx memoryOOMContext) HandleMemoryOOM() bool { return ctx.handle() }
+
+func TestMemoryBudgetCopyOutRetriesBeforeIO(t *testing.T) {
+	for _, vector := range []bool{false, true} {
+		t.Run(map[bool]string{false: "single-range", true: "vector"}[vector], func(t *testing.T) {
+			file, err := os.CreateTemp(t.TempDir(), "memory")
+			if err != nil {
+				t.Fatal(err)
+			}
+			mf, err := pgalloc.NewMemoryFile(file, pgalloc.MemoryFileOpts{MemoryLimit: 2 * hostarch.PageSize, DisableMemoryAccounting: true, DisableIMAWorkAround: true})
+			if err != nil {
+				file.Close()
+				t.Fatal(err)
+			}
+			defer mf.Destroy()
+			base := context.WithValue(contexttest.Context(t), pgalloc.CtxMemoryFile, mf)
+			mm := testMemoryManager(base, t)
+			defer mm.DecUsers(base)
+			addr, err := mm.MMap(base, memmap.MMapOpts{Length: 3 * hostarch.PageSize, Private: true, Perms: hostarch.ReadWrite, MaxPerms: hostarch.AnyAccess})
+			if err != nil {
+				t.Fatal(err)
+			}
+			if n, err := mm.CopyOut(base, addr, []byte{1}, usermem.IOOpts{IgnorePermissions: true}); n != 1 || err != nil {
+				t.Fatalf("initial write: n=%d, err=%v", n, err)
+			}
+			handled, copied := 0, 0
+			ctx := memoryOOMContext{Context: base, handle: func() bool {
+				handled++
+				// MUnmap takes both MM locks; the callback must hold neither.
+				if err := mm.MUnmap(base, addr, hostarch.PageSize); err != nil {
+					t.Fatal(err)
+				}
+				return true
+			}}
+			start := addr + hostarch.Addr(hostarch.PageSize)
+			copy := func(blocks safemem.BlockSeq) (uint64, error) { copied++; return safemem.ZeroSeq(blocks) }
+			var n int64
+			if vector {
+				middle := start + hostarch.Addr(hostarch.PageSize)
+				n, err = mm.withVecInternalMappings(ctx, hostarch.AddrRangeSeqFromSlice([]hostarch.AddrRange{{Start: start, End: middle}, {Start: middle, End: middle + hostarch.Addr(hostarch.PageSize)}}), hostarch.Write, true, copy)
+			} else {
+				n, err = mm.withInternalMappings(ctx, hostarch.AddrRange{Start: start, End: start + 2*hostarch.Addr(hostarch.PageSize)}, hostarch.Write, true, copy)
+			}
+			if err != nil || n != 2*hostarch.PageSize || handled != 1 || copied != 1 {
+				t.Fatalf("copy after OOM: n=%d, err=%v, handled=%d, copied=%d", n, err, handled, copied)
+			}
+			if _, err := mm.CopyOut(ctx, 0, []byte{1}, usermem.IOOpts{}); err != linuxerr.EFAULT || handled != 1 {
+				t.Fatalf("invalid address: err=%v, handled=%d", err, handled)
+			}
+		})
+	}
+}
+
 func TestMappedFileRanges(t *testing.T) {
 	target := &pgalloc.MemoryFile{}
 	other := &pgalloc.MemoryFile{}