[rm nixpkgs] make HEAD request to BinaryCache to ensure binary is cached

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,14 +4,19 @@
 package devpkg
 
 import (
+	"context"
 	"crypto/md5"
 	"encoding/hex"
 	"fmt"
 	"io"
+	"net/http"
 	"net/url"
 	"path/filepath"
 	"regexp"
 	"strings"
+	"sync"
+	"time"
+	"unicode"
 
 	"github.com/pkg/errors"
 	"github.com/samber/lo"
@@ -23,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.
@@ -53,6 +59,20 @@ type Package struct {
 	normalizedPackageAttributePathCache string // memoized value from normalizedPackageAttributePath()
 }
 
+// isNarInfoInCache checks if the .narinfo for this package is in the `BinaryCache`.
+// The key is the `Package.Raw` string.
+// This cannot be a field on the Package struct, because that struct
+// is constructed multiple times in a request (TODO: we could fix that).
+var isNarInfoInCache = struct {
+	status map[string]bool
+	lock   sync.RWMutex
+	// re-use httpClient to re-use the connection
+	httpClient http.Client
+}{
+	status:     map[string]bool{},
+	httpClient: http.Client{},
+}
+
 // PackageFromStrings constructs Package from the list of package names provided.
 // These names correspond to devbox packages from the devbox.json config.
 func PackageFromStrings(rawNames []string, l lock.Locker) []*Package {
@@ -172,6 +192,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
@@ -421,7 +442,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 {
@@ -462,37 +499,117 @@ 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) {
+	sysInfo, err := p.sysInfoIfExists()
+	if err != nil {
+		return false, err
+	}
+	return sysInfo != nil, nil
+}
+
+// sysInfoIfExists returns the system info for the user's system. If the sysInfo
+// is missing, then nil is returned
+// NOTE: this is called from multiple go-routines and needs to be concurrency safe.
+// Hence, we compute nix.Version, nix.System and lockfile.Resolve prior to calling this
+// function from within a goroutine.
+func (p *Package) sysInfoIfExists() (*lock.SystemInfo, error) {
 	if !featureflag.RemoveNixpkgs.Enabled() {
-		return false, nil
+		return nil, nil
 	}
 
 	if !p.isVersioned() {
-		return false, nil
+		return nil, nil
+	}
+
+	version, err := nix.Version()
+	if err != nil {
+		return nil, err
+	}
+
+	// enable for nix >= 2.17
+	if vercheck.SemverCompare(version, "2.17.0") < 0 {
+		return nil, err
 	}
 
 	entry, err := p.lockfile.Resolve(p.Raw)
 	if err != nil {
-		return false, err
+		return nil, err
 	}
 
+	userSystem := nix.System()
+
 	if entry.Systems == nil {
-		return false, nil
+		return nil, nil
 	}
 
 	// Check if the user's system's info is present in the lockfile
-	_, ok := entry.Systems[nix.System()]
+	sysInfo, ok := entry.Systems[userSystem]
 	if !ok {
+		return nil, nil
+	}
+	return sysInfo, nil
+}
+
+// IsInBinaryCache returns true if the package is in the binary cache.
+// ALERT: Callers must call FillNarInfoCache before calling this function.
+func (p *Package) IsInBinaryCache() (bool, error) {
+
+	if eligible, err := p.isEligibleForBinaryCache(); err != nil {
+		return false, err
+	} else if !eligible {
 		return false, nil
 	}
 
-	version, err := nix.Version()
+	// Check if the narinfo is present in the binary cache
+	isNarInfoInCache.lock.RLock()
+	exists, ok := isNarInfoInCache.status[p.Raw]
+	isNarInfoInCache.lock.RUnlock()
+	if !ok {
+		return false, errors.Errorf("narInfo cache miss: %v. call XYZ before invoking IsInBinaryCache", p.Raw)
+	}
+	return exists, nil
+}
+
+// fillNarInfoCache fills the cache value for the narinfo of this package,
+// if it is eligible for the binary cache.
+// NOTE: this must be concurrency safe.
+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 false, err
+		return err
+	} else if sysInfo == nil {
+		return errors.New(
+			"sysInfo is nil, but should not be because" +
+				" the package is eligible for binary cache",
+		)
 	}
 
-	// enable for nix >= 2.17
-	return vercheck.SemverCompare(version, "2.17.0") >= 0, nil
+	pathParts := newStorePathParts(sysInfo.StorePath)
+	reqURL := BinaryCache + "/" + pathParts.hash + ".narinfo"
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	req, err := http.NewRequestWithContext(ctx, http.MethodHead, reqURL, nil)
+	if err != nil {
+		return err
+	}
+	res, err := isNarInfoInCache.httpClient.Do(req)
+	if err != nil {
+		return err
+	}
+	// read the body fully, and close it to ensure the connection is reused.
+	_, _ = io.Copy(io.Discard, res.Body)
+	defer res.Body.Close()
+
+	isNarInfoInCache.lock.Lock()
+	isNarInfoInCache.status[p.Raw] = res.StatusCode == 200
+	isNarInfoInCache.lock.Unlock()
+	return nil
 }
 
 // InputAddressedPath is the input-addressed path in /nix/store
