Simply pgalloc.MemoryFile.Allocate() API.

Earlier we had another method for AllocateAndPopulate. This change moves the additional (optional) parameters to pgalloc.AllocOpts. * Makes `Reader` parameter optional. Not all callers have data to fill into the allocated pages. Avoids having to define no-op readers. * Avoid calling MapInternal() when `Reader == nil` and `Mode != AllocateAndWritePopulate`. * Avoid calling safemem.ReadFullToBlocks() when `Reader == nil`. * Move AllocationMode upgrade/downgrade logic to callsites. This makes Allocate() simpler. PiperOrigin-RevId: 547688619
2026-05-22 17:12:49 -07:00 · 2023-07-12 22:12:08 -07:00
parent 415845e2a1
commit dd3a7a1fb9
4 changed files with 179 additions and 147 deletions
@@ -307,10 +307,7 @@ func (rf *regularFile) Translate(ctx context.Context, required, optional memmap.
 		}
 		optional = required
 	}
-	pagesAlloced, cerr := rf.data.Fill(ctx, required, optional, rf.size.RacyLoad(), rf.inode.fs.mf, rf.memoryUsageKind, pgalloc.AllocateOnly, func(_ context.Context, dsts safemem.BlockSeq, _ uint64) (uint64, error) {
-		// Newly-allocated pages are zeroed, so we don't need to do anything.
-		return dsts.NumBytes(), nil
-	})
+	pagesAlloced, cerr := rf.data.Fill(ctx, required, optional, rf.size.RacyLoad(), rf.inode.fs.mf, rf.memoryUsageKind, pgalloc.AllocateOnly, nil /* r */)
 	// rf.data.Fill() may fail mid-way. We still want to account any pages that
 	// were allocated, irrespective of an error.
 	rf.inode.fs.adjustPageAcct(pagesToFill, pagesAlloced)
@@ -387,10 +384,17 @@ func (fd *regularFileFD) Allocate(ctx context.Context, mode, offset, length uint
 	}
 	// Given our definitions in pgalloc, fallocate(2) semantics imply that pages
 	// in the MemoryFile must be committed, in addition to being allocated.
-	pagesAlloced, err := f.data.Fill(ctx, required, required, newSize, f.inode.fs.mf, f.memoryUsageKind, pgalloc.AllocateAndCommit, func(_ context.Context, dsts safemem.BlockSeq, _ uint64) (uint64, error) {
-		// Newly-allocated pages are zeroed, so we don't need to do anything.
-		return dsts.NumBytes(), nil
-	})
+	allocMode := pgalloc.AllocateAndCommit
+	if !f.inode.fs.mf.IsDiskBacked() {
+		// Upgrade to AllocateAndWritePopulate for memory(shmem)-backed files. We
+		// take a more aggressive approach in populating pages for memory-backed
+		// MemoryFiles. shmem pages are subject to swap rather than disk writeback.
+		// They are not likely to be swapped before they are written to. Hence it
+		// is beneficial to populate (in addition to commit) shmem pages to avoid
+		// faulting page-by-page when these pages are written to in the future.
+		allocMode = pgalloc.AllocateAndWritePopulate
+	}
+	pagesAlloced, err := f.data.Fill(ctx, required, required, newSize, f.inode.fs.mf, f.memoryUsageKind, allocMode, nil /* r */)
 	// f.data.Fill() may fail mid-way. We still want to account any pages that
 	// were allocated, irrespective of an error.
 	f.inode.fs.adjustPageAcct(pagesToFill, pagesAlloced)
@@ -733,10 +737,21 @@ func (rw *regularFileReadWriter) WriteFromBlocks(srcs safemem.BlockSeq) (uint64,
 				goto exitLoop
 			}
 			gapMR.End = gapMR.Start + (hostarch.PageSize * pagesReserved)
