perf(vercheck): switch devbox versions before startup for devbox_version auto

internal/boxcli/midcobra/projectversion.go

@@ -0,0 +1,90 @@
+// Copyright 2024 Jetify Inc. and contributors. All rights reserved.
+// Use of this source code is governed by the license in the LICENSE file.
+
+package midcobra
+
+import (
+	"io"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+
+	"go.jetify.com/devbox/internal/devconfig"
+	"go.jetify.com/devbox/internal/vercheck"
+)
+
+// ProjectVersion returns a middleware that switches to the devbox version
+// required by the project's devbox_version before the command runs, when its
+// policy is "auto". Opening the project enforces devbox_version too, but only
+// after the command's other setup (such as checking the nix install), so
+// switching from there pays for devbox startup twice.
+//
+// The middleware only handles the common case: a command with a --config flag
+// that opens a single project. Commands it skips, such as devbox global and
+// --all-projects, still switch versions when they open a project.
+func ProjectVersion() Middleware {
+	return &projectVersionMiddleware{}
+}
+
+type projectVersionMiddleware struct{}
+
+var _ Middleware = (*projectVersionMiddleware)(nil)
+
+func (m *projectVersionMiddleware) preRun(cmd *cobra.Command, args []string) {
+	configPath, ok := projectConfigPath(cmd, args)
+	if !ok {
+		return
+	}
+
+	var cfg *devconfig.Config
+	var err error
+	if configPath == "" {
+		cfg, err = devconfig.Find(".")
+	} else {
+		cfg, err = devconfig.Open(configPath)
+	}
+	if err != nil {
+		// Let the command report the error when it opens the project.
+		return
+	}
+	vercheck.SwitchProjectVersion(cfg.Root.AbsRootPath, cfg.Root.DevboxVersion)
+}
+
+func (m *projectVersionMiddleware) postRun(*cobra.Command, []string, error) {}
+
+// projectConfigPath returns the value of the --config flag of the subcommand
+// that args run. It returns false if the subcommand doesn't open a single
+// project from its --config flag.
+func projectConfigPath(root *cobra.Command, args []string) (path string, ok bool) {
+	subcmd, flagArgs, err := root.Find(args)
+	if err != nil {
+		return "", false
+	}
+	// devbox global commands set --config to the global project in a
+	// PersistentPreRunE, so the flag's value isn't known yet.
+	for c := subcmd; c != nil; c = c.Parent() {
+		if c.Name() == "global" && c.Parent() == root {
+			return "", false
+		}
+	}
+	configFlag := subcmd.Flag("config")
+	if configFlag == nil || configFlag.Value.Type() != "string" {
+		return "", false
+	}
+
+	// Parse only the flags we need into a separate flag set so that the
+	// command's own flags are still unset when cobra parses them.
+	flags := pflag.NewFlagSet(subcmd.Name(), pflag.ContinueOnError)
+	flags.SetOutput(io.Discard)
+	flags.ParseErrorsAllowlist.UnknownFlags = true
+	flags.StringP("config", configFlag.Shorthand, configFlag.DefValue, "")
+	allProjects := false
+	if subcmd.Flag("all-projects") != nil {
+		flags.BoolVar(&allProjects, "all-projects", false, "")
+	}
+	if err := flags.Parse(flagArgs); err != nil || allProjects {
+		return "", false
+	}
+	path, err = flags.GetString("config")
+	return path, err == nil
+}

internal/boxcli/midcobra/projectversion_test.go

