From cc37e536cbf8f63d046271109d1904ccd6ac45ed Mon Sep 17 00:00:00 2001
From: Jing Chen <chjing@google.com>
Date: Mon, 25 Mar 2024 20:25:34 -0700
Subject: [PATCH] Increase the RLIMIT_MEMLOCK when TPUProxy is enabled.

PiperOrigin-RevId: 619045692
---
 runsc/boot/limits.go |  8 ++++++--
 runsc/boot/loader.go | 10 +++++-----
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/runsc/boot/limits.go b/runsc/boot/limits.go
index 13dccd60a..084fada34 100644
--- a/runsc/boot/limits.go
+++ b/runsc/boot/limits.go
@@ -113,12 +113,16 @@ func (d *defs) initDefaults() error {
 	return nil
 }
 
-func createLimitSet(spec *specs.Spec) (*limits.LimitSet, error) {
+func createLimitSet(spec *specs.Spec, enableTPUProxy bool) (*limits.LimitSet, error) {
 	ls, err := defaults.get()
 	if err != nil {
 		return nil, err
 	}
-
+	// Set RLIMIT_MEMLOCK's default value to unlimited when TPUProxy is enabled.
+	// The value will be overwritten if the exact rlimit is provided.
+	if enableTPUProxy {
+		ls.SetUnchecked(limits.MemoryLocked, limits.Limit{Cur: limits.Infinity, Max: limits.Infinity})
+	}
 	// Then apply overwrites on top of defaults.
 	for _, rl := range spec.Process.Rlimits {
 		lt, ok := limits.FromLinuxResourceName[rl.Type]
diff --git a/runsc/boot/loader.go b/runsc/boot/loader.go
index c92f75aa7..a9c214e41 100644
--- a/runsc/boot/loader.go
+++ b/runsc/boot/loader.go
@@ -514,7 +514,7 @@ func New(args Args) (*Loader, error) {
 	dogOpts.TaskTimeoutAction = args.Conf.WatchdogAction
 	dog := watchdog.New(k, dogOpts)
 
-	procArgs, err := createProcessArgs(args.ID, args.Spec, creds, k, k.RootPIDNamespace())
+	procArgs, err := createProcessArgs(args.ID, args.Spec, args.Conf, creds, k, k.RootPIDNamespace())
 	if err != nil {
 		return nil, fmt.Errorf("creating init process for root container: %w", err)
 	}
@@ -582,9 +582,9 @@ func New(args Args) (*Loader, error) {
 }
 
 // createProcessArgs creates args that can be used with kernel.CreateProcess.
-func createProcessArgs(id string, spec *specs.Spec, creds *auth.Credentials, k *kernel.Kernel, pidns *kernel.PIDNamespace) (kernel.CreateProcessArgs, error) {
+func createProcessArgs(id string, spec *specs.Spec, conf *config.Config, creds *auth.Credentials, k *kernel.Kernel, pidns *kernel.PIDNamespace) (kernel.CreateProcessArgs, error) {
 	// Create initial limits.
-	ls, err := createLimitSet(spec)
+	ls, err := createLimitSet(spec, specutils.TPUProxyIsEnabled(spec, conf))
 	if err != nil {
 		return kernel.CreateProcessArgs{}, fmt.Errorf("creating limits: %w", err)
 	}
@@ -899,7 +899,7 @@ func (l *Loader) startSubcontainer(spec *specs.Spec, conf *config.Config, cid st
 		nvidiaDriverVersion: l.root.nvidiaDriverVersion,
 	}
 	var err error
-	info.procArgs, err = createProcessArgs(cid, spec, creds, l.k, pidns)
+	info.procArgs, err = createProcessArgs(cid, spec, conf, creds, l.k, pidns)
 	if err != nil {
 		return fmt.Errorf("creating new process: %w", err)
 	}
@@ -1196,7 +1196,7 @@ func (l *Loader) executeAsync(args *control.ExecArgs) (kernel.ThreadID, error) {
 	}
 	args.PIDNamespace = tg.PIDNamespace()
 
-	args.Limits, err = createLimitSet(l.root.spec)
+	args.Limits, err = createLimitSet(l.root.spec, specutils.TPUProxyIsEnabled(l.root.spec, l.root.conf))
 	if err != nil {
 		return 0, fmt.Errorf("creating limits: %w", err)
 	}