From 70bcf5d91be0b904a89d0e60941f5d7ce07aa1da Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Mon, 8 Jul 2024 18:44:36 -0700
Subject: [PATCH] `profiletool merge`: Support walking through directories of
 profiles.

This allows the tool to recursively merge a directory full of N profiles,
as opposed to needing each profile to be explicitly listed on the command
line.

The intent is to run this as part of the build system for PGO builds, such
that individual profiles can exist in the repo but get merged at compile
time.

PiperOrigin-RevId: 650441896
---
 tools/profiletool/profiletool.go | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/tools/profiletool/profiletool.go b/tools/profiletool/profiletool.go
index bc1505261..a7ae15910 100644
--- a/tools/profiletool/profiletool.go
+++ b/tools/profiletool/profiletool.go
@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"path/filepath"
 
 	"github.com/google/pprof/profile"
 	"gvisor.dev/gvisor/pkg/log"
@@ -61,9 +62,30 @@ func mergeProfiles() error {
 	if err := mergeCmd.Parse(os.Args[2:]); err != nil {
 		return fmt.Errorf("invalid flags: %w", err)
 	}
-	profilePaths := mergeCmd.Args()
-	if len(profilePaths) < 2 {
-		return errors.New("must provide at least 2 profiles as positional arguments")
+	argPaths := mergeCmd.Args()
+	var profilePaths []string
+	for _, argPath := range argPaths {
+		st, err := os.Stat(argPath)
+		if err != nil {
+			return fmt.Errorf("cannot stat %q: %w", argPath, err)
+		}
+		if st.IsDir() {
+			filepath.Walk(argPath, func(path string, info os.FileInfo, err error) error {
+				if err != nil {
+					return fmt.Errorf("cannot walk %q: %w", path, err)
+				}
+				if info.IsDir() {
+					return nil
+				}
+				profilePaths = append(profilePaths, path)
+				return nil
+			})
+		} else {
+			profilePaths = append(profilePaths, argPath)
+		}
+	}
+	if len(profilePaths) == 0 {
+		return errors.New("no profiles (or directories containing profiles) specified as positional arguments")
 	}
 	profiles := make([]*profile.Profile, len(profilePaths))
 	for i, profilePath := range profilePaths {
@@ -73,7 +95,7 @@ func mergeProfiles() error {
 		}
 		prof, err := profile.Parse(profileFile)
 		if err != nil {
-			return fmt.Errorf("cannot parse %q: %w", profilePath, err)
+			return fmt.Errorf("cannot parse %q as a profile: %w", profilePath, err)
 		}
 		profileFile.Close()
 		profiles[i] = prof