From cf7a9aa08478e1a2dc2bac3f7e8ea168cd36bb8e Mon Sep 17 00:00:00 2001
From: "Jonathan A. Sternberg" <jonathan.sternberg@docker.com>
Date: Thu, 24 Oct 2024 09:56:27 -0500
Subject: [PATCH] pprof: take cpu and memory profiles by setting environment
 variables

When run in standalone mode, the environment variables
`DOCKER_BUILDX_CPU_PROFILE` and `DOCKER_BUILDX_MEM_PROFILE` will cause
profiles to be written by the CLI.

Signed-off-by: Jonathan A. Sternberg <jonathan.sternberg@docker.com>
---
 cmd/buildx/debug.go | 75 +++++++++++++++++++++++++++++++++++++++++++++
 cmd/buildx/main.go  | 18 +++++++----
 2 files changed, 87 insertions(+), 6 deletions(-)
 create mode 100644 cmd/buildx/debug.go

diff --git a/cmd/buildx/debug.go b/cmd/buildx/debug.go
new file mode 100644
index 00000000..49766c68
--- /dev/null
+++ b/cmd/buildx/debug.go
@@ -0,0 +1,75 @@
+package main
+
+import (
+	"context"
+	"os"
+	"runtime"
+	"runtime/pprof"
+
+	"github.com/moby/buildkit/util/bklog"
+	"github.com/sirupsen/logrus"
+)
+
+func setupDebugProfiles(ctx context.Context) (stop func()) {
+	var stopFuncs []func()
+	if fn := setupCPUProfile(ctx); fn != nil {
+		stopFuncs = append(stopFuncs, fn)
+	}
+	if fn := setupHeapProfile(ctx); fn != nil {
+		stopFuncs = append(stopFuncs, fn)
+	}
+	return func() {
+		for _, fn := range stopFuncs {
+			fn()
+		}
+	}
+}
+
+func setupCPUProfile(ctx context.Context) (stop func()) {
+	if cpuProfile := os.Getenv("BUILDX_CPU_PROFILE"); cpuProfile != "" {
+		f, err := os.Create(cpuProfile)
+		if err != nil {
+			bklog.G(ctx).Warn("could not create cpu profile", logrus.WithError(err))
+			return nil
+		}
+
+		if err := pprof.StartCPUProfile(f); err != nil {
+			bklog.G(ctx).Warn("could not start cpu profile", logrus.WithError(err))
+			_ = f.Close()
+			return nil
+		}
+
+		return func() {
+			pprof.StopCPUProfile()
+			if err := f.Close(); err != nil {
+				bklog.G(ctx).Warn("could not close file for cpu profile", logrus.WithError(err))
+			}
+		}
+	}
+	return nil
+}
+
+func setupHeapProfile(ctx context.Context) (stop func()) {
+	if heapProfile := os.Getenv("BUILDX_MEM_PROFILE"); heapProfile != "" {
+		// Memory profile is only created on stop.
+		return func() {
+			f, err := os.Create(heapProfile)
+			if err != nil {
+				bklog.G(ctx).Warn("could not create memory profile", logrus.WithError(err))
+				return
+			}
+
+			// get up-to-date statistics
+			runtime.GC()
+
+			if err := pprof.WriteHeapProfile(f); err != nil {
+				bklog.G(ctx).Warn("could not write memory profile", logrus.WithError(err))
+			}
+
+			if err := f.Close(); err != nil {
+				bklog.G(ctx).Warn("could not close file for memory profile", logrus.WithError(err))
+			}
+		}
+	}
+	return nil
+}
diff --git a/cmd/buildx/main.go b/cmd/buildx/main.go
index 62762f4c..db30e38d 100644
--- a/cmd/buildx/main.go
+++ b/cmd/buildx/main.go
@@ -73,6 +73,16 @@ func runPlugin(cmd *command.DockerCli) error {
 	})
 }
 
+func run(cmd *command.DockerCli) error {
+	stopProfiles := setupDebugProfiles(context.TODO())
+	defer stopProfiles()
+
+	if plugin.RunningStandalone() {
+		return runStandalone(cmd)
+	}
+	return runPlugin(cmd)
+}
+
 func main() {
 	cmd, err := command.NewDockerCli()
 	if err != nil {
@@ -80,15 +90,11 @@ func main() {
 		os.Exit(1)
 	}
 
-	if plugin.RunningStandalone() {
-		err = runStandalone(cmd)
-	} else {
-		err = runPlugin(cmd)
-	}
-	if err == nil {
+	if err = run(cmd); err == nil {
 		return
 	}
 
+	// Check the error from the run function above.
 	if sterr, ok := err.(cli.StatusError); ok {
 		if sterr.Status != "" {
 			fmt.Fprintln(cmd.Err(), sterr.Status)