add tests for prlimit

Signed-off-by: lifubang <[email protected]>

Author: lifubang

src/syscall/export_linux_test.go

@@ -11,6 +11,7 @@ import (
 var (
 	RawSyscallNoError = rawSyscallNoError
 	ForceClone3       = &forceClone3
+	Prlimit           = prlimit
 )
 
 const (

src/syscall/export_rlimit_test.go

@@ -6,6 +6,12 @@
 
 package syscall
 
+import "sync/atomic"
+
 func OrigRlimitNofile() *Rlimit {
 	return origRlimitNofile.Load()
 }
+
+func GetInternalOrigRlimitNofile() *atomic.Pointer[Rlimit] {
+	return &origRlimitNofile
+}

src/syscall/syscall_linux_test.go

@@ -654,3 +654,66 @@ func TestAllThreadsSyscallBlockedSyscall(t *testing.T) {
 	wr.Close()
 	wg.Wait()
 }
+
+func TestPrlimitSelf(t *testing.T) {
+	origLimit := syscall.OrigRlimitNofile()
+	origRlimitNofile := syscall.GetInternalOrigRlimitNofile()
+	defer origRlimitNofile.Store(origLimit)
+
+	origRlimitNofile.Store(&syscall.Rlimit{
+		Cur: 1024,
+		Max: 65536,
+	})
+
+	// Get the current nofile limit
+	var lim syscall.Rlimit
+	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &lim); err != nil {
+		t.Fatalf("Failed to get the current nofile limit: %v", err)
+	}
+	// Set current process's nofile limit through prlimit
+	if err := syscall.Prlimit(0, syscall.RLIMIT_NOFILE, &lim, nil); err != nil {
+		t.Fatalf("Prlimit self failed: %v", err)
+	}
+
+	rlimLater := origRlimitNofile.Load()
+	if rlimLater != nil {
+		t.Fatalf("origRlimitNofile expected to be nil, but got %v", rlimLater)
+	}
+}
+
+func TestPrlimitOtherProcess(t *testing.T) {
+	origLimit := syscall.OrigRlimitNofile()
+	origRlimitNofile := syscall.GetInternalOrigRlimitNofile()
+	defer origRlimitNofile.Store(origLimit)
+
+	origRlimitNofile.Store(&syscall.Rlimit{
+		Cur: 1024,
+		Max: 65536,
+	})
+	rlimOrig := origRlimitNofile.Load()
+
+	// Get the current nofile limit
+	var lim syscall.Rlimit
+	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &lim); err != nil {
+		t.Fatalf("Failed to get the current nofile limit: %v", err)
+	}
+
+	// Start a child process firstly,
+	// so we can use Prlimit to set it's nofile limit.
+	cmd := exec.Command("sleep", "infinity")
+	cmd.Start()
+	defer func() {
+		cmd.Process.Kill()
+		cmd.Process.Wait()
+	}()
+
+	// Set child process's nofile rlimit through prlimit
+	if err := syscall.Prlimit(cmd.Process.Pid, syscall.RLIMIT_NOFILE, &lim, nil); err != nil {
+		t.Fatalf("Prlimit(%d) failed: %v", cmd.Process.Pid, err)
+	}
+
+	rlimLater := origRlimitNofile.Load()
+	if rlimLater != rlimOrig {
+		t.Fatalf("origRlimitNofile expected to be %v, but got %v", rlimOrig, rlimLater)
+	}
+}