sift/internal/pdfpool/pdfpool_test.go

package pdfpool_test

// P5a task 04: the extraction pool. A few slots, a short queue, single-flight per key, and at
// most one job in flight per caller.

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"git.wntrmute.dev/kyle/sift/internal/pdfpool"
)

func waitUntil(t *testing.T, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(3 * time.Second)
	for !cond() {
		if time.Now().After(deadline) {
			t.Fatal("condition not reached within 3 s")
		}
		time.Sleep(5 * time.Millisecond)
	}
}

func TestSingleFlight(t *testing.T) {
	p := pdfpool.New(2, 4, time.Second)
	var calls atomic.Int32
	release := make(chan struct{})
	fn := func(ctx context.Context) (any, error) {
		calls.Add(1)
		<-release
		return "text", nil
	}
	var wg sync.WaitGroup
	results := make(chan any, 5)
	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			v, err := p.Do(context.Background(), "doc-1", "caller-"+string(rune('a'+i)), fn)
			if err != nil {
				t.Error(err)
			}
			results <- v
		}(i)
	}
	waitUntil(t, func() bool { return calls.Load() == 1 && p.Waiting("doc-1") == 4 })
	close(release)
	wg.Wait()
	close(results)
	for v := range results {
		if v != "text" {
			t.Errorf("a sharer got %v", v)
		}
	}
	if n := calls.Load(); n != 1 {
		t.Errorf("fn ran %d times for one key, want 1", n)
	}
	// Nothing is cached: the next call runs fn again.
	_, _ = p.Do(context.Background(), "doc-1", "x", func(ctx context.Context) (any, error) { calls.Add(1); return nil, nil })
	if n := calls.Load(); n != 2 {
		t.Errorf("after the flight ended, fn ran %d times in total, want 2", n)
	}
}

// Two slots and four waiting places: the seventh distinct job is refused at once.
func TestSlotsAndQueue(t *testing.T) {
	p := pdfpool.New(2, 4, 5*time.Second)
	release := make(chan struct{})
	var running atomic.Int32
	block := func(ctx context.Context) (any, error) {
		running.Add(1)
		defer running.Add(-1)
		<-release
		return nil, nil
	}
	var wg sync.WaitGroup
	for i := 0; i < 6; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			if _, err := p.Do(context.Background(), "doc-"+string(rune('a'+i)), "caller-"+string(rune('a'+i)), block); err != nil {
				t.Errorf("job %d: %v", i, err)
			}
		}(i)
	}
	waitUntil(t, func() bool { return running.Load() == 2 && p.Queued() == 4 })
	start := time.Now()
	_, err := p.Do(context.Background(), "doc-z", "caller-z", block)
	if !errors.Is(err, pdfpool.ErrBusy) {
		t.Errorf("seventh job: %v, want ErrBusy", err)
	}
	if d := time.Since(start); d > 200*time.Millisecond {
		t.Errorf("a full queue refused after %v; it must refuse at once", d)
	}
	if running.Load() > 2 {
		t.Errorf("%d jobs ran at once, want at most 2", running.Load())
	}
	close(release)
	wg.Wait()
}

// A queued job that waits longer than Wait is refused with ErrBusy and gives up its place.
func TestQueueWait(t *testing.T) {
	p := pdfpool.New(1, 1, 150*time.Millisecond)
	release := make(chan struct{})
	defer close(release)
	go func() {
		_, _ = p.Do(context.Background(), "a", "ca", func(ctx context.Context) (any, error) { <-release; return nil, nil })
	}()
	waitUntil(t, func() bool { return p.Running() == 1 })
	start := time.Now()
	_, err := p.Do(context.Background(), "b", "cb", func(ctx context.Context) (any, error) { return nil, nil })
	if !errors.Is(err, pdfpool.ErrBusy) {
		t.Fatalf("waiter: %v, want ErrBusy", err)
	}
	if d := time.Since(start); d < 100*time.Millisecond || d > time.Second {
		t.Errorf("waiter refused after %v, want about 150 ms", d)
	}
	if p.Queued() != 0 {
		t.Errorf("queued = %d after the waiter gave up, want 0", p.Queued())
	}
}

