Fix the sandbox memory usage via GetContainerMemoryUsage API.

The total(sandbox) memory usage using the GetContainerMemoryUsage API will
return incorrect usage when called before calling the API for each individual
containers in the sandbox. This is because the memory usage for the containers
cgroup is not updated while calculating the total usage. This CL fixes it by
updating the usage for every child cgroup, which will return the correct memory
usage for the parent cgroup.

PiperOrigin-RevId: 574300913

pkg/sentry/control/usage.go

@@ -88,7 +88,7 @@ func (u *Usage) UsageFD(opts *MemoryUsageFileOpts, out *MemoryUsageFile) error {
 func (u *Usage) Collect(opts *MemoryUsageOpts, out *MemoryUsage) error {
 	if opts.Full {
 		// Ensure everything is up to date.
-		if err := u.Kernel.MemoryFile().UpdateUsage(0); err != nil {
+		if err := u.Kernel.MemoryFile().UpdateUsage(nil); err != nil {
 			return err
 		}
 

pkg/sentry/fsimpl/cgroupfs/memory.go

@@ -121,13 +121,25 @@ type memoryCgroup struct {
 	*cgroupInode
 }
 
-func (memCg *memoryCgroup) collectMemoryUsage() uint64 {
-	_, totalBytes := usage.MemoryAccounting.CopyPerCg(memCg.ID())
-
+// Collects all the memory cgroup ids for the cgroup.
+func (memCg *memoryCgroup) collectMemCgIDs(memCgIDs map[uint32]struct{}) {
+	// Add ourselves.
+	memCgIDs[memCg.ID()] = struct{}{}
+	// Add our children.
 	memCg.forEachChildDir(func(d *dir) {
 		cg := memoryCgroup{d.cgi}
-		totalBytes += cg.collectMemoryUsage()
+		cg.collectMemCgIDs(memCgIDs)
 	})
+}
+
+// Returns the memory usage for all cgroup ids in memCgIDs.
+func getUsage(k *kernel.Kernel, memCgIDs map[uint32]struct{}) uint64 {
+	k.MemoryFile().UpdateUsage(memCgIDs)
+	var totalBytes uint64
+	for id := range memCgIDs {
+		_, bytes := usage.MemoryAccounting.CopyPerCg(id)
+		totalBytes += bytes
+	}
 	return totalBytes
 }
 
@@ -139,10 +151,10 @@ type memoryUsageInBytesData struct {
 // Generate implements vfs.DynamicBytesSource.Generate.
 func (d *memoryUsageInBytesData) Generate(ctx context.Context, buf *bytes.Buffer) error {
 	k := kernel.KernelFromContext(ctx)
-	mf := k.MemoryFile()
-	mf.UpdateUsage(d.memCg.ID())
 
-	totalBytes := d.memCg.collectMemoryUsage()
+	memCgIDs := make(map[uint32]struct{})
+	d.memCg.collectMemCgIDs(memCgIDs)
+	totalBytes := getUsage(k, memCgIDs)
 	fmt.Fprintf(buf, "%d\n", totalBytes)
 	return nil
 }

pkg/sentry/fsimpl/proc/tasks_files.go

@@ -268,7 +268,7 @@ var _ dynamicInode = (*meminfoData)(nil)
 // Generate implements vfs.DynamicBytesSource.Generate.
 func (*meminfoData) Generate(ctx context.Context, buf *bytes.Buffer) error {
 	mf := kernel.KernelFromContext(ctx).MemoryFile()
-	_ = mf.UpdateUsage(0) // Best effort
+	_ = mf.UpdateUsage(nil) // Best effort
 	snapshot, totalUsage := usage.MemoryAccounting.Copy()
 	totalSize := usage.TotalMemory(mf.TotalSize(), totalUsage)
 	anon := snapshot.Anonymous + snapshot.Tmpfs

pkg/sentry/pgalloc/pgalloc.go

@@ -1088,10 +1088,10 @@ func (f *MemoryFile) ShouldCacheEvictable() bool {
 }
 
 // UpdateUsage ensures that the memory usage statistics in
-// usage.MemoryAccounting are up to date. If forceScan is true, the
-// UsageScanDuration is ignored and the memory file is scanned to get the
-// memory usage.
-func (f *MemoryFile) UpdateUsage(memCgID uint32) error {
+// usage.MemoryAccounting are up to date. If memCgIDs is nil, all the pages
+// will be scanned. Else only the pages which belong to the memory cgroup ids
+// in memCgIDs will be scanned and the memory usage will be updated.
+func (f *MemoryFile) UpdateUsage(memCgIDs map[uint32]struct{}) error {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 
@@ -1120,10 +1120,10 @@ func (f *MemoryFile) UpdateUsage(memCgID uint32) error {
 		return nil
 	}
 
-	if memCgID == 0 {
+	if memCgIDs == nil {
 		f.usageLast = time.Now()
 	}
-	err = f.updateUsageLocked(currentUsage, memCgID, mincore)
+	err = f.updateUsageLocked(currentUsage, memCgIDs, mincore)
 	log.Debugf("UpdateUsage: currentUsage=%d, usageExpected=%d, usageSwapped=%d.",
 		currentUsage, f.usageExpected, f.usageSwapped)
 	log.Debugf("UpdateUsage: took %v.", time.Since(f.usageLast))
@@ -1136,7 +1136,7 @@ func (f *MemoryFile) UpdateUsage(memCgID uint32) error {
 //
 // Precondition: f.mu must be held; it may be unlocked and reacquired.
 // +checklocks:f.mu
-func (f *MemoryFile) updateUsageLocked(currentUsage uint64, memCgID uint32, checkCommitted func(bs []byte, committed []byte) error) error {
+func (f *MemoryFile) updateUsageLocked(currentUsage uint64, memCgIDs map[uint32]struct{}, checkCommitted func(bs []byte, committed []byte) error) error {
 	// Track if anything changed to elide the merge. In the common case, we
 	// expect all segments to be committed and no merge to occur.
 	changedAny := false
@@ -1183,10 +1183,12 @@ func (f *MemoryFile) updateUsageLocked(currentUsage uint64, memCgID uint32, chec
 		// Scan the pages of the given memCgID only. This will avoid scanning the
 		// whole memory file when the memory usage is required only for a specific
 		// cgroup. The total memory usage of all cgroups can be obtained when the
-		// memCgID is passed as zero.
-		if memCgID != 0 && seg.ValuePtr().memCgID != memCgID {
-			seg = seg.NextSegment()
-			continue
+		// memCgIDs is nil.
+		if memCgIDs != nil {
+			if _, ok := memCgIDs[seg.ValuePtr().memCgID]; !ok {
+				seg = seg.NextSegment()
+				continue
+			}
 		}
 
 		// Get the range for this segment. As we touch slices, the

pkg/sentry/pgalloc/save_restore.go

@@ -50,7 +50,7 @@ func (f *MemoryFile) SaveTo(ctx context.Context, w wire.Writer) error {
 	// Ensure that all pages that contain data have knownCommitted set, since
 	// we only store knownCommitted pages below.
 	zeroPage := make([]byte, hostarch.PageSize)
-	err := f.updateUsageLocked(0, 0, func(bs []byte, committed []byte) error {
+	err := f.updateUsageLocked(0, nil, func(bs []byte, committed []byte) error {
 		for pgoff := 0; pgoff < len(bs); pgoff += hostarch.PageSize {
 			i := pgoff / hostarch.PageSize
 			pg := bs[pgoff : pgoff+hostarch.PageSize]

runsc/boot/events.go

@@ -100,7 +100,7 @@ func (cm *containerManager) Event(cid *string, out *EventOut) error {
 
 	// Memory usage.
 	mem := cm.l.k.MemoryFile()
-	_ = mem.UpdateUsage(0) // best effort to update.
+	_ = mem.UpdateUsage(nil) // best effort to update.
 	_, totalUsage := usage.MemoryAccounting.Copy()
 	switch containers := cm.l.containerCount(); containers {
 	case 0: