diff --git a/pkg/sentry/fsutil/frame_ref_set.go b/pkg/sentry/fsutil/frame_ref_set.go index a808894df..a0035a423 100644 --- a/pkg/sentry/fsutil/frame_ref_set.go +++ b/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) diff --git a/pkg/sentry/pgalloc/pgalloc.go b/pkg/sentry/pgalloc/pgalloc.go index 1b36b7338..b175241d2 100644 --- a/pkg/sentry/pgalloc/pgalloc.go +++ b/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)) } diff --git a/pkg/sentry/pgalloc/save_restore.go b/pkg/sentry/pgalloc/save_restore.go index b289e2617..657fe6a21 100644 --- a/pkg/sentry/pgalloc/save_restore.go +++ b/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 diff --git a/pkg/sentry/usage/BUILD b/pkg/sentry/usage/BUILD index e142a7698..23fc563fc 100644 --- a/pkg/sentry/usage/BUILD +++ b/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", diff --git a/pkg/sentry/usage/memory.go b/pkg/sentry/usage/memory.go index 88c5a0f94..5bec6cb1b 100644 --- a/pkg/sentry/usage/memory.go +++ b/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