From 20ab5d0e450af1af503d1ee6e4d2988010533b6f Mon Sep 17 00:00:00 2001
From: coffeegoddd <dustin@dolthub.com>
Date: Thu, 15 Feb 2024 14:03:46 -0800
Subject: [PATCH] /go/performance/utils/dolt_builder: allow profile builds

---
 go/performance/utils/dolt_builder/cmd/main.go | 15 +++++++++--
 go/performance/utils/dolt_builder/run.go      | 27 ++++++++++++++-----
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/go/performance/utils/dolt_builder/cmd/main.go b/go/performance/utils/dolt_builder/cmd/main.go
index e38eab30f3f..532892d8dd5 100644
--- a/go/performance/utils/dolt_builder/cmd/main.go
+++ b/go/performance/utils/dolt_builder/cmd/main.go
@@ -16,6 +16,7 @@ package main
 
 import (
 	"context"
+	"flag"
 	"fmt"
 	"log"
 	"os"
@@ -23,8 +24,14 @@ import (
 	builder "github.com/dolthub/dolt/go/performance/utils/dolt_builder"
 )
 
+var profile = flag.String("profile", "", "path to profile used during build")
+
 func main() {
+	flag.Parse()
 	commitList := os.Args[1:]
+	if *profile != "" {
+		commitList = commitList[2:]
+	}
 	if len(commitList) < 1 {
 		helpStr := "dolt-builder takes Dolt commit shas or tags as arguments\n" +
 			"and builds corresponding binaries to a path specified\n" +
@@ -32,12 +39,16 @@ func main() {
 			"If DOLT_BIN is not set, ./doltBin will be used\n" +
 			"usage: dolt-builder dccba46 4bad226 ...\n" +
 			"usage: dolt-builder v0.19.0 v0.22.6 ...\n" +
-			"set DEBUG=1 to run in debug mode\n"
+			"set DEBUG=1 to run in debug mode\n" +
+			"use the -profile flag to supply a pprof profile\n" +
+			"which will be used to create a PGO build\n" +
+			"usage: dolt-builder -profile /path/to/profile v1.33.0\n" +
+			"only one version may be specified when supplying a profile\n"
 		fmt.Print(helpStr)
 		os.Exit(2)
 	}
 
-	err := builder.Run(context.Background(), commitList)
+	err := builder.Run(context.Background(), commitList, *profile)
 	if err != nil {
 		log.Fatal(err)
 	}
diff --git a/go/performance/utils/dolt_builder/run.go b/go/performance/utils/dolt_builder/run.go
index 5e91675b6e2..344dc2c9fec 100644
--- a/go/performance/utils/dolt_builder/run.go
+++ b/go/performance/utils/dolt_builder/run.go
@@ -16,6 +16,7 @@ package dolt_builder
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"os/signal"
@@ -29,7 +30,11 @@ import (
 
 const envDoltBin = "DOLT_BIN"
 
-func Run(parentCtx context.Context, commitList []string) error {
+func Run(parentCtx context.Context, commitList []string, profilePath string) error {
+	if profilePath != "" && len(commitList) > 1 {
+		return errors.New("cannot build more that one binary when a profile is supplied")
+	}
+
 	doltBin, err := getDoltBin()
 	if err != nil {
 		return err
@@ -79,7 +84,7 @@ func Run(parentCtx context.Context, commitList []string) error {
 	for _, commit := range commitList {
 		commit := commit // https://golang.org/doc/faq#closures_and_goroutines
 		g.Go(func() error {
-			return buildBinaries(ctx, tempDir, repoDir, doltBin, commit)
+			return buildBinaries(ctx, tempDir, repoDir, doltBin, profilePath, commit)
 		})
 	}
 
@@ -128,7 +133,7 @@ func getDoltBin() (string, error) {
 }
 
 // buildBinaries builds a dolt binary at the given commit and stores it in the doltBin
-func buildBinaries(ctx context.Context, tempDir, repoDir, doltBinDir, commit string) error {
+func buildBinaries(ctx context.Context, tempDir, repoDir, doltBinDir, profilePath, commit string) error {
 	checkoutDir := filepath.Join(tempDir, commit)
 	if err := os.MkdirAll(checkoutDir, os.ModePerm); err != nil {
 		return err
@@ -144,7 +149,7 @@ func buildBinaries(ctx context.Context, tempDir, repoDir, doltBinDir, commit str
 		return err
 	}
 
-	command, err := goBuild(ctx, checkoutDir, commitDir)
+	command, err := goBuild(ctx, checkoutDir, commitDir, profilePath)
 	if err != nil {
 		return err
 	}
@@ -153,14 +158,24 @@ func buildBinaries(ctx context.Context, tempDir, repoDir, doltBinDir, commit str
 }
 
 // goBuild builds the dolt binary and returns the filename
-func goBuild(ctx context.Context, source, dest string) (string, error) {
+func goBuild(ctx context.Context, source, dest, profilePath string) (string, error) {
 	goDir := filepath.Join(source, "go")
 	doltFileName := "dolt"
 	if runtime.GOOS == "windows" {
 		doltFileName = "dolt.exe"
 	}
+
+	args := make([]string, 0)
+	args = append(args, "build")
+
+	if profilePath != "" {
+		args = append(args, fmt.Sprintf("-pgo=%s", profilePath))
+	}
+
 	toBuild := filepath.Join(dest, doltFileName)
-	build := ExecCommand(ctx, "go", "build", "-o", toBuild, filepath.Join(goDir, "cmd", "dolt"))
+	args = append(args, "-o", toBuild, filepath.Join(goDir, "cmd", "dolt"))
+
+	build := ExecCommand(ctx, "go", args...)
 	build.Dir = goDir
 	err := build.Run()
 	if err != nil {