Merge pull request #1866 from NitroCao/fix/optimize-openfiles

fix: optimize performance of process.OpenFiles()

internal/common/readlink_linux.go

@@ -0,0 +1,53 @@
+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)
+
+	n, err := ignoringEINTR2(func() (int, error) {
+		return fixCount(syscall.Readlink(name, *b))
+	})
+	if err != nil {
+		bufferPool.Put(b)
+		return "", &os.PathError{Op: "readlink", Path: name, Err: err}
+	}
+
+	result := string((*b)[:n])
+	bufferPool.Put(b)
+	return result, nil
+}

process/process_linux.go

@@ -372,15 +372,7 @@ func (p *Process) ChildrenWithContext(ctx context.Context) ([]*Process, error) {
 
 func (p *Process) OpenFilesWithContext(ctx context.Context) ([]OpenFilesStat, error) {
 	_, ofs, err := p.fillFromfdWithContext(ctx)
-	if err != nil {
-		return nil, err
-	}
-	ret := make([]OpenFilesStat, len(ofs))
-	for i, o := range ofs {
-		ret[i] = *o
-	}
-
-	return ret, nil
+	return ofs, err
 }
 
 func (p *Process) ConnectionsWithContext(ctx context.Context) ([]net.ConnectionStat, error) {
@@ -629,17 +621,17 @@ func (p *Process) fillFromfdListWithContext(ctx context.Context) (string, []stri
 }
 
 // Get num_fds from /proc/(pid)/fd
-func (p *Process) fillFromfdWithContext(ctx context.Context) (int32, []*OpenFilesStat, error) {
+func (p *Process) fillFromfdWithContext(ctx context.Context) (int32, []OpenFilesStat, error) {
 	statPath, fnames, err := p.fillFromfdListWithContext(ctx)
 	if err != nil {
 		return 0, nil, err
 	}
 	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
 		}
@@ -647,7 +639,7 @@ func (p *Process) fillFromfdWithContext(ctx context.Context) (int32, []*OpenFile
 		if err != nil {
 			return numFDs, openfiles, err
 		}
-		o := &OpenFilesStat{
+		o := OpenFilesStat{
 			Path: path,
 			Fd:   t,
 		}

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)

process/testdata/linux/1/fd/0

@@ -0,0 +1 @@
+/foo

process/testdata/linux/1/fd/1

@@ -0,0 +1 @@
+/bar

process/testdata/linux/1/fd/2

@@ -0,0 +1 @@
+/baz