perf(devbox): skip the nix recompute when only env, init_hook or scripts change

internal/devbox/devbox.go

@@ -208,12 +208,29 @@ func (d *Devbox) Config() *devconfig.Config {
 	return d.cfg
 }
 
+// ConfigHash hashes every input that goes into the devbox environment: the
+// whole config, the packages and the plugin configs.
 func (d *Devbox) ConfigHash() (string, error) {
 	h, err := d.cfg.Hash()
 	if err != nil {
 		return "", err
 	}
+	return d.hashWithPackagesAndPlugins(h)
+}
 
+// NixConfigHash is like ConfigHash, but leaves out config fields that are only
+// read after the Nix environment is built (env, init_hook, scripts, etc.). If
+// it hasn't changed, the generated flake, the print-dev-env output and the
+// nix profile are still valid. See configfile.ConfigFile.NixHash.
+func (d *Devbox) NixConfigHash() (string, error) {
+	h, err := d.cfg.NixHash()
+	if err != nil {
+		return "", err
+	}
+	return d.hashWithPackagesAndPlugins(h)
+}
+
+func (d *Devbox) hashWithPackagesAndPlugins(h string) (string, error) {
 	buf := bytes.Buffer{}
 	buf.WriteString(h)
 	for _, pkg := range d.AllPackages() {
@@ -888,11 +905,19 @@ func (d *Devbox) ensureStateIsUpToDateAndComputeEnv(
 ) (map[string]string, error) {
 	defer debug.FunctionTimer().End()
 
-	upToDate, err := d.lockfile.IsUpToDateAndInstalled(isFishShell())
+	status, err := d.lockfile.StateStatus(isFishShell())
 	if err != nil {
 		return nil, err
 	}
-	if !upToDate {
+	if envOpts.SkipRecompute && status == lock.StateNonNixStale {
+		// Refreshing the non-Nix state doesn't run Nix and is cheap, so do it
+		// even when the caller asked to skip the (expensive) recompute.
+		if err := d.refreshNonNixState(ctx); err != nil {
+			return nil, err
+		}
+		status = lock.StateUpToDate
+	}
+	if status != lock.StateUpToDate {
 		if envOpts.Hooks.OnStaleState != nil {
 			envOpts.Hooks.OnStaleState()
 		}

internal/devbox/packages.go

@@ -309,10 +309,11 @@ func (d *Devbox) ensureStateIsUpToDate(ctx context.Context, mode installMode) er
 	defer trace.StartRegion(ctx, "devboxEnsureStateIsUpToDate").End()
 	defer debug.FunctionTimer().End()
 
-	upToDate, err := d.lockfile.IsUpToDateAndInstalled(isFishShell())
+	status, err := d.lockfile.StateStatus(isFishShell())
 	if err != nil {
 		return err
 	}
+	upToDate := status == lock.StateUpToDate
 
 	// if mode is install or uninstall, then we need to compute some state
 	// like updating the flake or installing packages locally, so must continue
@@ -322,6 +323,15 @@ func (d *Devbox) ensureStateIsUpToDate(ctx context.Context, mode installMode) er
 		if upToDate {
 			return nil
 		}
+		// Only env, init_hook, scripts, etc. changed. The Nix environment
+		// is still valid, so skip installing packages and recomputing it.
+		if status == lock.StateNonNixStale {
+			if err := d.refreshNonNixState(ctx); err != nil {
+				return err
+			}
+			d.warnIfEnvIsStale(ctx)
+			return nil
+		}
 		ux.Finfof(d.stderr, "Ensuring packages are installed.\n")
 	}
 
@@ -348,19 +358,64 @@ func (d *Devbox) ensureStateIsUpToDate(ctx context.Context, mode installMode) er
 		}
 	}
 
+	if !upToDate {
+		d.warnIfEnvIsStale(ctx)
+	}
+
+	return d.updateLockfile(mode, recomputeState)
+}
+
+// warnIfEnvIsStale warns that the current shell's environment is out of date.
+// It should be called after the project state changed.
+func (d *Devbox) warnIfEnvIsStale(ctx context.Context) {
 	// If we're in a devbox shell (global or project), then the environment might
 	// be out of date after the user installs something. If have direnv active
 	// it should reload automatically so we don't need to refresh.
-	if d.IsEnvEnabled() && !upToDate && !d.IsDirenvActive() {
+	if d.IsEnvEnabled() && !d.IsDirenvActive() {
 		ux.FHidableWarning(
 			ctx,
 			d.stderr,
 			StateOutOfDateMessage,
 			d.RefreshAliasOrCommand(),
 		)
 	}
+}
 
-	return d.updateLockfile(mode, recomputeState)
+// refreshNonNixState brings the project state up to date when only inputs
+// that don't affect the Nix environment changed (lock.StateNonNixStale). The
+// generated flake, the print-dev-env cache and the nix profile are still
+// valid, so this skips all Nix work. It rewrites the scripts and init hook
+// files and saves the new state hashes. Env vars, env_from and aliases need no
+// work here because they are read from the config every time the environment
+// is computed.
+func (d *Devbox) refreshNonNixState(ctx context.Context) error {
+	defer trace.StartRegion(ctx, "devboxRefreshNonNixState").End()
+	defer debug.FunctionTimer().End()
+	slog.Debug("only non-nix inputs changed, skipping nix environment recompute")
+
+	if err := shellgen.WriteScriptsToFiles(d); err != nil {
+		return err
+	}
+	return d.saveStateHashFile()
+}
+
+// saveStateHashFile records the current inputs in .devbox/state.json so that
+// later commands know the state is up to date.
+func (d *Devbox) saveStateHashFile() error {
+	configHash, err := d.ConfigHash()
+	if err != nil {
+		return err
+	}
+	nixConfigHash, err := d.NixConfigHash()
+	if err != nil {
+		return err
+	}
+	return lock.UpdateAndSaveStateHashFile(lock.UpdateStateHashFileArgs{
+		ProjectDir:    d.projectDir,
+		ConfigHash:    configHash,
+		NixConfigHash: nixConfigHash,
+		IsFish:        isFishShell(),
+	})
 }
 
 // updateLockfile will ensure devbox.lock is up to date with the current state of the project.update
@@ -394,15 +449,7 @@ func (d *Devbox) updateLockfile(mode installMode, recomputeState bool) error {
 	// If not, we leave the local.lock in a stale state, so that state is recomputed
 	// on the next ensureStateIsUpToDate call with mode=ensure.
 	if recomputeState {
-		configHash, err := d.ConfigHash()
-		if err != nil {
-			return err
-		}
-		return lock.UpdateAndSaveStateHashFile(lock.UpdateStateHashFileArgs{
-			ProjectDir: d.projectDir,
-			ConfigHash: configHash,
-			IsFish:     isFishShell(),
-		})
+		return d.saveStateHashFile()
 	}
 	return nil
 }

internal/devbox/statehash_test.go

@@ -0,0 +1,164 @@
+// Copyright 2024 Jetify Inc. and contributors. All rights reserved.
+// Use of this source code is governed by the license in the LICENSE file.
+
+package devbox
+
+import (
+	"context"
+	"errors"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"go.jetify.com/devbox/internal/devbox/devopt"
+	"go.jetify.com/devbox/internal/lock"
+	"go.jetify.com/devbox/internal/nix"
+	"go.jetify.com/devbox/internal/shellgen"
+)
+
+// cacheOnlyNix serves print-dev-env from the cache and fails if devbox tries
+// to run `nix print-dev-env`.
+type cacheOnlyNix struct {
+	uncachedCalls int
+}
+
+func (n *cacheOnlyNix) PrintDevEnv(_ context.Context, args *nix.PrintDevEnvArgs) (*nix.PrintDevEnvOut, error) {
+	if !args.UsePrintDevEnvCache {
+		n.uncachedCalls++
+		return nil, errors.New("unexpected uncached nix print-dev-env call")
+	}
+	return &nix.PrintDevEnvOut{Variables: map[string]nix.Variable{
+		"PATH": {Type: "exported", Value: "/nix/store/fake/bin"},
+	}}, nil
+}
+
+const stateTestConfig = `{
+  "packages": ["github:NixOS/nixpkgs/nixpkgs-unstable#hello"],
+  "env": {"MY_VAR": "old"},
+  "shell": {
+    "init_hook": ["echo old-hook"],
+    "scripts": {"old-script": "echo old"}
+  }
+}`
+
+func openStateTestProject(t *testing.T, dir, config string) *Devbox {
+	t.Helper()
+	require.NoError(t, os.WriteFile(filepath.Join(dir, "devbox.json"), []byte(config), 0o644))
+	d, err := Open(&devopt.Opts{Dir: dir, Stderr: os.Stderr})
+	require.NoError(t, err)
+	return d
+}
+
+func stateStatus(t *testing.T, d *Devbox) lock.StateStatus {
+	t.Helper()
+	status, err := d.lockfile.StateStatus(isFishShell())
+	require.NoError(t, err)
+	return status
+}
+
+func TestEnsureStateNonNixChangeSkipsNix(t *testing.T) {
+	dir := t.TempDir()
+	box := openStateTestProject(t, dir, stateTestConfig)
+
+	// Simulate a previous full recompute.
+	require.NoError(t, os.MkdirAll(filepath.Join(dir, ".devbox"), 0o755))
+	require.NoError(t, box.saveStateHashFile())
+	require.Equal(t, lock.StateUpToDate, stateStatus(t, box))
+
+	// Edit env, the init hook and the scripts.
+	box = openStateTestProject(t, dir, `{
+  "packages": ["github:NixOS/nixpkgs/nixpkgs-unstable#hello"],
+  "env": {"MY_VAR": "new"},
+  "shell": {
+    "init_hook": ["echo new-hook"],
+    "scripts": {"new-script": "echo new"}
+  }
+}`)
+	fakeNix := &cacheOnlyNix{}
+	box.nix = fakeNix
+	require.Equal(t, lock.StateNonNixStale, stateStatus(t, box))
+
+	ctx := t.Context()
+	require.NoError(t, box.ensureStateIsUpToDate(ctx, ensure))
+	assert.Equal(t, 0, fakeNix.uncachedCalls, "ensureStateIsUpToDate should not run nix")
+	assert.Equal(t, lock.StateUpToDate, stateStatus(t, box))
+
+	hooks, err := os.ReadFile(shellgen.ScriptPath(dir, shellgen.HooksFilename))
+	require.NoError(t, err)
+	assert.Contains(t, string(hooks), "new-hook")
+	assert.FileExists(t, shellgen.ScriptPath(dir, "new-script"))
+	assert.NoFileExists(t, shellgen.ScriptPath(dir, "old-script"))
+
+	env, err := box.computeEnv(ctx, true /*usePrintDevEnvCache*/, devopt.EnvOptions{})
+	require.NoError(t, err)
+	assert.Equal(t, "new", env["MY_VAR"])
+	assert.Equal(t, 0, fakeNix.uncachedCalls)
+}
+
+func TestSkipRecomputeRefreshesNonNixState(t *testing.T) {
+	dir := t.TempDir()
+	box := openStateTestProject(t, dir, stateTestConfig)
+	require.NoError(t, os.MkdirAll(filepath.Join(dir, ".devbox"), 0o755))
+	require.NoError(t, box.saveStateHashFile())
+
+	box = openStateTestProject(t, dir, `{
+  "packages": ["github:NixOS/nixpkgs/nixpkgs-unstable#hello"],
+  "env": {"MY_VAR": "new"},
+  "shell": {"init_hook": ["echo new-hook"]}
+}`)
+	fakeNix := &cacheOnlyNix{}
+	box.nix = fakeNix
+
+	staleCalls := 0
+	env, err := box.ensureStateIsUpToDateAndComputeEnv(t.Context(), devopt.EnvOptions{
+		SkipRecompute: true,
+		Hooks:         devopt.LifecycleHooks{OnStaleState: func() { staleCalls++ }},
+	})
+	require.NoError(t, err)
+	assert.Equal(t, 0, staleCalls, "non-nix changes should be refreshed, not reported as stale")
+	assert.Equal(t, 0, fakeNix.uncachedCalls)
+	assert.Equal(t, "new", env["MY_VAR"])
+	assert.Equal(t, lock.StateUpToDate, stateStatus(t, box))
+
+	hooks, err := os.ReadFile(shellgen.ScriptPath(dir, shellgen.HooksFilename))
+	require.NoError(t, err)
+	assert.Contains(t, string(hooks), "new-hook")
+}
+
+func TestStateStatusNixChanges(t *testing.T) {
+	tests := []struct {
+		name   string
+		config string
+	}{
+		{"packages", `{
+  "packages": ["github:NixOS/nixpkgs/nixpkgs-unstable#cowsay"],
+  "env": {"MY_VAR": "old"},
+  "shell": {"init_hook": ["echo old-hook"], "scripts": {"old-script": "echo old"}}
+}`},
+		{"packages and env", `{
+  "packages": ["github:NixOS/nixpkgs/nixpkgs-unstable#hello", "github:NixOS/nixpkgs/nixpkgs-unstable#cowsay"],
+  "env": {"MY_VAR": "new"},
+  "shell": {"init_hook": ["echo old-hook"], "scripts": {"old-script": "echo old"}}
+}`},
+		{"include", `{
+  "packages": ["github:NixOS/nixpkgs/nixpkgs-unstable#hello"],
+  "env": {"MY_VAR": "old"},
+  "shell": {"init_hook": ["echo old-hook"], "scripts": {"old-script": "echo old"}},
+  "include": ["plugin:nginx"]
+}`},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			dir := t.TempDir()
+			box := openStateTestProject(t, dir, stateTestConfig)
+			require.NoError(t, os.MkdirAll(filepath.Join(dir, ".devbox"), 0o755))
+			require.NoError(t, box.saveStateHashFile())
+
+			box = openStateTestProject(t, dir, tt.config)
+			assert.Equal(t, lock.StateStale, stateStatus(t, box))
+		})
+	}
+}

