[perf] Improve install and ensure perf (#2076)

## Summary

This improves install and ensure perf by making the following changes:

* Batch all store path existence checks. (previously we did one by one,
each `nix ls` call is about 50ms)
* If lockfile has store paths we use them (instead of trying to get them
from nix)
* Fix regression in `FillNarInfoCache`

## How was it tested?

Installed large devbox.json (30+ packages) and used
`DEVBOX_PRINT_EXEC_TIME` and observed almost all "prep" time disappear.
prep time is the work devbox does before actually building the packages.

Author: mikeland73

internal/devbox/packages.go

@@ -539,6 +539,7 @@ func (d *Devbox) packagesToInstallInStore(ctx context.Context, mode installMode)
 
 	// Second, check which packages are not in the nix store
 	packagesToInstall := []*devpkg.Package{}
+	storePathsForPackage := map[*devpkg.Package][]string{}
 	for _, pkg := range packages {
 		installables, err := pkg.Installables()
 		if err != nil {
@@ -549,16 +550,35 @@ func (d *Devbox) packagesToInstallInStore(ctx context.Context, mode installMode)
 				packagesToInstall = append(packagesToInstall, pkg)
 				continue
 			}
-			storePaths, err := nix.StorePathsFromInstallable(ctx, installable, pkg.HasAllowInsecure())
+
+			resolvedStorePaths, err := pkg.GetResolvedStorePaths()
 			if err != nil {
-				return nil, packageInstallErrorHandler(err, pkg, installable)
+				return nil, err
+			}
+			if resolvedStorePaths != nil {
+				storePathsForPackage[pkg] = append(storePathsForPackage[pkg], resolvedStorePaths...)
+				continue
 			}
-			isInStore, err := nix.StorePathsAreInStore(ctx, storePaths)
+
+			storePathsForPackage[pkg], err = nix.StorePathsFromInstallable(
+				ctx, installable, pkg.HasAllowInsecure())
 			if err != nil {
-				return nil, err
+				return nil, packageInstallErrorHandler(err, pkg, installable)
 			}
-			if !isInStore {
+		}
+	}
+
+	// Batch this for perf
+	storePathMap, err := nix.StorePathsAreInStore(ctx, lo.Flatten(lo.Values(storePathsForPackage)))
+	if err != nil {
+		return nil, err
+	}
+
+	for pkg, storePaths := range storePathsForPackage {
+		for _, storePath := range storePaths {
+			if !storePathMap[storePath] {
 				packagesToInstall = append(packagesToInstall, pkg)
+				break
 			}
 		}
 	}

internal/devpkg/narinfo_cache.go

@@ -132,7 +132,6 @@ func (p *Package) areExpectedOutputsInCacheOnce(outputName string) (bool, error)
 func (p *Package) fetchNarInfoStatusOnce(
 	outputName string,
 ) (map[string]string, error) {
-	defer debug.FunctionTimer().End()
 	ctx := context.TODO()
 
 	outputToCache := map[string]string{}

internal/devpkg/package.go

@@ -715,3 +715,22 @@ func (p *Package) GetOutputsWithCache() ([]Output, error) {
 	}
 	return outputs, nil
 }
+
+// GetResolvedStorePaths returns the store paths that are resolved (in lockfile)
+func (p *Package) GetResolvedStorePaths() ([]string, error) {
+	names, err := p.outputs.GetNames(p)
+	if err != nil {
+		return nil, err
+	}
+	storePaths := []string{}
+	for _, name := range names {
+		outputs, err := p.outputsForOutputName(name)
+		if err != nil {
+			return nil, err
+		}
+		for _, output := range outputs {
+			storePaths = append(storePaths, output.Path)
+		}
+	}
+	return storePaths, nil
+}

internal/nix/store.go

@@ -48,49 +48,77 @@ func StorePathsFromInstallable(ctx context.Context, installable string, allowIns
 
 		return nil, err
 	}
-	return parseStorePathFromInstallableOutput(installable, resultBytes)
+
+	validPaths, err := parseStorePathFromInstallableOutput(resultBytes)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse path-info for %s: %w", installable, err)
+	}
+
+	return maps.Keys(validPaths), nil
 }
 
