Pass memory cgroup id in pgalloc.Allocate()

For memory accounting per task, the memory cgroup id of the task is required.
This CL retrieves the cgroup id from the task context and stores it in the
AllocOpts struct which will later be used to account for memory usage.

PiperOrigin-RevId: 546996195

pkg/sentry/fsimpl/iouringfs/iouringfs.go

@@ -124,8 +124,8 @@ func New(ctx context.Context, vfsObj *vfs.VirtualFilesystem, entries uint32, par
 	ringsBufferSize = uint64(hostarch.Addr(ringsBufferSize).MustRoundUp())
 
 	mf := mfp.MemoryFile()
-
-	rbfr, err := mf.Allocate(ringsBufferSize, pgalloc.AllocOpts{Kind: usage.Anonymous})
+	memCgID := pgalloc.MemoryCgroupIDFromContext(ctx)
+	rbfr, err := mf.Allocate(ringsBufferSize, pgalloc.AllocOpts{Kind: usage.Anonymous, MemCgID: memCgID})
 	if err != nil {
 		return nil, linuxerr.ENOMEM
 	}
@@ -133,7 +133,7 @@ func New(ctx context.Context, vfsObj *vfs.VirtualFilesystem, entries uint32, par
 	// Allocate enough space to store the given number of submission queue entries.
 	sqEntriesSize := uint64(numSqEntries * uint32((*linux.IOUringSqe)(nil).SizeBytes()))
 	sqEntriesSize = uint64(hostarch.Addr(sqEntriesSize).MustRoundUp())
-	sqefr, err := mf.Allocate(sqEntriesSize, pgalloc.AllocOpts{Kind: usage.Anonymous})
+	sqefr, err := mf.Allocate(sqEntriesSize, pgalloc.AllocOpts{Kind: usage.Anonymous, MemCgID: memCgID})
 	if err != nil {
 		return nil, linuxerr.ENOMEM
 	}

pkg/sentry/fsimpl/tmpfs/regular_file.go

@@ -427,7 +427,7 @@ func (fd *regularFileFD) PRead(ctx context.Context, dst usermem.IOSequence, offs
 		return 0, nil
 	}
 	f := fd.inode().impl.(*regularFile)
-	rw := getRegularFileReadWriter(f, offset)
+	rw := getRegularFileReadWriter(f, offset, 0)
 	n, err := dst.CopyOutFrom(ctx, rw)
 	putRegularFileReadWriter(rw)
 	fd.inode().touchAtime(fd.vfsfd.Mount())
@@ -489,7 +489,7 @@ func (fd *regularFileFD) pwrite(ctx context.Context, src usermem.IOSequence, off
 	src = src.TakeFirst64(srclen)
 
 	// Perform the write.
-	rw := getRegularFileReadWriter(f, offset)
+	rw := getRegularFileReadWriter(f, offset, pgalloc.MemoryCgroupIDFromContext(ctx))
 	n, err := src.CopyInTo(ctx, rw)
 
 	f.inode.touchCMtimeLocked()
@@ -559,6 +559,10 @@ type regularFileReadWriter struct {
 	// Offset into the file to read/write at. Note that this may be
 	// different from the FD offset if PRead/PWrite is used.
 	off uint64
+
+	// memCgID is the memory cgroup ID used for accounting the allocated
+	// pages.
+	memCgID uint32
 }
 
 var regularFileReadWriterPool = sync.Pool{
@@ -567,10 +571,11 @@ var regularFileReadWriterPool = sync.Pool{
 	},
 }
 
-func getRegularFileReadWriter(file *regularFile, offset int64) *regularFileReadWriter {
+func getRegularFileReadWriter(file *regularFile, offset int64, memCgID uint32) *regularFileReadWriter {
 	rw := regularFileReadWriterPool.Get().(*regularFileReadWriter)
 	rw.file = file
 	rw.off = uint64(offset)
+	rw.memCgID = memCgID
 	return rw
 }
 
@@ -728,7 +733,7 @@ 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, pgalloc.AllocateAndWritePopulate, safemem.ReaderFunc(func(dsts safemem.BlockSeq) (uint64, error) {
+			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
 			}))

pkg/sentry/fsutil/file_range_set.go

@@ -114,6 +114,7 @@ func (frs *FileRangeSet) PagesToFill(required, optional memmap.MappableRange) ui
 func (frs *FileRangeSet) Fill(ctx context.Context, required, optional memmap.MappableRange, fileSize uint64, mf *pgalloc.MemoryFile, kind usage.MemoryKind, allocMode pgalloc.AllocationMode, readAt func(ctx context.Context, dsts safemem.BlockSeq, offset uint64) (uint64, error)) (uint64, error) {
 	gap := frs.LowerBoundGap(required.Start)
 	var pagesAlloced uint64
+	memCgID := pgalloc.MemoryCgroupIDFromContext(ctx)
 	for gap.Ok() && gap.Start() < required.End {
 		if gap.Range().Length() == 0 {
 			gap = gap.NextGap()
@@ -122,7 +123,7 @@ 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, allocMode, safemem.ReaderFunc(func(dsts safemem.BlockSeq) (uint64, error) {
+		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) {

pkg/sentry/kernel/kcov.go

@@ -242,7 +242,11 @@ func (kcov *Kcov) ConfigureMMap(ctx context.Context, opts *memmap.MMapOpts) erro
 
 	if kcov.mappable == nil {
 		// Set up the kcov area.
-		fr, err := kcov.mfp.MemoryFile().Allocate(kcov.size*8, pgalloc.AllocOpts{Kind: usage.Anonymous})
+		opts := pgalloc.AllocOpts{
+			Kind:    usage.Anonymous,
+			MemCgID: pgalloc.MemoryCgroupIDFromContext(ctx),
+		}
+		fr, err := kcov.mfp.MemoryFile().Allocate(kcov.size*8, opts)
 		if err != nil {
 			return err
 		}

pkg/sentry/kernel/kernel.go

@@ -774,6 +774,8 @@ func (ctx *createProcessContext) Value(key any) any {
 		return ctx.kernel.RealtimeClock()
 	case limits.CtxLimits:
 		return ctx.args.Limits
+	case pgalloc.CtxMemoryCgroupID:
+		return ctx.getMemoryCgroupID()
 	case pgalloc.CtxMemoryFile:
 		return ctx.kernel.mf
 	case pgalloc.CtxMemoryFileProvider:
@@ -793,6 +795,17 @@ func (ctx *createProcessContext) Value(key any) any {
 	}
 }
 
+func (ctx *createProcessContext) getMemoryCgroupID() uint32 {
+	for cg := range ctx.args.InitialCgroups {
+		for _, ctl := range cg.Controllers() {
+			if ctl.Type() == CgroupControllerMemory {
+				return cg.ID()
+			}
+		}
+	}
+	return InvalidCgroupID
+}
+
 // CreateProcess creates a new task in a new thread group with the given
 // options. The new task has no parent and is in the root PID namespace.
 //
@@ -1656,6 +1669,7 @@ func (k *Kernel) ReleaseCgroupHierarchy(hid uint32) {
 		for cg := range t.cgroups {
 			if cg.HierarchyID() == hid {
 				cg.Leave(t)
+				t.resetMemCgID(cg)
 				delete(t.cgroups, cg)
 				releasedCGs = append(releasedCGs, cg)
 				// A task can't be part of multiple cgroups from the same

pkg/sentry/kernel/shm/shm.go

@@ -211,7 +211,7 @@ func (r *Registry) newShmLocked(ctx context.Context, pid int32, key ipc.Key, cre
 	}
 
 	effectiveSize := uint64(hostarch.Addr(size).MustRoundUp())
-	fr, err := mfp.MemoryFile().Allocate(effectiveSize, pgalloc.AllocOpts{Kind: usage.Anonymous})
+	fr, err := mfp.MemoryFile().Allocate(effectiveSize, pgalloc.AllocOpts{Kind: usage.Anonymous, MemCgID: pgalloc.MemoryCgroupIDFromContext(ctx)})
 	if err != nil {
 		return nil, err
 	}

pkg/sentry/kernel/task.go

@@ -588,6 +588,9 @@ type Task struct {
 	// +checklocks:mu
 	cgroups map[Cgroup]struct{}
 
+	// memCgID is the memory cgroup id.
+	memCgID atomicbitops.Uint32
+
 	// userCounters is a pointer to a set of user counters.
 	//
 	// The userCounters pointer is exclusive to the task goroutine, but the

pkg/sentry/kernel/task_cgroup.go

@@ -51,6 +51,25 @@ func (t *Task) EnterInitialCgroups(parent *Task, initCgroups map[Cgroup]struct{}
 		// Since t isn't in any cgroup yet, we can skip the check against
 		// existing cgroups.
 		c.Enter(t)
+		t.setMemCgID(c)
+	}
+}
+
+// TODO(b/277772401): setMemCgIDLocked should be called after adding support for
+// task migration for cgroup memory controllers.
+func (t *Task) setMemCgID(cg Cgroup) {
+	for _, ctl := range cg.Controllers() {
+		if ctl.Type() == CgroupControllerMemory {
+			t.memCgID.Store(cg.ID())
+		}
+	}
+}
+
+func (t *Task) resetMemCgID(cg Cgroup) {
+	for _, ctl := range cg.Controllers() {
+		if ctl.Type() == CgroupControllerMemory {
+			t.memCgID.Store(0)
+		}
 	}
 }
 
@@ -82,6 +101,7 @@ func (t *Task) enterCgroupLocked(c Cgroup) {
 	c.IncRef()
 	t.cgroups[c] = struct{}{}
 	c.Enter(t)
+	t.setMemCgID(c)
 }
 
 // +checklocks:t.mu
@@ -101,6 +121,7 @@ func (t *Task) LeaveCgroups() {
 	for c := range cgs {
 		c.Leave(t)
 	}
+	t.memCgID.Store(0)
 	t.mu.Unlock()
 	t.tg.pidns.owner.mu.Unlock()
 

pkg/sentry/kernel/task_context.go

@@ -111,6 +111,8 @@ func (t *Task) contextValue(key any, isTaskGoroutine bool) any {
 		return func(sig linux.Signal) error {
 			return t.SendSignal(SignalInfoNoInfo(sig, t, t))
 		}
+	case pgalloc.CtxMemoryCgroupID:
+		return t.memCgID.Load()
 	case pgalloc.CtxMemoryFile:
 		return t.k.mf
 	case pgalloc.CtxMemoryFileProvider:

pkg/sentry/mm/aio_context.go

@@ -252,8 +252,8 @@ type aioMappable struct {
 
 var aioRingBufferSize = uint64(hostarch.Addr(linux.AIORingSize).MustRoundUp())
 
-func newAIOMappable(mfp pgalloc.MemoryFileProvider) (*aioMappable, error) {
-	fr, err := mfp.MemoryFile().Allocate(aioRingBufferSize, pgalloc.AllocOpts{Kind: usage.Anonymous})
+func newAIOMappable(ctx context.Context, mfp pgalloc.MemoryFileProvider) (*aioMappable, error) {
+	fr, err := mfp.MemoryFile().Allocate(aioRingBufferSize, pgalloc.AllocOpts{Kind: usage.Anonymous, MemCgID: pgalloc.MemoryCgroupIDFromContext(ctx)})
 	if err != nil {
 		return nil, err
 	}
@@ -368,7 +368,7 @@ func (mm *MemoryManager) NewAIOContext(ctx context.Context, events uint32) (uint
 	// libaio peeks inside looking for a magic number. This function allocates
 	// a page per context and keeps it set to zeroes to ensure it will not
 	// match AIO_RING_MAGIC and make libaio happy.
-	m, err := newAIOMappable(mm.mfp)
+	m, err := newAIOMappable(ctx, mm.mfp)
 	if err != nil {
 		return 0, err
 	}

pkg/sentry/mm/pma.go

@@ -205,7 +205,8 @@ func (mm *MemoryManager) getPMAsInternalLocked(ctx context.Context, vseg vmaIter
 		}
 	}
 
-	opts := pgalloc.AllocOpts{Kind: usage.Anonymous, Dir: pgalloc.BottomUp}
+	memCgID := pgalloc.MemoryCgroupIDFromContext(ctx)
+	opts := pgalloc.AllocOpts{Kind: usage.Anonymous, Dir: pgalloc.BottomUp, MemCgID: memCgID}
 	vma := vseg.ValuePtr()
 	if uintptr(ar.Start) < atomic.LoadUintptr(&vma.lastFault) {
 		// Detect cases where memory is accessed downwards and change memory file
@@ -373,7 +374,7 @@ 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, pgalloc.AllocateAndWritePopulate, &safemem.BlockSeqReader{mm.internalMappingsLocked(pseg, copyAR)})
+					fr, err := mf.AllocateAndFill(uint64(copyAR.Length()), usage.Anonymous, memCgID, pgalloc.AllocateAndWritePopulate, &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

pkg/sentry/pgalloc/context.go

@@ -27,6 +27,9 @@ const (
 
 	// CtxMemoryFileProvider is a Context.Value key for a MemoryFileProvider.
 	CtxMemoryFileProvider
+
+	// CtxMemoryCgroupID is the memory cgroup id which the task belongs to.
+	CtxMemoryCgroupID
 )
 
 // MemoryFileFromContext returns the MemoryFile used by ctx, or nil if no such
@@ -46,3 +49,12 @@ func MemoryFileProviderFromContext(ctx context.Context) MemoryFileProvider {
 	}
 	return nil
 }
+
+// MemoryCgroupIDFromContext returns the memory cgroup id of the ctx, or
+// zero if the ctx does not belong to any memory cgroup.
+func MemoryCgroupIDFromContext(ctx context.Context) uint32 {
+	if v := ctx.Value(CtxMemoryCgroupID); v != nil {
+		return v.(uint32)
+	}
+	return 0
+}

pkg/sentry/pgalloc/pgalloc.go

@@ -421,6 +421,9 @@ func (f *MemoryFile) Destroy() {
 type AllocOpts struct {
 	Kind usage.MemoryKind
 	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
 }
 
 // Allocate returns a range of initially-zeroed pages of the given length with
@@ -625,7 +628,7 @@ const (
 // Preconditions:
 //   - length > 0.
 //   - length must be page-aligned.
-func (f *MemoryFile) AllocateAndFill(length uint64, kind usage.MemoryKind, allocMode AllocationMode, r safemem.Reader) (memmap.FileRange, error) {
+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
@@ -643,7 +646,7 @@ func (f *MemoryFile) AllocateAndFill(length uint64, kind usage.MemoryKind, alloc
 		// and we also additionally incur useless disk writebacks.
 		allocMode = AllocateOnly
 	}
-	fr, err := f.Allocate(length, AllocOpts{Kind: kind})
+	fr, err := f.Allocate(length, AllocOpts{Kind: kind, MemCgID: memCgID})
 	if err != nil {
 		return memmap.FileRange{}, err
 	}