From ca459dce1bc05dd6c03364a0f282934215f9622b Mon Sep 17 00:00:00 2001
From: ttyS3 <ttys3.rust@gmail.com>
Date: Sat, 25 Dec 2021 22:51:14 +0800
Subject: [PATCH] fix: fixup memory usage for cgroup v2

---
 common/pkg/cgroups/cgroups.go | 23 +++++++++++++++++++++++
 common/pkg/cgroups/memory.go  | 24 ++++++++++++------------
 2 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/common/pkg/cgroups/cgroups.go b/common/pkg/cgroups/cgroups.go
index d8eac537e0..0bf275f385 100644
--- a/common/pkg/cgroups/cgroups.go
+++ b/common/pkg/cgroups/cgroups.go
@@ -365,6 +365,29 @@ func readFileAsUint64(path string) (uint64, error) {
 	return ret, nil
 }
 
+func readFileByKeyAsUint64(path, key string) (uint64, error) {
+	content, err := ioutil.ReadFile(path)
+	if err != nil {
+		return 0, err
+	}
+	for _, line := range strings.Split(string(content), "\n") {
+		fields := strings.SplitN(line, " ", 2)
+		if fields[0] == key {
+			v := cleanString(string(fields[1]))
+			if v == "max" {
+				return math.MaxUint64, nil
+			}
+			ret, err := strconv.ParseUint(v, 10, 64)
+			if err != nil {
+				return ret, errors.Wrapf(err, "parse %s from %s", v, path)
+			}
+			return ret, nil
+		}
+	}
+
+	return 0, fmt.Errorf("no key named %s from %s", key, path)
+}
+
 // New creates a new cgroup control
 func New(path string, resources *spec.LinuxResources) (*CgroupControl, error) {
 	cgroup2, err := IsCgroup2UnifiedMode()
diff --git a/common/pkg/cgroups/memory.go b/common/pkg/cgroups/memory.go
index b3991f7e3e..10d65893c6 100644
--- a/common/pkg/cgroups/memory.go
+++ b/common/pkg/cgroups/memory.go
@@ -7,8 +7,7 @@ import (
 	spec "github.com/opencontainers/runtime-spec/specs-go"
 )
 
-type memHandler struct {
-}
+type memHandler struct{}
 
 func getMemoryHandler() *memHandler {
 	return &memHandler{}
@@ -41,22 +40,23 @@ func (c *memHandler) Stat(ctr *CgroupControl, m *Metrics) error {
 	usage := MemoryUsage{}
 
 	var memoryRoot string
-	filenames := map[string]string{}
+	var limitFilename string
 
 	if ctr.cgroup2 {
 		memoryRoot = filepath.Join(cgroupRoot, ctr.path)
-		filenames["usage"] = "memory.current"
-		filenames["limit"] = "memory.max"
+		limitFilename = "memory.max"
+		if usage.Usage, err = readFileByKeyAsUint64(filepath.Join(memoryRoot, "memory.stat"), "anon"); err != nil {
+			return err
+		}
 	} else {
 		memoryRoot = ctr.getCgroupv1Path(Memory)
-		filenames["usage"] = "memory.usage_in_bytes"
-		filenames["limit"] = "memory.limit_in_bytes"
+		limitFilename = "memory.limit_in_bytes"
+		if usage.Usage, err = readFileAsUint64(filepath.Join(memoryRoot, "memory.usage_in_bytes")); err != nil {
+			return err
+		}
 	}
-	usage.Usage, err = readFileAsUint64(filepath.Join(memoryRoot, filenames["usage"]))
-	if err != nil {
-		return err
-	}
-	usage.Limit, err = readFileAsUint64(filepath.Join(memoryRoot, filenames["limit"]))
+
+	usage.Limit, err = readFileAsUint64(filepath.Join(memoryRoot, limitFilename))
 	if err != nil {
 		return err
 	}