func TestCallerBusy(t *testing.T) {
	p := pdfpool.New(2, 4, time.Second)
	release := make(chan struct{})
	go func() {
		_, _ = p.Do(context.Background(), "doc-1", "straylight", func(ctx context.Context) (any, error) { <-release; return nil, nil })
	}()
	waitUntil(t, func() bool { return p.Running() == 1 })
	if _, err := p.Do(context.Background(), "doc-2", "straylight", func(ctx context.Context) (any, error) { return nil, nil }); !errors.Is(err, pdfpool.ErrCallerBusy) {
		t.Errorf("second job from the same caller: %v, want ErrCallerBusy", err)
	}
	// The same caller asking for the same document again joins the flight instead.
	done := make(chan error, 1)
	go func() {
		_, err := p.Do(context.Background(), "doc-1", "straylight", func(ctx context.Context) (any, error) { return nil, nil })
		done <- err
	}()
	waitUntil(t, func() bool { return p.Waiting("doc-1") == 1 })
	close(release)
	if err := <-done; err != nil {
		t.Errorf("joining its own flight: %v", err)
	}
	// Once the job is done the caller may start another.
	if _, err := p.Do(context.Background(), "doc-2", "straylight", func(ctx context.Context) (any, error) { return nil, nil }); err != nil {
		t.Errorf("after the first job ended: %v", err)
	}
}

func TestErrorsAndPanicsReachEverySharer(t *testing.T) {
	p := pdfpool.New(2, 4, time.Second)
	boom := errors.New("boom")
	if _, err := p.Do(context.Background(), "k", "c", func(ctx context.Context) (any, error) { return nil, boom }); !errors.Is(err, boom) {
		t.Errorf("fn error: %v", err)
	}
	_, err := p.Do(context.Background(), "k", "c", func(ctx context.Context) (any, error) { panic("kaboom") })
	if err == nil {
		t.Fatal("a panicking fn returned no error")
	}
	if p.Running() != 0 {
		t.Errorf("running = %d after a panic, want 0 (the slot leaked)", p.Running())
	}
	if _, err := p.Do(context.Background(), "k2", "c", func(ctx context.Context) (any, error) { return 1, nil }); err != nil {
		t.Errorf("the caller is still marked busy after its job panicked: %v", err)
	}
}

func TestContextCancelWhileQueued(t *testing.T) {
	p := pdfpool.New(1, 2, 5*time.Second)
	release := make(chan struct{})
	defer close(release)
	go func() {
		_, _ = p.Do(context.Background(), "a", "ca", func(ctx context.Context) (any, error) { <-release; return nil, nil })
	}()
	waitUntil(t, func() bool { return p.Running() == 1 })
	ctx, cancel := context.WithCancel(context.Background())
	errc := make(chan error, 1)
	go func() {
		_, err := p.Do(ctx, "b", "cb", func(ctx context.Context) (any, error) { return nil, nil })
		errc <- err
	}()
	waitUntil(t, func() bool { return p.Queued() == 1 })
	cancel()
	if err := <-errc; !errors.Is(err, context.Canceled) {
		t.Errorf("cancelled waiter: %v, want context.Canceled", err)
	}
	if p.Queued() != 0 {
		t.Errorf("queued = %d after cancel, want 0", p.Queued())
	}
}

// The job's own context is not the first caller's: if the caller that started a shared flight
// goes away, the sharers still get the result.
func TestFlightOutlivesItsStarter(t *testing.T) {
	p := pdfpool.New(1, 2, 5*time.Second)
	release := make(chan struct{})
	ctx, cancel := context.WithCancel(context.Background())
	jobCtxErr := make(chan error, 1)
	go func() {
		_, _ = p.Do(ctx, "doc", "first", func(jctx context.Context) (any, error) {
			<-release
			jobCtxErr <- jctx.Err()
			return "done", nil
		})
	}()
	waitUntil(t, func() bool { return p.Running() == 1 })
	res := make(chan any, 1)
	go func() {
		v, _ := p.Do(context.Background(), "doc", "second", nil)
		res <- v
	}()
	waitUntil(t, func() bool { return p.Waiting("doc") == 1 })
	cancel()
	close(release)
	if v := <-res; v != "done" {
		t.Errorf("sharer got %v after the starter left, want done", v)
	}
	if e := <-jobCtxErr; e != nil {
		t.Errorf("the job's context was cancelled with its starter: %v", e)
	}
}