diff --git a/cgroup2/manager.go b/cgroup2/manager.go index 3536fc97..4ff27fa6 100644 --- a/cgroup2/manager.go +++ b/cgroup2/manager.go @@ -51,6 +51,22 @@ const ( cpuQuotaPeriodUSecSupportedVersion = 242 ) +// StatMask represents which controller stats to collect. +type StatMask uint64 + +const ( + StatPids StatMask = 1 << iota + StatCPU + StatMemory + StatMemoryEvents + StatIO + StatRdma + StatHugetlb + + // StatAll collects all available stats (default behavior of Stat). + StatAll = StatPids | StatCPU | StatMemory | StatMemoryEvents | StatIO | StatRdma | StatHugetlb +) + var ( canDelegate bool @@ -558,41 +574,62 @@ func (c *Manager) MoveTo(destination *Manager) error { return nil } +// Stat returns all cgroup stats. +// This is equivalent to calling StatFiltered(StatAll). func (c *Manager) Stat() (*stats.Metrics, error) { + return c.StatFiltered(StatAll) +} + +// StatFiltered returns cgroup stats for the specified controllers. +func (c *Manager) StatFiltered(mask StatMask) (*stats.Metrics, error) { var metrics stats.Metrics var err error - metrics.Pids = &stats.PidsStat{ - Current: getStatFileContentUint64(filepath.Join(c.path, "pids.current")), - Limit: getStatFileContentUint64(filepath.Join(c.path, "pids.max")), + if mask&StatPids != 0 { + metrics.Pids = &stats.PidsStat{ + Current: getStatFileContentUint64(filepath.Join(c.path, "pids.current")), + Limit: getStatFileContentUint64(filepath.Join(c.path, "pids.max")), + } } - metrics.CPU, err = readCPUStats(c.path) - if err != nil { - return nil, err + if mask&StatCPU != 0 { + metrics.CPU, err = readCPUStats(c.path) + if err != nil { + return nil, err + } } - metrics.Memory, err = readMemoryStats(c.path) - if err != nil { - return nil, err + if mask&StatMemory != 0 { + metrics.Memory, err = readMemoryStats(c.path) + if err != nil { + return nil, err + } } - metrics.MemoryEvents, err = readMemoryEvents(c.path) - if err != nil { - return nil, err + if mask&StatMemoryEvents != 0 { + metrics.MemoryEvents, err = readMemoryEvents(c.path) + if err != nil { + return nil, err + } } - metrics.Io = &stats.IOStat{ - Usage: readIoStats(c.path), - PSI: getStatPSIFromFile(filepath.Join(c.path, "io.pressure")), + if mask&StatIO != 0 { + metrics.Io = &stats.IOStat{ + Usage: readIoStats(c.path), + PSI: getStatPSIFromFile(filepath.Join(c.path, "io.pressure")), + } } - metrics.Rdma = &stats.RdmaStat{ - Current: rdmaStats(filepath.Join(c.path, "rdma.current")), - Limit: rdmaStats(filepath.Join(c.path, "rdma.max")), + if mask&StatRdma != 0 { + metrics.Rdma = &stats.RdmaStat{ + Current: rdmaStats(filepath.Join(c.path, "rdma.current")), + Limit: rdmaStats(filepath.Join(c.path, "rdma.max")), + } } - metrics.Hugetlb = readHugeTlbStats(c.path) + if mask&StatHugetlb != 0 { + metrics.Hugetlb = readHugeTlbStats(c.path) + } return &metrics, nil } diff --git a/cgroup2/manager_test.go b/cgroup2/manager_test.go index 4902e576..5ee6971c 100644 --- a/cgroup2/manager_test.go +++ b/cgroup2/manager_test.go @@ -463,6 +463,146 @@ func BenchmarkStat(b *testing.B) { } } +func TestStatFiltered(t *testing.T) { + checkCgroupMode(t) + group := "/stat-filtered-test-cg" + groupPath := fmt.Sprintf("%s-%d", group, os.Getpid()) + c, err := NewManager(defaultCgroup2Path, groupPath, &Resources{}) + require.NoError(t, err, "failed to init new cgroup manager") + t.Cleanup(func() { + _ = c.Delete() + }) + + t.Run("StatAll", func(t *testing.T) { + statAll, err := c.StatFiltered(StatAll) + require.NoError(t, err) + + assert.NotNil(t, statAll.Pids) + assert.NotNil(t, statAll.CPU) + assert.NotNil(t, statAll.Memory) + assert.NotNil(t, statAll.Io) + assert.NotNil(t, statAll.Rdma) + }) + + t.Run("CPUOnly", func(t *testing.T) { + stats, err := c.StatFiltered(StatCPU) + require.NoError(t, err) + + assert.NotNil(t, stats.CPU, "CPU stats should be populated") + assert.Nil(t, stats.Pids, "Pids stats should be nil") + assert.Nil(t, stats.Memory, "Memory stats should be nil") + assert.Nil(t, stats.MemoryEvents, "MemoryEvents should be nil") + assert.Nil(t, stats.Io, "IO stats should be nil") + assert.Nil(t, stats.Rdma, "RDMA stats should be nil") + assert.Nil(t, stats.Hugetlb, "Hugetlb stats should be nil") + }) + + t.Run("MemoryOnly", func(t *testing.T) { + stats, err := c.StatFiltered(StatMemory) + require.NoError(t, err) + + assert.NotNil(t, stats.Memory, "Memory stats should be populated") + assert.Nil(t, stats.Pids, "Pids stats should be nil") + assert.Nil(t, stats.CPU, "CPU stats should be nil") + assert.Nil(t, stats.MemoryEvents, "MemoryEvents should be nil") + assert.Nil(t, stats.Io, "IO stats should be nil") + }) + + t.Run("MemoryEventsOnly", func(t *testing.T) { + stats, err := c.StatFiltered(StatMemoryEvents) + require.NoError(t, err) + + assert.NotNil(t, stats.MemoryEvents, "MemoryEvents should be populated") + assert.Nil(t, stats.Memory, "Memory stats should be nil") + assert.Nil(t, stats.CPU, "CPU stats should be nil") + }) + + t.Run("CPUAndMemory", func(t *testing.T) { + stats, err := c.StatFiltered(StatCPU | StatMemory) + require.NoError(t, err) + + assert.NotNil(t, stats.CPU, "CPU stats should be populated") + assert.NotNil(t, stats.Memory, "Memory stats should be populated") + assert.Nil(t, stats.Pids, "Pids stats should be nil") + assert.Nil(t, stats.MemoryEvents, "MemoryEvents should be nil") + assert.Nil(t, stats.Io, "IO stats should be nil") + }) + + t.Run("PidsOnly", func(t *testing.T) { + stats, err := c.StatFiltered(StatPids) + require.NoError(t, err) + + assert.NotNil(t, stats.Pids, "Pids stats should be populated") + assert.Nil(t, stats.CPU, "CPU stats should be nil") + assert.Nil(t, stats.Memory, "Memory stats should be nil") + }) + + t.Run("IOOnly", func(t *testing.T) { + stats, err := c.StatFiltered(StatIO) + require.NoError(t, err) + + assert.NotNil(t, stats.Io, "IO stats should be populated") + assert.Nil(t, stats.CPU, "CPU stats should be nil") + assert.Nil(t, stats.Memory, "Memory stats should be nil") + }) + + t.Run("ZeroMask", func(t *testing.T) { + stats, err := c.StatFiltered(0) + require.NoError(t, err) + + assert.Nil(t, stats.Pids, "Pids stats should be nil") + assert.Nil(t, stats.CPU, "CPU stats should be nil") + assert.Nil(t, stats.Memory, "Memory stats should be nil") + assert.Nil(t, stats.MemoryEvents, "MemoryEvents should be nil") + assert.Nil(t, stats.Io, "IO stats should be nil") + assert.Nil(t, stats.Rdma, "RDMA stats should be nil") + assert.Nil(t, stats.Hugetlb, "Hugetlb stats should be nil") + }) +} + +func BenchmarkStatFiltered(b *testing.B) { + checkCgroupMode(b) + group := "/stat-filtered-bench-cg" + groupPath := fmt.Sprintf("%s-%d", group, os.Getpid()) + c, err := NewManager(defaultCgroup2Path, groupPath, &Resources{}) + require.NoErrorf(b, err, "failed to init new cgroup manager") + b.Cleanup(func() { + _ = c.Delete() + }) + + b.Run("StatAll", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _, err := c.StatFiltered(StatAll) + require.NoError(b, err) + } + }) + + b.Run("CPUOnly", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _, err := c.StatFiltered(StatCPU) + require.NoError(b, err) + } + }) + + b.Run("MemoryOnly", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _, err := c.StatFiltered(StatMemory) + require.NoError(b, err) + } + }) + + b.Run("CPUAndMemory", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _, err := c.StatFiltered(StatCPU | StatMemory) + require.NoError(b, err) + } + }) +} + func toPtr[T any](v T) *T { return &v }