From 81edea074ad69e86db59e3903557ba4c2c6ef9bc Mon Sep 17 00:00:00 2001
From: NitroCao <jaycecao520@gmail.com>
Date: Fri, 6 Jun 2025 19:04:25 +0800
Subject: [PATCH] fix: optimize performance of process.OpenFiles()

---
 internal/common/readlink_linux.go | 51 ++++++++++++++++++++++++++++++++++++
 process/process_linux.go          |  4 +--
 process/process_linux_test.go     | 55 +++++++++++++++++++++++++++++++++++++++
 process/testdata/linux/1/fd/0     |  1 +
 process/testdata/linux/1/fd/1     |  1 +
 process/testdata/linux/1/fd/2     |  1 +
 6 files changed, 111 insertions(+), 2 deletions(-)
 create mode 100644 internal/common/readlink_linux.go
 create mode 120000 process/testdata/linux/1/fd/0
 create mode 120000 process/testdata/linux/1/fd/1
 create mode 120000 process/testdata/linux/1/fd/2

diff --git a/internal/common/readlink_linux.go b/internal/common/readlink_linux.go
new file mode 100644
index 0000000..35ed026
--- /dev/null
+++ b/internal/common/readlink_linux.go
@@ -0,0 +1,51 @@
+package common
+
+import (
+	"errors"
+	"os"
+	"sync"
+	"syscall"
+)
+
+var bufferPool = sync.Pool{
+	New: func() any {
+		b := make([]byte, syscall.PathMax)
+		return &b
+	},
+}
+
+// The following three functions are copied from stdlib.
+
+// ignoringEINTR2 is ignoringEINTR, but returning an additional value.
+func ignoringEINTR2[T any](fn func() (T, error)) (T, error) {
+	for {
+		v, err := fn()
+		if !errors.Is(err, syscall.EINTR) {
+			return v, err
+		}
+	}
+}
+
+// Many functions in package syscall return a count of -1 instead of 0.
+// Using fixCount(call()) instead of call() corrects the count.
+func fixCount(n int, err error) (int, error) {
+	if n < 0 {
+		n = 0
+	}
+	return n, err
+}
+
+// Readlink behaves like os.Readlink but caches the buffer passed to syscall.Readlink.
+func Readlink(name string) (string, error) {
+	b := bufferPool.Get().(*[]byte)
+	defer bufferPool.Put(b)
+
+	n, err := ignoringEINTR2(func() (int, error) {
+		return fixCount(syscall.Readlink(name, *b))
+	})
+	if err != nil {
+		return "", &os.PathError{Op: "readlink", Path: name, Err: err}
+	}
+
+	return string((*b)[:n]), nil
+}
diff --git a/process/process_linux.go b/process/process_linux.go
index f8ddbb6..7b10e34 100644
--- a/process/process_linux.go
+++ b/process/process_linux.go
@@ -636,10 +636,10 @@ func (p *Process) fillFromfdWithContext(ctx context.Context) (int32, []*OpenFile
 	}
 	numFDs := int32(len(fnames))
 
-	var openfiles []*OpenFilesStat
+	openfiles := make([]*OpenFilesStat, 0, numFDs)
 	for _, fd := range fnames {
 		fpath := filepath.Join(statPath, fd)
-		path, err := os.Readlink(fpath)
+		path, err := common.Readlink(fpath)
 		if err != nil {
 			continue
 		}
diff --git a/process/process_linux_test.go b/process/process_linux_test.go
index 8c1c203..4a0812d 100644
--- a/process/process_linux_test.go
+++ b/process/process_linux_test.go
@@ -15,6 +15,61 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestFillFromfdWithContext(t *testing.T) {
+	type expect struct {
+		numFDs    int32
+		openFiles []*OpenFilesStat
+		err       error
+	}
+	type testCase struct {
+		name     string
+		pid      int32
+		expected *expect
+	}
+
+	cases := []testCase{
+		{
+			name: "good path",
+			pid:  1,
+			expected: &expect{
+				numFDs: 3,
+				openFiles: []*OpenFilesStat{
+					{
+						Path: "/foo",
+						Fd:   0,
+					},
+					{
+						Path: "/bar",
+						Fd:   1,
+					},
+					{
+						Path: "/baz",
+						Fd:   2,
+					},
+				},
+			},
+		},
+	}
+
+	t.Setenv("HOST_PROC", "testdata/linux")
+	for _, tt := range cases {
+		t.Run(tt.name, func(t *testing.T) {
+			p, err := NewProcess(tt.pid)
+			require.NoError(t, err)
+			numFDs, openFiles, err := p.fillFromfdWithContext(context.TODO())
+			if tt.expected.err != nil {
+				assert.ErrorContains(t, err, tt.expected.err.Error())
+				return
+			}
+
+			//nolint:testifylint // false positive
+			assert.NoError(t, err)
+			assert.Equal(t, tt.expected.numFDs, numFDs)
+			assert.ElementsMatch(t, tt.expected.openFiles, openFiles)
+		})
+	}
+}
+
 func TestSplitProcStat(t *testing.T) {
 	expectedFieldsNum := 53
 	statLineContent := make([]string, expectedFieldsNum-1)
diff --git a/process/testdata/linux/1/fd/0 b/process/testdata/linux/1/fd/0
new file mode 120000
index 0000000..7370610
--- /dev/null
+++ b/process/testdata/linux/1/fd/0
@@ -0,0 +1 @@
+/foo
\ No newline at end of file
diff --git a/process/testdata/linux/1/fd/1 b/process/testdata/linux/1/fd/1
new file mode 120000
index 0000000..f77af99
--- /dev/null
+++ b/process/testdata/linux/1/fd/1
@@ -0,0 +1 @@
+/bar
\ No newline at end of file
diff --git a/process/testdata/linux/1/fd/2 b/process/testdata/linux/1/fd/2
new file mode 120000
index 0000000..7aa88ef
--- /dev/null
+++ b/process/testdata/linux/1/fd/2
@@ -0,0 +1 @@
+/baz
\ No newline at end of file