memory: support cgroup with systemd (tikv#7627)

close tikv#7628 Signed-off-by: bufferflies <[email protected]> Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
lhy1024 · Jan 4, 2024 · 754095a · 754095a
1 parent 61b15c6
commit 754095a
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/cmd/pd-server/main.go b/cmd/pd-server/main.go
@@ -30,6 +30,7 @@ import (
 	resource_manager "github.com/tikv/pd/pkg/mcs/resourcemanager/server"
 	scheduling "github.com/tikv/pd/pkg/mcs/scheduling/server"
 	tso "github.com/tikv/pd/pkg/mcs/tso/server"
+	"github.com/tikv/pd/pkg/memory"
 	"github.com/tikv/pd/pkg/schedule/schedulers"
 	"github.com/tikv/pd/pkg/swaggerserver"
 	"github.com/tikv/pd/pkg/utils/configutil"
@@ -229,7 +230,7 @@ func start(cmd *cobra.Command, args []string, services ...string) {
 	}
 	// Flushing any buffered log entries
 	defer log.Sync()
-
+	memory.InitMemoryHook()
 	if len(services) != 0 {
 		versioninfo.Log(server.APIServiceMode)
 	} else {

diff --git a/pkg/memory/meminfo.go b/pkg/memory/meminfo.go
@@ -52,9 +52,13 @@ func MemTotalNormal() (uint64, error) {
 	if time.Since(t) < 60*time.Second {
 		return total, nil
 	}
+	return totalMem()
+}
+
+func totalMem() (uint64, error) {
 	v, err := mem.VirtualMemory()
 	if err != nil {
-		return v.Total, err
+		return 0, err
 	}
 	memLimit.set(v.Total, time.Now())
 	return v.Total, nil
@@ -182,6 +186,38 @@ func init() {
 	mustNil(err)
 }
 
+// InitMemoryHook initializes the memory hook.
+// It is to solve the problem that tidb cannot read cgroup in the systemd.
+// so if we are not in the container, we compare the cgroup memory limit and the physical memory,
+// the cgroup memory limit is smaller, we use the cgroup memory hook.
+// ref https://github.com/pingcap/tidb/pull/48096/
+func InitMemoryHook() {
+	if cgroup.InContainer() {
+		log.Info("use cgroup memory hook because pd is in the container")
+		return
+	}
+	cgroupValue, err := cgroup.GetMemoryLimit()
+	if err != nil {
+		return
+	}
+	physicalValue, err := totalMem()
+	if err != nil {
+		return
+	}
+	if physicalValue > cgroupValue && cgroupValue != 0 {
+		MemTotal = MemTotalCGroup
+		MemUsed = MemUsedCGroup
+		sysutil.RegisterGetMemoryCapacity(MemTotalCGroup)
+		log.Info("use cgroup memory hook", zap.Int64("cgroupMemorySize", int64(cgroupValue)), zap.Int64("physicalMemorySize", int64(physicalValue)))
+	} else {
+		log.Info("use physical memory hook", zap.Int64("cgroupMemorySize", int64(cgroupValue)), zap.Int64("physicalMemorySize", int64(physicalValue)))
+	}
+	_, err = MemTotal()
+	mustNil(err)
+	_, err = MemUsed()
+	mustNil(err)
+}
+
 // InstanceMemUsed returns the memory usage of this process
 func InstanceMemUsed() (uint64, error) {
 	used, t := serverMemUsage.get()