From 746264387fe1f3d0daa17d05ee7a1751f526dde8 Mon Sep 17 00:00:00 2001
From: Song Zhang <zhangsong34@huawei.com>
Date: Thu, 21 Dec 2023 18:57:19 +0800
Subject: [PATCH] containerd: disable Transparent HugePage for shim process if 
 SHIM_DISABLE_THP is set

Signed-off-by: Song Zhang <zhangsong34@huawei.com>
---
 runtime/v1/shim/client/client.go              | 13 +++++++
 runtime/v1/shim/reaper.go                     | 13 +++++++
 sys/thp.go                                    | 34 +++++++++++++++++++
 sys/thp_amd64.go                              |  3 ++
 sys/thp_arm64.go                              |  3 ++
 .../github.com/containerd/go-runc/monitor.go  | 20 +++++++++--
 6 files changed, 83 insertions(+), 3 deletions(-)
 create mode 100644 sys/thp.go
 create mode 100644 sys/thp_amd64.go
 create mode 100644 sys/thp_arm64.go

diff --git a/runtime/v1/shim/client/client.go b/runtime/v1/shim/client/client.go
index 4c134c2..29a89e7 100644
--- a/runtime/v1/shim/client/client.go
+++ b/runtime/v1/shim/client/client.go
@@ -105,9 +105,22 @@ func WithStart(binary, address, daemonAddress, cgroup string, debug bool, exitHa
 		if err != nil {
 			return nil, nil, err
 		}
+		// Set THP disabled for shim process
+		if err := sys.SetTHPState(1, false); err != nil {
+			if err := sys.SetTHPState(0, true); err != nil {
+				return nil, nil, err
+			}
+			return nil, nil, err
+		}
 		if err := cmd.Start(); err != nil {
+			if err := sys.SetTHPState(0, true); err != nil {
+				return nil, nil, err
+			}
 			return nil, nil, errors.Wrapf(err, "failed to start shim")
 		}
+		if err := sys.SetTHPState(0, true); err != nil {
+			return nil, nil, err
+		}
 		defer func() {
 			if err != nil {
 				cmd.Process.Kill()
diff --git a/runtime/v1/shim/reaper.go b/runtime/v1/shim/reaper.go
index f5f8096..03759a6 100644
--- a/runtime/v1/shim/reaper.go
+++ b/runtime/v1/shim/reaper.go
@@ -76,9 +76,22 @@ type Monitor struct {
 
 // Start starts the command a registers the process with the reaper
 func (m *Monitor) Start(c *exec.Cmd) (chan runc.Exit, error) {
+	// Set THP enabled for subprocess.
+	if err := sys.SetTHPState(0, false); err != nil {
+		if err := sys.SetTHPState(1, true); err != nil {
+			return nil, err
+		}
+		return nil, err
+	}
 	ec := m.Subscribe()
 	if err := c.Start(); err != nil {
 		m.Unsubscribe(ec)
+		if err := sys.SetTHPState(1, true); err != nil {
+			return nil, err
+		}
+		return nil, err
+	}
+	if err := sys.SetTHPState(1, true); err != nil {
 		return nil, err
 	}
 	return ec, nil
diff --git a/sys/thp.go b/sys/thp.go
new file mode 100644
index 0000000..3bbbd00
--- /dev/null
+++ b/sys/thp.go
@@ -0,0 +1,34 @@
+package sys
+
+import (
+	"os"
+	"runtime"
+	"syscall"
+
+	"github.com/sirupsen/logrus"
+)
+
+const (
+	PR_SET_THP_DISABLE = 41
+)
+
+func SetTHPState(flag int, resume bool) error {
+	logrus.Debug("start to set THP")
+	if os.Getenv("SHIM_DISABLE_THP") != "1" {
+		logrus.Debug("skip set THP")
+		return nil
+	}
+
+	if resume {
+		defer runtime.UnlockOSThread()
+	} else {
+		runtime.LockOSThread()
+	}
+
+	_, _, errno := syscall.RawSyscall6(uintptr(PRCTL_SYSCALL), uintptr(PR_SET_THP_DISABLE), uintptr(flag), 0, 0, 0, 0)
+	if errno != 0 {
+		logrus.Errorf("disable THP failed: %v", errno)
+		return errno
+	}
+	return nil
+}
diff --git a/sys/thp_amd64.go b/sys/thp_amd64.go
new file mode 100644
index 0000000..e1e977e
--- /dev/null
+++ b/sys/thp_amd64.go
@@ -0,0 +1,3 @@
+package sys
+
+const PRCTL_SYSCALL = 157
diff --git a/sys/thp_arm64.go b/sys/thp_arm64.go
new file mode 100644
index 0000000..a6db8d6
--- /dev/null
+++ b/sys/thp_arm64.go
@@ -0,0 +1,3 @@
+package sys
+
+const PRCTL_SYSCALL = 167
diff --git a/vendor/github.com/containerd/go-runc/monitor.go b/vendor/github.com/containerd/go-runc/monitor.go
index bb8bbab..a9ac837 100644
--- a/vendor/github.com/containerd/go-runc/monitor.go
+++ b/vendor/github.com/containerd/go-runc/monitor.go
@@ -17,14 +17,15 @@
 package runc
 
 import (
-	"os/exec"
-	"syscall"
-	"time"
 	"io/ioutil"
+	"os/exec"
 	"path/filepath"
 	"strconv"
 	"strings"
+	"syscall"
+	"time"
 
+	"github.com/containerd/containerd/sys"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 )
@@ -54,7 +55,20 @@ type defaultMonitor struct {
 }
 
 func (m *defaultMonitor) Start(c *exec.Cmd) (chan Exit, error) {
+	// Set THP enabled for runc process.
+	if err := sys.SetTHPState(0, false); err != nil {
+		if err := sys.SetTHPState(1, true); err != nil {
+			return nil, err
+		}
+		return nil, err
+	}
 	if err := c.Start(); err != nil {
+		if err := sys.SetTHPState(1, true); err != nil {
+			return nil, err
+		}
+		return nil, err
+	}
+	if err := sys.SetTHPState(1, true); err != nil {
 		return nil, err
 	}
 	ec := make(chan Exit, 1)
-- 
2.33.0