Make changes for CgroupsReadControlFile method to get memory usage per cgroup.

- Adds methods to Enter, Leave and Migrate the tasks for memory controller.
- Update the memory cgroup id when the task enters and leaves the cgroups.

PiperOrigin-RevId: 555568875

pkg/sentry/fsimpl/cgroupfs/memory.go

@@ -31,13 +31,15 @@ import (
 // +stateify savable
 type memoryController struct {
 	controllerCommon
-	controllerStateless
 	controllerNoResource
 
 	limitBytes            atomicbitops.Int64
 	softLimitBytes        atomicbitops.Int64
 	moveChargeAtImmigrate atomicbitops.Int64
 	pressureLevel         int64
+
+	// memCg is the memory cgroup for this controller.
+	memCg *memoryCgroup
 }
 
 var _ controller = (*memoryController)(nil)
@@ -79,26 +81,68 @@ func (c *memoryController) Clone() controller {
 }
 
 // AddControlFiles implements controller.AddControlFiles.
-func (c *memoryController) AddControlFiles(ctx context.Context, creds *auth.Credentials, _ *cgroupInode, contents map[string]kernfs.Inode) {
-	contents["memory.usage_in_bytes"] = c.fs.newControllerFile(ctx, creds, &memoryUsageInBytesData{}, true)
+func (c *memoryController) AddControlFiles(ctx context.Context, creds *auth.Credentials, cg *cgroupInode, contents map[string]kernfs.Inode) {
+	c.memCg = &memoryCgroup{cg}
+	contents["memory.usage_in_bytes"] = c.fs.newControllerFile(ctx, creds, &memoryUsageInBytesData{memCg: &memoryCgroup{cg}}, true)
 	contents["memory.limit_in_bytes"] = c.fs.newStubControllerFile(ctx, creds, &c.limitBytes, true)
 	contents["memory.soft_limit_in_bytes"] = c.fs.newStubControllerFile(ctx, creds, &c.softLimitBytes, true)
 	contents["memory.move_charge_at_immigrate"] = c.fs.newStubControllerFile(ctx, creds, &c.moveChargeAtImmigrate, true)
 	contents["memory.pressure_level"] = c.fs.newStaticControllerFile(ctx, creds, linux.FileMode(0644), fmt.Sprintf("%d\n", c.pressureLevel))
 }
 
+// Enter implements controller.Enter.
+func (c *memoryController) Enter(t *kernel.Task) {
+	// Update the new cgroup id for the task.
+	t.SetMemCgID(c.memCg.ID())
+}
+
+// Leave implements controller.Leave.
+func (c *memoryController) Leave(t *kernel.Task) {
+	// Update the cgroup id for the task to zero.
+	t.SetMemCgID(0)
+}
+
+// PrepareMigrate implements controller.PrepareMigrate.
+func (c *memoryController) PrepareMigrate(t *kernel.Task, src controller) error {
+	return nil
+}
+
+// CommitMigrate implements controller.CommitMigrate.
+func (c *memoryController) CommitMigrate(t *kernel.Task, src controller) {
+	// Start tracking t at dst by updating the memCgID.
+	t.SetMemCgID(c.memCg.ID())
+}
+
+// AbortMigrate implements controller.AbortMigrate.
+func (c *memoryController) AbortMigrate(t *kernel.Task, src controller) {}
+
 // +stateify savable
-type memoryUsageInBytesData struct{}
+type memoryCgroup struct {
+	*cgroupInode
+}
+
+func (memCg *memoryCgroup) collectMemoryUsage() uint64 {
+	_, totalBytes := usage.MemoryAccounting.CopyPerCg(memCg.ID())
+
+	memCg.forEachChildDir(func(d *dir) {
+		cg := memoryCgroup{d.cgi}
+		totalBytes += cg.collectMemoryUsage()
+	})
+	return totalBytes
+}
+
+// +stateify savable
+type memoryUsageInBytesData struct {
+	memCg *memoryCgroup
+}
 
 // Generate implements vfs.DynamicBytesSource.Generate.
 func (d *memoryUsageInBytesData) Generate(ctx context.Context, buf *bytes.Buffer) error {
-	// TODO(b/183151557): This is a giant hack, we're using system-wide
-	// accounting since we know there is only one cgroup.
 	k := kernel.KernelFromContext(ctx)
 	mf := k.MemoryFile()
 	mf.UpdateUsage()
-	_, totalBytes := usage.MemoryAccounting.Copy()
 
+	totalBytes := d.memCg.collectMemoryUsage()
 	fmt.Fprintf(buf, "%d\n", totalBytes)
 	return nil
 }

pkg/sentry/kernel/kernel.go

@@ -1703,7 +1703,7 @@ func (k *Kernel) ReleaseCgroupHierarchy(hid uint32) {
 		for cg := range t.cgroups {
 			if cg.HierarchyID() == hid {
 				cg.Leave(t)
-				t.resetMemCgID(cg)
+				t.ResetMemCgIDFromCgroup(cg)
 				delete(t.cgroups, cg)
 				releasedCGs = append(releasedCGs, cg)
 				// A task can't be part of multiple cgroups from the same

pkg/sentry/kernel/task_cgroup.go

@@ -51,24 +51,32 @@ func (t *Task) EnterInitialCgroups(parent *Task, initCgroups map[Cgroup]struct{}
 		// Since t isn't in any cgroup yet, we can skip the check against
 		// existing cgroups.
 		c.Enter(t)
-		t.setMemCgID(c)
+		t.SetMemCgIDFromCgroup(c)
 	}
 }
 
-// TODO(b/277772401): setMemCgIDLocked should be called after adding support for
-// task migration for cgroup memory controllers.
-func (t *Task) setMemCgID(cg Cgroup) {
+// SetMemCgID sets the given memory cgroup id to the task.
+func (t *Task) SetMemCgID(memCgID uint32) {
+	t.memCgID.Store(memCgID)
+}
+
+// SetMemCgIDFromCgroup sets the id of the given memory cgroup to the task.
+func (t *Task) SetMemCgIDFromCgroup(cg Cgroup) {
 	for _, ctl := range cg.Controllers() {
 		if ctl.Type() == CgroupControllerMemory {
-			t.memCgID.Store(cg.ID())
+			t.SetMemCgID(cg.ID())
+			return
 		}
 	}
 }
 
-func (t *Task) resetMemCgID(cg Cgroup) {
+// ResetMemCgIDFromCgroup sets the memory cgroup id to zero, if the task has
+// a memory cgroup.
+func (t *Task) ResetMemCgIDFromCgroup(cg Cgroup) {
 	for _, ctl := range cg.Controllers() {
 		if ctl.Type() == CgroupControllerMemory {
-			t.memCgID.Store(0)
+			t.SetMemCgID(0)
+			return
 		}
 	}
 }
@@ -101,7 +109,7 @@ func (t *Task) enterCgroupLocked(c Cgroup) {
 	c.IncRef()
 	t.cgroups[c] = struct{}{}
 	c.Enter(t)
-	t.setMemCgID(c)
+	t.SetMemCgIDFromCgroup(c)
 }
 
 // +checklocks:t.mu
@@ -121,7 +129,7 @@ func (t *Task) LeaveCgroups() {
 	for c := range cgs {
 		c.Leave(t)
 	}
-	t.memCgID.Store(0)
+	t.SetMemCgID(0)
 	t.mu.Unlock()
 	t.tg.pidns.owner.mu.Unlock()
 

pkg/sentry/usage/memory.go

@@ -323,9 +323,15 @@ func (m *MemoryLocked) Total() uint64 {
 func (m *MemoryLocked) TotalPerCg(memCgID uint32) uint64 {
 	m.mu.Lock()
 	defer m.mu.Unlock()
+
+	// Total memory usage including the sentry memory.
+	if memCgID == 0 {
+		return m.totalLocked()
+	}
+	// Memory usage for all cgroups except sentry memory.
 	ms, ok := m.MemCgIDToMemStats[memCgID]
 	if !ok {
-		panic(fmt.Sprintf("invalid memory cgroup id: %v", memCgID))
+		return 0
 	}
 	return ms.totalLocked()
 }
@@ -345,9 +351,15 @@ func (m *MemoryLocked) Copy() (MemoryStats, uint64) {
 func (m *MemoryLocked) CopyPerCg(memCgID uint32) (MemoryStats, uint64) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
+
+	// Total memory usage including the sentry memory.
+	if memCgID == 0 {
+		return m.copyLocked(), m.totalLocked()
+	}
+	// Memory usage for all cgroups except sentry memory.
 	ms, ok := m.MemCgIDToMemStats[memCgID]
 	if !ok {
-		panic(fmt.Sprintf("invalid memory cgroup id: %v", memCgID))
+		return MemoryStats{}, 0
 	}
 	return ms.copyLocked(), ms.totalLocked()
 }

test/syscalls/linux/cgroup.cc

@@ -651,8 +651,9 @@ TEST(MemoryCgroup, MemoryUsageInBytes) {
 
   Mounter m(ASSERT_NO_ERRNO_AND_VALUE(TempPath::CreateDir()));
   Cgroup c = ASSERT_NO_ERRNO_AND_VALUE(m.MountCgroupfs("memory"));
-  EXPECT_THAT(c.ReadIntegerControlFile("memory.usage_in_bytes"),
-              IsPosixErrorOkAndHolds(Gt(0)));
+  const uint64_t usage = ASSERT_NO_ERRNO_AND_VALUE(
+      c.ReadIntegerControlFile("memory.usage_in_bytes"));
+  EXPECT_GE(usage, 0);
 }
 
 TEST(CPUCgroup, ControlFilesHaveDefaultValues) {