[perf] Add flag to fix missing store paths (#2102)

mikeland73 · web-flow · commit 3bb13c162e10 · 2024-06-05T18:48:59.000-07:00
diff --git a/internal/boxcli/install.go b/internal/boxcli/install.go
@@ -11,10 +11,17 @@ import (
 
 	"go.jetpack.io/devbox/internal/devbox"
 	"go.jetpack.io/devbox/internal/devbox/devopt"
+	"go.jetpack.io/devbox/internal/devpkg"
+	"go.jetpack.io/devbox/internal/ux"
 )
 
+type installCmdFlags struct {
+	runCmdFlags
+	tidyLockfile bool
+}
+
 func installCmd() *cobra.Command {
-	flags := runCmdFlags{}
+	flags := installCmdFlags{}
 	command := &cobra.Command{
 		Use:     "install",
 		Short:   "Install all packages mentioned in devbox.json",
@@ -26,11 +33,16 @@ func installCmd() *cobra.Command {
 	}
 
 	flags.config.register(command)
+	command.Flags().BoolVar(
+		&flags.tidyLockfile, "tidy-lockfile", false,
+		"Fix missing store paths in the devbox.lock file.",
+		// Could potentially do more in the future.
+	)
 
 	return command
 }
 
-func installCmdFunc(cmd *cobra.Command, flags runCmdFlags) error {
+func installCmdFunc(cmd *cobra.Command, flags installCmdFlags) error {
 	// Check the directory exists.
 	box, err := devbox.Open(&devopt.Opts{
 		Dir:         flags.config.path,
@@ -40,9 +52,18 @@ func installCmdFunc(cmd *cobra.Command, flags runCmdFlags) error {
 	if err != nil {
 		return errors.WithStack(err)
 	}
-	if err = box.Install(cmd.Context()); err != nil {
+	ctx := cmd.Context()
+	if flags.tidyLockfile {
+		ctx = ux.HideMessage(ctx, devpkg.MissingStorePathsWarning)
+	}
+	if err = box.Install(ctx); err != nil {
 		return errors.WithStack(err)
 	}
+	if flags.tidyLockfile {
+		if err = box.FixMissingStorePaths(ctx); err != nil {
+			return errors.WithStack(err)
+		}
+	}
 	fmt.Fprintln(cmd.ErrOrStderr(), "Finished installing packages.")
 	return nil
 }
diff --git a/internal/boxcli/pull.go b/internal/boxcli/pull.go
@@ -106,7 +106,9 @@ func pullCmdFunc(cmd *cobra.Command, url string, flags *pullCmdFlags) error {
 
 	return installCmdFunc(
 		cmd,
-		runCmdFlags{config: configFlags{pathFlag: pathFlag{path: flags.config.path}}},
+		installCmdFlags{
+			runCmdFlags: runCmdFlags{config: configFlags{pathFlag: pathFlag{path: flags.config.path}}},
+		},
 	)
 }
 
diff --git a/internal/devbox/devbox.go b/internal/devbox/devbox.go
@@ -856,6 +856,7 @@ func (d *Devbox) configureProcessCompose(ctx context.Context, processComposeOpts
 // some additional processing. The computeEnv environment won't necessarily
 // represent the final "devbox run" or "devbox shell" environments.
 func (d *Devbox) computeEnv(ctx context.Context, usePrintDevEnvCache bool) (map[string]string, error) {
+	defer debug.FunctionTimer().End()
 	defer trace.StartRegion(ctx, "devboxComputeEnv").End()
 
 	// Append variables from current env if --pure is not passed
diff --git a/internal/devbox/nixprofile.go b/internal/devbox/nixprofile.go
@@ -16,6 +16,7 @@ import (
 //
 // It also removes any packages from the nix profile that are no longer in the buildInputs.
 func (d *Devbox) syncNixProfileFromFlake(ctx context.Context) error {
+	defer debug.FunctionTimer().End()
 	// Get the computed Devbox environment from the generated flake
 	env, err := d.computeEnv(ctx, false /*usePrintDevEnvCache*/)
 	if err != nil {
diff --git a/internal/devbox/packages.go b/internal/devbox/packages.go
@@ -335,6 +335,7 @@ func (d *Devbox) updateLockfile(recomputeState bool) error {
 // - the generated flake
 // - the nix-profile
 func (d *Devbox) recomputeState(ctx context.Context) error {
+	defer debug.FunctionTimer().End()
 	if err := shellgen.GenerateForPrintEnv(ctx, d); err != nil {
 		return err
 	}
@@ -643,3 +644,49 @@ func (d *Devbox) moveAllowInsecureFromLockfile(writer io.Writer, lockfile *lock.
 
 	return nil
 }
+
+func (d *Devbox) FixMissingStorePaths(ctx context.Context) error {
+	packages := d.InstallablePackages()
+	for _, pkg := range packages {
+		if pkg.IsRunX() {
+			continue
+		}
+		existingStorePaths, err := pkg.GetResolvedStorePaths()
+		if err != nil {
+			return err
+		}
+
+		if len(existingStorePaths) > 0 {
+			continue
+		}
+
+		installables, err := pkg.Installables()
+		if err != nil {
+			return err
+		}
+
+		outputs := []lock.Output{}
+		for _, installable := range installables {
+			storePaths, err := nix.StorePathsFromInstallable(ctx, installable, pkg.HasAllowInsecure())
+			if err != nil {
+				return err
+			}
+			if len(storePaths) == 0 {
+				return fmt.Errorf("no store paths found for package %s", pkg.Raw)
+			}
+			for _, storePath := range storePaths {
+				parts := nix.NewStorePathParts(storePath)
+				outputs = append(outputs, lock.Output{
+					Path: storePath,
+					Name: parts.Output,
+					// Ugh, not sure this is true, but it's more true than not.
+					Default: true,
+				})
+			}
+		}
+		if err = d.lockfile.SetOutputsForPackage(pkg.Raw, outputs); err != nil {
+			return err
+		}
+	}
+	return d.lockfile.Save()
+}
diff --git a/internal/devbox/update.go b/internal/devbox/update.go
@@ -75,6 +75,12 @@ func (d *Devbox) Update(ctx context.Context, opts devopt.UpdateOpts) error {
 	// It will return an error if .devbox/gen/flake is missing
 	// TODO: Remove this if it's not needed.
 	_ = nix.FlakeUpdate(shellgen.FlakePath(d))
+
+	// fix any missing store paths.
+	if err = d.FixMissingStorePaths(ctx); err != nil {
+		return errors.WithStack(err)
+	}
+
 	return plugin.Update()
 }
 
diff --git a/internal/devpkg/package.go b/internal/devpkg/package.go
@@ -736,19 +736,19 @@ func (p *Package) GetResolvedStorePaths() ([]string, error) {
 	return storePaths, nil
 }
 
+const MissingStorePathsWarning = "Outputs for %s are not in lockfile. To fix this issue and improve performance, please run " +
+	"`devbox install --tidy-lockfile`\n"
+
 func (p *Package) GetStorePaths(ctx context.Context, w io.Writer) ([]string, error) {
 	storePathsForPackage, err := p.GetResolvedStorePaths()
 	if err != nil || len(storePathsForPackage) > 0 {
 		return storePathsForPackage, err
 	}
 
-	// No fast path, we need to query nix.
-	// TODO we should give people the option to add paths to lockfile.
-	ux.Fwarning(
-		w,
-		"Outputs for %s are not in lockfile. Fetching store paths from nix, this may take a while\n",
-		p.Raw,
-	)
+	if p.IsDevboxPackage {
+		// No fast path, we need to query nix.
+		ux.FHidableWarning(ctx, w, MissingStorePathsWarning, p.Raw)
+	}
 
 	installables, err := p.Installables()
 	if err != nil {
diff --git a/internal/lock/lockfile.go b/internal/lock/lockfile.go
@@ -14,6 +14,7 @@ import (
 	"github.com/samber/lo"
 	"go.jetpack.io/devbox/internal/cachehash"
 	"go.jetpack.io/devbox/internal/devpkg/pkgtype"
+	"go.jetpack.io/devbox/internal/nix"
 	"go.jetpack.io/devbox/internal/searcher"
 	"go.jetpack.io/pkg/runx/impl/types"
 
@@ -199,6 +200,21 @@ func (f *File) IsUpToDateAndInstalled(isFish bool) (bool, error) {
 	})
 }
 
+func (f *File) SetOutputsForPackage(pkg string, outputs []Output) error {
+	p, err := f.Resolve(pkg)
+	if err != nil {
+		return err
+	}
+	if p.Systems == nil {
+		p.Systems = map[string]*SystemInfo{}
+	}
+	if p.Systems[nix.System()] == nil {
+		p.Systems[nix.System()] = &SystemInfo{}
+	}
+	p.Systems[nix.System()].Outputs = outputs
+	return f.Save()
+}
+
 func (f *File) isDirty() (bool, error) {
 	currentHash, err := cachehash.JSON(f)
 	if err != nil {
diff --git a/internal/nix/nixprofile/profile.go b/internal/nix/nixprofile/profile.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/pkg/errors"
 	"github.com/samber/lo"
+	"go.jetpack.io/devbox/internal/debug"
 	"go.jetpack.io/devbox/internal/devpkg"
 	"go.jetpack.io/devbox/internal/lock"
 	"go.jetpack.io/devbox/internal/nix"
@@ -23,6 +24,7 @@ func ProfileListItems(
 	writer io.Writer,
 	profileDir string,
 ) ([]*NixProfileListItem, error) {
+	defer debug.FunctionTimer().End()
 	output, err := nix.ProfileList(writer, profileDir, true /*useJSON*/)
 	if err != nil {
 		// fallback to legacy profile list
diff --git a/internal/nix/profiles.go b/internal/nix/profiles.go
@@ -38,6 +38,7 @@ type ProfileInstallArgs struct {
 }
 
 func ProfileInstall(ctx context.Context, args *ProfileInstallArgs) error {
+	defer debug.FunctionTimer().End()
 	if !IsInsecureAllowed() && PackageIsInsecure(args.Installable) {
 		knownVulnerabilities := PackageKnownVulnerabilities(args.Installable)
 		errString := fmt.Sprintf("Package %s is insecure. \n\n", args.Installable)
@@ -78,6 +79,7 @@ func ProfileInstall(ctx context.Context, args *ProfileInstallArgs) error {
 // ProfileRemove removes packages from a profile.
 // WARNING, don't use indexes, they are not supported by nix 2.20+
 func ProfileRemove(profilePath string, packageNames ...string) error {
+	defer debug.FunctionTimer().End()
 	cmd := command(
 		append([]string{
 			"profile", "remove",
diff --git a/internal/nix/storepath.go b/internal/nix/storepath.go
@@ -13,23 +13,23 @@ type StorePathParts struct {
 	Hash    string
 	Name    string
 	Version string
+	Output  string
 }
 
 // NewStorePathParts splits a Nix store path into its hash, name and version
 // components in the same way that Nix does.
 //
 // See https://nixos.org/manual/nix/stable/language/builtins.html#builtins-parseDrvName
-//
-// TODO: store paths can also have `-{output}` suffixes, which need to be handled below.
 func NewStorePathParts(path string) StorePathParts {
 	path = strings.TrimPrefix(path, "/nix/store/")
-	// path is now <hash>-<name>-<version
+	// path is now <hash>-<name>-<version>[-output]
 
 	hash, name := path[:32], path[33:]
 	dashIndex := 0
 	for i, r := range name {
 		if dashIndex != 0 && !unicode.IsLetter(r) {
-			return StorePathParts{Hash: hash, Name: name[:dashIndex], Version: name[i:]}
+			version, output, _ := strings.Cut(name[i:], "-")
+			return StorePathParts{Hash: hash, Name: name[:dashIndex], Version: version, Output: output}
 		}
 		dashIndex = 0
 		if r == '-' {
diff --git a/internal/nix/storepath_test.go b/internal/nix/storepath_test.go
@@ -35,6 +35,16 @@ func TestStorePathParts(t *testing.T) {
 				Name: "coreutils",
 			},
 		},
+		// With output
+		{
+			storePath: "/nix/store/0z1zq1zq1zq1zq1zq1zq1zq1zq1zq1zq-foo-1.0.0-bar",
+			expected: StorePathParts{
+				Hash:    "0z1zq1zq1zq1zq1zq1zq1zq1zq1zq1zq",
+				Name:    "foo",
+				Version: "1.0.0",
+				Output:  "bar",
+			},
+		},
 	}
 
 	for _, testCase := range testCases {
diff --git a/internal/ux/messages.go b/internal/ux/messages.go
@@ -4,6 +4,7 @@
 package ux
 
 import (
+	"context"
 	"fmt"
 	"io"
 
@@ -29,3 +30,24 @@ func Ferror(w io.Writer, format string, a ...any) {
 	color.New(color.FgHiRed).Fprint(w, "Error: ")
 	fmt.Fprintf(w, format, a...)
 }
+
+// Hidable messages allow the use of context to disable a message. Messages can be hidden
+// by their format string.
+
+type ctxKey string
+
+func HideMessage(ctx context.Context, format string) context.Context {
+	return context.WithValue(ctx, ctxKey(format), true)
+}
+
+func FHidableWarning(ctx context.Context, w io.Writer, format string, a ...any) {
+	if isHidden(ctx, format) {
+		return
+	}
+	Fwarning(w, format, a...)
+}
+
+func isHidden(ctx context.Context, format string) bool {
+	isHidden, _ := ctx.Value(ctxKey(format)).(bool)
+	return isHidden
+}

Original file line number	Diff line number	Diff line change
`@@ -106,7 +106,9 @@ func pullCmdFunc(cmd cobra.Command, url string, flags pullCmdFlags) error {`
`106`	`106`
`107`	`107`	`return installCmdFunc(`
`108`	`108`	`cmd,`
`109`		`- runCmdFlags{config: configFlags{pathFlag: pathFlag{path: flags.config.path}}},`
	`109`	`+ installCmdFlags{`
	`110`	`+ runCmdFlags: runCmdFlags{config: configFlags{pathFlag: pathFlag{path: flags.config.path}}},`
	`111`	`+ },`
`110`	`112`	`)`
`111`	`113`	`}`
`112`	`114`