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{}
}
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{}{}}
}
// 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)
}
}
}