From 9819d8265f92e0797df646f9a7056bb979d0005b Mon Sep 17 00:00:00 2001
From: Michael Pratt <mpratt@google.com>
Date: Wed, 31 Jan 2024 17:56:47 -0500
Subject: [PATCH] cmd/compile: avoid reading entire PGO profile just to check
 the header

isPreProfileFile reads the entire file into memory just to check the
first few bytes, and then throws it all away. We can avoid this by just
peeking at the beginning.

For #58102.

Change-Id: Id2c2844e5e44a2f3a9c7cdb9a027d94d26bdf71d
Reviewed-on: https://go-review.googlesource.com/c/go/+/560035
Reviewed-by: Cherry Mui <cherryyz@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Auto-Submit: Michael Pratt <mpratt@google.com>
---
 src/cmd/compile/internal/pgo/irgraph.go | 71 ++++++++++++-------------
 1 file changed, 33 insertions(+), 38 deletions(-)

diff --git a/src/cmd/compile/internal/pgo/irgraph.go b/src/cmd/compile/internal/pgo/irgraph.go
index 588d681779..224f14368f 100644
--- a/src/cmd/compile/internal/pgo/irgraph.go
+++ b/src/cmd/compile/internal/pgo/irgraph.go
@@ -49,7 +49,7 @@ import (
 	"errors"
 	"fmt"
 	"internal/profile"
-	"io/ioutil"
+	"io"
 	"os"
 	"sort"
 	"strconv"
@@ -145,51 +145,52 @@ type Profile struct {
 
 var wantHdr = "GO PREPROFILE V1\n"
 
-func isPreProfileFile(filename string) (bool, error) {
-	content, err := ioutil.ReadFile(filename)
-	if err != nil {
-		return false, err
+func isPreProfileFile(r *bufio.Reader) (bool, error) {
+	hdr, err := r.Peek(len(wantHdr))
+	if err == io.EOF {
+		// Empty file.
+		return false, nil
+	} else if err != nil {
+		return false, fmt.Errorf("error reading profile header: %w", err)
 	}
 
-	/* check the header */
-	fileContent := string(content)
-	if strings.HasPrefix(fileContent, wantHdr) {
-		return true, nil
-	}
-	return false, nil
+	return string(hdr) == wantHdr, nil
 }
 
 // New generates a profile-graph from the profile or pre-processed profile.
 func New(profileFile string) (*Profile, error) {
-	var profile *Profile
-	var err error
-	isPreProf, err := isPreProfileFile(profileFile)
+	f, err := os.Open(profileFile)
 	if err != nil {
 		return nil, fmt.Errorf("error opening profile: %w", err)
 	}
-	if !isPreProf {
-		profile, err = processProto(profileFile)
-		if err != nil {
-			return nil, fmt.Errorf("error processing pprof PGO profile: %w", err)
-		}
-	} else {
-		profile, err = processPreprof(profileFile)
+	defer f.Close()
+
+	r := bufio.NewReader(f)
+
+	isPreProf, err := isPreProfileFile(r)
+	if err != nil {
+		return nil, fmt.Errorf("error processing profile header: %w", err)
+	}
+
+	if isPreProf {
+		profile, err := processPreprof(r)
 		if err != nil {
 			return nil, fmt.Errorf("error processing preprocessed PGO profile: %w", err)
 		}
+		return profile, nil
+	}
+
+	profile, err := processProto(r)
+	if err != nil {
+		return nil, fmt.Errorf("error processing pprof PGO profile: %w", err)
 	}
 	return profile, nil
 
 }
 
 // processProto generates a profile-graph from the profile.
-func processProto(profileFile string) (*Profile, error) {
-	f, err := os.Open(profileFile)
-	if err != nil {
-		return nil, fmt.Errorf("error opening profile: %w", err)
-	}
-	defer f.Close()
-	p, err := profile.Parse(f)
+func processProto(r io.Reader) (*Profile, error) {
+	p, err := profile.Parse(r)
 	if errors.Is(err, profile.ErrNoData) {
 		// Treat a completely empty file the same as a profile with no
 		// samples: nothing to do.
@@ -242,8 +243,8 @@ func processProto(profileFile string) (*Profile, error) {
 }
 
 // processPreprof generates a profile-graph from the pre-procesed profile.
-func processPreprof(preprofileFile string) (*Profile, error) {
-	namedEdgeMap, totalWeight, err := createNamedEdgeMapFromPreprocess(preprofileFile)
+func processPreprof(r io.Reader) (*Profile, error) {
+	namedEdgeMap, totalWeight, err := createNamedEdgeMapFromPreprocess(r)
 	if err != nil {
 		return nil, err
 	}
@@ -297,14 +298,8 @@ func postProcessNamedEdgeMap(weight map[NamedCallEdge]int64, weightVal int64) (e
 
 // restore NodeMap information from a preprocessed profile.
 // The reader can refer to the format of preprocessed profile in cmd/preprofile/main.go.
-func createNamedEdgeMapFromPreprocess(preprofileFile string) (edgeMap NamedEdgeMap, totalWeight int64, err error) {
-	readFile, err := os.Open(preprofileFile)
-	if err != nil {
-		return NamedEdgeMap{}, 0, fmt.Errorf("error opening preprocessed profile: %w", err)
-	}
-	defer readFile.Close()
-
-	fileScanner := bufio.NewScanner(readFile)
+func createNamedEdgeMapFromPreprocess(r io.Reader) (edgeMap NamedEdgeMap, totalWeight int64, err error) {
+	fileScanner := bufio.NewScanner(r)
 	fileScanner.Split(bufio.ScanLines)
 	weight := make(map[NamedCallEdge]int64)
 
-- 
2.48.1