From 0794b5eb8173a33ab68cbdc516f3e9d0ba781720 Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Wed, 10 Jan 2024 14:36:24 +0800
Subject: [PATCH] memory: support cgroup with systemd  (#7627) (#7666)

close tikv/pd#7628

Signed-off-by: bufferflies <1045931706@qq.com>

Co-authored-by: bufferflies <1045931706@qq.com>
---
 cmd/pd-server/main.go |  3 ++-
 pkg/memory/meminfo.go | 38 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/cmd/pd-server/main.go b/cmd/pd-server/main.go
index c09b323d0d1..d0f4e458412 100644
--- 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
index 5a81cf57d3d..64505e2e5f3 100644
--- 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()