-			fr, err := rw.file.inode.fs.mf.AllocateAndFill(gapMR.Length(), rw.file.memoryUsageKind, rw.memCgID, pgalloc.AllocateAndWritePopulate, safemem.ReaderFunc(func(dsts safemem.BlockSeq) (uint64, error) {
-				// No-op here. The write to dsts will happen in the next iteration.
-				return dsts.NumBytes(), nil
-			}))
+			allocMode := pgalloc.AllocateAndWritePopulate
+			if rw.file.inode.fs.mf.IsDiskBacked() {
+				// Don't populate pages for disk-backed files. Benchmarking showed that
+				// disk-backed pages are likely to be written back to disk before we
+				// can write to them. The pages fault again on write anyways. In total,
+				// prepopulating disk-backed pages deteriorates performance as it fails
+				// to eliminate future page faults and we also additionally incur
+				// useless disk writebacks.
+				allocMode = pgalloc.AllocateOnly
+			}
+			fr, err := rw.file.inode.fs.mf.Allocate(gapMR.Length(), pgalloc.AllocOpts{
+				Kind:    rw.file.memoryUsageKind,
+				Mode:    allocMode,
+				MemCgID: rw.memCgID,
+			})
 			if err != nil {
 				retErr = err
 				rw.file.inode.fs.unaccountPages(pagesReserved)
@@ -123,45 +123,53 @@ func (frs *FileRangeSet) Fill(ctx context.Context, required, optional memmap.Map
 		gr := gap.Range().Intersect(optional)

 		// Read data into the gap.
-		fr, err := mf.AllocateAndFill(gr.Length(), kind, memCgID, allocMode, safemem.ReaderFunc(func(dsts safemem.BlockSeq) (uint64, error) {
-			var done uint64
-			for !dsts.IsEmpty() {
-				n, err := func() (uint64, error) {
-					off := gr.Start + done
-					if off >= fileSize {
-						return 0, io.EOF
-					}
-					if off+dsts.NumBytes() > fileSize {
-						rd := fileSize - off
-						n, err := readAt(ctx, dsts.TakeFirst64(rd), off)
-						if n == rd && err == nil {
-							return n, io.EOF
+		opts := pgalloc.AllocOpts{
+			Kind:    kind,
+			Mode:    allocMode,
+			MemCgID: memCgID,
+		}
+		if readAt != nil {
+			opts.Reader = safemem.ReaderFunc(func(dsts safemem.BlockSeq) (uint64, error) {
+				var done uint64
+				for !dsts.IsEmpty() {
+					n, err := func() (uint64, error) {
+						off := gr.Start + done
+						if off >= fileSize {
+							return 0, io.EOF
 						}
-						return n, err
-					}
-					return readAt(ctx, dsts, off)
-				}()
-				done += n
-				dsts = dsts.DropFirst64(n)
-				if err != nil {
-					if err == io.EOF {
-						// MemoryFile.AllocateAndFill truncates down to a page
-						// boundary, but FileRangeSet.Fill is supposed to
-						// zero-fill to the end of the page in this case.
-						donepgaddr, ok := hostarch.Addr(done).RoundUp()
-						if donepg := uint64(donepgaddr); ok && donepg != done {
-							dsts.DropFirst64(donepg - done)
-							done = donepg
-							if dsts.IsEmpty() {
-								return done, nil
+						if off+dsts.NumBytes() > fileSize {
+							rd := fileSize - off
+							n, err := readAt(ctx, dsts.TakeFirst64(rd), off)
+							if n == rd && err == nil {
+								return n, io.EOF
+							}
+							return n, err
+						}
+						return readAt(ctx, dsts, off)
+					}()
+					done += n
+					dsts = dsts.DropFirst64(n)
+					if err != nil {
+						if err == io.EOF {
+							// MemoryFile.AllocateAndFill truncates down to a page
+							// boundary, but FileRangeSet.Fill is supposed to
+							// zero-fill to the end of the page in this case.
+							donepgaddr, ok := hostarch.Addr(done).RoundUp()
+							if donepg := uint64(donepgaddr); ok && donepg != done {
+								dsts.DropFirst64(donepg - done)
+								done = donepg
+								if dsts.IsEmpty() {
+									return done, nil
+								}
 							}
 						}
+						return done, err
 					}
-					return done, err
 				}
-			}
-			return done, nil
-		}))
+				return done, nil
+			})
+		}
+		fr, err := mf.Allocate(gr.Length(), opts)

 		// Store anything we managed to read into the cache.
 		if done := fr.Length(); done != 0 {
@@ -374,7 +374,12 @@ func (mm *MemoryManager) getPMAsInternalLocked(ctx context.Context, vseg vmaIter
 						return pstart, pseg.PrevGap(), err
 					}
 					// Copy contents.
-					fr, err := mf.AllocateAndFill(uint64(copyAR.Length()), usage.Anonymous, memCgID, pgalloc.AllocateAndWritePopulate, &safemem.BlockSeqReader{mm.internalMappingsLocked(pseg, copyAR)})
+					fr, err := mf.Allocate(uint64(copyAR.Length()), pgalloc.AllocOpts{
+						Kind:    usage.Anonymous,
+						Mode:    pgalloc.AllocateAndWritePopulate,
+						MemCgID: memCgID,
+						Reader:  &safemem.BlockSeqReader{mm.internalMappingsLocked(pseg, copyAR)},
+					})
 					if _, ok := err.(safecopy.BusError); ok {
 						// If we got SIGBUS during the copy, deliver SIGBUS to
 						// userspace (instead of SIGSEGV) if we're breaking
@@ -417,13 +417,53 @@ func (f *MemoryFile) Destroy() {
 	f.reclaimCond.Signal()
 }

+// AllocationMode provides a way to inform the pgalloc API how to allocate
+// memory and pages on the host.
+// A page will exist in one of the following incremental states:
+//  1. Allocated: A page is allocated if it was returned by Allocate() and its
+//     reference count hasn't dropped to 0 since then.
+//  2. Committed: As described in MemoryFile documentation above, a page is
+//     committed if the host kernel is spending resources to store its
+//     contents. A committed page is implicitly allocated.
+//  3. Populated: A page is populated for reading/writing in a page table
+//     hierarchy if it has a page table entry that permits reading/writing
+//     respectively. A populated page is implicitly committed, since the page
+//     table entry needs a physical page to point to, but not vice versa.
+type AllocationMode int
+
+const (
+	// AllocateOnly indicates that pages need to only be allocated.
+	AllocateOnly AllocationMode = iota
+	// AllocateAndCommit indicates that pages need to be committed, in addition
+	// to being allocated.
+	AllocateAndCommit
+	// AllocateAndWritePopulate indicates that writable pages should ideally be
+	// populated in the page table, in addition to being allocated. This is a
+	// suggestion, not a requirement.
+	AllocateAndWritePopulate
+)
+
 // AllocOpts are options used in MemoryFile.Allocate.
 type AllocOpts struct {
+	// Kind is the memory kind to be used for accounting.
 	Kind usage.MemoryKind
-	Dir  Direction
+	// Dir indicates the direction in which offsets are allocated.
+	Dir Direction
 	// MemCgID is the memory cgroup ID and the zero value indicates that
 	// the memory will not be accounted to any cgroup.
 	MemCgID uint32
+	// Mode allows the callers to select how the pages are allocated in the
+	// MemoryFile. Callers that will fill the allocated memory by writing to it
+	// should pass AllocateAndWritePopulate to avoid faulting page-by-page. Callers
+	// that will fill the allocated memory by invoking host system calls should
+	// pass AllocateOnly.
+	Mode AllocationMode
+	// If Reader is provided, the allocated memory is filled by calling
+	// ReadToBlocks() repeatedly until either length bytes are read or a non-nil
+	// error is returned. It returns the allocated memory, truncated down to the
+	// nearest page. If this is shorter than length bytes due to an error
+	// returned by ReadToBlocks(), it returns the partially filled fr and error.
+	Reader safemem.Reader
 }

 // Allocate returns a range of initially-zeroed pages of the given length with
@@ -434,6 +474,63 @@ type AllocOpts struct {
 //
 // Preconditions: length must be page-aligned and non-zero.
 func (f *MemoryFile) Allocate(length uint64, opts AllocOpts) (memmap.FileRange, error) {
+	fr, err := f.allocate(length, &opts)
+	if err != nil {
+		return memmap.FileRange{}, err
+	}
+	var dsts safemem.BlockSeq
+	switch opts.Mode {
+	case AllocateOnly: // Allocation is handled above. Nothing more to do.
+	case AllocateAndCommit:
+		if err := f.commitFile(fr); err != nil {
+			f.DecRef(fr)
+			return memmap.FileRange{}, err
+		}
+	case AllocateAndWritePopulate:
+		dsts, err = f.MapInternal(fr, hostarch.Write)
+		if err != nil {
+			f.DecRef(fr)
+			return memmap.FileRange{}, err
+		}
+		if canPopulate() {
+			rem := dsts
+			for {
+				if !tryPopulate(rem.Head()) {
+					break
+				}
+				rem = rem.Tail()
+				if rem.IsEmpty() {
+					break
+				}
+			}
+		}
+	default:
+		panic(fmt.Sprintf("unknown allocation mode: %d", opts.Mode))
+	}
+	if opts.Reader != nil {
+		if dsts.IsEmpty() {
+			dsts, err = f.MapInternal(fr, hostarch.Write)
+			if err != nil {
+				f.DecRef(fr)
+				return memmap.FileRange{}, err
+			}
+		}
+		n, err := safemem.ReadFullToBlocks(opts.Reader, dsts)
+		un := uint64(hostarch.Addr(n).RoundDown())
+		if un < length {
+			// Free unused memory and update fr to contain only the memory that is
+			// still allocated.
+			f.DecRef(memmap.FileRange{fr.Start + un, fr.End})
+			fr.End = fr.Start + un
+		}
+		if err != nil {
+			return fr, err
+		}
+	}
+	return fr, nil
+}
+
+func (f *MemoryFile) allocate(length uint64, opts *AllocOpts) (memmap.FileRange, error) {
 	if length == 0 || length%hostarch.PageSize != 0 {
 		panic(fmt.Sprintf("invalid allocation length: %#x", length))
 	}
@@ -492,7 +589,7 @@ func (f *MemoryFile) Allocate(length uint64, opts AllocOpts) (memmap.FileRange,
 // then forwards. This heuristic has important consequence for how sequential
 // mappings can be merged in the host VMAs, given that addresses for both
 // application and sentry mappings are allocated top-down (from higher to
-// lower addresses). The file is also grown expoentially in order to create
+// lower addresses). The file is also grown exponentially in order to create
 // space for mappings to be allocated downwards.
 //
 // Precondition: alignment must be a power of 2.
@@ -586,104 +683,6 @@ func findAvailableRangeBottomUp(usage *usageSet, length, alignment uint64) (memm
 	panic(fmt.Sprintf("NextLargeEnoughGap didn't return a gap at the end, length: %d", length))
 }

-// AllocationMode provides a way to inform the pgalloc API how to allocate
-// memory and pages on the host.
-// A page will exist in one of the following incremental states:
-//  1. Allocated: A page is allocated if it was returned by Allocate() and its
-//     reference count hasn't dropped to 0 since then.
-//  2. Committed: As described in MemoryFile documentation above, a page is
-//     committed if the host kernel is spending resources to store its
-//     contents. A committed page is implicitly allocated.
-//  3. Populated: A page is populated for reading/writing in a page table
-//     hierarchy if it has a page table entry that permits reading/writing
-//     respectively. A populated page is implicitly committed, since the page
-//     table entry needs a physical page to point to, but not vice versa.
-type AllocationMode int
-
-const (
-	// AllocateOnly indicates that pages need to only be allocated.
-	AllocateOnly AllocationMode = iota
-	// AllocateAndCommit indicates that pages need to be committed, in addition
-	// to being allocated.
-	AllocateAndCommit
-	// AllocateAndWritePopulate indicates that writable pages should ideally be
-	// populated in the page table, in addition to being allocated. This is a
-	// suggestion, not a requirement.
-	AllocateAndWritePopulate
-)
-
-// AllocateAndFill allocates memory of the given kind and fills it by calling
-// r.ReadToBlocks() repeatedly until either length bytes are read or a non-nil
-// error is returned. It returns the memory filled by r, truncated down to the
-// nearest page. If this is shorter than length bytes due to an error returned
-// by r.ReadToBlocks(), it returns that error.
-//
-// allocMode allows the callers to select how the pages are allocated in the
-// MemoryFile. Callers that will fill the allocated memory by writing to it
-// should pass AllocateAndWritePopulate to avoid faulting page-by-page. Callers
-// that will fill the allocated memory by invoking host system calls should
-// pass AllocateOnly. Note that the mode may be upgraded in certain scenarios
-// for performance. See implementation for more details.
-//
-// Preconditions:
-//   - length > 0.
-//   - length must be page-aligned.
-func (f *MemoryFile) AllocateAndFill(length uint64, kind usage.MemoryKind, memCgID uint32, allocMode AllocationMode, r safemem.Reader) (memmap.FileRange, error) {
-	if !f.opts.DiskBackedFile && allocMode == AllocateAndCommit {
-		// Upgrade to AllocateAndWritePopulate for memory(shmem)-backed files. We
-		// take a more aggressive approach in populating pages for memory-backed
-		// MemoryFiles. shmem pages are subject to swap rather than disk writeback.
-		// They are not likely to be swapped before they are written to. Hence it
-		// is beneficial to populate (in addition to commit) shmem pages to avoid
-		// faulting page-by-page when these pages are written to in the future.
-		allocMode = AllocateAndWritePopulate
-	} else if f.opts.DiskBackedFile && allocMode == AllocateAndWritePopulate {
-		// Don't populate pages for disk-backed files.
-		// Benchmarking showed that prepopulated pages are likely to be written
-		// back to disk before the application can write to them. The pages will
-		// fault again on write anyways. In total, prepopulating disk-backed pages
-		// deteriorates performance as it fails to eliminate future page faults
-		// and we also additionally incur useless disk writebacks.
-		allocMode = AllocateOnly
-	}
-	fr, err := f.Allocate(length, AllocOpts{Kind: kind, MemCgID: memCgID})
-	if err != nil {
-		return memmap.FileRange{}, err
-	}
-	if allocMode == AllocateAndCommit {
-		if err := f.commitFile(fr); err != nil {
-			f.DecRef(fr)
-			return memmap.FileRange{}, err
-		}
-	}
-	dsts, err := f.MapInternal(fr, hostarch.Write)
-	if err != nil {
-		f.DecRef(fr)
-		return memmap.FileRange{}, err
-	}
-	if allocMode == AllocateAndWritePopulate && canPopulate() {
-		rem := dsts
-		for {
-			if !tryPopulate(rem.Head()) {
-				break
-			}
-			rem = rem.Tail()
-			if rem.IsEmpty() {
-				break
-			}
-		}
-	}
-	n, err := safemem.ReadFullToBlocks(r, dsts)
-	un := uint64(hostarch.Addr(n).RoundDown())
-	if un < length {
-		// Free unused memory and update fr to contain only the memory that is
-		// still allocated.
-		f.DecRef(memmap.FileRange{fr.Start + un, fr.End})
-		fr.End = fr.Start + un
-	}
-	return fr, err
-}
-
 var mlockDisabled atomicbitops.Uint32
 var madvPopulateWriteDisabled atomicbitops.Uint32

@@ -1290,6 +1289,11 @@ func (f *MemoryFile) FD() int {
 	return int(f.file.Fd())
 }

+// IsDiskBacked returns true if f is backed by a file on disk.
+func (f *MemoryFile) IsDiskBacked() bool {
+	return f.opts.DiskBackedFile
+}
+
 // String implements fmt.Stringer.String.
 //
 // Note that because f.String locks f.mu, calling f.String internally