add FillNarInfoCache to make HEAD requests in parallel

Author: savil

go.mod

@@ -45,13 +45,17 @@ require (
 	gopkg.in/yaml.v3 v3.0.1
 )
 
-require github.com/joho/godotenv v1.5.1
+require (
+	github.com/joho/godotenv v1.5.1
+	golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f
+)
 
 require (
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
 	github.com/buger/jsonparser v1.1.1 // indirect
 	github.com/kr/pretty v0.1.0 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
+	golang.org/x/net v0.8.0 // indirect
 )
 
 require (

go.sum

@@ -231,6 +231,8 @@ golang.org/x/mod v0.11.0 h1:bUO06HqtnRcc/7l71XBe4WcqTZ+3AH1J59zWDDwLKgU=
 golang.org/x/mod v0.11.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
 golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
 golang.org/x/net v0.8.0 h1:Zrh2ngAOFYneWTAIAPethzeaQLuHwhuBkuV6ZiRnUaQ=
+golang.org/x/net v0.8.0/go.mod h1:QVkue5JL9kW//ek3r6jTKnTFis1tRmNAW2P1shuFdJc=
+golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f h1:wMNYb4v58l5UBM7MYRLPG6ZhfOqbKu7X5eyFl8ZhKvA=
 golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190204203706-41f3e6584952/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=

internal/devpkg/package.go

@@ -4,6 +4,7 @@
 package devpkg
 
 import (
+	"context"
 	"crypto/md5"
 	"encoding/hex"
 	"fmt"
@@ -27,6 +28,7 @@ import (
 	"go.jetpack.io/devbox/internal/nix"
 	"go.jetpack.io/devbox/internal/vercheck"
 	"go.jetpack.io/devbox/plugins"
+	"golang.org/x/sync/errgroup"
 )
 
 // Package represents a "package" added to the devbox.json config.
@@ -182,6 +184,7 @@ func (p *Package) IsInstallable() bool {
 // Installable for this package. Installable is a nix concept defined here:
 // https://nixos.org/manual/nix/stable/command-ref/new-cli/nix.html#installables
 func (p *Package) Installable() (string, error) {
+
 	inCache, err := p.IsInBinaryCache()
 	if err != nil {
 		return "", err
@@ -435,7 +438,23 @@ func (p *Package) LegacyToVersioned() string {
 	return p.Raw + "@latest"
 }
 
-func (p *Package) EnsureNixpkgsPrefetched(w io.Writer) error {
+// ensureNixpkgsPrefetched will prefetch flake for the nixpkgs registry for the package.
+// This is an internal method, and should not be called directly.
+func EnsureNixpkgsPrefetched(ctx context.Context, w io.Writer, pkgs []*Package) error {
+	if err := FillNarInfoCache(ctx, pkgs...); err != nil {
+		return err
+	}
+	for _, input := range pkgs {
+		if err := input.ensureNixpkgsPrefetched(w); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// ensureNixpkgsPrefetched should be called via the public EnsureNixpkgsPrefetched.
+// See function comment there.
+func (p *Package) ensureNixpkgsPrefetched(w io.Writer) error {
 
 	inCache, err := p.IsInBinaryCache()
 	if err != nil {
@@ -476,7 +495,7 @@ func (p *Package) HashFromNixPkgsURL() string {
 // It is used as FromStore in builtins.fetchClosure.
 const BinaryCache = "https://cache.nixos.org"
 
-func (p *Package) IsInBinaryCache() (bool, error) {
+func (p *Package) isEligibleForBinaryCache() (bool, error) {
 	if !featureflag.RemoveNixpkgs.Enabled() {
 		return false, nil
 	}
@@ -485,40 +504,84 @@ func (p *Package) IsInBinaryCache() (bool, error) {
 		return false, nil
 	}
 
-	entry, err := p.lockfile.Resolve(p.Raw)
+	sysInfo, err := p.sysInfoIfExists()
 	if err != nil {
 		return false, err
+	} else if sysInfo == nil {
+		return false, nil
 	}
 
-	userSystem, err := nix.System()
+	version, err := nix.Version()
 	if err != nil {
 		return false, err
 	}
 
+	// enable for nix >= 2.17
+	return vercheck.SemverCompare(version, "2.17.0") >= 0, nil
+}
+
+// sysInfoIfExists returns the system info for the user's system. If the sysInfo
+// is missing, then nil is returned
+func (p *Package) sysInfoIfExists() (*lock.SystemInfo, error) {
+
+	entry, err := p.lockfile.Resolve(p.Raw)
+	if err != nil {
+		return nil, err
+	}
+
+	userSystem, err := nix.System()
+	if err != nil {
+		return nil, err
+	}
+
 	if entry.Systems == nil {
-		return false, nil
+		return nil, nil
 	}
 
 	// Check if the user's system's info is present in the lockfile
 	sysInfo, ok := entry.Systems[userSystem]
 	if !ok {
-		return false, nil
+		return nil, nil
 	}
+	return sysInfo, nil
+}
 
-	version, err := nix.Version()
-	if err != nil {
-		return false, err
-	}
+// IsInBinaryCache returns true if the package is in the binary cache.
+// Callers should call FillNarInfoCache before calling this function.
+func (p *Package) IsInBinaryCache() (bool, error) {
 
-	// enable for nix >= 2.17
-	if vercheck.SemverCompare(version, "2.17.0") < 0 {
+	if eligible, err := p.isEligibleForBinaryCache(); err != nil {
+		return false, err
+	} else if !eligible {
 		return false, nil
 	}
 
 	// Check if the narinfo is present in the binary cache
-	if exists, ok := isNarInfoInCache[p.Raw]; ok {
-		fmt.Printf("narInfo cache hit: %v\n", exists)
-		return exists, nil
+	exists, ok := isNarInfoInCache[p.Raw]
+	if !ok {
+		return false, errors.Errorf("narInfo cache miss: %v. call XYZ before invoking IsInBinaryCache", p.Raw)
+	}
+	fmt.Printf("narInfo cache hit: %v\n", exists)
+	return exists, nil
+}
+
+// fillNarInfoCache fills the cache value for the narinfo of this package,
+// if it is eligible for the binary cache.
+func (p *Package) fillNarInfoCache() error {
+	if eligible, err := p.isEligibleForBinaryCache(); err != nil {
+		return err
+	} else if !eligible {
+		return nil
+	}
+
+	sysInfo, err := p.sysInfoIfExists()
+	if err != nil {
+		return err
+	} else if sysInfo == nil {
+		return errors.New(
+			"sysInfo is nil, but should not be because" +
+				" the package is eligible for binary cache",
+		)
 	}
 
 	httpClient := &http.Client{
@@ -530,14 +593,14 @@ func (p *Package) IsInBinaryCache() (bool, error) {
 	if err != nil {
 		if os.IsTimeout(err) {
 			isNarInfoInCache[p.Raw] = false // set this to avoid re-tries
-			return false, nil
+			return nil
 		}
-		return false, err
+		return err
 	}
 	fmt.Printf("res status code %d\n", res.StatusCode)
 
 	isNarInfoInCache[p.Raw] = res.StatusCode == 200
-	return isNarInfoInCache[p.Raw], nil
+	return nil
 }
 
 // InputAddressedPath is the input-addressed path in /nix/store
@@ -628,3 +691,24 @@ func newStorePathParts(path string) *storePathParts {
 func (p *storePathParts) Equal(other *storePathParts) bool {
 	return p.hash == other.hash && p.name == other.name && p.version == other.version
 }
+
+// FillNarInfoCache checks the remote binary cache for the narinfo of each
+// package in the list, and caches the result.
+// Callers of IsInBinaryCache must call this function first.
+func FillNarInfoCache(ctx context.Context, packages ...*Package) error {
+	g, ctx := errgroup.WithContext(ctx)
+
+	for _, p := range packages {
+		// If the package's NarInfo status is already known, skip it
+		if _, ok := isNarInfoInCache[p.Raw]; ok {
+			continue
+		}
+		g.Go(func() error {
+			return p.fillNarInfoCache()
+		})
+	}
+	if err := g.Wait(); err != nil {
+		return err
+	}
+	return nil
+}

internal/impl/devbox.go

@@ -706,7 +706,7 @@ func (d *Devbox) StartProcessManager(
 	processComposePath, err := utilityLookPath("process-compose")
 	if err != nil {
 		fmt.Fprintln(d.writer, "Installing process-compose. This may take a minute but will only happen once.")
-		if err = d.addDevboxUtilityPackage("github:F1bonacc1/process-compose/v0.43.1"); err != nil {
+		if err = d.addDevboxUtilityPackage(ctx, "github:F1bonacc1/process-compose/v0.43.1"); err != nil {
 			return err
 		}
 

internal/impl/packages.go

@@ -42,6 +42,12 @@ func (d *Devbox) Add(ctx context.Context, platforms, excludePlatforms []string,
 	// replace it.
 	pkgs := devpkg.PackageFromStrings(lo.Uniq(pkgsNames), d.lockfile)
 
+	// Fill in narinfo cache for all packages, even if the package-names are bogus
+	// (we'll just not use the result later)
+	if err := devpkg.FillNarInfoCache(ctx, pkgs...); err != nil {
+		return err
+	}
+
 	// addedPackageNames keeps track of the possibly transformed (versioned)
 	// names of added packages (even if they are already in config). We use this
 	// to know the exact name to mark as allowed insecure later on.
@@ -295,10 +301,8 @@ func (d *Devbox) addPackagesToProfile(ctx context.Context, mode installMode) err
 	// If packages are in profile but nixpkgs has been purged, the experience
 	// will be poor when we try to run print-dev-env. So we ensure nixpkgs is
 	// prefetched for all relevant packages (those not in binary cache).
-	for _, input := range pkgs {
-		if err := input.EnsureNixpkgsPrefetched(d.writer); err != nil {
-			return err
-		}
+	if err := devpkg.EnsureNixpkgsPrefetched(ctx, d.writer, pkgs); err != nil {
+		return err
 	}
 
 	var msg string
@@ -321,7 +325,7 @@ func (d *Devbox) addPackagesToProfile(ctx context.Context, mode installMode) err
 
 		stepMsg := fmt.Sprintf("[%d/%d] %s", stepNum, total, pkg)
 
-		if err := nixprofile.ProfileInstall(&nixprofile.ProfileInstallArgs{
+		if err := nixprofile.ProfileInstall(ctx, &nixprofile.ProfileInstallArgs{
 			CustomStepMessage: stepMsg,
 			Lockfile:          d.lockfile,
 			Package:           pkg.Raw,
@@ -343,7 +347,12 @@ func (d *Devbox) removePackagesFromProfile(ctx context.Context, pkgs []string) e
 		return err
 	}
 
-	for _, input := range devpkg.PackageFromStrings(pkgs, d.lockfile) {
+	packages := devpkg.PackageFromStrings(pkgs, d.lockfile)
+	if err := devpkg.FillNarInfoCache(ctx, packages...); err != nil {
+		return err
+	}
+
+	for _, input := range packages {
 		index, err := nixprofile.ProfileListIndex(&nixprofile.ProfileListIndexArgs{
 			Lockfile:   d.lockfile,
 			Writer:     d.writer,
@@ -414,6 +423,9 @@ func (d *Devbox) pendingPackagesForInstallation(ctx context.Context) ([]*devpkg.
 	if err != nil {
 		return nil, err
 	}
+	if err := devpkg.FillNarInfoCache(ctx, packages...); err != nil {
+		return nil, err
+	}
 	for _, pkg := range packages {
 		_, err := nixprofile.ProfileListIndex(&nixprofile.ProfileListIndexArgs{
 			List:       list,

internal/impl/update.go

@@ -49,6 +49,10 @@ func (d *Devbox) Update(ctx context.Context, pkgs ...string) error {
 		}
 	}
 
+	if devpkg.FillNarInfoCache(ctx, pendingPackagesToUpdate...); err != nil {
+		return err
+	}
+
 	for _, pkg := range pendingPackagesToUpdate {
 		if _, _, isVersioned := searcher.ParseVersionedPackage(pkg.Raw); !isVersioned {
 			if err = d.attemptToUpgradeFlake(pkg); err != nil {

internal/impl/util.go

@@ -4,6 +4,7 @@
 package impl
 
 import (
+	"context"
 	"io/fs"
 	"os"
 	"path/filepath"
@@ -18,13 +19,13 @@ import (
 // It's used to install applications devbox might need, like process-compose
 // This is an alternative to a global install which would modify a user's
 // environment.
-func (d *Devbox) addDevboxUtilityPackage(pkg string) error {
+func (d *Devbox) addDevboxUtilityPackage(ctx context.Context, pkg string) error {
 	profilePath, err := utilityNixProfilePath()
 	if err != nil {
 		return err
 	}
 
-	return nixprofile.ProfileInstall(&nixprofile.ProfileInstallArgs{
+	return nixprofile.ProfileInstall(ctx, &nixprofile.ProfileInstallArgs{
 		Lockfile:    d.lockfile,
 		Package:     pkg,
 		ProfilePath: profilePath,

internal/nix/nixprofile/profile.go

@@ -5,6 +5,7 @@ package nixprofile
 
 import (
 	"bufio"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -250,9 +251,15 @@ type ProfileInstallArgs struct {
 }
 
 // ProfileInstall calls nix profile install with default profile
-func ProfileInstall(args *ProfileInstallArgs) error {
+func ProfileInstall(ctx context.Context, args *ProfileInstallArgs) error {
 	input := devpkg.PackageFromString(args.Package, args.Lockfile)
 
+	// Fill in the narinfo cache for the input package. It's okay to call this for a single package
+	// because installing is a slow operation anyway.
+	if err := devpkg.FillNarInfoCache(ctx, input); err != nil {
+		return err
+	}
+
 	inCache, err := input.IsInBinaryCache()
 	if err != nil {
 		return err

internal/shellgen/flake_input.go

@@ -77,13 +77,18 @@ func (f *flakeInput) BuildInputs() ([]string, error) {
 // i.e. have a commit hash and always resolve to the same package/version.
 // Note: inputs returned by this function include plugin packages. (php only for now)
 // It's not entirely clear we always want to add plugin packages to the top level
-func flakeInputs(ctx context.Context, packages []*devpkg.Package) []*flakeInput {
+func flakeInputs(ctx context.Context, packages []*devpkg.Package) ([]*flakeInput, error) {
 	defer trace.StartRegion(ctx, "flakeInputs").End()
 
 	// Use the verbose name flakeInputs to distinguish from `inputs`
 	// which refer to `nix.Input` in most of the codebase.
 	flakeInputs := map[string]*flakeInput{}
 
+	// Fill the NarInfo Cache so we can check IsInBinaryCache() for each package, below.
+	if err := devpkg.FillNarInfoCache(ctx, packages...); err != nil {
+		return nil, err
+	}
+
 	packages = lo.Filter(packages, func(item *devpkg.Package, _ int) bool {
 		// Include packages (like local or remote flakes) that cannot be
 		// fetched from a Binary Cache Store.
@@ -115,5 +120,5 @@ func flakeInputs(ctx context.Context, packages []*devpkg.Package) []*flakeInput
 		}
 	}
 
-	return goutil.PickByKeysSorted(flakeInputs, order)
+	return goutil.PickByKeysSorted(flakeInputs, order), nil
 }

internal/shellgen/flake_plan.go

@@ -38,7 +38,10 @@ func newFlakePlan(ctx context.Context, devbox devboxer) (*flakePlan, error) {
 	if err != nil {
 		return nil, err
 	}
-	flakeInputs := flakeInputs(ctx, packages)
+	flakeInputs, err := flakeInputs(ctx, packages)
+	if err != nil {
+		return nil, err
+	}
 	nixpkgsInfo := getNixpkgsInfo(devbox.Config().NixPkgsCommitHash())
 
 	// This is an optimization. Try to reuse the nixpkgs info from the flake