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 <[email protected]>

Author: jsternberg

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
+}

cmd/buildx/main.go

@@ -40,6 +40,9 @@ func init() {
 }
 
 func runStandalone(cmd *command.DockerCli) error {
+	stopProfiles := setupDebugProfiles(context.TODO())
+	defer stopProfiles()
+
 	if err := cmd.Initialize(cliflags.NewClientOptions()); err != nil {
 		return err
 	}
@@ -65,6 +68,9 @@ func flushMetrics(cmd *command.DockerCli) {
 }
 
 func runPlugin(cmd *command.DockerCli) error {
+	stopProfiles := setupDebugProfiles(context.TODO())
+	defer stopProfiles()
+
 	rootCmd := commands.NewRootCmd("buildx", true, cmd)
 	return plugin.RunPlugin(cmd, rootCmd, manager.Metadata{
 		SchemaVersion: "0.1.0",