100 lines
2.5 KiB
Go
100 lines
2.5 KiB
Go
|
package processing
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"sync"
|
|||
|
|
|
|||
|
|
"github.com/google/uuid"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// DefaultProcessingWorkers is the in-process bound for concurrent ClaimNext+ProcessJob.
|
|||
|
|
// ClaimNext uses FOR UPDATE SKIP LOCKED so each worker gets a distinct pending job.
|
|||
|
|
const DefaultProcessingWorkers = 2
|
|||
|
|
|
|||
|
|
// MaxProcessingWorkers caps in-process job parallelism (OpenAI RPM + DB pool).
|
|||
|
|
const MaxProcessingWorkers = 8
|
|||
|
|
|
|||
|
|
// ClampProcessingWorkers bounds n to [1, MaxProcessingWorkers].
|
|||
|
|
func ClampProcessingWorkers(n int) int {
|
|||
|
|
if n < 1 {
|
|||
|
|
return 1
|
|||
|
|
}
|
|||
|
|
if n > MaxProcessingWorkers {
|
|||
|
|
return MaxProcessingWorkers
|
|||
|
|
}
|
|||
|
|
return n
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// JobSlots limits concurrent ProcessJob goroutines. Safe for multi-job parallelism
|
|||
|
|
// because ClaimNext is SKIP LOCKED. Not for same-job item parallelism (completion
|
|||
|
|
// protocol assumes a single ProcessJob owns final status).
|
|||
|
|
type JobSlots struct {
|
|||
|
|
Workers int
|
|||
|
|
sem chan struct{}
|
|||
|
|
wg sync.WaitGroup
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// NewJobSlots creates a bounded slot set for concurrent processing jobs.
|
|||
|
|
func NewJobSlots(workers int) *JobSlots {
|
|||
|
|
w := ClampProcessingWorkers(workers)
|
|||
|
|
return &JobSlots{
|
|||
|
|
Workers: w,
|
|||
|
|
sem: make(chan struct{}, w),
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Wait blocks until all in-flight ProcessJob goroutines finish.
|
|||
|
|
func (s *JobSlots) Wait() {
|
|||
|
|
s.wg.Wait()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TryStart claims one free slot (non-blocking). claim must be SKIP LOCKED–safe.
|
|||
|
|
// If claim fails, the slot is released. On success, process runs in a new goroutine.
|
|||
|
|
func (s *JobSlots) TryStart(
|
|||
|
|
ctx context.Context,
|
|||
|
|
claim func(context.Context) (uuid.UUID, error),
|
|||
|
|
process func(context.Context, uuid.UUID) error,
|
|||
|
|
onDone func(jobID uuid.UUID, err error),
|
|||
|
|
) (started bool, claimErr error) {
|
|||
|
|
select {
|
|||
|
|
case s.sem <- struct{}{}:
|
|||
|
|
default:
|
|||
|
|
return false, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
jobID, err := claim(ctx)
|
|||
|
|
if err != nil {
|
|||
|
|
<-s.sem
|
|||
|
|
return false, err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
s.wg.Add(1)
|
|||
|
|
go func(id uuid.UUID) {
|
|||
|
|
defer s.wg.Done()
|
|||
|
|
defer func() { <-s.sem }()
|
|||
|
|
procErr := process(ctx, id)
|
|||
|
|
if onDone != nil {
|
|||
|
|
onDone(id, procErr)
|
|||
|
|
}
|
|||
|
|
}(jobID)
|
|||
|
|
return true, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Fill starts jobs until all free slots are occupied or claim returns an error
|
|||
|
|
// (including pgx.ErrNoRows when the queue is empty). Each tick should call Fill
|
|||
|
|
// once so ClaimNext fills up to Workers concurrent ProcessJob goroutines.
|
|||
|
|
func (s *JobSlots) Fill(
|
|||
|
|
ctx context.Context,
|
|||
|
|
claim func(context.Context) (uuid.UUID, error),
|
|||
|
|
process func(context.Context, uuid.UUID) error,
|
|||
|
|
onDone func(jobID uuid.UUID, err error),
|
|||
|
|
) (started int, lastErr error) {
|
|||
|
|
for {
|
|||
|
|
ok, err := s.TryStart(ctx, claim, process, onDone)
|
|||
|
|
if !ok {
|
|||
|
|
return started, err
|
|||
|
|
}
|
|||
|
|
started++
|
|||
|
|
}
|
|||
|
|
}
|