Cg2: Add the ability to filter stats

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
 }

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
 }