pdfpool: a few extraction slots, a short queue, single-flight, one job per caller

Generic in-process scheduler for the expensive PDF-extraction child process. A pool runs at most `slots` jobs, keeps at most `queue` waiting, and refuses a waiter past `wait`. Do() applies six rules in order: join an existing running or queued job (no slot, allowed while the caller holds another, fn may be nil); refuse a caller that already has a job in flight; take a free slot, else queue (first-come-first-served, else ErrBusy); run fn under context.WithoutCancel so the flight outlives a caller that leaves; broadcast the result to every sharer with no caching and a recovered panic returned as an error. Running counts only jobs holding a slot; Queued counts the line; Waiting(key) counts joiners. Co-Authored-By: ornith-1.5-35b-a3b <ornith-1.5-35b-a3b@llama.cpp.invalid>
2026-09-25 23:32:16 -07:00
parent b93b7d5816
commit bb9a6544d1
3 changed files with 525 additions and 0 deletions
@@ -0,0 +1,298 @@
+// Package pdfpool is a small in-process scheduler for expensive, stateless work: a PDF
+// extraction that spawns a child process and can consume a few hundred MB. It bounds how many
+// jobs run at once, keeps a short waiting line, coalesces identical requests into one job, and
+// lets a single caller hold at most one job. It knows nothing about PDFs.
+package pdfpool
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+)
+
+var (
+	ErrBusy       = errors.New("pdf workers busy") // queue full, or waited longer than Wait
+	ErrCallerBusy = errors.New("busy")             // this caller already has a job in flight
+)
+
+// Pool runs at most `slots` jobs at once, keeps at most `queue` jobs waiting, and refuses a
+// waiter that has waited longer than the deadline set at construction.
+type Pool struct {
+	slots int
+	queue int
+	wait  time.Duration
+
+	mu         sync.Mutex
+	jobs       map[string]*job // key -> job currently in flight (running or queued)
+	callerJobs map[string]*job // caller -> the job they started
+	waiters    []*waiter       // FIFO line of jobs waiting for a slot
+	running    int             // jobs currently holding a slot
+}
+
+// job is one extraction in flight for a key. Its result is broadcast to every sharer.
+type job struct {
+	key    string
+	caller string
+	fn     func(context.Context) (any, error)
+
+	pool *Pool
+
+	mu    sync.Mutex
+	done  chan struct{} // closed once the result is ready
+	res   any
+	err   error
+	count int // sharers waiting on this job, not counting the starter
+}
+
+func (j *job) run(ctx context.Context) (any, error) {
+	v, err := j.invoke(context.WithoutCancel(ctx))
+	j.pool.finish(j, v, err)
+	return v, err
+}
+
+func (j *job) invoke(ctx context.Context) (v any, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			v, err = nil, fmt.Errorf("pdfpool: job panicked: %v", r)
+		}
+	}()
+	return j.fn(ctx)
+}
+
+// finish broadcasts the result, frees the slot and the caller, and hands the slot to the next
+// waiter. It is called exactly once per job.
+func (p *Pool) finish(j *job, res any, err error) {
+	p.mu.Lock()
+	p.finishLocked(j, res, err)
+	p.mu.Unlock()
+}
+
+// finishLocked releases every sharer, frees the slot and the caller, and grants the freed slot,
+// assuming p.mu is held.
+func (p *Pool) finishLocked(j *job, res any, err error) {
+	j.broadcast(res, err)
+	if cur, ok := p.jobs[j.key]; ok && cur == j {
+		delete(p.jobs, j.key)
+	}
+	if cur, ok := p.callerJobs[j.caller]; ok && cur == j {
+		delete(p.callerJobs, j.caller)
+	}
+	p.running--
+	p.grant()
+}
+
+// broadcast records the result and releases every sharer waiting on the job.
+func (j *job) broadcast(res any, err error) {
+	j.mu.Lock()
+	j.res = res
+	j.err = err
+	close(j.done)
+	j.mu.Unlock()
+}
+
+// abandon releases the sharers of a queued job that never took a slot and drops it from the pool.
+// It leaves the slot count untouched, since the job never held one. Must be called with p.mu held.
+func (p *Pool) abandon(j *job, err error) {
+	j.broadcast(nil, err)
+	if cur, ok := p.jobs[j.key]; ok && cur == j {
+		delete(p.jobs, j.key)
+	}
+	if cur, ok := p.callerJobs[j.caller]; ok && cur == j {
+		delete(p.callerJobs, j.caller)
+	}
+}
+
+// startLocked registers a new job for key under caller. It does not take a slot; a caller must
+// call takeSlotLocked to have the job run, or grant/abandon to settle it. Must be called with
+// p.mu held.
+func (p *Pool) startLocked(key, caller string, fn func(context.Context) (any, error)) *job {
+	j := &job{
+		key:    key,
+		caller: caller,
+		fn:     fn,
+		pool:   p,
+		done:   make(chan struct{}),
+	}
+	p.jobs[key] = j
+	p.callerJobs[caller] = j
+	return j
+}
+
+// takeSlotLocked registers a new job and has it hold a slot now. Must be called with p.mu held.
+func (p *Pool) takeSlotLocked(key, caller string, fn func(context.Context) (any, error)) *job {
+	j := p.startLocked(key, caller, fn)
+	p.running++
+	return j
+}
+
+// grant starts the first still-waiting job in line while slots are free. Must be called with
+// p.mu held.
+func (p *Pool) grant() {
+	for len(p.waiters) > 0 {
+		w := p.waiters[0]
+		if w.claimStart() {
+			p.waiters = p.waiters[1:]
+			p.running++
+			return
+		}
+		// The waiter gave up its place; drop it.
+		p.waiters = p.waiters[1:]
+	}
+}
+
+// removeWaiter drops w from the line. Must be called with p.mu held.
+func (p *Pool) removeWaiter(w *waiter) {
+	for i := range p.waiters {
+		if p.waiters[i] == w {
+			p.waiters = append(p.waiters[:i], p.waiters[i+1:]...)
+			return
+		}
+	}
+}
+
+// waiter is a job waiting in line for a slot.
+type waiter struct {
+	job   *job
+	start chan struct{} // closed when the job is granted a slot
+	mu    sync.Mutex
+	done  bool // claimStart or claimExit has won
+}
+
+func (w *waiter) claimStart() bool {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.done {
+		return false
+	}
+	w.done = true
+	close(w.start)
+	return true
+}
+
+func (w *waiter) claimExit() bool {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.done {
+		return false
+	}
+	w.done = true
+	return true
+}
+
+// New returns a pool with `slots` jobs running at once, at most `queue` waiting, and a waiter
+// refused after `wait`.
+func New(slots, queue int, wait time.Duration) *Pool {
+	return &Pool{
+		slots:      slots,
+		queue:      queue,
+		wait:       wait,
+		jobs:       make(map[string]*job),
+		callerJobs: make(map[string]*job),
+	}
+}
+
+// Do runs fn under the pool and returns its result.
+func (p *Pool) Do(ctx context.Context, key, caller string, fn func(context.Context) (any, error)) (any, error) {
+	p.mu.Lock()
+
+	// 1. Join: a job for this key is already running or queued. Sharing takes no slot and is
+	// allowed even if the caller holds another job.
+	if j, ok := p.jobs[key]; ok {
+		j.mu.Lock()
+		j.count++
+		j.mu.Unlock()
+		p.mu.Unlock()
+		defer func() {
+			j.mu.Lock()
+			j.count--
+			j.mu.Unlock()
+		}()
+		return j.wait(ctx)
+	}
+
+	// 2. Caller: this caller already has a job in flight for another key.
+	if _, busy := p.callerJobs[caller]; busy {
+		p.mu.Unlock()
+		return nil, ErrCallerBusy
+	}
+
+	// 3. Slot: start at once, else line up, else refuse.
+	if p.running < p.slots {
+		j := p.takeSlotLocked(key, caller, fn)
+		p.mu.Unlock()
+		return j.run(ctx)
+	}
+	if len(p.waiters) >= p.queue {
+		p.mu.Unlock()
+		return nil, ErrBusy
+	}
+
+	w := &waiter{job: p.startLocked(key, caller, fn), start: make(chan struct{})}
+	p.waiters = append(p.waiters, w)
+	timer := time.NewTimer(p.wait)
+	defer timer.Stop()
+	p.mu.Unlock()
+
+	select {
+	case <-w.start:
+		return w.job.run(ctx)
+	case <-timer.C:
+		if w.claimExit() {
+			j := w.job
+			p.mu.Lock()
+			p.removeWaiter(w)
+			p.abandon(j, ErrBusy)
+			p.mu.Unlock()
+			return nil, ErrBusy
+		}
+		// The slot was granted while the timer fired; run the job the grant started.
+		return w.job.run(ctx)
+	case <-ctx.Done():
+		if w.claimExit() {
+			j := w.job
+			p.mu.Lock()
+			p.removeWaiter(w)
+			p.abandon(j, ctx.Err())
+			p.mu.Unlock()
+			return nil, ctx.Err()
+		}
+		return w.job.run(ctx)
+	}
+}
+
+func (j *job) wait(ctx context.Context) (any, error) {
+	select {
+	case <-j.done:
+		j.mu.Lock()
+		defer j.mu.Unlock()
+		return j.res, j.err
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}
+
+func (p *Pool) Running() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.running
+}
+
+func (p *Pool) Queued() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return len(p.waiters)
+}
+
+func (p *Pool) Waiting(key string) int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	j := p.jobs[key]
+	if j == nil {
+		return 0
+	}
+	j.mu.Lock()
+	defer j.mu.Unlock()
+	return j.count
+}
@@ -0,0 +1,226 @@
+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)
+	}
+}