[rm nixpkgs] make HEAD request to BinaryCache to ensure binary is cached (#1318)

## Summary

This PR makes a HTTP HEAD request to
`https://cache.nixos.org/<hash>.narinfo` to verify that the binary is
indeed cached.

If not, we fallback to the legacy path that is dependent on downloading
nixpkgs.

**Implementation note**
In `IsInBinaryCache`, I wanted to insert the HEAD request. However, this
function can be called in a loop for all packages, which is bad for
perf.

Instead, in this design, a new function `FillNarInfoCache(packageList)`
makes the HEAD requests concurrently, storing results in an in-memory
cache.

We _require_ that callers invoke `FillNarInfoCache(packageList)` prior
to invoking `IsInBinaryCache`. Failing this, `IsInBinaryCache` will
return an error. Not thrilled with this design, since it may lead to
inadvertent errors, but this code is pretty core and our test coverage
should expose any codepaths that call `IsInBinaryCache` without
`FillNarInfoCache`.




## How was it tested?

I could do `devbox shell` in a project that has `hello` package. I
inserted some print statements to verify that the HEAD
response was 200 and that the caching for the boolean value of
isNarInfoInCache was working.

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,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
 		}