internal/devconfig/config.go

@@ -423,6 +423,26 @@ func (c *Config) Hash() (string, error) {
 	return cachehash.Bytes(data), nil
 }
 
+// NixHash is like Hash but ignores the root config's fields that don't affect
+// the Nix environment (see configfile.ConfigFile.NixHash). Included configs
+// (plugins) are hashed in full because they can add packages and create files.
+func (c *Config) NixHash() (string, error) {
+	data := []byte{}
+	for _, i := range c.included {
+		hash, err := i.Hash()
+		if err != nil {
+			return "", err
+		}
+		data = append(data, hash...)
+	}
+	hash, err := c.Root.NixHash()
+	if err != nil {
+		return "", err
+	}
+	data = append(data, hash...)
+	return cachehash.Bytes(data), nil
+}
+
 func (c *Config) IsJetifyCloudEnvFrom() bool {
 	for _, i := range c.included {
 		if i.IsJetifyCloudEnvFrom() {

internal/devconfig/config_test.go

@@ -713,7 +713,8 @@ type testLockProject struct {
 	dir string
 }
 
-func (p *testLockProject) ConfigHash() (string, error) { return "", nil }
+func (p *testLockProject) ConfigHash() (string, error)    { return "", nil }
+func (p *testLockProject) NixConfigHash() (string, error) { return "", nil }
 
 func (p *testLockProject) Stdenv() flake.Ref                                        { return flake.Ref{} }
 func (p *testLockProject) AllPackageNamesIncludingRemovedTriggerPackages() []string { return nil }