depot/go/nix/nixbuild/peerresolver.go

// SPDX-FileCopyrightText: 2023 Luke Granger-Brown <depot@lukegb.com>
//
// SPDX-License-Identifier: Apache-2.0

package nixbuild

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"path"

	"hg.lukegb.com/lukegb/depot/go/nix/nar"
	"hg.lukegb.com/lukegb/depot/go/nix/nar/narinfo"
	"hg.lukegb.com/lukegb/depot/go/nix/nixdrv"
	"hg.lukegb.com/lukegb/depot/go/nix/nixpool"
	"hg.lukegb.com/lukegb/depot/go/nix/nixstore"
)

type PeerFetcher struct {
	Pool *nixpool.Pool
}

var _ Fetcher = ((*PeerFetcher)(nil))
var _ DerivationFetcher = ((*PeerFetcher)(nil))

func (r PeerFetcher) FetchNarInfo(ctx context.Context, path string) (*narinfo.NarInfo, error) {
	d, err := r.Pool.Get()
	if err != nil {
		return nil, err
	}
	defer r.Pool.Put(d)

	return d.NARInfo(path)
}

type poolReturningReadCloser struct {
	io.ReadCloser
	pool *nixpool.Pool
	d    *nixstore.Daemon
}

func (prc *poolReturningReadCloser) Close() error {
	err := prc.ReadCloser.Close()
	prc.pool.Put(prc.d)
	return err
}

func (r PeerFetcher) FetchNar(ctx context.Context, ni *narinfo.NarInfo) (io.ReadCloser, error) {
	d, err := r.Pool.Get()
	if err != nil {
		return nil, err
	}

	rc, err := d.NARFromPath(ni.StorePath)
	if err != nil {
		return nil, err
	}
	return &poolReturningReadCloser{
		ReadCloser: rc,
		pool:       r.Pool,
		d:          d,
	}, nil
}

func (r PeerFetcher) FetchDerivationForPath(ctx context.Context, path string) (*nixdrv.BasicDerivation, string, error) {
	d, err := r.Pool.Get()
	if err != nil {
		return nil, "", err
	}

	derivers, err := d.QueryValidDerivers(path)
	r.Pool.Put(d)
	if err != nil {
		return nil, "", err
	}
	if len(derivers) == 0 {
		return nil, "", fmt.Errorf("don't know how to build %v", path)
	}

	// Do we need to do something smarter than take the first derivation?
	drvPath := derivers[0]
	drv, err := r.LoadDerivation(ctx, drvPath)
	if err != nil {
		return nil, "", fmt.Errorf("fetching derivation %v for %v: %w", drvPath, path, err)
	}
	bd, err := drv.ToBasicDerivation(ctx, r)
	if err != nil {
		return nil, "", fmt.Errorf("converting %v to a basic derivation: %w", drvPath, err)
	}
	//drvBase := strings.TrimSuffix(drvPath, filepath.Ext(drvPath))
	return bd, drvPath, nil
}

func (r PeerFetcher) LoadDerivation(ctx context.Context, drvPath string) (*nixdrv.Derivation, error) {
	ni, err := r.FetchNarInfo(ctx, drvPath)
	if err != nil {
		return nil, fmt.Errorf("fetching narinfo for %v: %w", drvPath, err)
	}
	rc, err := r.FetchNar(ctx, ni)
	if err != nil {
		return nil, fmt.Errorf("fetching nar for %v: %w", drvPath, err)
	}
	defer rc.Close()

	imfs := nar.NewInmemoryFS()
	fn := path.Base(ni.StorePath)
	if err := nar.Unpack(rc, imfs, fn); err != nil {
		return nil, fmt.Errorf("unpacking nar for %v: %w", drvPath, err)
	}
	dent, ok := imfs.Dirent[fn]
	if !ok {
		return nil, fmt.Errorf("unpacking nar for %v yielded no file", drvPath)
	} else if dent.Content == nil {
		return nil, fmt.Errorf("unpacking nar for %v yielded non-file", drvPath)
	}

	return nixdrv.Load(bufio.NewReader(bytes.NewReader(dent.Content)))
}

type PeerResolver struct {
	PeerFetcher
	Priority int
}

var _ Resolver = ((*PeerResolver)(nil))

func (r PeerResolver) RelativePriority() int { return r.Priority }

type PeerTarget struct {
	PeerFetcher
	ActivityTracker *nixstore.ActivityTracker
}

var _ Target = ((*PeerTarget)(nil))

func (t *PeerTarget) EnsurePaths(ctx context.Context, paths []string) ([]string, error) {
	d, err := t.Pool.Get()
	if err != nil {
		return nil, err
	}
	defer t.Pool.Put(d)

	var ensuredPaths []string
	for _, p := range paths {
		if err := d.EnsurePath(t.ActivityTracker, p); err != nil {
			var nv nixstore.PathNotValidError
			if !errors.As(err, &nv) {
				return nil, fmt.Errorf("ensuring %v: %w", p, err)
			}
		} else {
			ensuredPaths = append(ensuredPaths, p)
		}
	}
	return ensuredPaths, nil
}

func (t *PeerTarget) PutNar(ctx context.Context, ni *narinfo.NarInfo, w io.Reader) error {
	d, err := t.Pool.Get()
	if err != nil {
		return err
	}
	defer t.Pool.Put(d)

	return d.AddToStoreNar(ni, w)
}

func (t *PeerTarget) ValidPaths(ctx context.Context, paths []string) ([]string, error) {
	d, err := t.Pool.Get()
	if err != nil {
		return nil, err
	}
	defer t.Pool.Put(d)

	return d.ValidPaths(paths)
}