-// StorePathsAreInStore returns true if the store path is in the store
-// It relies on `nix store ls` to check if the store path is in the store
-func StorePathsAreInStore(ctx context.Context, storePaths []string) (bool, error) {
+// StorePathsAreInStore a map of store paths to whether they are in the store.
+func StorePathsAreInStore(ctx context.Context, storePaths []string) (map[string]bool, error) {
 	defer debug.FunctionTimer().End()
+	if len(storePaths) == 0 {
+		return map[string]bool{}, nil
+	}
+	args := append([]string{"path-info", "--offline", "--json"}, storePaths...)
+	cmd := commandContext(ctx, args...)
+	debug.Log("Running cmd %s", cmd)
+	output, err := cmd.Output()
+	if err != nil {
+		return nil, err
+	}
+
+	validPaths, err := parseStorePathFromInstallableOutput(output)
+	if err != nil {
+		return nil, err
+	}
+
+	result := map[string]bool{}
 	for _, storePath := range storePaths {
-		cmd := commandContext(ctx, "store", "ls", storePath)
-		debug.Log("Running cmd %s", cmd)
-		if err := cmd.Run(); err != nil {
-			if exitErr := (&exec.ExitError{}); errors.As(err, &exitErr) {
-				return false, nil
-			}
-			return false, err
-		}
+		_, ok := validPaths[storePath]
+		result[storePath] = ok
 	}
-	return true, nil
+
+	return result, nil
+}
+
+// Older nix versions (like 2.17) are an array of objects that contain path and valid fields
+type pathInfoLegacy struct {
+	Path  string `json:"path"`
+	Valid bool   `json:"valid"`
 }
 
 // parseStorePathFromInstallableOutput parses the output of `nix store path-from-installable --json`
 // This function is decomposed out of StorePathFromInstallable to make it testable.
-func parseStorePathFromInstallableOutput(installable string, output []byte) ([]string, error) {
-	// Newer nix versions (like 2.20)
+func parseStorePathFromInstallableOutput(output []byte) (map[string]any, error) {
+	// Newer nix versions (like 2.20) have output of the form
+	// {"<store-path>": {}}
+	// if a store path is used as an installable, the keys will be present even if invalid but
+	// the values will be null.
 	var out1 map[string]any
 	if err := json.Unmarshal(output, &out1); err == nil {
-		return maps.Keys(out1), nil
+		maps.DeleteFunc(out1, func(k string, v any) bool {
+			return v == nil
+		})
+		return out1, nil
 	}
 
-	// Older nix versions (like 2.17)
-	var out2 []struct {
-		Path  string `json:"path"`
-		Valid bool   `json:"valid"`
-	}
+	var out2 []pathInfoLegacy
+
 	if err := json.Unmarshal(output, &out2); err == nil {
-		res := []string{}
+		res := map[string]any{}
 		for _, outValue := range out2 {
-			res = append(res, outValue.Path)
+			if outValue.Valid {
+				res[outValue.Path] = true
+			}
 		}
 		return res, nil
 	}
 
-	return nil, fmt.Errorf("failed to parse store path from installable (%s) output: %s", installable, output)
+	return nil, fmt.Errorf("failed to parse path-info output: %s", output)
 }
 
 // DaemonError reports an unsuccessful attempt to connect to the Nix daemon.

internal/nix/store_test.go

@@ -3,6 +3,8 @@ package nix
 import (
 	"slices"
 	"testing"
+
+	"golang.org/x/exp/maps"
 )
 
 func TestParseStorePathFromInstallableOutput(t *testing.T) {
@@ -19,18 +21,18 @@ func TestParseStorePathFromInstallableOutput(t *testing.T) {
 		},
 		{
 			name:     "go-basic-nix-2-17-0",
-			input:    `[{"path":"/nix/store/fgkl3qk8p5hnd07b0dhzfky3ys5gxjmq-go-1.22.0","valid":false}]`,
+			input:    `[{"path":"/nix/store/fgkl3qk8p5hnd07b0dhzfky3ys5gxjmq-go-1.22.0","valid":true}]`,
 			expected: []string{"/nix/store/fgkl3qk8p5hnd07b0dhzfky3ys5gxjmq-go-1.22.0"},
 		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			actual, err := parseStorePathFromInstallableOutput(tc.name, []byte(tc.input))
+			actual, err := parseStorePathFromInstallableOutput([]byte(tc.input))
 			if err != nil {
 				t.Errorf("Expected no error but got error: %s", err)
 			}
-			if !slices.Equal(tc.expected, actual) {
+			if !slices.Equal(tc.expected, maps.Keys(actual)) {
 				t.Errorf("Expected store path %s but got %s", tc.expected, actual)
 			}
 		})