cmd: add memprofile flags and tests

These are only available when building with the DEBUG=1 make target.

Author: lzap

HACKING.md

@@ -29,6 +29,12 @@ Build by running:
 $ go build ./cmd/image-builder/
 ```
 
+To include optional **profiling** support (`--memprofile`, `--memprofile-goroutine`, `--memprofile-rate`), set **`DEBUG`** to any non-empty value when invoking `make build`, for example:
+
+```console
+$ DEBUG=1 make build
+```
+
 ## Unit tests
 
 Run the unit tests via:

Makefile

@@ -99,9 +99,12 @@ $(BUILDDIR)/%/:
 # keep in sync with:
 # https://github.com/containers/podman/blob/2981262215f563461d449b9841741339f4d9a894/Makefile#L51
 TAGS := containers_image_openpgp,exclude_graphdriver_btrfs,exclude_graphdriver_devicemapper
+ifneq ($(DEBUG),)
+TAGS := $(TAGS),profiling
+endif
 
 .PHONY: build
-build: $(BUILDDIR)/bin/  ## build the binary from source
+build: $(BUILDDIR)/bin/  ## build the binary from source (set DEBUG=1 to include extra build tags)
 	go build -tags="$(TAGS)" -ldflags="-X main.version=${VERSION}" -o $<image-builder ./cmd/image-builder/
 	# Note that this is only needed for the bib container to detect if qemu-user is available
 	for arch in amd64 arm64; do \

cmd/image-builder/main.go

@@ -596,6 +596,7 @@ func normalizeRootArgs(_ *pflag.FlagSet, name string) pflag.NormalizedName {
 func run() error {
 	// Initialize console logger (stderr, no prefix)
 	log.SetFlags(0)
+	memProfileResetForProcessStart()
 
 	rootCmd := &cobra.Command{
 		Use:   "image-builder",
@@ -634,6 +635,8 @@ operating systems like Fedora, CentOS and RHEL with easy customizations support.
 	rootCmd.PersistentFlags().StringArray("force-repo", nil, `Override the base repositories during build (these will not be part of the final image)`)
 	rootCmd.PersistentFlags().String("output-dir", "", `Put output into the specified directory`)
 	rootCmd.PersistentFlags().BoolP("verbose", "v", false, `Switch to verbose mode (more logging on stderr and verbose progress)`)
+	registerMemProfileFlags(rootCmd)
+	rootCmd.PersistentPreRun = memProfilePersistentPreRun
 
 	rootCmd.SetOut(osStdout)
 	rootCmd.SetErr(osStderr)
@@ -764,7 +767,11 @@ operating systems like Fedora, CentOS and RHEL with easy customizations support.
 		ilog.SetDefault(log.New(os.Stderr, "", 0))
 	}
 
-	return rootCmd.Execute()
+	execErr := rootCmd.Execute()
+	if flushErr := memProfileFlush(); flushErr != nil {
+		return errors.Join(execErr, flushErr)
+	}
+	return execErr
 }
 
 func main() {

cmd/image-builder/memprofile.go

@@ -0,0 +1,89 @@
+//go:build profiling
+
+// See @cmd/image-builder/memprofile_stub.go for API documentation shared with the non-profiling build.
+
+package main
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"runtime"
+	"runtime/pprof"
+
+	"github.com/spf13/cobra"
+)
+
+// Go default allocation sampling rate (see runtime.MemProfileRate).
+const defaultMemProfileRate = 512 * 1024
+
+var (
+	memProfilePath          string
+	memProfileGoroutinePath string
+	memProfileRateOpt       int = -1
+)
+
+func memProfileResetForProcessStart() {
+	runtime.MemProfileRate = 0
+}
+
+func registerMemProfileFlags(root *cobra.Command) {
+	f := root.PersistentFlags()
+	f.StringVar(&memProfilePath, "memprofile", "", "write a heap memory profile in pprof format when the program exits (use with go tool pprof)")
+	f.StringVar(&memProfileGoroutinePath, "memprofile-goroutine", "", "write a goroutine profile in pprof format on exit (optional; use with go tool pprof)")
+	f.IntVar(&memProfileRateOpt, "memprofile-rate", -1, "when --memprofile is set, sets runtime.MemProfileRate for allocation sampling (-1 uses Go's default 524288; 0 samples only live heap at exit with minimal runtime overhead)")
+}
+
+func memProfilePersistentPreRun(_ *cobra.Command, _ []string) {
+	runtime.MemProfileRate = 0
+	if memProfilePath != "" {
+		if memProfileRateOpt >= 0 {
+			runtime.MemProfileRate = memProfileRateOpt
+		} else {
+			runtime.MemProfileRate = defaultMemProfileRate
+		}
+	}
+}
+
+func memProfileFlush() error {
+	var errs []error
+	if memProfilePath != "" {
+		if err := writeHeapProfile(memProfilePath); err != nil {
+			errs = append(errs, err)
+		}
+	}
+	if memProfileGoroutinePath != "" {
+		if err := writeGoroutineProfile(memProfileGoroutinePath); err != nil {
+			errs = append(errs, err)
+		}
+	}
+	return errors.Join(errs...)
+}
+
+func writeHeapProfile(path string) error {
+	f, err := os.Create(path)
+	if err != nil {
+		return fmt.Errorf("memprofile: create heap profile %q: %w", path, err)
+	}
+	defer f.Close()
+	if err := pprof.WriteHeapProfile(f); err != nil {
+		return fmt.Errorf("memprofile: write heap profile %q: %w", path, err)
+	}
+	return nil
+}
+
+func writeGoroutineProfile(path string) error {
+	f, err := os.Create(path)
+	if err != nil {
+		return fmt.Errorf("memprofile: create goroutine profile %q: %w", path, err)
+	}
+	defer f.Close()
+	prof := pprof.Lookup("goroutine")
+	if prof == nil {
+		return fmt.Errorf("memprofile: goroutine profile unavailable")
+	}
+	if err := prof.WriteTo(f, 0); err != nil {
+		return fmt.Errorf("memprofile: write goroutine profile %q: %w", path, err)
+	}
+	return nil
+}

cmd/image-builder/memprofile_stub.go

@@ -0,0 +1,23 @@
+//go:build !profiling
+
+// Memory profiling hooks: this file is the no-op implementation when the binary is built
+// without -tags profiling. See memprofile.go for the profiling implementation.
+
+package main
+
+import "github.com/spf13/cobra"
+
+// memProfileResetForProcessStart is called at the start of run() before cobra runs; it keeps
+// default allocation sampling behavior. With profiling, it forces sampling off until flags apply.
+func memProfileResetForProcessStart() {}
+
+// registerMemProfileFlags would attach --memprofile* flags to the root command; no flags are
+// registered without the profiling build tag.
+func registerMemProfileFlags(_ *cobra.Command) {}
+
+// memProfilePersistentPreRun would set runtime.MemProfileRate from --memprofile* flags; it does
+// nothing without the profiling build tag.
+func memProfilePersistentPreRun(_ *cobra.Command, _ []string) {}
+
+// memProfileFlush would write heap and/or goroutine profiles on exit; it always returns nil here.
+func memProfileFlush() error { return nil }

cmd/image-builder/memprofile_test.go

@@ -0,0 +1,42 @@
+//go:build profiling
+
+package main_test
+
+import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	main "github.com/osbuild/image-builder-cli/cmd/image-builder"
+	testrepos "github.com/osbuild/images/test/data/repositories"
+)
+
+func TestMemProfileWritesHeapAndGoroutineFiles(t *testing.T) {
+	restore := main.MockNewRepoRegistry(testrepos.New)
+	defer restore()
+
+	dir := t.TempDir()
+	heapPath := filepath.Join(dir, "heap.pprof")
+	gPath := filepath.Join(dir, "goroutine.pprof")
+
+	restore = main.MockOsArgs([]string{"list", "--format=json", "--memprofile", heapPath, "--memprofile-goroutine", gPath})
+	defer restore()
+
+	var fakeStdout bytes.Buffer
+	restore = main.MockOsStdout(&fakeStdout)
+	defer restore()
+
+	err := main.Run()
+	require.NoError(t, err)
+
+	st, err := os.Stat(heapPath)
+	require.NoError(t, err)
+	require.Greater(t, st.Size(), int64(0))
+
+	stg, err := os.Stat(gPath)
+	require.NoError(t, err)
+	require.Greater(t, stg.Size(), int64(0))
+}