@@ -0,0 +1,72 @@
+package midcobra
+
+import (
+	"testing"
+
+	"github.com/spf13/cobra"
+)
+
+func testRootCmd() *cobra.Command {
+	noop := func(*cobra.Command, []string) {}
+	withConfig := func(cmd *cobra.Command) *cobra.Command {
+		cmd.Flags().StringP("config", "c", "/default", "")
+		return cmd
+	}
+
+	root := &cobra.Command{Use: "devbox"}
+	root.PersistentFlags().BoolP("quiet", "q", false, "")
+
+	shell := withConfig(&cobra.Command{Use: "shell", Run: noop})
+	shell.Flags().Bool("pure", false, "")
+	shell.Flags().StringSlice("env", nil, "")
+	run := withConfig(&cobra.Command{Use: "run", Run: noop})
+	run.Flags().Bool("all-projects", false, "")
+	global := &cobra.Command{Use: "global"}
+	global.AddCommand(withConfig(&cobra.Command{Use: "shellenv", Run: noop}))
+	root.AddCommand(shell, run, global, &cobra.Command{Use: "version", Run: noop})
+	return root
+}
+
+func TestProjectConfigPath(t *testing.T) {
+	tests := []struct {
+		args     []string
+		wantPath string
+		wantOK   bool
+	}{
+		{args: []string{"shell"}, wantPath: "/default", wantOK: true},
+		{args: []string{"shell", "--config", "/a"}, wantPath: "/a", wantOK: true},
+		{args: []string{"shell", "--config=/a"}, wantPath: "/a", wantOK: true},
+		{args: []string{"shell", "-c", "/a"}, wantPath: "/a", wantOK: true},
+		{args: []string{"-q", "shell", "--pure", "--env", "A=1", "-c", "/a"}, wantPath: "/a", wantOK: true},
+		{args: []string{"shell", "--unknown", "value", "-c", "/a"}, wantPath: "/a", wantOK: true},
+		{args: []string{"run", "-c", "/a", "--", "script", "-c", "/b"}, wantPath: "/a", wantOK: true},
+		{args: []string{"run", "--all-projects", "--", "script"}, wantOK: false},
+		{args: []string{"global", "shellenv"}, wantOK: false},
+		{args: []string{"version"}, wantOK: false},
+		{args: []string{"nope"}, wantOK: false},
+	}
+	for _, tc := range tests {
+		gotPath, gotOK := projectConfigPath(testRootCmd(), tc.args)
+		if gotPath != tc.wantPath || gotOK != tc.wantOK {
+			t.Errorf("projectConfigPath(%q) = %q, %t; want %q, %t",
+				tc.args, gotPath, gotOK, tc.wantPath, tc.wantOK)
+		}
+	}
+}
+
+// TestProjectConfigPathLeavesFlagsUnset checks that the middleware doesn't
+// parse the command's own flags, which cobra parses later.
+func TestProjectConfigPathLeavesFlagsUnset(t *testing.T) {
+	root := testRootCmd()
+	projectConfigPath(root, []string{"shell", "-c", "/a", "--env", "A=1"})
+
+	shell, _, err := root.Find([]string{"shell"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, name := range []string{"config", "env"} {
+		if shell.Flag(name).Changed {
+			t.Errorf("flag --%s was changed", name)
+		}
+	}
+}

internal/boxcli/root.go

@@ -102,6 +102,9 @@ func Execute(ctx context.Context, args []string) int {
 	defer debug.Recover()
 	rootCmd := RootCmd()
 	exe := midcobra.New(rootCmd)
+	// Runs first because it might re-run the command with another devbox
+	// version, which makes the work done by any earlier middleware wasted.
+	exe.AddMiddleware(midcobra.ProjectVersion())
 	exe.AddMiddleware(traceMiddleware)
 	exe.AddMiddleware(midcobra.Telemetry())
 	exe.AddMiddleware(debugMiddleware)

internal/vercheck/project.go

@@ -8,6 +8,8 @@ import (
 	"io"
 	"log/slog"
 	"os"
+	"path/filepath"
+	"runtime"
 	"strings"
 	"syscall"
 
@@ -43,13 +45,9 @@ func CheckProjectVersion(w io.Writer, configPath string, required *configfile.De
 		return nil
 	}
 
-	policy := required.Policy()
-	if env := os.Getenv(envir.DevboxVersionPolicy); env != "" {
-		p, err := configfile.ParseVersionPolicy(env)
-		if err != nil {
-			return usererr.New("Invalid %s: %v", envir.DevboxVersionPolicy, err)
-		}
-		policy = p
+	policy, err := projectPolicy(required)
+	if err != nil {
+		return err
 	}
 	if policy == configfile.VersionPolicyOff {
 		return nil
@@ -59,12 +57,7 @@ func CheckProjectVersion(w io.Writer, configPath string, required *configfile.De
 	if err != nil {
 		return err
 	}
-	current, err := semver.NewVersion(currentDevboxVersion)
-	if err != nil {
-		slog.Debug("skipping devbox_version check: can't parse running version", "version", currentDevboxVersion, "err", err)
-		return nil
-	}
-	if constraint.Check(current) {
+	if satisfies(constraint) {
 		return nil
 	}
 
@@ -79,9 +72,69 @@ func CheckProjectVersion(w io.Writer, configPath string, required *configfile.De
 	}
 }
 
+// SwitchProjectVersion is a fast path for the "auto" policy that runs before
+// the CLI does any other work, so that switching versions doesn't pay for
+// devbox startup twice. It re-runs the current command with the required
+// version (and doesn't return) only when CheckProjectVersion would. In every
+// other case it returns without printing anything, leaving any warnings or
+// errors to CheckProjectVersion when the project is opened.
+func SwitchProjectVersion(configPath string, required *configfile.DevboxVersion) {
+	if required == nil || isDevBuild {
+		return
+	}
+	if policy, err := projectPolicy(required); err != nil || policy != configfile.VersionPolicyAuto {
+		return
+	}
+	constraint, err := required.Constraint()
+	if err != nil || satisfies(constraint) {
+		return
+	}
+	// CheckProjectVersion warns about this case, so leave it to that.
+	if userChosenVersion() != "" {
+		return
+	}
+	// CheckProjectVersion reports any error when the project is opened.
+	_ = switchVersion(io.Discard, configPath, required, "")
+}
+
+// projectPolicy returns the policy for required, which the
+// DEVBOX_VERSION_POLICY environment variable overrides.
+func projectPolicy(required *configfile.DevboxVersion) (configfile.VersionPolicy, error) {
+	env := os.Getenv(envir.DevboxVersionPolicy)
+	if env == "" {
+		return required.Policy(), nil
+	}
+	p, err := configfile.ParseVersionPolicy(env)
+	if err != nil {
+		return "", usererr.New("Invalid %s: %v", envir.DevboxVersionPolicy, err)
+	}
+	return p, nil
+}
+
+// satisfies reports whether the running devbox version satisfies constraint.
+// It's true when the running version can't be parsed, which skips the check.
+func satisfies(constraint *semver.Constraints) bool {
+	current, err := semver.NewVersion(currentDevboxVersion)
+	if err != nil {
+		slog.Debug("skipping devbox_version check: can't parse running version", "version", currentDevboxVersion, "err", err)
+		return true
+	}
+	return constraint.Check(current)
+}
+
+// userChosenVersion returns DEVBOX_USE_VERSION if the user set it, or an
+// empty string if it's unset or the auto policy set it.
+func userChosenVersion() string {
+	if useVersion := os.Getenv(envir.DevboxUseVersion); useVersion != os.Getenv(autoVersionEnvName) {
+		return useVersion
+	}
+	return ""
+}
+
 // switchVersion re-runs the current command with the exact version required
-// by devbox_version. It uses the devbox launcher, which downloads the version
-// if needed and runs the version named by DEVBOX_USE_VERSION.
+// by devbox_version. It runs that version's binary directly if the launcher
+// already downloaded it. Otherwise it uses the launcher, which downloads the
+// version and runs the version named by DEVBOX_USE_VERSION.
 func switchVersion(w io.Writer, configPath string, required *configfile.DevboxVersion, msg string) error {
 	target, ok := required.ExactVersion()
 	if !ok {
@@ -94,15 +147,14 @@ func switchVersion(w io.Writer, configPath string, required *configfile.DevboxVe
 	}
 
 	// Respect a version the user chose explicitly.
-	autoVersion := os.Getenv(autoVersionEnvName)
-	if useVersion := os.Getenv(envir.DevboxUseVersion); useVersion != "" && useVersion != autoVersion {
+	if useVersion := userChosenVersion(); useVersion != "" {
 		return warnOnce(w, configPath, fmt.Sprintf(
 			"%s\nNot switching versions automatically because %s=%s is set.",
 			msg, envir.DevboxUseVersion, useVersion,
 		))
 	}
 
-	if autoVersion == target {
+	if os.Getenv(autoVersionEnvName) == target {
 		return usererr.New(
 			"%s\nDevbox tried to switch to version %s automatically, but the launcher ran version %s instead.",
 			msg, target, currentDevboxVersion,
@@ -118,18 +170,44 @@ func switchVersion(w io.Writer, configPath string, required *configfile.DevboxVe
 		)
 	}
 
-	slog.Debug("switching devbox version for devbox_version", "from", currentDevboxVersion, "to", target, "launcher", launcher)
+	// Running the downloaded binary directly skips the launcher's startup
+	// cost. It still gets the launcher's environment variables because this
+	// process inherited them.
+	bin := installedBinary(target)
+	if bin == "" {
+		bin = launcher
+	}
+	slog.Debug("switching devbox version for devbox_version", "from", currentDevboxVersion, "to", target, "bin", bin)
 	env := withEnv(os.Environ(), map[string]string{
 		envir.DevboxUseVersion: target,
 		autoVersionEnvName:     target,
 	})
-	args := append([]string{launcher}, os.Args[1:]...)
-	if err := execFunc(launcher, args, env); err != nil {
-		return usererr.WithUserMessage(err, "Failed to switch to devbox version %s using the launcher at %s.", target, launcher)
+	args := append([]string{bin}, os.Args[1:]...)
+	if err := execFunc(bin, args, env); err != nil {
+		return usererr.WithUserMessage(err, "Failed to switch to devbox version %s using %s.", target, bin)
 	}
 	return nil
 }
 
+// installedBinary returns the path to the devbox binary for version if the
+// launcher already downloaded it, or an empty string otherwise.
+//
+// Note: keep this in sync with INSTALL_DIR and install_flow in launch.sh,
+// which use $HOME/.cache even when XDG_CACHE_HOME is set.
+func installedBinary(version string) string {
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return ""
+	}
+	platform := runtime.GOOS + "_" + runtime.GOARCH
+	bin := filepath.Join(home, ".cache", "devbox", "bin", version+"_"+platform, "devbox")
+	fi, err := os.Stat(bin)
+	if err != nil || !fi.Mode().IsRegular() || fi.Mode().Perm()&0o111 == 0 {
+		return ""
+	}
+	return bin
+}
+
 // warnOnce prints msg as a warning unless it was already printed for the
 // project at configPath by this process or a parent devbox process.
 func warnOnce(w io.Writer, configPath, msg string) error {