Add GetMemoryUsage() API to the gVisor sandbox.

Adds GetMemoryUsage() API to get the memory usage of the containers.

PiperOrigin-RevId: 557923509

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(); err != nil {
+		if err := u.Kernel.MemoryFile().UpdateUsage(0); err != nil {
 			return err
 		}
 

pkg/sentry/fsimpl/cgroupfs/memory.go

@@ -140,7 +140,7 @@ type memoryUsageInBytesData struct {
 func (d *memoryUsageInBytesData) Generate(ctx context.Context, buf *bytes.Buffer) error {
 	k := kernel.KernelFromContext(ctx)
 	mf := k.MemoryFile()
-	mf.UpdateUsage()
+	mf.UpdateUsage(d.memCg.ID())
 
 	totalBytes := d.memCg.collectMemoryUsage()
 	fmt.Fprintf(buf, "%d\n", totalBytes)

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() // Best effort
+	_ = mf.UpdateUsage(0) // Best effort
 	snapshot, totalUsage := usage.MemoryAccounting.Copy()
 	totalSize := usage.TotalMemory(mf.TotalSize(), totalUsage)
 	anon := snapshot.Anonymous + snapshot.Tmpfs

pkg/sentry/pgalloc/pgalloc.go

@@ -1088,8 +1088,10 @@ func (f *MemoryFile) ShouldCacheEvictable() bool {
 }
 
 // UpdateUsage ensures that the memory usage statistics in
-// usage.MemoryAccounting are up to date.
-func (f *MemoryFile) UpdateUsage() error {
+// 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 {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 
@@ -1111,14 +1113,17 @@ func (f *MemoryFile) UpdateUsage() error {
 		log.Debugf("UpdateUsage: skipped with usageSwapped!=0.")
 		return nil
 	}
+
 	// Linux updates usage values at CONFIG_HZ.
 	if scanningAfter := time.Now().Sub(f.usageLast).Milliseconds(); scanningAfter < time.Second.Milliseconds()/linux.CLOCKS_PER_SEC {
 		log.Debugf("UpdateUsage: skipped because previous scan happened %d ms back", scanningAfter)
 		return nil
 	}
 
-	f.usageLast = time.Now()
-	err = f.updateUsageLocked(currentUsage, mincore)
+	if memCgID == 0 {
+		f.usageLast = time.Now()
+	}
+	err = f.updateUsageLocked(currentUsage, memCgID, mincore)
 	log.Debugf("UpdateUsage: currentUsage=%d, usageExpected=%d, usageSwapped=%d.",
 		currentUsage, f.usageExpected, f.usageSwapped)
 	log.Debugf("UpdateUsage: took %v.", time.Since(f.usageLast))
@@ -1131,7 +1136,7 @@ func (f *MemoryFile) UpdateUsage() error {
 //
 // Precondition: f.mu must be held; it may be unlocked and reacquired.
 // +checklocks:f.mu
-func (f *MemoryFile) updateUsageLocked(currentUsage uint64, checkCommitted func(bs []byte, committed []byte) error) error {
+func (f *MemoryFile) updateUsageLocked(currentUsage uint64, memCgID uint32, 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
@@ -1175,6 +1180,15 @@ func (f *MemoryFile) updateUsageLocked(currentUsage uint64, checkCommitted func(
 			continue
 		}
 
+		// 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
+		}
+
 		// Get the range for this segment. As we touch slices, the
 		// Start value will be walked along.
 		r := seg.Range()

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, func(bs []byte, committed []byte) error {
+	err := f.updateUsageLocked(0, 0, 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() // best effort to update.
+	_ = mem.UpdateUsage(0) // best effort to update.
 	_, totalUsage := usage.MemoryAccounting.Copy()
 	switch containers := cm.l.containerCount(); containers {
 	case 0: