From f8c4846f214db2cb577992b76f5ab0c992ee553a Mon Sep 17 00:00:00 2001 From: Nayana Bidari Date: Mon, 10 Jul 2023 15:14:24 -0700 Subject: [PATCH] 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 | 6 +++--- pkg/sentry/fsimpl/tmpfs/regular_file.go | 13 +++++++++---- pkg/sentry/fsutil/file_range_set.go | 3 ++- pkg/sentry/kernel/kcov.go | 6 +++++- pkg/sentry/kernel/kernel.go | 14 ++++++++++++++ pkg/sentry/kernel/shm/shm.go | 2 +- pkg/sentry/kernel/task.go | 3 +++ pkg/sentry/kernel/task_cgroup.go | 21 +++++++++++++++++++++ pkg/sentry/kernel/task_context.go | 2 ++ pkg/sentry/mm/aio_context.go | 6 +++--- pkg/sentry/mm/pma.go | 5 +++-- pkg/sentry/pgalloc/context.go | 12 ++++++++++++ pkg/sentry/pgalloc/pgalloc.go | 7 +++++-- 13 files changed, 83 insertions(+), 17 deletions(-) diff --git a/pkg/sentry/fsimpl/iouringfs/iouringfs.go b/pkg/sentry/fsimpl/iouringfs/iouringfs.go index df9108e23..a5edd44e5 100644 --- a/pkg/sentry/fsimpl/iouringfs/iouringfs.go +++ b/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 } diff --git a/pkg/sentry/fsimpl/tmpfs/regular_file.go b/pkg/sentry/fsimpl/tmpfs/regular_file.go index dc50af169..91639db17 100644 --- a/pkg/sentry/fsimpl/tmpfs/regular_file.go +++ b/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 })) diff --git a/pkg/sentry/fsutil/file_range_set.go b/pkg/sentry/fsutil/file_range_set.go index 46736fbd0..af7840a96 100644 --- a/pkg/sentry/fsutil/file_range_set.go +++ b/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) { diff --git a/pkg/sentry/kernel/kcov.go b/pkg/sentry/kernel/kcov.go index c974d807f..27c3515e0 100644 --- a/pkg/sentry/kernel/kcov.go +++ b/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 } diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go index 4e7ee69c3..d9b1eb4c8 100644 --- a/pkg/sentry/kernel/kernel.go +++ b/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 diff --git a/pkg/sentry/kernel/shm/shm.go b/pkg/sentry/kernel/shm/shm.go index 33e0ceafa..ef62a9c30 100644 --- a/pkg/sentry/kernel/shm/shm.go +++ b/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 } diff --git a/pkg/sentry/kernel/task.go b/pkg/sentry/kernel/task.go index 97516c6c6..c29b76123 100644 --- a/pkg/sentry/kernel/task.go +++ b/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 diff --git a/pkg/sentry/kernel/task_cgroup.go b/pkg/sentry/kernel/task_cgroup.go index ca577061d..970540ee2 100644 --- a/pkg/sentry/kernel/task_cgroup.go +++ b/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() diff --git a/pkg/sentry/kernel/task_context.go b/pkg/sentry/kernel/task_context.go index 2357eaad1..f57d79434 100644 --- a/pkg/sentry/kernel/task_context.go +++ b/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: diff --git a/pkg/sentry/mm/aio_context.go b/pkg/sentry/mm/aio_context.go index dd1879810..d80ae11d7 100644 --- a/pkg/sentry/mm/aio_context.go +++ b/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 } diff --git a/pkg/sentry/mm/pma.go b/pkg/sentry/mm/pma.go index 8ddf42db7..572ecd734 100644 --- a/pkg/sentry/mm/pma.go +++ b/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 diff --git a/pkg/sentry/pgalloc/context.go b/pkg/sentry/pgalloc/context.go index d25215418..5350e1f5b 100644 --- a/pkg/sentry/pgalloc/context.go +++ b/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 +} diff --git a/pkg/sentry/pgalloc/pgalloc.go b/pkg/sentry/pgalloc/pgalloc.go index 732804800..2bbb63c10 100644 --- a/pkg/sentry/pgalloc/pgalloc.go +++ b/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 }