sift/internal/pdfpool/pdfpool.go

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