From 4bc9e37b0f5cbb28af8679de511c14ba5ad3d7e0 Mon Sep 17 00:00:00 2001
From: Antoine Toulme <antoine@lunar-ocean.com>
Date: Fri, 25 Aug 2023 13:35:07 -0700
Subject: [PATCH] faster file read

---
 internal/common/common.go       | 23 +++++++++++
 internal/common/common_linux.go | 73 ++++++++++++++++++---------------
 process/process_test.go         |  7 ++++
 3 files changed, 71 insertions(+), 32 deletions(-)

diff --git a/internal/common/common.go b/internal/common/common.go
index 9bfece362..bad8d1f06 100644
--- a/internal/common/common.go
+++ b/internal/common/common.go
@@ -114,6 +114,29 @@ func ReadLines(filename string) ([]string, error) {
 	return ReadLinesOffsetN(filename, 0, -1)
 }
 
+// ReadLine reads a file and returns the first occurrence of a line that is prefixed with prefix.
+func ReadLine(filename string, prefix string) (string, error) {
+	f, err := os.Open(filename)
+	if err != nil {
+		return "", err
+	}
+	defer f.Close()
+	r := bufio.NewReader(f)
+	for {
+		line, err := r.ReadString('\n')
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+		}
+		if strings.HasPrefix(line, prefix) {
+			return line, nil
+		}
+	}
+
+	return "", nil
+}
+
 // ReadLinesOffsetN reads contents from file and splits them by new line.
 // The offset tells at which line number to start.
 // The count determines the number of lines to read (starting from offset):
diff --git a/internal/common/common_linux.go b/internal/common/common_linux.go
index b58edbeb0..607d963a9 100644
--- a/internal/common/common_linux.go
+++ b/internal/common/common_linux.go
@@ -62,17 +62,38 @@ func BootTimeWithContext(ctx context.Context) (uint64, error) {
 		return 0, err
 	}
 
-	statFile := "stat"
+	useStatFile := true
 	if system == "lxc" && role == "guest" {
 		// if lxc, /proc/uptime is used.
-		statFile = "uptime"
+		useStatFile = false
 	} else if system == "docker" && role == "guest" {
 		// also docker, guest
-		statFile = "uptime"
+		useStatFile = false
 	}
 
-	filename := HostProcWithContext(ctx, statFile)
-	lines, err := ReadLines(filename)
+	if useStatFile {
+		return readBootTimeStat(ctx)
+	} else {
+		filename := HostProcWithContext(ctx, "uptime")
+		lines, err := ReadLines(filename)
+		if err != nil {
+			return handleBootTimeFileReadErr(err)
+		}
+		if len(lines) != 1 {
+			return 0, fmt.Errorf("wrong uptime format")
+		}
+		f := strings.Fields(lines[0])
+		b, err := strconv.ParseFloat(f[0], 64)
+		if err != nil {
+			return 0, err
+		}
+		currentTime := float64(time.Now().UnixNano()) / float64(time.Second)
+		t := currentTime - b
+		return uint64(t), nil
+	}
+}
+
+func handleBootTimeFileReadErr(err error) (uint64, error) {
 	if os.IsPermission(err) {
 		var info syscall.Sysinfo_t
 		err := syscall.Sysinfo(&info)
@@ -81,42 +102,30 @@ func BootTimeWithContext(ctx context.Context) (uint64, error) {
 		}
 
 		currentTime := time.Now().UnixNano() / int64(time.Second)
-		t := currentTime - int64(info.Uptime)
+		t := currentTime - info.Uptime
 		return uint64(t), nil
 	}
+	return 0, err
+}
+
+func readBootTimeStat(ctx context.Context) (uint64, error) {
+	filename := HostProcWithContext(ctx, "stat")
+	line, err := ReadLine(filename, "btime")
 	if err != nil {
-		return 0, err
+		return handleBootTimeFileReadErr(err)
 	}
-
-	if statFile == "stat" {
-		for _, line := range lines {
-			if strings.HasPrefix(line, "btime") {
-				f := strings.Fields(line)
-				if len(f) != 2 {
-					return 0, fmt.Errorf("wrong btime format")
-				}
-				b, err := strconv.ParseInt(f[1], 10, 64)
-				if err != nil {
-					return 0, err
-				}
-				t := uint64(b)
-				return t, nil
-			}
-		}
-	} else if statFile == "uptime" {
-		if len(lines) != 1 {
-			return 0, fmt.Errorf("wrong uptime format")
+	if strings.HasPrefix(line, "btime") {
+		f := strings.Fields(line)
+		if len(f) != 2 {
+			return 0, fmt.Errorf("wrong btime format")
 		}
-		f := strings.Fields(lines[0])
-		b, err := strconv.ParseFloat(f[0], 64)
+		b, err := strconv.ParseInt(f[1], 10, 64)
 		if err != nil {
 			return 0, err
 		}
-		currentTime := float64(time.Now().UnixNano()) / float64(time.Second)
-		t := currentTime - b
-		return uint64(t), nil
+		t := uint64(b)
+		return t, nil
 	}
-
 	return 0, fmt.Errorf("could not find btime")
 }
 
diff --git a/process/process_test.go b/process/process_test.go
index 9281c93c3..ae269606b 100644
--- a/process/process_test.go
+++ b/process/process_test.go
@@ -21,6 +21,7 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/shirou/gopsutil/v3/internal/common"
+	"github.com/stretchr/testify/require"
 )
 
 var mu sync.Mutex
@@ -862,3 +863,9 @@ func BenchmarkProcessPpid(b *testing.B) {
 		p.Ppid()
 	}
 }
+
+func BenchmarkProcesses(b *testing.B) {
+	ps, err := Processes()
+	require.NoError(b, err)
+	require.Greater(b, len(ps), 0)
+}