From aff516812158f00b768f7c5358a31e634018b347 Mon Sep 17 00:00:00 2001 From: Nayana Bidari Date: Thu, 20 Jul 2023 09:37:28 -0700 Subject: [PATCH] Plumb memory cgroup id in memmap.IncRef. Update the memmap IncRef method to pass memory cgroup id and store it in the FrameRefSet which will be used for memory accounting. During DecRef, the memCgID from the FrameRefSet will be retrieved and passed to MemoryLocked.Dec to remove the memory from the cgroup. PiperOrigin-RevId: 549656411 --- pkg/sentry/devices/nvproxy/frontend_mmap.go | 2 +- pkg/sentry/devices/nvproxy/uvm_mmap.go | 2 +- pkg/sentry/fsimpl/gofer/regular_file.go | 4 +- pkg/sentry/fsimpl/gofer/special_file.go | 4 +- pkg/sentry/fsimpl/kernfs/mmap_util.go | 4 +- pkg/sentry/fsutil/BUILD | 2 +- pkg/sentry/fsutil/frame_ref_set.go | 54 ++++++++++++--------- pkg/sentry/memmap/memmap.go | 6 ++- pkg/sentry/mm/lifecycle.go | 3 +- pkg/sentry/mm/pma.go | 11 +++-- pkg/sentry/pgalloc/pgalloc.go | 2 +- 11 files changed, 54 insertions(+), 40 deletions(-) diff --git a/pkg/sentry/devices/nvproxy/frontend_mmap.go b/pkg/sentry/devices/nvproxy/frontend_mmap.go index b57f13924..95b8ad6cb 100644 --- a/pkg/sentry/devices/nvproxy/frontend_mmap.go +++ b/pkg/sentry/devices/nvproxy/frontend_mmap.go @@ -65,7 +65,7 @@ type frontendFDMemmapFile struct { } // IncRef implements memmap.File.IncRef. -func (mf *frontendFDMemmapFile) IncRef(fr memmap.FileRange) { +func (mf *frontendFDMemmapFile) IncRef(fr memmap.FileRange, memCgID uint32) { } // DecRef implements memmap.File.DecRef. diff --git a/pkg/sentry/devices/nvproxy/uvm_mmap.go b/pkg/sentry/devices/nvproxy/uvm_mmap.go index dad7c9ff9..c0d06941e 100644 --- a/pkg/sentry/devices/nvproxy/uvm_mmap.go +++ b/pkg/sentry/devices/nvproxy/uvm_mmap.go @@ -67,7 +67,7 @@ type uvmFDMemmapFile struct { } // IncRef implements memmap.File.IncRef. -func (mf *uvmFDMemmapFile) IncRef(fr memmap.FileRange) { +func (mf *uvmFDMemmapFile) IncRef(fr memmap.FileRange, memCgID uint32) { } // DecRef implements memmap.File.DecRef. diff --git a/pkg/sentry/fsimpl/gofer/regular_file.go b/pkg/sentry/fsimpl/gofer/regular_file.go index 30ee20fee..ad1bd7ae3 100644 --- a/pkg/sentry/fsimpl/gofer/regular_file.go +++ b/pkg/sentry/fsimpl/gofer/regular_file.go @@ -931,9 +931,9 @@ type dentryPlatformFile struct { } // IncRef implements memmap.File.IncRef. -func (d *dentryPlatformFile) IncRef(fr memmap.FileRange) { +func (d *dentryPlatformFile) IncRef(fr memmap.FileRange, memCgID uint32) { d.dataMu.Lock() - d.fdRefs.IncRefAndAccount(fr) + d.fdRefs.IncRefAndAccount(fr, memCgID) d.dataMu.Unlock() } diff --git a/pkg/sentry/fsimpl/gofer/special_file.go b/pkg/sentry/fsimpl/gofer/special_file.go index 5292c5447..ebdecbfbc 100644 --- a/pkg/sentry/fsimpl/gofer/special_file.go +++ b/pkg/sentry/fsimpl/gofer/special_file.go @@ -485,10 +485,10 @@ func (fd *specialFileFD) InvalidateUnsavable(ctx context.Context) error { } // IncRef implements memmap.File.IncRef. -func (fd *specialFileFD) IncRef(fr memmap.FileRange) { +func (fd *specialFileFD) IncRef(fr memmap.FileRange, memCgID uint32) { fd.fileRefsMu.Lock() defer fd.fileRefsMu.Unlock() - fd.fileRefs.IncRefAndAccount(fr) + fd.fileRefs.IncRefAndAccount(fr, memCgID) } // DecRef implements memmap.File.DecRef. diff --git a/pkg/sentry/fsimpl/kernfs/mmap_util.go b/pkg/sentry/fsimpl/kernfs/mmap_util.go index 269f9fb9e..acb3991b7 100644 --- a/pkg/sentry/fsimpl/kernfs/mmap_util.go +++ b/pkg/sentry/fsimpl/kernfs/mmap_util.go @@ -52,9 +52,9 @@ type inodePlatformFile struct { var _ memmap.File = (*inodePlatformFile)(nil) // IncRef implements memmap.File.IncRef. -func (i *inodePlatformFile) IncRef(fr memmap.FileRange) { +func (i *inodePlatformFile) IncRef(fr memmap.FileRange, memCgID uint32) { i.fdRefsMu.Lock() - i.fdRefs.IncRefAndAccount(fr) + i.fdRefs.IncRefAndAccount(fr, memCgID) i.fdRefsMu.Unlock() } diff --git a/pkg/sentry/fsutil/BUILD b/pkg/sentry/fsutil/BUILD index 8dff960c6..3591067d0 100644 --- a/pkg/sentry/fsutil/BUILD +++ b/pkg/sentry/fsutil/BUILD @@ -50,7 +50,7 @@ go_template_instance( types = { "Key": "uint64", "Range": "memmap.FileRange", - "Value": "uint64", + "Value": "FrameRefSegInfo", "Functions": "FrameRefSetFunctions", }, ) diff --git a/pkg/sentry/fsutil/frame_ref_set.go b/pkg/sentry/fsutil/frame_ref_set.go index a0035a423..6f95cffcb 100644 --- a/pkg/sentry/fsutil/frame_ref_set.go +++ b/pkg/sentry/fsutil/frame_ref_set.go @@ -21,6 +21,15 @@ import ( "gvisor.dev/gvisor/pkg/sentry/usage" ) +// FrameRefSegInfo holds reference count and memory cgroup id of the segment. +type FrameRefSegInfo struct { + // refs indicates the reference count of the segment. + refs uint64 + // memCgID is the memory cgroup id of the first task which touches the + // segment. This will not be changed over the lifetime of the segment. + memCgID uint32 +} + // FrameRefSetFunctions implements segment.Functions for FrameRefSet. type FrameRefSetFunctions struct{} @@ -35,39 +44,41 @@ func (FrameRefSetFunctions) MaxKey() uint64 { } // ClearValue implements segment.Functions.ClearValue. -func (FrameRefSetFunctions) ClearValue(val *uint64) { +func (FrameRefSetFunctions) ClearValue(val *FrameRefSegInfo) { } // Merge implements segment.Functions.Merge. -func (FrameRefSetFunctions) Merge(_ memmap.FileRange, val1 uint64, _ memmap.FileRange, val2 uint64) (uint64, bool) { +func (FrameRefSetFunctions) Merge(_ memmap.FileRange, val1 FrameRefSegInfo, _ memmap.FileRange, val2 FrameRefSegInfo) (FrameRefSegInfo, bool) { if val1 != val2 { - return 0, false + return FrameRefSegInfo{}, false } return val1, true } // Split implements segment.Functions.Split. -func (FrameRefSetFunctions) Split(_ memmap.FileRange, val uint64, _ uint64) (uint64, uint64) { +func (FrameRefSetFunctions) Split(_ memmap.FileRange, val FrameRefSegInfo, _ uint64) (FrameRefSegInfo, FrameRefSegInfo) { return val, val } // IncRefAndAccount adds a reference on the range fr. All newly inserted segments -// are accounted as host page cache memory mappings. -func (refs *FrameRefSet) IncRefAndAccount(fr memmap.FileRange) { - seg, gap := refs.Find(fr.Start) +// are accounted as host page cache memory mappings. The new segments will be +// associated with the memCgID, if the segment already exists then the memCgID +// will not be changed. +func (frSet *FrameRefSet) IncRefAndAccount(fr memmap.FileRange, memCgID uint32) { + seg, gap := frSet.Find(fr.Start) for { switch { case seg.Ok() && seg.Start() < fr.End: - seg = refs.Isolate(seg, fr) - seg.SetValue(seg.Value() + 1) + seg = frSet.Isolate(seg, fr) + seg.ValuePtr().refs++ seg, gap = seg.NextNonEmpty() case gap.Ok() && gap.Start() < fr.End: newRange := gap.Range().Intersect(fr) - // TODO(b/277772401): Get memCgID from memmap.File.IncRef method. - usage.MemoryAccounting.Inc(newRange.Length(), usage.Mapped, 0) - seg, gap = refs.InsertWithoutMerging(gap, newRange, 1).NextNonEmpty() + usage.MemoryAccounting.Inc(newRange.Length(), usage.Mapped, memCgID) + frInfo := FrameRefSegInfo{refs: 1, memCgID: memCgID} + seg, gap = frSet.InsertWithoutMerging(gap, newRange, frInfo).NextNonEmpty() default: - refs.MergeAdjacent(fr) + frSet.MergeAdjacent(fr) return } } @@ -75,19 +86,18 @@ func (refs *FrameRefSet) IncRefAndAccount(fr memmap.FileRange) { // DecRefAndAccount removes a reference on the range fr and untracks segments // that are removed from memory accounting. -func (refs *FrameRefSet) DecRefAndAccount(fr memmap.FileRange) { - seg := refs.FindSegment(fr.Start) +func (frSet *FrameRefSet) DecRefAndAccount(fr memmap.FileRange) { + seg := frSet.FindSegment(fr.Start) for seg.Ok() && seg.Start() < fr.End { - seg = refs.Isolate(seg, fr) - if old := seg.Value(); old == 1 { - // TODO(b/277772401): Get memCgID from memmap.File.DecRef method. - usage.MemoryAccounting.Dec(seg.Range().Length(), usage.Mapped, 0) - seg = refs.Remove(seg).NextSegment() + seg = frSet.Isolate(seg, fr) + if old := seg.ValuePtr().refs; old == 1 { + usage.MemoryAccounting.Dec(seg.Range().Length(), usage.Mapped, seg.ValuePtr().memCgID) + seg = frSet.Remove(seg).NextSegment() } else { - seg.SetValue(old - 1) + seg.ValuePtr().refs-- seg = seg.NextSegment() } } - refs.MergeAdjacent(fr) + frSet.MergeAdjacent(fr) } diff --git a/pkg/sentry/memmap/memmap.go b/pkg/sentry/memmap/memmap.go index f433f3829..5537fd639 100644 --- a/pkg/sentry/memmap/memmap.go +++ b/pkg/sentry/memmap/memmap.go @@ -386,7 +386,9 @@ type MMapOpts struct { type File interface { // All pages in a File are reference-counted. - // IncRef increments the reference count on all pages in fr. + // IncRef increments the reference count on all pages in fr and + // associates each page with a memCgID (memory cgroup id) to which it + // belongs. memCgID will not be changed if the page already exists. // // Preconditions: // * fr.Start and fr.End must be page-aligned. @@ -394,7 +396,7 @@ type File interface { // * At least one reference must be held on all pages in fr. (The File // interface does not provide a way to acquire an initial reference; // implementors may define mechanisms for doing so.) - IncRef(fr FileRange) + IncRef(fr FileRange, memCgID uint32) // DecRef decrements the reference count on all pages in fr. // diff --git a/pkg/sentry/mm/lifecycle.go b/pkg/sentry/mm/lifecycle.go index 088a77da8..4ed7d38f3 100644 --- a/pkg/sentry/mm/lifecycle.go +++ b/pkg/sentry/mm/lifecycle.go @@ -149,6 +149,7 @@ func (mm *MemoryManager) Fork(ctx context.Context) (*MemoryManager, error) { srcvseg := mm.vmas.FirstSegment() dstpgap := mm2.pmas.FirstGap() var unmapAR hostarch.AddrRange + memCgID := pgalloc.MemoryCgroupIDFromContext(ctx) for srcpseg := mm.pmas.FirstSegment(); srcpseg.Ok(); srcpseg = srcpseg.NextSegment() { pma := srcpseg.ValuePtr() if !pma.private { @@ -198,7 +199,7 @@ func (mm *MemoryManager) Fork(ctx context.Context) (*MemoryManager, error) { } fr := srcpseg.fileRange() mm2.incPrivateRef(fr) - srcpseg.ValuePtr().file.IncRef(fr) + srcpseg.ValuePtr().file.IncRef(fr, memCgID) addrRange := srcpseg.Range() mm2.addRSSLocked(addrRange) dstpgap = mm2.pmas.Insert(dstpgap, addrRange, *pma).NextGap() diff --git a/pkg/sentry/mm/pma.go b/pkg/sentry/mm/pma.go index e353f9609..c0f16698c 100644 --- a/pkg/sentry/mm/pma.go +++ b/pkg/sentry/mm/pma.go @@ -252,7 +252,7 @@ func (mm *MemoryManager) getPMAsInternalLocked(ctx context.Context, vseg vmaIter } mm.addRSSLocked(allocAR) mm.incPrivateRef(fr) - mf.IncRef(fr) + mf.IncRef(fr, memCgID) pseg, pgap = mm.pmas.Insert(pgap, allocAR, pma{ file: mf, off: fr.Start, @@ -308,7 +308,7 @@ func (mm *MemoryManager) getPMAsInternalLocked(ctx context.Context, vseg vmaIter newpma.needCOW = true } mm.addRSSLocked(newpmaAR) - t.File.IncRef(t.FileRange()) + t.File.IncRef(t.FileRange(), memCgID) // This is valid because memmap.Mappable.Translate is // required to return Translations in increasing // Translation.Source order. @@ -410,7 +410,7 @@ func (mm *MemoryManager) getPMAsInternalLocked(ctx context.Context, vseg vmaIter } oldpma.file.DecRef(pseg.fileRange()) mm.incPrivateRef(fr) - mf.IncRef(fr) + mf.IncRef(fr, memCgID) oldpma.file = mf oldpma.off = fr.Start oldpma.translatePerms = hostarch.AnyAccess @@ -490,7 +490,7 @@ func (mm *MemoryManager) getPMAsInternalLocked(ctx context.Context, vseg vmaIter newpma.maxPerms.Write = false newpma.needCOW = true } - t.File.IncRef(t.FileRange()) + t.File.IncRef(t.FileRange(), memCgID) pseg = mm.pmas.Insert(pgap, newpmaAR, newpma) pgap = pseg.NextGap() } @@ -710,13 +710,14 @@ func (mm *MemoryManager) Pin(ctx context.Context, ar hostarch.AddrRange, at host ar.End = pendaddr } + memCgID := pgalloc.MemoryCgroupIDFromContext(ctx) // Gather pmas. var prs []PinnedRange for pseg.Ok() && pseg.Start() < ar.End { psar := pseg.Range().Intersect(ar) f := pseg.ValuePtr().file fr := pseg.fileRangeOf(psar) - f.IncRef(fr) + f.IncRef(fr, memCgID) prs = append(prs, PinnedRange{ Source: psar, File: f, diff --git a/pkg/sentry/pgalloc/pgalloc.go b/pkg/sentry/pgalloc/pgalloc.go index b175241d2..4195a4c77 100644 --- a/pkg/sentry/pgalloc/pgalloc.go +++ b/pkg/sentry/pgalloc/pgalloc.go @@ -867,7 +867,7 @@ func (f *MemoryFile) markDecommitted(fr memmap.FileRange) { } // IncRef implements memmap.File.IncRef. -func (f *MemoryFile) IncRef(fr memmap.FileRange) { +func (f *MemoryFile) IncRef(fr memmap.FileRange, memCgID uint32) { if !fr.WellFormed() || fr.Length() == 0 || fr.Start%hostarch.PageSize != 0 || fr.End%hostarch.PageSize != 0 { panic(fmt.Sprintf("invalid range: %v", fr)) }