diff --git a/pkg/sentry/fsimpl/cgroupfs/memory.go b/pkg/sentry/fsimpl/cgroupfs/memory.go index 2b238837b..9296adb4d 100644 --- a/pkg/sentry/fsimpl/cgroupfs/memory.go +++ b/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 } diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go index 6bf133fbb..04d241996 100644 --- a/pkg/sentry/kernel/kernel.go +++ b/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 diff --git a/pkg/sentry/kernel/task_cgroup.go b/pkg/sentry/kernel/task_cgroup.go index 970540ee2..82c0dd54b 100644 --- a/pkg/sentry/kernel/task_cgroup.go +++ b/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() diff --git a/pkg/sentry/usage/memory.go b/pkg/sentry/usage/memory.go index 5bec6cb1b..c451e08e8 100644 --- a/pkg/sentry/usage/memory.go +++ b/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() } diff --git a/test/syscalls/linux/cgroup.cc b/test/syscalls/linux/cgroup.cc index db7ee4084..fe09b29cd 100644 --- a/test/syscalls/linux/cgroup.cc +++ b/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) {