From 367c64c9a22d02cc439bc2cde4046c6b8c4f912d Mon Sep 17 00:00:00 2001
From: Jake Coffman <jake@jakecoffman.com>
Date: Wed, 3 Apr 2024 09:38:37 -0500
Subject: [PATCH] add flag to request and gather a flamegraph

---
 cmd/dependabot/internal/cmd/root.go   |  1 +
 cmd/dependabot/internal/cmd/update.go |  2 ++
 internal/infra/run.go                 | 33 ++++++++++++++++++++++++++-
 testdata/scripts/flamegraph.txt       | 32 ++++++++++++++++++++++++++
 4 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 testdata/scripts/flamegraph.txt

diff --git a/cmd/dependabot/internal/cmd/root.go b/cmd/dependabot/internal/cmd/root.go
index 42e08b5..b8fe736 100644
--- a/cmd/dependabot/internal/cmd/root.go
+++ b/cmd/dependabot/internal/cmd/root.go
@@ -15,6 +15,7 @@ type SharedFlags struct {
 	file                string
 	cache               string
 	debugging           bool
+	flamegraph          bool
 	proxyCertPath       string
 	collectorConfigPath string
 	extraHosts          []string
diff --git a/cmd/dependabot/internal/cmd/update.go b/cmd/dependabot/internal/cmd/update.go
index 8fa2e1e..188f821 100644
--- a/cmd/dependabot/internal/cmd/update.go
+++ b/cmd/dependabot/internal/cmd/update.go
@@ -86,6 +86,7 @@ func NewUpdateCommand() *cobra.Command {
 				CollectorImage:      collectorImage,
 				Creds:               input.Credentials,
 				Debug:               flags.debugging,
+				Flamegraph:          flags.flamegraph,
 				Expected:            nil, // update subcommand doesn't use expectations
 				ExtraHosts:          flags.extraHosts,
 				InputName:           flags.file,
@@ -126,6 +127,7 @@ func NewUpdateCommand() *cobra.Command {
 	cmd.Flags().StringVar(&flags.collectorConfigPath, "collector-config", "", "path to an OpenTelemetry collector config file")
 	cmd.Flags().BoolVar(&flags.pullImages, "pull", true, "pull the image if it isn't present")
 	cmd.Flags().BoolVar(&flags.debugging, "debug", false, "run an interactive shell inside the updater")
+	cmd.Flags().BoolVar(&flags.flamegraph, "flamegraph", false, "generate a flamegraph and other metrics")
 	cmd.Flags().StringArrayVarP(&flags.volumes, "volume", "v", nil, "mount volumes in Docker")
 	cmd.Flags().StringArrayVar(&flags.extraHosts, "extra-hosts", nil, "Docker extra hosts setting on the proxy")
 	cmd.Flags().DurationVarP(&flags.timeout, "timeout", "t", 0, "max time to run an update")
diff --git a/internal/infra/run.go b/internal/infra/run.go
index 5d28601..a04982a 100644
--- a/internal/infra/run.go
+++ b/internal/infra/run.go
@@ -1,6 +1,7 @@
 package infra
 
 import (
+	"archive/tar"
 	"context"
 	"encoding/base64"
 	"encoding/json"
@@ -48,6 +49,8 @@ type RunParams struct {
 	PullImages bool
 	// run an interactive shell?
 	Debug bool
+	// generate performance metrics?
+	Flamegraph bool
 	// Volumes are used to mount directories in Docker
 	Volumes []string
 	// Timeout specifies an optional maximum duration the CLI will run an update.
@@ -408,10 +411,17 @@ func runContainers(ctx context.Context, params RunParams) (err error) {
 			return err
 		}
 	} else {
+		env := userEnv(prox.url, params.ApiUrl)
+		if params.Flamegraph {
+			env = append(env, "FLAMEGRAPH=1")
+		}
 		const cmd = "update-ca-certificates && bin/run fetch_files && bin/run update_files"
-		if err := updater.RunCmd(ctx, cmd, dependabot, userEnv(prox.url, params.ApiUrl)...); err != nil {
+		if err := updater.RunCmd(ctx, cmd, dependabot, env...); err != nil {
 			return err
 		}
+		if params.Flamegraph {
+			getFromContainer(ctx, cli, updater.containerID, "/tmp/dependabot-flamegraph.html")
+		}
 		// If the exit code is non-zero, error when using the `update` subcommand, but not the `test` subcommand.
 		if params.Expected == nil && *updater.ExitCode != 0 {
 			return fmt.Errorf("updater exited with code %d", *updater.ExitCode)
@@ -421,6 +431,27 @@ func runContainers(ctx context.Context, params RunParams) (err error) {
 	return nil
 }
 
+func getFromContainer(ctx context.Context, cli *client.Client, containerID, srcPath string) {
+	reader, _, err := cli.CopyFromContainer(ctx, containerID, srcPath)
+	if err != nil {
+		log.Println("Failed to get from container:", err)
+		return
+	}
+	defer reader.Close()
+	outFile, err := os.Create("flamegraph.html")
+	if err != nil {
+		log.Println("Failed to create file while getting from container:", err)
+		return
+	}
+	defer outFile.Close()
+	tarReader := tar.NewReader(reader)
+	tarReader.Next()
+	_, err = io.Copy(outFile, tarReader)
+	if err != nil {
+		log.Printf("Failed copy while getting from container %v: %v\n", srcPath, err)
+	}
+}
+
 func putCloneDir(ctx context.Context, cli *client.Client, updater *Updater, dir string) error {
 	// Docker won't create the directory, so we have to do it first.
 	const cmd = "mkdir -p " + guestRepoDir
diff --git a/testdata/scripts/flamegraph.txt b/testdata/scripts/flamegraph.txt
new file mode 100644
index 0000000..b2f24ef
--- /dev/null
+++ b/testdata/scripts/flamegraph.txt
@@ -0,0 +1,32 @@
+# Build the dummy Dockerfile
+exec docker build -qt flamegraph-updater .
+
+# Run the dependabot command
+dependabot update go_modules dependabot/cli --updater-image flamegraph-updater --flamegraph
+
+# There should be a flamegraph file in the current directory
+exists flamegraph.html
+
+exec docker rmi -f flamegraph-updater
+
+-- Dockerfile --
+FROM ubuntu:22.04
+
+RUN useradd dependabot
+
+COPY --chown=dependabot --chmod=755 update-ca-certificates /usr/bin/update-ca-certificates
+COPY --chown=dependabot --chmod=755 run bin/run
+
+-- update-ca-certificates --
+#!/usr/bin/env bash
+
+echo "Updated those certificates for ya"
+
+-- run --
+#!/usr/bin/env bash
+
+# if there's an environment variable "FLAMEGRAPH" set, create a fake flamegraph in tmp
+if [ -n "$FLAMEGRAPH" ]; then
+  echo "Creating flamegraph"
+  echo "fake flamegraph" > /tmp/dependabot-flamegraph.html
+fi