package archive
import (
"context"
"errors"
"os"
"sync"
"sync/atomic"
"testing"
"time"
)
// memObjects is a bucket in memory that counts what it was asked.
type memObjects struct {
where string
mu sync.Mutex
objects map[string][]byte
heads atomic.Int32
uploads atomic.Int32
// refuse fails every upload; gate, when set, holds each upload
// until it is closed, and uploading hears each one begin.
refuse bool
uploadErr error
gate chan struct{}
uploading chan string
}
func newMemObjects(where string) *memObjects {
return &memObjects{where: where, objects: map[string][]byte{}, uploading: make(chan string, 16)}
}
func (m *memObjects) Where() string { return m.where }
func (m *memObjects) Has(_ context.Context, key string) (bool, error) {
m.heads.Add(1)
m.mu.Lock()
defer m.mu.Unlock()
_, ok := m.objects[key]
return ok, nil
}
func (m *memObjects) Upload(ctx context.Context, key, path string) error {
m.uploads.Add(1)
m.uploading <- key
if m.gate != nil {
select {
case <-m.gate:
case <-ctx.Done():
return ctx.Err()
}
}
if m.uploadErr != nil {
return m.uploadErr
}
if m.refuse {
return errors.New("PUT: 500")
}
b, err := os.ReadFile(path)
if err != nil {
return err
}
m.mu.Lock()
defer m.mu.Unlock()
m.objects[key] = b
return nil
}
func (m *memObjects) has(key string) bool {
m.mu.Lock()
defer m.mu.Unlock()
_, ok := m.objects[key]
return ok
}
func writes(body string) func(string) error {
return func(dst string) error { return os.WriteFile(dst, []byte(body), 0o644) }
}
// Many readers asking for the same archive at once cost one build and
// one upload; the rest wait for it and are answered by it.
func TestBucketedBuildsOnceUnderConcurrency(t *testing.T) {
b := NewBucketed()
objs := newMemObjects("e/b")
var builds atomic.Int32
release := make(chan struct{})
build := func(dst string) error {
builds.Add(1)
<-release
return os.WriteFile(dst, []byte("tarball"), 0o644)
}
const n = 16
var wg sync.WaitGroup
var ran atomic.Int32
errs := make(chan error, n)
for range n {
wg.Add(1)
go func() {
defer wg.Done()
res, err := b.Ensure(context.Background(), objs, Job{Scope: "archive/r/", Key: "archive/r/h.tar.gz", Build: build})
if res.Built {
ran.Add(1)
}
errs <- err
}()
}
time.Sleep(50 * time.Millisecond)
close(release)
wg.Wait()
close(errs)
for err := range errs {
if err != nil {
t.Fatalf("FAIL: %v", err)
}
}
if builds.Load() != 1 || objs.uploads.Load() != 1 || ran.Load() != 1 {
t.Fatalf("FAIL: builds=%d uploads=%d built-reports=%d, want 1 each", builds.Load(), objs.uploads.Load(), ran.Load())
}
if string(objs.objects["archive/r/h.tar.gz"]) != "tarball" {
t.Fatalf("FAIL: object = %q", objs.objects["archive/r/h.tar.gz"])
}
t.Logf("ok: %d readers, 1 build, 1 upload", n)
}
// A key once seen in the bucket is not asked about again, and one the
// bucket already holds is never built.
func TestBucketedRemembersAndSkipsWhatTheBucketHolds(t *testing.T) {
b := NewBucketed()
objs := newMemObjects("e/b")
objs.objects["archive/r/h.zip"] = []byte("old")
build := func(string) error { t.Fatal("FAIL: built an archive the bucket holds"); return nil }
for range 3 {
res, err := b.Ensure(context.Background(), objs, Job{Key: "archive/r/h.zip", Build: build})
if err != nil || res.Built {
t.Fatalf("FAIL: res=%+v err=%v", res, err)
}
}
if objs.heads.Load() != 1 {
t.Fatalf("FAIL: %d HEADs, want 1: a positive answer is remembered", objs.heads.Load())
}
other := newMemObjects("e/other")
var built atomic.Int32
if _, err := b.Ensure(context.Background(), other, Job{Key: "archive/r/h.zip", Build: func(dst string) error {
built.Add(1)
return writes("x")(dst)
}}); err != nil || built.Load() != 1 {
t.Fatalf("FAIL: rebinding to another bucket did not build there: built=%d err=%v", built.Load(), err)
}
t.Log("ok")
}
// A remembered key is asked about again once it is an hour old, and
// rebuilt if the bucket no longer has it (a lifecycle rule, a person).
func TestBucketedRechecksAnHourOldMemory(t *testing.T) {
now := time.Date(2026, 10, 3, 0, 0, 0, 0, time.UTC)
nowFunc = func() time.Time { return now }
t.Cleanup(func() { nowFunc = time.Now })
b := NewBucketed()
objs := newMemObjects("e/b")
ctx := context.Background()
job := Job{Key: "k", Build: writes("v")}
if _, err := b.Ensure(ctx, objs, job); err != nil {
t.Fatal(err)
}
heads := objs.heads.Load()
now = now.Add(59 * time.Minute)
b.Ensure(ctx, objs, job)
if objs.heads.Load() != heads {
t.Fatal("FAIL: asked the bucket again within the hour")
}
objs.mu.Lock()
delete(objs.objects, "k")
objs.mu.Unlock()
now = now.Add(2 * time.Minute)
res, err := b.Ensure(ctx, objs, job)
if err != nil || !res.Built || !objs.has("k") {
t.Fatalf("FAIL: an hour-old memory of a vanished object was trusted: res=%+v err=%v", res, err)
}
t.Log("ok: rechecked after an hour, rebuilt what was gone")
}
// A failed upload hands the built file to Spill, so the reader is
// served without a second build; the bucket is then left alone for a
// minute, and every request in that minute is told so at once.
func TestBucketedUploadFailureSpillsAndRestsTheBucket(t *testing.T) {
now := time.Date(2026, 10, 3, 0, 0, 0, 0, time.UTC)
nowFunc = func() time.Time { return now }
t.Cleanup(func() { nowFunc = time.Now })
b := NewBucketed()
objs := newMemObjects("e/b")
objs.refuse = true
var spilled string
job := Job{Key: "k", Build: writes("built"), Spill: func(path string) bool {
body, _ := os.ReadFile(path)
spilled = string(body)
return false
}}
_, err := b.Ensure(context.Background(), objs, job)
var ue *UploadError
if !errors.As(err, &ue) || spilled != "built" {
t.Fatalf("FAIL: err=%v spilled=%q, want an UploadError and the built file spilled", err, spilled)
}
heads := objs.heads.Load()
if _, err := b.Ensure(context.Background(), objs, Job{Key: "k2", Build: writes("x")}); !errors.Is(err, ErrBucketDown) {
t.Fatalf("FAIL: within the minute: %v, want ErrBucketDown", err)
}
if objs.heads.Load() != heads || objs.uploads.Load() != 1 {
t.Fatal("FAIL: the resting bucket was asked again")
}
objs.refuse = false
now = now.Add(61 * time.Second)
if res, err := b.Ensure(context.Background(), objs, Job{Key: "k2", Build: writes("x")}); err != nil || !res.Built {
t.Fatalf("FAIL: after the minute: res=%+v err=%v", res, err)
}
t.Log("ok")
}
// A build that panics fails its waiters instead of hanging them.
func TestBucketedPanickingBuildReleasesWaiters(t *testing.T) {
b := NewBucketed()
objs := newMemObjects("e/b")
started := make(chan struct{})
release := make(chan struct{})
job := Job{Key: "k", Build: func(string) error {
close(started)
<-release
panic("zit fell over")
}}
leader := make(chan error, 1)
go func() { _, err := b.Ensure(context.Background(), objs, job); leader <- err }()
<-started
waiter := make(chan error, 1)
go func() {
_, err := b.Ensure(context.Background(), objs, Job{Key: "k", Build: func(string) error { return errors.New("second build") }})
waiter <- err
}()
time.Sleep(20 * time.Millisecond)
close(release)
for name, ch := range map[string]chan error{"leader": leader, "waiter": waiter} {
select {
case err := <-ch:
if err == nil {
t.Fatalf("FAIL: the %s saw success from a panicking build", name)
}
case <-time.After(5 * time.Second):
t.Fatalf("FAIL: the %s hung on a panicking build", name)
}
}
t.Log("ok")
}
// Forgetting a scope while a build for it is in flight stops that build
// from uploading, and its key is not remembered: a deleted repository's
// late build must not reappear in the bucket.
func TestBucketedForgottenScopeRefusesALateUpload(t *testing.T) {
b := NewBucketed()
objs := newMemObjects("e/b")
started := make(chan struct{})
release := make(chan struct{})
var spilled atomic.Bool
job := Job{Scope: "archive/a/", Key: "archive/a/h.zip", Build: func(dst string) error {
close(started)
<-release
return writes("late")(dst)
}, Spill: func(string) bool { spilled.Store(true); return false }}
done := make(chan error, 1)
go func() { _, err := b.Ensure(context.Background(), objs, job); done <- err }()
<-started
b.Forget("archive/a/")
close(release)
err := <-done
var ue *UploadError
if !errors.As(err, &ue) || !errors.Is(err, ErrForgotten) {
t.Fatalf("FAIL: err=%v, want an UploadError for a forgotten scope", err)
}
if objs.uploads.Load() != 0 || objs.has("archive/a/h.zip") {
t.Fatal("FAIL: a forgotten scope's build was uploaded")
}
if !spilled.Load() {
t.Fatal("FAIL: the built file was not offered to the reader who waited for it")
}
// A new repository under the same name starts clean and works.
if res, err := b.Ensure(context.Background(), objs, Job{Scope: "archive/a/", Gen: b.Gen("archive/a/"), Key: "archive/a/h.zip", Build: writes("new")}); err != nil || !res.Built {
t.Fatalf("FAIL: after Forget: res=%+v err=%v", res, err)
}
t.Log("ok")
}
// Forget waits for an upload already under way, so a sweep of the
// prefix that follows it sees that object rather than missing it.
func TestBucketedForgetWaitsForAnUploadInFlight(t *testing.T) {
b := NewBucketed()
objs := newMemObjects("e/b")
objs.gate = make(chan struct{})
done := make(chan error, 1)
go func() {
_, err := b.Ensure(context.Background(), objs, Job{Scope: "archive/a/", Key: "archive/a/h.zip", Build: writes("x")})
done <- err
}()
<-objs.uploading
forgot := make(chan struct{})
go func() { b.Forget("archive/a/"); close(forgot) }()
select {
case <-forgot:
t.Fatal("FAIL: Forget returned while an upload into the scope was in flight")
case <-time.After(50 * time.Millisecond):
}
close(objs.gate)
<-forgot
if err := <-done; err != nil || !objs.has("archive/a/h.zip") {
t.Fatalf("FAIL: err=%v", err)
}
t.Log("ok: Forget waited; the object it waited for is there to be swept")
}
// Forget drops what is remembered under a prefix, so a deleted
// repository's archives are asked about again rather than assumed.
func TestBucketedForgetsAPrefix(t *testing.T) {
b := NewBucketed()
objs := newMemObjects("e/b")
objs.objects["archive/a/h.zip"] = []byte("a")
objs.objects["archive/ab/h.zip"] = []byte("ab")
ctx := context.Background()
none := func(string) error { return errors.New("no build") }
b.Ensure(ctx, objs, Job{Scope: "archive/a/", Key: "archive/a/h.zip", Build: none})
b.Ensure(ctx, objs, Job{Scope: "archive/ab/", Key: "archive/ab/h.zip", Build: none})
b.Forget("archive/a/")
delete(objs.objects, "archive/a/h.zip")
if _, err := b.Ensure(ctx, objs, Job{Scope: "archive/a/", Key: "archive/a/h.zip", Build: none}); err == nil {
t.Fatal("FAIL: a forgotten key was still assumed present")
}
before := objs.heads.Load()
if _, err := b.Ensure(ctx, objs, Job{Scope: "archive/ab/", Key: "archive/ab/h.zip", Build: none}); err != nil || objs.heads.Load() != before {
t.Fatalf("FAIL: Forget(archive/a/) touched archive/ab/: err=%v", err)
}
t.Log("ok")
}
// A failed build uploads nothing and is not remembered; the next
// reader builds again.
func TestBucketedFailedBuildIsNotRemembered(t *testing.T) {
b := NewBucketed()
objs := newMemObjects("e/b")
_, err := b.Ensure(context.Background(), objs, Job{Key: "k", Build: func(string) error { return errors.New("boom") }})
var be *BuildError
if !errors.As(err, &be) {
t.Fatalf("FAIL: err=%v, want a BuildError", err)
}
if objs.uploads.Load() != 0 {
t.Fatal("FAIL: a failed build was uploaded")
}
res, err := b.Ensure(context.Background(), objs, Job{Key: "k", Build: writes("ok")})
if err != nil || !res.Built || res.Size != 2 {
t.Fatalf("FAIL: res=%+v err=%v", res, err)
}
t.Log("ok")
}
// A reader who leaves while the archive it asked for is uploading does
// not abort the upload others wait on, and does not rest the bucket.
func TestBucketedLeaderLeavingMidUploadKeepsTheBucket(t *testing.T) {
b := NewBucketed()
objs := newMemObjects("e/b")
objs.gate = make(chan struct{})
ctx, cancel := context.WithCancel(context.Background())
done := make(chan error, 1)
go func() { _, err := b.Ensure(ctx, objs, Job{Key: "k", Build: writes("v")}); done <- err }()
<-objs.uploading
cancel()
time.Sleep(20 * time.Millisecond)
close(objs.gate)
<-done
if !objs.has("k") {
t.Fatal("FAIL: the upload was abandoned with the reader who left")
}
res, err := b.Ensure(context.Background(), objs, Job{Key: "k", Build: writes("v")})
if err != nil || res.Built {
t.Fatalf("FAIL: the next request: res=%+v err=%v, want the held archive", res, err)
}
t.Log("ok")
}
// An upload that fails only because its context ended says nothing
// about the bucket, so the bucket is not rested.
func TestBucketedContextFailureDoesNotRestTheBucket(t *testing.T) {
for _, cause := range []error{context.Canceled, context.DeadlineExceeded} {
b := NewBucketed()
objs := newMemObjects("e/b")
objs.uploadErr = cause
if _, err := b.Ensure(context.Background(), objs, Job{Key: "k", Build: writes("v")}); err == nil {
t.Fatalf("FAIL: %v: a failed upload reported success", cause)
}
objs.uploadErr = nil
if _, err := b.Ensure(context.Background(), objs, Job{Key: "k2", Build: writes("v")}); err != nil {
t.Fatalf("FAIL: %v rested the bucket: %v", cause, err)
}
}
t.Log("ok")
}
// The generation is the caller's, read before its read check: a
// request that began before a Forget is refused even if its build
// starts after it.
func TestBucketedGenerationIsTheCallers(t *testing.T) {
b := NewBucketed()
objs := newMemObjects("e/b")
gen := b.Gen("archive/a/")
b.Forget("archive/a/")
_, err := b.Ensure(context.Background(), objs, Job{Scope: "archive/a/", Gen: gen, Key: "archive/a/h.zip", Build: writes("late")})
if !errors.Is(err, ErrForgotten) || objs.has("archive/a/h.zip") {
t.Fatalf("FAIL: err=%v: a request from before the Forget uploaded", err)
}
t.Log("ok")
}