Merge pull request #66367 from cpuguy83/speedup_pidof

getPids - don't recursively traverse every dir in /proc

pkg/util/procfs/procfs_linux.go

@@ -21,6 +21,7 @@ package procfs
 import (
 	"bytes"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"os"
 	"path"
@@ -105,47 +106,60 @@ func PidOf(name string) ([]int, error) {
 
 func getPids(re *regexp.Regexp) []int {
 	pids := []int{}
-	filepath.Walk("/proc", func(path string, info os.FileInfo, err error) error {
+
+	dirFD, err := os.Open("/proc")
+	if err != nil {
+		return nil
+	}
+	defer dirFD.Close()
+
+	for {
+		// Read a small number at a time in case there are many entries, we don't want to
+		// allocate a lot here.
+		ls, err := dirFD.Readdir(10)
+		if err == io.EOF {
+			break
+		}
 		if err != nil {
-			// We should continue processing other directories/files
 			return nil
 		}
-		base := filepath.Base(path)
-		// Traverse only the directories we are interested in
-		if info.IsDir() && path != "/proc" {
+
+		for _, entry := range ls {
+			if !entry.IsDir() {
+				continue
+			}
+
 			// If the directory is not a number (i.e. not a PID), skip it
-			if _, err := strconv.Atoi(base); err != nil {
-				return filepath.SkipDir
+			pid, err := strconv.Atoi(entry.Name())
+			if err != nil {
+				continue
+			}
+
+			cmdline, err := ioutil.ReadFile(filepath.Join("/proc", entry.Name(), "cmdline"))
+			if err != nil {
+				glog.V(4).Infof("Error reading file %s: %+v", filepath.Join("/proc", entry.Name(), "cmdline"), err)
+				continue
+			}
+
+			// The bytes we read have '\0' as a separator for the command line
+			parts := bytes.SplitN(cmdline, []byte{0}, 2)
+			if len(parts) == 0 {
+				continue
+			}
+			// Split the command line itself we are interested in just the first part
+			exe := strings.FieldsFunc(string(parts[0]), func(c rune) bool {
+				return unicode.IsSpace(c) || c == ':'
+			})
+			if len(exe) == 0 {
+				continue
+			}
+			// Check if the name of the executable is what we are looking for
+			if re.MatchString(exe[0]) {
+				// Grab the PID from the directory path
+				pids = append(pids, pid)
 			}
 		}
-		if base != "cmdline" {
-			return nil
-		}
-		cmdline, err := ioutil.ReadFile(path)
-		if err != nil {
-			glog.V(4).Infof("Error reading file %s: %+v", path, err)
-			return nil
-		}
-		// The bytes we read have '\0' as a separator for the command line
-		parts := bytes.SplitN(cmdline, []byte{0}, 2)
-		if len(parts) == 0 {
-			return nil
-		}
-		// Split the command line itself we are interested in just the first part
-		exe := strings.FieldsFunc(string(parts[0]), func(c rune) bool {
-			return unicode.IsSpace(c) || c == ':'
-		})
-		if len(exe) == 0 {
-			return nil
-		}
-		// Check if the name of the executable is what we are looking for
-		if re.MatchString(exe[0]) {
-			dirname := filepath.Base(filepath.Dir(path))
-			// Grab the PID from the directory path
-			pid, _ := strconv.Atoi(dirname)
-			pids = append(pids, pid)
-		}
-		return nil
-	})
+	}
+
 	return pids
 }

pkg/util/procfs/procfs_linux_test.go

@@ -23,6 +23,7 @@ import (
 	"os"
 	"os/signal"
 	"path/filepath"
+	"regexp"
 	"runtime"
 	"syscall"
 	"testing"
@@ -95,3 +96,21 @@ func TestPKill(t *testing.T) {
 		t.Fatalf("timeout waiting for %v", sig)
 	}
 }
+
+func BenchmarkGetPids(b *testing.B) {
+	if runtime.GOOS == "darwin" || runtime.GOOS == "windows" {
+		b.Skipf("not supported on GOOS=%s", runtime.GOOS)
+	}
+
+	re, err := regexp.Compile("(^|/)" + filepath.Base(os.Args[0]) + "$")
+	assert.Empty(b, err)
+
+	for i := 0; i < b.N; i++ {
+		pids := getPids(re)
+
+		b.StopTimer()
+		assert.NotZero(b, pids)
+		assert.Contains(b, pids, os.Getpid())
+		b.StartTimer()
+	}
+}