From 29d1c24611d84fbc5e006e3aac1e6cacba3008f9 Mon Sep 17 00:00:00 2001
From: Mike Landau <mikeland86@gmail.com>
Date: Mon, 18 Sep 2023 14:36:26 -0700
Subject: [PATCH] [perf] Profile possibly slow functions

---
 internal/boxcli/root.go      |  3 ++-
 internal/debug/time.go       | 48 ++++++++++++++++++++++++++++++------
 internal/impl/devbox.go      |  1 +
 internal/impl/packages.go    |  2 ++
 internal/shellgen/scripts.go |  1 +
 internal/wrapnix/wrapper.go  |  1 +
 6 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/internal/boxcli/root.go b/internal/boxcli/root.go
index 2138f6244c3..83520f579b7 100644
--- a/internal/boxcli/root.go
+++ b/internal/boxcli/root.go
@@ -106,6 +106,7 @@ func Execute(ctx context.Context, args []string) int {
 }
 
 func Main() {
+	timer := debug.Timer(strings.Join(os.Args, " "))
 	ctx := context.Background()
 	if strings.HasSuffix(os.Args[0], "ssh") ||
 		strings.HasSuffix(os.Args[0], "scp") {
@@ -126,7 +127,7 @@ func Main() {
 	code := Execute(ctx, os.Args[1:])
 	// Run out here instead of as a middleware so we can capture any time we spend
 	// in middlewares as well.
-	debug.PrintExecutionTime()
+	timer.End()
 	os.Exit(code)
 }
 
diff --git a/internal/debug/time.go b/internal/debug/time.go
index d2a9b284d65..e05fcc6bacd 100644
--- a/internal/debug/time.go
+++ b/internal/debug/time.go
@@ -6,22 +6,54 @@ package debug
 import (
 	"fmt"
 	"os"
+	"runtime"
 	"strconv"
 	"strings"
 	"time"
 )
 
+var timerEnabled, _ = strconv.ParseBool(os.Getenv(devboxPrintExecTime))
+
 const devboxPrintExecTime = "DEVBOX_PRINT_EXEC_TIME"
 
-var start = time.Now()
+var headerPrinted = false
+
+type timer struct {
+	name string
+	time time.Time
+}
+
+func Timer(name string) *timer {
+	if !timerEnabled {
+		return nil
+	}
+	return &timer{
+		name: name,
+		time: time.Now(),
+	}
+}
+
+func FunctionTimer() *timer {
+	if !timerEnabled {
+		return nil
+	}
+	pc := make([]uintptr, 15)
+	n := runtime.Callers(2, pc)
+	frames := runtime.CallersFrames(pc[:n])
+	frame, _ := frames.Next()
+	parts := strings.Split(frame.Function, ".")
+	return Timer(parts[len(parts)-1])
+}
 
-func PrintExecutionTime() {
-	if enabled, _ := strconv.ParseBool(os.Getenv(devboxPrintExecTime)); !enabled {
+func (t *timer) End() {
+	if t == nil {
 		return
 	}
-	fmt.Fprintf(
-		os.Stderr,
-		"\"%s\" took %s\n", strings.Join(os.Args, " "),
-		time.Since(start),
-	)
+	if !headerPrinted {
+		fmt.Fprintln(os.Stderr, "\nExec times over 1ms:")
+		headerPrinted = true
+	}
+	if time.Since(t.time) >= time.Millisecond {
+		fmt.Fprintf(os.Stderr, "\"%s\" took %s\n", t.name, time.Since(t.time))
+	}
 }
diff --git a/internal/impl/devbox.go b/internal/impl/devbox.go
index 94e91942095..4775fc0bf4b 100644
--- a/internal/impl/devbox.go
+++ b/internal/impl/devbox.go
@@ -923,6 +923,7 @@ var nixEnvCache map[string]string
 // Note that this is in-memory cache of the final environment, and not the same
 // as the nix print-dev-env cache which is stored in a file.
 func (d *Devbox) nixEnv(ctx context.Context) (map[string]string, error) {
+	defer debug.FunctionTimer().End()
 	if nixEnvCache != nil {
 		return nixEnvCache, nil
 	}
diff --git a/internal/impl/packages.go b/internal/impl/packages.go
index ba03cf54052..bdd738c8db1 100644
--- a/internal/impl/packages.go
+++ b/internal/impl/packages.go
@@ -203,6 +203,7 @@ const (
 // what operations are happening, because this function may take time to execute.
 func (d *Devbox) ensurePackagesAreInstalled(ctx context.Context, mode installMode) error {
 	defer trace.StartRegion(ctx, "ensurePackages").End()
+	defer debug.FunctionTimer().End()
 
 	if upToDate, err := d.lockfile.IsUpToDateAndInstalled(); err != nil || upToDate {
 		return err
@@ -266,6 +267,7 @@ func (d *Devbox) profilePath() (string, error) {
 // syncPackagesToProfile ensures that all packages in devbox.json exist in the nix profile,
 // and no more.
 func (d *Devbox) syncPackagesToProfile(ctx context.Context, mode installMode) error {
+	defer debug.FunctionTimer().End()
 	// TODO: we can probably merge these two operations to be faster and minimize chances of
 	// the devbox.json and nix profile falling out of sync.
 	if err := d.addPackagesToProfile(ctx, mode); err != nil {
diff --git a/internal/shellgen/scripts.go b/internal/shellgen/scripts.go
index 9078ccdcbe6..ea5b33c0430 100644
--- a/internal/shellgen/scripts.go
+++ b/internal/shellgen/scripts.go
@@ -32,6 +32,7 @@ type devboxer interface {
 // WriteScriptsToFiles writes scripts defined in devbox.json into files inside .devbox/gen/scripts.
 // Scripts (and hooks) are persisted so that we can easily call them from devbox run (inside or outside shell).
 func WriteScriptsToFiles(devbox devboxer) error {
+	defer debug.FunctionTimer().End()
 	err := os.MkdirAll(filepath.Join(devbox.ProjectDir(), scriptsDir), 0755) // Ensure directory exists.
 	if err != nil {
 		return errors.WithStack(err)
diff --git a/internal/wrapnix/wrapper.go b/internal/wrapnix/wrapper.go
index 7ec1c9a8999..d60f90710a7 100644
--- a/internal/wrapnix/wrapper.go
+++ b/internal/wrapnix/wrapper.go
@@ -38,6 +38,7 @@ var devboxSymlinkDir = xdg.CacheSubpath(filepath.Join("devbox", "bin", "current"
 
 // CreateWrappers creates wrappers for all the executables in nix paths
 func CreateWrappers(ctx context.Context, devbox devboxer) error {
+	defer debug.FunctionTimer().End()
 	shellEnvHash, err := devbox.ShellEnvHash(ctx)
 	if err != nil {
 		return err