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>
This commit is contained in:
ornith-1.5-35b-a3b
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user