fix: read cgroupv2 memory usage from the cgroup that owns memory.max

cgroupv2.go

@@ -156,7 +156,17 @@ func (s cgroupV2Statter) cpuTotal() (total float64, err error) {
 	return quotaUs / periodUs, nil
 }
 
-func (s cgroupV2Statter) memoryMaxBytes() (*float64, error) {
+// memoryLimit is the effective memory.max and the cgroup that sets it.
+// owner is kept so usage can be read from the same cgroup the limit is
+// enforced on.
+type memoryLimit struct {
+	owner *cgroupV2Statter
+	bytes float64
+}
+
+// memoryLimit returns the nearest cgroup, starting at s and walking up,
+// that sets memory.max. It returns nil when no cgroup sets a limit.
+func (s *cgroupV2Statter) memoryLimit() (*memoryLimit, error) {
 	memoryMaxPath := filepath.Join(s.path, cgroupV2MemoryMaxBytes)
 
 	maxUsageBytes, err := readInt64(s.fs, memoryMaxPath)
@@ -168,12 +178,26 @@ func (s cgroupV2Statter) memoryMaxBytes() (*float64, error) {
 		// If the value is not a valid integer _or_ the memory max file
 		// does not exist, we can assume that the limit is 'max'.
 		// Try the parent or return nil (no limit).
-		return getFromParentOrFallback(&s, (*cgroupV2Statter).memoryMaxBytes, func() (*float64, error) {
+		return getFromParentOrFallback(s, (*cgroupV2Statter).memoryLimit, func() (*memoryLimit, error) {
 			return nil, nil
 		})
 	}
 
-	return ptr.To(float64(maxUsageBytes)), nil
+	return &memoryLimit{owner: s, bytes: float64(maxUsageBytes)}, nil
+}
+
+func (s cgroupV2Statter) memoryUsedBytes() (int64, error) {
+	currUsageBytes, err := s.memoryCurrentBytes()
+	if err != nil {
+		return 0, xerrors.Errorf("read memory usage: %w", err)
+	}
+
+	inactiveFileBytes, err := s.memoryInactiveFileBytes()
+	if err != nil {
+		return 0, xerrors.Errorf("read memory stats: %w", err)
+	}
+
+	return currUsageBytes - inactiveFileBytes, nil
 }
 
 func (s cgroupV2Statter) memoryCurrentBytes() (int64, error) {
@@ -218,22 +242,37 @@ func (s cgroupV2Statter) memory(p Prefix) (*Result, error) {
 		Unit:   "B",
 		Prefix: p,
 	}
-	if total, err := s.memoryMaxBytes(); err != nil {
+	limit, err := s.memoryLimit()
+	if err != nil {
 		return nil, xerrors.Errorf("read memory total: %w", err)
-	} else {
-		r.Total = total
 	}
 
-	currUsageBytes, err := s.memoryCurrentBytes()
-	if err != nil {
-		return nil, xerrors.Errorf("read memory usage: %w", err)
+	// Read usage from the cgroup that owns the limit so that used and
+	// total describe the same domain. Otherwise sibling cgroups under
+	// the limit owner (for example /docker/<id> next to /init.scope in
+	// docker-in-docker) would count against the limit but not the usage.
+	if limit != nil {
+		r.Total = ptr.To(limit.bytes)
+
+		if limit.owner.path != s.path {
+			if used, err := limit.owner.memoryUsedBytes(); err != nil {
+				if !errors.Is(err, fs.ErrNotExist) {
+					return nil, xerrors.Errorf("read limit owner memory: %w", err)
+				}
+				// The limit owner does not expose usage files, so fall
+				// back to the process's own cgroup below.
+			} else {
+				r.Used = float64(used)
+				return r, nil
+			}
+		}
 	}
 
-	inactiveFileBytes, err := s.memoryInactiveFileBytes()
+	used, err := s.memoryUsedBytes()
 	if err != nil {
-		return nil, xerrors.Errorf("read memory stats: %w", err)
+		return nil, err
 	}
 
-	r.Used = float64(currUsageBytes - inactiveFileBytes)
+	r.Used = float64(used)
 	return r, nil
 }

stat_internal_test.go

@@ -507,6 +507,25 @@ func TestStatter(t *testing.T) {
 				assert.Equal(t, "B", mem.Unit)
 			})
 
+			t.Run("Memory/UsageFromLimitOwner", func(t *testing.T) {
+				t.Parallel()
+
+				fs := initFS(t, fsContainerCgroupV2DockerInDocker)
+				s, err := New(WithFS(fs), withNoWait, withIsCgroupV2(true))
+				require.NoError(t, err)
+
+				mem, err := s.ContainerMemory(PrefixDefault)
+				require.NoError(t, err)
+
+				require.NotNil(t, mem)
+				// Usage comes from the root that owns memory.max, which
+				// includes the sibling /docker/<id> cgroup.
+				assert.Equal(t, 805306368.0, mem.Used)
+				require.NotNil(t, mem.Total)
+				assert.Equal(t, 1073741824.0, *mem.Total)
+				assert.Equal(t, "B", mem.Unit)
+			})
+
 			t.Run("Memory/StatInParent", func(t *testing.T) {
 				t.Parallel()
 
@@ -835,6 +854,31 @@ sysboxfs /proc/sys sysboxfs rw,nosuid,nodev,noexec,relatime 0 0`,
 		filepath.Join(cgroupRootPath, "init.scope", cgroupV2MemoryStat):       "inactive_file 268435456",
 		filepath.Join(cgroupRootPath, "init.scope", cgroupV2MemoryUsageBytes): "536870912",
 	}
+	// fsContainerCgroupV2DockerInDocker simulates a docker-in-docker workspace.
+	// The agent runs in /init.scope and a nested container runs in the
+	// sibling /docker/<id>. Only the container root sets memory.max, so
+	// usage must be read from the root to include both children.
+	fsContainerCgroupV2DockerInDocker = map[string]string{
+		procOneCgroup:  "0::/init.scope",
+		procSelfCgroup: "0::/init.scope",
+		procMounts: `overlay / overlay rw,relatime,lowerdir=/some/path:/some/path,upperdir=/some/path:/some/path,workdir=/some/path:/some/path 0 0
+proc /proc/sys proc ro,nosuid,nodev,noexec,relatime 0 0`,
+		sysCgroupType: "domain",
+
+		filepath.Join(cgroupRootPath, cgroupV2CPUMax):           "max 100000",
+		filepath.Join(cgroupRootPath, cgroupV2MemoryMaxBytes):   "1073741824",
+		filepath.Join(cgroupRootPath, cgroupV2MemoryUsageBytes): "939524096",
+		filepath.Join(cgroupRootPath, cgroupV2MemoryStat):       "inactive_file 134217728",
+
+		filepath.Join(cgroupRootPath, "init.scope", cgroupV2CPUStat):          "usage_usec 0",
+		filepath.Join(cgroupRootPath, "init.scope", cgroupV2MemoryMaxBytes):   "max",
+		filepath.Join(cgroupRootPath, "init.scope", cgroupV2MemoryUsageBytes): "67108864",
+		filepath.Join(cgroupRootPath, "init.scope", cgroupV2MemoryStat):       "inactive_file 0",
+
+		filepath.Join(cgroupRootPath, cgroupV2Path, cgroupV2MemoryMaxBytes):   "max",
+		filepath.Join(cgroupRootPath, cgroupV2Path, cgroupV2MemoryUsageBytes): "872415232",
+		filepath.Join(cgroupRootPath, cgroupV2Path, cgroupV2MemoryStat):       "inactive_file 134217728",
+	}
 	// Variant where child has memory.stat but NOT memory.current (should inherit from root)
 	fsContainerCgroupV2KubernetesMissingMemoryCurrent = map[string]string{
 		procOneCgroup:  "0::/",