all: add tracing (#750)

Add a hidden `--trace` flag that turns on tracing to measure execution
time. When set, a trace file is written (by default to trace.out) which
can be inspected with `go tool trace trace.out`.

internal/boxcli/midcobra/midcobra.go

@@ -55,15 +55,16 @@ func (ex *midcobraExecutable) AddMiddleware(mids ...Middleware) {
 
 func (ex *midcobraExecutable) Execute(ctx context.Context, args []string) int {
 	// Ensure cobra uses the same arguments
-	ex.cmd.SetArgs(args)
+	ex.cmd.SetContext(ctx)
+	_ = ex.cmd.ParseFlags(args)
 
 	// Run the 'pre' hooks
 	for _, m := range ex.middlewares {
 		m.preRun(ex.cmd, args)
 	}
 
 	// Execute the cobra command:
-	err := ex.cmd.ExecuteContext(ctx)
+	err := ex.cmd.Execute()
 
 	var postRunErr error
 	var userExecErr *usererr.ExitError
@@ -76,8 +77,8 @@ func (ex *midcobraExecutable) Execute(ctx context.Context, args []string) int {
 	// run even if the command resulted in an error. This is useful when we still want to clean up
 	// before the program exists or we want to log something. The error, if any, gets passed
 	// to the post hook.
-	for _, m := range ex.middlewares {
-		m.postRun(ex.cmd, args, postRunErr)
+	for i := len(ex.middlewares) - 1; i >= 0; i-- {
+		ex.middlewares[i].postRun(ex.cmd, args, postRunErr)
 	}
 
 	if err != nil {

internal/boxcli/midcobra/telemetry.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"runtime/trace"
 	"strings"
 	"time"
 
@@ -55,13 +56,19 @@ func (m *telemetryMiddleware) withExecutionID(execID string) Middleware {
 
 func (m *telemetryMiddleware) preRun(cmd *cobra.Command, args []string) {
 	m.startTime = telemetry.CommandStartTime()
-	if !m.disabled {
-		sentry := telemetry.NewSentry(m.opts.SentryDSN)
-		sentry.Init(m.opts.AppName, m.opts.AppVersion, m.executionID)
+
+	ctx := cmd.Context()
+	defer trace.StartRegion(ctx, "telemetryPreRun").End()
+	if m.disabled {
+		trace.Log(ctx, "telemetry", "telemetry is disabled")
+		return
 	}
+	sentry := telemetry.NewSentry(m.opts.SentryDSN)
+	sentry.Init(m.opts.AppName, m.opts.AppVersion, m.executionID)
 }
 
 func (m *telemetryMiddleware) postRun(cmd *cobra.Command, args []string, runErr error) {
+	defer trace.StartRegion(cmd.Context(), "telemetryPostRun").End()
 	if m.disabled {
 		return
 	}

internal/boxcli/midcobra/trace.go

@@ -0,0 +1,67 @@
+// Copyright 2022 Jetpack Technologies Inc and contributors. All rights reserved.
+// Use of this source code is governed by the license in the LICENSE file.
+
+package midcobra
+
+import (
+	"context"
+	"os"
+	"runtime/trace"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+)
+
+type TraceMiddleware struct {
+	executionID string // uuid
+	tracef      *os.File
+	flag        *pflag.Flag
+	task        *trace.Task
+}
+
+var _ Middleware = (*DebugMiddleware)(nil)
+
+func (t *TraceMiddleware) AttachToFlag(flags *pflag.FlagSet, flagName string) {
+	flags.String(flagName, "", "write a trace to a file")
+	t.flag = flags.Lookup(flagName)
+	t.flag.Hidden = true
+	t.flag.NoOptDefVal = "trace.out"
+}
+
+func (t *TraceMiddleware) preRun(cmd *cobra.Command, args []string) {
+	if t == nil {
+		return
+	}
+	path := t.flag.Value.String()
+	if path == "" {
+		return
+	}
+	var err error
+	t.tracef, err = os.Create(path)
+	if err != nil {
+		panic("error enabling tracing: " + err.Error())
+	}
+	if err := trace.Start(t.tracef); err != nil {
+		panic("error enabling tracing: " + err.Error())
+	}
+
+	var ctx context.Context
+	ctx, t.task = trace.NewTask(cmd.Context(), "cliCommand")
+	cmd.SetContext(ctx)
+}
+
+func (t *TraceMiddleware) postRun(cmd *cobra.Command, args []string, runErr error) {
+	if t.tracef == nil {
+		return
+	}
+	t.task.End()
+	trace.Stop()
+	if err := t.tracef.Close(); err != nil {
+		panic("error closing trace file: " + err.Error())
+	}
+}
+
+func (t *TraceMiddleware) withExecutionID(execID string) Middleware {
+	t.executionID = execID
+	return t
+}

internal/boxcli/root.go

@@ -15,7 +15,10 @@ import (
 	"go.jetpack.io/devbox/internal/debug"
 )
 
-var debugMiddleware *midcobra.DebugMiddleware = &midcobra.DebugMiddleware{}
+var (
+	debugMiddleware *midcobra.DebugMiddleware = &midcobra.DebugMiddleware{}
+	traceMiddleware *midcobra.TraceMiddleware = &midcobra.TraceMiddleware{}
+)
 
 type rootCmdFlags struct {
 	quiet bool
@@ -58,13 +61,15 @@ func RootCmd() *cobra.Command {
 	command.PersistentFlags().BoolVarP(
 		&flags.quiet, "quiet", "q", false, "suppresses logs")
 	debugMiddleware.AttachToFlag(command.PersistentFlags(), "debug")
+	traceMiddleware.AttachToFlag(command.PersistentFlags(), "trace")
 
 	return command
 }
 
 func Execute(ctx context.Context, args []string) int {
 	defer debug.Recover()
 	exe := midcobra.New(RootCmd())
+	exe.AddMiddleware(traceMiddleware)
 	exe.AddMiddleware(midcobra.Telemetry())
 	exe.AddMiddleware(debugMiddleware)
 	return exe.Execute(ctx, args)

internal/impl/devbox.go

@@ -11,6 +11,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"runtime/trace"
 	"strconv"
 	"strings"
 
@@ -138,7 +139,10 @@ func (d *Devbox) Generate() error {
 }
 
 func (d *Devbox) Shell() error {
-	if err := d.ensurePackagesAreInstalled(ensure); err != nil {
+	ctx, task := trace.NewTask(context.Background(), "devboxShell")
+	defer task.End()
+
+	if err := d.ensurePackagesAreInstalled(ctx, ensure); err != nil {
 		return err
 	}
 	fmt.Fprintln(d.writer, "Starting a devbox shell...")
@@ -153,7 +157,7 @@ func (d *Devbox) Shell() error {
 		return err
 	}
 
-	env, err := d.computeNixEnv()
+	env, err := d.computeNixEnv(ctx)
 	if err != nil {
 		return err
 	}
@@ -183,15 +187,18 @@ func (d *Devbox) Shell() error {
 }
 
 func (d *Devbox) RunScript(cmdName string, cmdArgs []string) error {
-	if err := d.ensurePackagesAreInstalled(ensure); err != nil {
+	ctx, task := trace.NewTask(context.Background(), "devboxRun")
+	defer task.End()
+
+	if err := d.ensurePackagesAreInstalled(ctx, ensure); err != nil {
 		return err
 	}
 
 	if err := d.writeScriptsToFiles(); err != nil {
 		return err
 	}
 
-	env, err := d.computeNixEnv()
+	env, err := d.computeNixEnv(ctx)
 	if err != nil {
 		return err
 	}
@@ -228,7 +235,10 @@ func (d *Devbox) ListScripts() []string {
 }
 
 func (d *Devbox) PrintEnv() (string, error) {
-	envs, err := d.computeNixEnv()
+	ctx, task := trace.NewTask(context.Background(), "devboxPrintEnv")
+	defer task.End()
+
+	envs, err := d.computeNixEnv(ctx)
 	if err != nil {
 		return "", err
 	}
@@ -448,7 +458,9 @@ func (d *Devbox) generateShellFiles() error {
 // Note that the shellrc.tmpl template (which sources this environment) does
 // some additional processing. The computeNixEnv environment won't necessarily
 // represent the final "devbox run" or "devbox shell" environments.
-func (d *Devbox) computeNixEnv() (map[string]string, error) {
+func (d *Devbox) computeNixEnv(ctx context.Context) (map[string]string, error) {
+	defer trace.StartRegion(ctx, "computeNixEnv").End()
+
 	currentEnv := os.Environ()
 	env := make(map[string]string, len(currentEnv))
 	for _, kv := range currentEnv {
@@ -471,7 +483,7 @@ func (d *Devbox) computeNixEnv() (map[string]string, error) {
 		originalPath = currentEnvPath
 	}
 
-	vaf, err := nix.PrintDevEnv(d.nixShellFilePath(), d.nixFlakesFilePath())
+	vaf, err := nix.PrintDevEnv(ctx, d.nixShellFilePath(), d.nixFlakesFilePath())
 	if err != nil {
 		return nil, err
 	}

internal/impl/packages.go

@@ -1,10 +1,12 @@
 package impl
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/exec"
 	"path/filepath"
+	"runtime/trace"
 	"sort"
 	"strings"
 
@@ -24,6 +26,9 @@ import (
 
 // Add adds the `pkgs` to the config (i.e. devbox.json) and nix profile for this devbox project
 func (d *Devbox) Add(pkgs ...string) error {
+	ctx, task := trace.NewTask(context.Background(), "devboxAdd")
+	defer task.End()
+
 	original := d.cfg.RawPackages
 	// Check packages are valid before adding.
 	for _, pkg := range pkgs {
@@ -45,7 +50,7 @@ func (d *Devbox) Add(pkgs ...string) error {
 	}
 
 	d.pluginManager.ApplyOptions(plugin.WithAddMode())
-	if err := d.ensurePackagesAreInstalled(install); err != nil {
+	if err := d.ensurePackagesAreInstalled(ctx, install); err != nil {
 		// if error installing, revert devbox.json
 		// This is not perfect because there may be more than 1 package being
 		// installed and we don't know which one failed. But it's better than
@@ -77,6 +82,8 @@ func (d *Devbox) Add(pkgs ...string) error {
 
 // Remove removes the `pkgs` from the config (i.e. devbox.json) and nix profile for this devbox project
 func (d *Devbox) Remove(pkgs ...string) error {
+	ctx, task := trace.NewTask(context.Background(), "devboxRemove")
+	defer task.End()
 
 	// First, save which packages are being uninstalled. Do this before we modify d.cfg.RawPackages below.
 	uninstalledPackages := lo.Intersect(d.cfg.RawPackages, pkgs)
@@ -99,11 +106,11 @@ func (d *Devbox) Remove(pkgs ...string) error {
 		return err
 	}
 
-	if err := d.removePackagesFromProfile(uninstalledPackages); err != nil {
+	if err := d.removePackagesFromProfile(ctx, uninstalledPackages); err != nil {
 		return err
 	}
 
-	if err := d.ensurePackagesAreInstalled(uninstall); err != nil {
+	if err := d.ensurePackagesAreInstalled(ctx, uninstall); err != nil {
 		return err
 	}
 
@@ -122,7 +129,9 @@ const (
 // ensurePackagesAreInstalled ensures that the nix profile has the packages specified
 // in the config (devbox.json). The `mode` is used for user messaging to explain
 // what operations are happening, because this function may take time to execute.
-func (d *Devbox) ensurePackagesAreInstalled(mode installMode) error {
+func (d *Devbox) ensurePackagesAreInstalled(ctx context.Context, mode installMode) error {
+	defer trace.StartRegion(ctx, "ensurePackages").End()
+
 	if err := d.generateShellFiles(); err != nil {
 		return err
 	}
@@ -131,7 +140,7 @@ func (d *Devbox) ensurePackagesAreInstalled(mode installMode) error {
 	}
 
 	if featureflag.Flakes.Enabled() {
-		if err := d.addPackagesToProfile(mode); err != nil {
+		if err := d.addPackagesToProfile(ctx, mode); err != nil {
 			return err
 		}
 
@@ -145,7 +154,7 @@ func (d *Devbox) ensurePackagesAreInstalled(mode installMode) error {
 		}
 
 		// We need to re-install the packages
-		if err := d.installNixProfile(); err != nil {
+		if err := d.installNixProfile(ctx); err != nil {
 			fmt.Fprintln(d.writer)
 			return errors.Wrap(err, "apply Nix derivation")
 		}
@@ -206,7 +215,9 @@ func (d *Devbox) printPackageUpdateMessage(
 
 // installNixProfile installs or uninstalls packages to or from this
 // devbox's Nix profile so that it matches what's in development.nix
-func (d *Devbox) installNixProfile() (err error) {
+func (d *Devbox) installNixProfile(ctx context.Context) (err error) {
+	defer trace.StartRegion(ctx, "installNixProfile").End()
+
 	profileDir, err := d.profilePath()
 	if err != nil {
 		return err
@@ -255,15 +266,17 @@ func (d *Devbox) profilePath() (string, error) {
 // addPackagesToProfile inspects the packages in devbox.json, checks which of them
 // are missing from the nix profile, and then installs each package individually into the
 // nix profile.
-func (d *Devbox) addPackagesToProfile(mode installMode) error {
+func (d *Devbox) addPackagesToProfile(ctx context.Context, mode installMode) error {
+	defer trace.StartRegion(ctx, "addNixProfilePkgs").End()
+
 	if featureflag.Flakes.Disabled() {
 		return nil
 	}
 	if mode == uninstall {
 		return nil
 	}
 
-	pkgs, err := d.pendingPackagesForInstallation()
+	pkgs, err := d.pendingPackagesForInstallation(ctx)
 	if err != nil {
 		return err
 	}
@@ -317,7 +330,9 @@ func (d *Devbox) addPackagesToProfile(mode installMode) error {
 	return nil
 }
 
-func (d *Devbox) removePackagesFromProfile(pkgs []string) error {
+func (d *Devbox) removePackagesFromProfile(ctx context.Context, pkgs []string) error {
+	defer trace.StartRegion(ctx, "removeNixProfilePkgs").End()
+
 	if !featureflag.Flakes.Enabled() {
 		return nil
 	}
@@ -367,7 +382,9 @@ func (d *Devbox) removePackagesFromProfile(pkgs []string) error {
 	return nil
 }
 
-func (d *Devbox) pendingPackagesForInstallation() ([]string, error) {
+func (d *Devbox) pendingPackagesForInstallation(ctx context.Context) ([]string, error) {
+	defer trace.StartRegion(ctx, "pendingPackages").End()
+
 	if featureflag.Flakes.Disabled() {
 		return nil, errors.New("Not implemented for legacy non-flakes devbox")
 	}