@@ -542,3 +659,78 @@ func (p *Package) EnsureUninstallableIsInLockfile() error {
 	_, err := p.lockfile.Resolve(p.Raw)
 	return err
 }
+
+// storePath are the constituent parts of
+// /nix/store/<hash>-<name>-<version>
+//
+// This is a helper struct for analyzing the string representation
+type storePathParts struct {
+	hash    string
+	name    string
+	version 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
+func newStorePathParts(path string) storePathParts {
+	path = strings.TrimPrefix(path, "/nix/store/")
+	// path is now <hash>-<name>-<version
+
+	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:]}
+		}
+		dashIndex = 0
+		if r == '-' {
+			dashIndex = i
+		}
+	}
+	return storePathParts{hash: hash, name: name}
+}
+
+// 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 {
+
+	// Pre-compute values read in fillNarInfoCache
+	// so they can be read from multiple go-routines without locks
+	_, err := nix.Version()
+	if err != nil {
+		return err
+	}
+	_ = nix.System()
+	for _, p := range packages {
+		_, err := p.lockfile.Resolve(p.Raw)
+		if err != nil {
+			return err
+		}
+	}
+
+	group, _ := errgroup.WithContext(ctx)
+	for _, p := range packages {
+		// If the package's NarInfo status is already known, skip it
+		isNarInfoInCache.lock.RLock()
+		_, ok := isNarInfoInCache.status[p.Raw]
+		isNarInfoInCache.lock.RUnlock()
+		if ok {
+			continue
+		}
+		pkg := p // copy the loop variable since its used in a closure below
+		group.Go(func() error {
+			err := pkg.fillNarInfoCache()
+			if err != nil {
+				// default to false if there was an error, so we don't re-try
+				isNarInfoInCache.lock.Lock()
+				isNarInfoInCache.status[pkg.Raw] = false
+				isNarInfoInCache.lock.Unlock()
+			}
+			return err
+		})
+	}
+	return group.Wait()
+}

internal/devpkg/package_test.go

@@ -183,3 +183,46 @@ func TestHashFromNixPkgsURL(t *testing.T) {
 		}
 	}
 }
+
+func TestStorePathParts(t *testing.T) {
+	testCases := []struct {
+		storePath string
+		expected  storePathParts
+	}{
+		// simple case:
+		{
+			storePath: "/nix/store/cvrn84c1hshv2wcds7n1rhydi6lacqns-gnumake-4.4.1",
+			expected: storePathParts{
+				hash:    "cvrn84c1hshv2wcds7n1rhydi6lacqns",
+				name:    "gnumake",
+				version: "4.4.1",
+			},
+		},
+		// the package name can have dashes:
+		{
+			storePath: "/nix/store/q2xdxsswjqmqcbax81pmazm367s7jzyb-cctools-binutils-darwin-wrapper-973.0.1",
+			expected: storePathParts{
+				hash:    "q2xdxsswjqmqcbax81pmazm367s7jzyb",
+				name:    "cctools-binutils-darwin-wrapper",
+				version: "973.0.1",
+			},
+		},
+		// version is optional. This is an artificial example I constructed
+		{
+			storePath: "/nix/store/gfxwrd5nggc68pjj3g3jhlldim9rpg0p-coreutils",
+			expected: storePathParts{
+				hash: "gfxwrd5nggc68pjj3g3jhlldim9rpg0p",
+				name: "coreutils",
+			},
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.storePath, func(t *testing.T) {
+			parts := newStorePathParts(testCase.storePath)
+			if parts != testCase.expected {
+				t.Errorf("Expected %v, got %v", testCase.expected, parts)
+			}
+		})
+	}
+}

internal/impl/devbox.go

@@ -723,7 +723,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
 		}