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)