Skip to content
quarry Artifact [ea03246c13]
Login

Artifact [ea03246c13]

Artifact ea03246c13f89ed5c7ab405c417850c4bcc82ccba676134f2c7b6aa72bb72af6:


package archive

import (
	"context"
	"errors"
	"os"
	"strings"
	"sync"
)

// Objects is the bucket an archive is kept in, as much of it as
// Bucketed needs.
type Objects interface {
	// Where names the bucket. A key known present in one bucket says
	// nothing about another bound later under the same repository.
	Where() string
	Has(ctx context.Context, key string) (bool, error)
	// Upload sends the file at path to key, streaming it.
	Upload(ctx context.Context, key, path string) error
}

// Bucketed puts each archive in a bucket once and remembers that it is
// there. The bucket is the cache: the forge builds an archive on the
// first request for it and every later reader is sent to the bucket,
// so the compression is paid once per check-in and container rather
// than once per reader, and the bytes never pass through the forge
// again.
//
// Keys are content-addressed, so an object is never stale and a
// positive answer can be remembered for the life of the process. A
// negative one is not: another node may have built it since.
type Bucketed struct {
	mu      sync.Mutex
	flights map[place]*flight
	known   map[place]struct{}
	swept   map[string]bool
}

type place struct{ where, key string }

type flight struct {
	done chan struct{}
	res  Result
	err  error
}

// Result says what one Ensure did, for the line the caller logs.
type Result struct {
	// Built is true for the one call that ran the build; every other
	// caller, waiting on it or arriving after, sees false.
	Built bool
	Size  int64
}

func NewBucketed() *Bucketed {
	return &Bucketed{flights: map[place]*flight{}, known: map[place]struct{}{}, swept: map[string]bool{}}
}

// Ensure makes key present in objs, running build (into a temporary
// file, then uploaded) only if the bucket does not already hold it.
//
// Concurrent callers for the same key share one build: the first runs
// it and the rest wait for its answer. A leader whose own request went
// away hands the build on rather than failing every reader behind it.
func (b *Bucketed) Ensure(ctx context.Context, objs Objects, key string, build func(dst string) error) (Result, error) {
	p := place{objs.Where(), key}
	for {
		b.mu.Lock()
		if _, ok := b.known[p]; ok {
			b.mu.Unlock()
			return Result{}, nil
		}
		if f, ok := b.flights[p]; ok {
			b.mu.Unlock()
			select {
			case <-f.done:
			case <-ctx.Done():
				return Result{}, ctx.Err()
			}
			if errors.Is(f.err, context.Canceled) || errors.Is(f.err, context.DeadlineExceeded) {
				continue
			}
			return Result{}, f.err
		}
		f := &flight{done: make(chan struct{})}
		b.flights[p] = f
		b.mu.Unlock()

		f.res, f.err = b.fill(ctx, objs, key, build)

		b.mu.Lock()
		delete(b.flights, p)
		if f.err == nil {
			b.known[p] = struct{}{}
		}
		b.mu.Unlock()
		close(f.done)
		return f.res, f.err
	}
}

func (b *Bucketed) fill(ctx context.Context, objs Objects, key string, build func(dst string) error) (Result, error) {
	have, err := objs.Has(ctx, key)
	if err != nil {
		return Result{}, err
	}
	if have {
		return Result{}, nil
	}
	tmp, err := os.CreateTemp("", "quarry-archive-*")
	if err != nil {
		return Result{}, err
	}
	path := tmp.Name()
	tmp.Close()
	defer os.Remove(path)
	if err := build(path); err != nil {
		return Result{}, &BuildError{err}
	}
	fi, err := os.Stat(path)
	if err != nil {
		return Result{}, err
	}
	if err := objs.Upload(ctx, key, path); err != nil {
		return Result{}, err
	}
	return Result{Built: true, Size: fi.Size()}, nil
}

// BuildError is a failure of the build itself, as opposed to the
// bucket's: a caller may serve the archive another way when the bucket
// fails, and should not when the build did.
type BuildError struct{ Err error }

func (e *BuildError) Error() string { return e.Err.Error() }
func (e *BuildError) Unwrap() error { return e.Err }

// Forget drops every remembered key under prefix, in every bucket, so
// the next request asks the bucket again instead of assuming.
func (b *Bucketed) Forget(prefix string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	for p := range b.known {
		if strings.HasPrefix(p.key, prefix) {
			delete(b.known, p)
		}
	}
}

// FirstSweep reports true once per name for the life of the process,
// so a sweep of what older builds left runs once rather than per build.
func (b *Bucketed) FirstSweep(name string) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.swept[name] {
		return false
	}
	b.swept[name] = true
	return true
}