Increment/decrement memory accounted per cgroup.

- Adds a new field in the usageInfo to store the memory cgroup id.
- Creates a map of cgroup ids and memory stats to track the memory per cgroup
in MemoryLocked struct.
- Introduces new methods to increment, decrement, move, copy and get the total
memory usage per cgroup.

PiperOrigin-RevId: 549148091

pkg/sentry/fsutil/frame_ref_set.go

@@ -63,7 +63,8 @@ func (refs *FrameRefSet) IncRefAndAccount(fr memmap.FileRange) {
 			seg, gap = seg.NextNonEmpty()
 		case gap.Ok() && gap.Start() < fr.End:
 			newRange := gap.Range().Intersect(fr)
-			usage.MemoryAccounting.Inc(newRange.Length(), usage.Mapped)
+			// 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()
 		default:
 			refs.MergeAdjacent(fr)
@@ -80,7 +81,8 @@ func (refs *FrameRefSet) DecRefAndAccount(fr memmap.FileRange) {
 	for seg.Ok() && seg.Start() < fr.End {
 		seg = refs.Isolate(seg, fr)
 		if old := seg.Value(); old == 1 {
-			usage.MemoryAccounting.Dec(seg.Range().Length(), usage.Mapped)
+			// TODO(b/277772401): Get memCgID from memmap.File.DecRef method.
+			usage.MemoryAccounting.Dec(seg.Range().Length(), usage.Mapped, 0)
 			seg = refs.Remove(seg).NextSegment()
 		} else {
 			seg.SetValue(old - 1)

pkg/sentry/pgalloc/pgalloc.go

@@ -257,6 +257,9 @@ type usageInfo struct {
 	knownCommitted bool
 
 	refs uint64
+
+	// memCgID is the memory cgroup id to which this page is committed.
+	memCgID uint32
 }
 
 // canCommit returns true if the tracked region can be committed.
@@ -574,8 +577,9 @@ func (f *MemoryFile) allocate(length uint64, opts *AllocOpts) (memmap.FileRange,
 	}
 	// Mark selected pages as in use.
 	if !f.usage.Add(fr, usageInfo{
-		kind: opts.Kind,
-		refs: 1,
+		kind:    opts.Kind,
+		refs:    1,
+		memCgID: opts.MemCgID,
 	}) {
 		panic(fmt.Sprintf("allocating %v: failed to insert into usage set:\n%v", fr, &f.usage))
 	}
@@ -850,10 +854,11 @@ func (f *MemoryFile) markDecommitted(fr memmap.FileRange) {
 		if val.knownCommitted {
 			// Drop the usageExpected appropriately.
 			amount := seg.Range().Length()
-			usage.MemoryAccounting.Dec(amount, val.kind)
+			usage.MemoryAccounting.Dec(amount, val.kind, val.memCgID)
 			f.usageExpected -= amount
 			val.knownCommitted = false
 		}
+		val.memCgID = 0
 	})
 	if gap.Ok() {
 		panic(fmt.Sprintf("Decommit(%v): attempted to decommit unallocated pages %v:\n%v", fr, gap.Range(), &f.usage))
@@ -904,7 +909,7 @@ func (f *MemoryFile) DecRef(fr memmap.FileRange) {
 			// Reclassify memory as System, until it's freed by the reclaim
 			// goroutine.
 			if val.knownCommitted {
-				usage.MemoryAccounting.Move(seg.Range().Length(), usage.System, val.kind)
+				usage.MemoryAccounting.Move(seg.Range().Length(), usage.System, val.kind, val.memCgID)
 			}
 			val.kind = usage.System
 		}
@@ -1144,9 +1149,9 @@ func (f *MemoryFile) updateUsageLocked(currentUsage uint64, checkCommitted func(
 			// that have been swapped.
 			newUsageSwapped := currentUsage - f.usageExpected
 			if f.usageSwapped < newUsageSwapped {
-				usage.MemoryAccounting.Inc(newUsageSwapped-f.usageSwapped, usage.System)
+				usage.MemoryAccounting.Inc(newUsageSwapped-f.usageSwapped, usage.System, 0)
 			} else {
-				usage.MemoryAccounting.Dec(f.usageSwapped-newUsageSwapped, usage.System)
+				usage.MemoryAccounting.Dec(f.usageSwapped-newUsageSwapped, usage.System, 0)
 			}
 			f.usageSwapped = newUsageSwapped
 		} else if f.usageSwapped != 0 {
@@ -1154,7 +1159,7 @@ func (f *MemoryFile) updateUsageLocked(currentUsage uint64, checkCommitted func(
 			// That's fine, we probably caught a race where pages were
 			// being committed while the below loop was running. Just
 			// report the higher number that we found and ignore swap.
-			usage.MemoryAccounting.Dec(f.usageSwapped, usage.System)
+			usage.MemoryAccounting.Dec(f.usageSwapped, usage.System, 0)
 			f.usageSwapped = 0
 		}
 	}()
@@ -1229,7 +1234,7 @@ func (f *MemoryFile) updateUsageLocked(currentUsage uint64, checkCommitted func(
 							seg = f.usage.Isolate(seg, committedFR)
 							seg.ValuePtr().knownCommitted = true
 							amount := seg.Range().Length()
-							usage.MemoryAccounting.Inc(amount, seg.ValuePtr().kind)
+							usage.MemoryAccounting.Inc(amount, seg.ValuePtr().kind, seg.ValuePtr().memCgID)
 							f.usageExpected += amount
 							changedAny = true
 						}
@@ -1433,6 +1438,7 @@ func (f *MemoryFile) markReclaimed(fr memmap.FileRange) {
 		kind:           usage.System,
 		knownCommitted: false,
 		refs:           0,
+		memCgID:        0,
 	}); got != want {
 		panic(fmt.Sprintf("reclaimed pages %v in segment %v has incorrect state %v, wanted %v:\n%v", fr, seg.Range(), got, want, &f.usage))
 	}

pkg/sentry/pgalloc/save_restore.go

@@ -189,7 +189,7 @@ func (f *MemoryFile) LoadFrom(ctx context.Context, r wire.Reader) error {
 		// Update accounting for restored pages. We need to do this here since
 		// these segments are marked as "known committed", and will be skipped
 		// over on accounting scans.
-		usage.MemoryAccounting.Inc(seg.End()-seg.Start(), seg.Value().kind)
+		usage.MemoryAccounting.Inc(seg.End()-seg.Start(), seg.Value().kind, seg.Value().memCgID)
 	}
 
 	return nil

pkg/sentry/usage/BUILD

@@ -1,12 +1,12 @@
 load("//tools:defs.bzl", "go_library")
-load("//pkg/sync/locking:locking.bzl", "declare_rwmutex")
+load("//pkg/sync/locking:locking.bzl", "declare_mutex")
 
 package(
     default_applicable_licenses = ["//:license"],
     licenses = ["notice"],
 )
 
-declare_rwmutex(
+declare_mutex(
     name = "memory_mutex",
     out = "memory_mutex.go",
     package = "usage",

pkg/sentry/usage/memory.go

@@ -84,11 +84,81 @@ type memoryStats struct {
 	Anonymous atomicbitops.Uint64
 	PageCache atomicbitops.Uint64
 	Tmpfs     atomicbitops.Uint64
-	// Lazily updated based on the value in RTMapped.
 	Mapped    atomicbitops.Uint64
 	Ramdiskfs atomicbitops.Uint64
 }
 
+// incLocked adds a usage of 'val' bytes from memory category 'kind'.
+//
+// Precondition: must be called when locked.
+func (ms *memoryStats) incLocked(val uint64, kind MemoryKind) {
+	switch kind {
+	case System:
+		ms.System.Add(val)
+	case Anonymous:
+		ms.Anonymous.Add(val)
+	case PageCache:
+		ms.PageCache.Add(val)
+	case Mapped:
+		ms.Mapped.Add(val)
+	case Tmpfs:
+		ms.Tmpfs.Add(val)
+	case Ramdiskfs:
+		ms.Ramdiskfs.Add(val)
+	default:
+		panic(fmt.Sprintf("invalid memory kind: %v", kind))
+	}
+}
+
+// decLocked removes a usage of 'val' bytes from memory category 'kind'.
+//
+// Precondition: must be called when locked.
+func (ms *memoryStats) decLocked(val uint64, kind MemoryKind) {
+	switch kind {
+	case System:
+		ms.System.Add(^(val - 1))
+	case Anonymous:
+		ms.Anonymous.Add(^(val - 1))
+	case PageCache:
+		ms.PageCache.Add(^(val - 1))
+	case Mapped:
+		ms.Mapped.Add(^(val - 1))
+	case Tmpfs:
+		ms.Tmpfs.Add(^(val - 1))
+	case Ramdiskfs:
+		ms.Ramdiskfs.Add(^(val - 1))
+	default:
+		panic(fmt.Sprintf("invalid memory kind: %v", kind))
+	}
+}
+
+// totalLocked returns a total usage.
+//
+// Precondition: must be called when locked.
+func (ms *memoryStats) totalLocked() (total uint64) {
+	total += ms.System.RacyLoad()
+	total += ms.Anonymous.RacyLoad()
+	total += ms.PageCache.RacyLoad()
+	total += ms.Mapped.RacyLoad()
+	total += ms.Tmpfs.RacyLoad()
+	total += ms.Ramdiskfs.RacyLoad()
+	return
+}
+
+// copyLocked returns a copy of the structure.
+//
+// Precondition: must be called when locked.
+func (ms *memoryStats) copyLocked() MemoryStats {
+	return MemoryStats{
+		System:    ms.System.RacyLoad(),
+		Anonymous: ms.Anonymous.RacyLoad(),
+		PageCache: ms.PageCache.RacyLoad(),
+		Tmpfs:     ms.Tmpfs.RacyLoad(),
+		Mapped:    ms.Mapped.RacyLoad(),
+		Ramdiskfs: ms.Ramdiskfs.RacyLoad(),
+	}
+}
+
 // MemoryStats tracks application memory usage in bytes. All fields correspond
 // to the memory category with the same name.
 type MemoryStats struct {
@@ -117,7 +187,7 @@ type RTMemoryStats struct {
 
 // MemoryLocked is Memory with access methods.
 type MemoryLocked struct {
-	mu memoryRWMutex
+	mu memoryMutex
 	// memoryStats records the memory stats.
 	memoryStats
 	// RTMemoryStats records the memory stats that need to be exposed through
@@ -125,6 +195,8 @@ type MemoryLocked struct {
 	*RTMemoryStats
 	// File is the backing file storing the memory stats.
 	File *os.File
+	// MemCgIDToMemStats is the map of cgroup ids to memory stats.
+	MemCgIDToMemStats map[uint32]*memoryStats
 }
 
 // Init initializes global 'MemoryAccounting'.
@@ -148,8 +220,9 @@ func Init() error {
 	}
 
 	MemoryAccounting = &MemoryLocked{
-		File:          file,
-		RTMemoryStats: RTMemoryStatsPointer(mmap),
+		File:              file,
+		RTMemoryStats:     RTMemoryStatsPointer(mmap),
+		MemCgIDToMemStats: make(map[uint32]*memoryStats),
 	}
 	return nil
 }
@@ -161,85 +234,78 @@ func Init() error {
 // resident.
 var MemoryAccounting *MemoryLocked
 
-func (m *MemoryLocked) incLocked(val uint64, kind MemoryKind) {
-	switch kind {
-	case System:
-		m.System.Add(val)
-	case Anonymous:
-		m.Anonymous.Add(val)
-	case PageCache:
-		m.PageCache.Add(val)
-	case Mapped:
-		m.RTMapped.Add(val)
-	case Tmpfs:
-		m.Tmpfs.Add(val)
-	case Ramdiskfs:
-		m.Ramdiskfs.Add(val)
-	default:
-		panic(fmt.Sprintf("invalid memory kind: %v", kind))
+func (m *MemoryLocked) incLockedPerCg(val uint64, kind MemoryKind, memCgID uint32) {
+	if _, ok := m.MemCgIDToMemStats[memCgID]; !ok {
+		m.MemCgIDToMemStats[memCgID] = &memoryStats{}
 	}
+
+	ms := m.MemCgIDToMemStats[memCgID]
+	ms.incLocked(val, kind)
 }
 
-// Inc adds an additional usage of 'val' bytes to memory category 'kind'.
+// Inc adds an additional usage of 'val' bytes to memory category 'kind' for a
+// cgroup with id 'memCgID'. If 'memCgID' is zero, the memory is accounted only
+// for the total memory usage.
 //
 // This method is thread-safe.
-func (m *MemoryLocked) Inc(val uint64, kind MemoryKind) {
-	m.mu.RLock()
+func (m *MemoryLocked) Inc(val uint64, kind MemoryKind, memCgID uint32) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
 	m.incLocked(val, kind)
-	m.mu.RUnlock()
-}
+	if memCgID != 0 {
+		m.incLockedPerCg(val, kind, memCgID)
+	}
 
-func (m *MemoryLocked) decLocked(val uint64, kind MemoryKind) {
-	switch kind {
-	case System:
-		m.System.Add(^(val - 1))
-	case Anonymous:
-		m.Anonymous.Add(^(val - 1))
-	case PageCache:
-		m.PageCache.Add(^(val - 1))
-	case Mapped:
-		m.RTMapped.Add(^(val - 1))
-	case Tmpfs:
-		m.Tmpfs.Add(^(val - 1))
-	case Ramdiskfs:
-		m.Ramdiskfs.Add(^(val - 1))
-	default:
-		panic(fmt.Sprintf("invalid memory kind: %v", kind))
+	// If the memory category is 'Mapped', update RTMapped.
+	if kind == Mapped {
+		m.RTMapped.Add(val)
 	}
 }
 
-// Dec remove a usage of 'val' bytes from memory category 'kind'.
-//
-// This method is thread-safe.
-func (m *MemoryLocked) Dec(val uint64, kind MemoryKind) {
-	m.mu.RLock()
-	m.decLocked(val, kind)
-	m.mu.RUnlock()
+func (m *MemoryLocked) decLockedPerCg(val uint64, kind MemoryKind, memCgID uint32) {
+	if _, ok := m.MemCgIDToMemStats[memCgID]; !ok {
+		panic(fmt.Sprintf("invalid memory cgroup id: %v", memCgID))
+	}
+
+	ms := m.MemCgIDToMemStats[memCgID]
+	ms.decLocked(val, kind)
 }
 
-// Move moves a usage of 'val' bytes from 'from' to 'to'.
+// Dec removes a usage of 'val' bytes from memory category 'kind' for a cgroup
+// with id 'memCgID'. If 'memCgID' is zero, the memory is removed only from the
+// total usage.
 //
 // This method is thread-safe.
-func (m *MemoryLocked) Move(val uint64, to MemoryKind, from MemoryKind) {
-	m.mu.RLock()
-	// Just call decLocked and incLocked directly. We held the RLock to
+func (m *MemoryLocked) Dec(val uint64, kind MemoryKind, memCgID uint32) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.decLocked(val, kind)
+	if memCgID != 0 {
+		m.decLockedPerCg(val, kind, memCgID)
+	}
+
+	// If the memory category is 'Mapped', update RTMapped.
+	if kind == Mapped {
+		m.RTMapped.Add(^(val - 1))
+	}
+}
+
+// Move moves a usage of 'val' bytes from 'from' to 'to' for a cgroup with
+// id 'memCgID'.
+//
+// This method is thread-safe.
+func (m *MemoryLocked) Move(val uint64, to MemoryKind, from MemoryKind, memCgID uint32) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	// Just call decLocked and incLocked directly. We held the Lock to
 	// protect against concurrent callers to Total().
 	m.decLocked(val, from)
 	m.incLocked(val, to)
-	m.mu.RUnlock()
-}
 
-// totalLocked returns a total usage.
-//
-// Precondition: must be called when locked.
-func (m *MemoryLocked) totalLocked() (total uint64) {
-	total += m.System.Load()
-	total += m.Anonymous.Load()
-	total += m.PageCache.Load()
-	total += m.RTMapped.Load()
-	total += m.Tmpfs.Load()
-	total += m.Ramdiskfs.Load()
-	return
+	if memCgID != 0 {
+		m.decLockedPerCg(val, from, memCgID)
+		m.incLockedPerCg(val, to, memCgID)
+	}
 }
 
 // Total returns a total memory usage.
@@ -251,21 +317,39 @@ func (m *MemoryLocked) Total() uint64 {
 	return m.totalLocked()
 }
 
+// TotalPerCg returns a total memory usage for a cgroup.
+//
+// This method is thread-safe.
+func (m *MemoryLocked) TotalPerCg(memCgID uint32) uint64 {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	ms, ok := m.MemCgIDToMemStats[memCgID]
+	if !ok {
+		panic(fmt.Sprintf("invalid memory cgroup id: %v", memCgID))
+	}
+	return ms.totalLocked()
+}
+
 // Copy returns a copy of the structure with a total.
 //
 // This method is thread-safe.
 func (m *MemoryLocked) Copy() (MemoryStats, uint64) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
-	ms := MemoryStats{
-		System:    m.System.RacyLoad(),
-		Anonymous: m.Anonymous.RacyLoad(),
-		PageCache: m.PageCache.RacyLoad(),
-		Tmpfs:     m.Tmpfs.RacyLoad(),
-		Mapped:    m.RTMapped.RacyLoad(),
-		Ramdiskfs: m.Ramdiskfs.RacyLoad(),
+	return m.copyLocked(), m.totalLocked()
+}
+
+// CopyPerCg returns a copy of the structure with a total for a cgroup.
+//
+// This method is thread-safe.
+func (m *MemoryLocked) CopyPerCg(memCgID uint32) (MemoryStats, uint64) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	ms, ok := m.MemCgIDToMemStats[memCgID]
+	if !ok {
+		panic(fmt.Sprintf("invalid memory cgroup id: %v", memCgID))
 	}
-	return ms, m.totalLocked()
+	return ms.copyLocked(), ms.totalLocked()
 }
 
 // These options control how much total memory the is reported to the