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

@@ -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.

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.

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()
 }
 

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.

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()
 }
 

pkg/sentry/fsutil/BUILD

@@ -50,7 +50,7 @@ go_template_instance(
     types = {
         "Key": "uint64",
         "Range": "memmap.FileRange",
-        "Value": "uint64",
+        "Value": "FrameRefSegInfo",
         "Functions": "FrameRefSetFunctions",
     },
 )

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)
 }

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.
 	//

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()

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,

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))
 	}