Something went wrong. Try again.
forked from charmbracelet/soft-serve
Something went wrong. Try again.
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495package sync
import ( "context" "sync"
"golang.org/x/sync/semaphore")
// WorkPool is a pool of work to be done.type WorkPool struct { workers int work sync.Map sem *semaphore.Weighted ctx context.Context logger func(string, ...interface{})}
// WorkPoolOption is a function that configures a WorkPool.type WorkPoolOption func(*WorkPool)
// WithWorkPoolLogger sets the logger to use.func WithWorkPoolLogger(logger func(string, ...interface{})) WorkPoolOption { return func(wq *WorkPool) { wq.logger = logger }}
// NewWorkPool creates a new work pool. The workers argument specifies the// number of concurrent workers to run the work.// The queue will chunk the work into batches of workers size.func NewWorkPool(ctx context.Context, workers int, opts ...WorkPoolOption) *WorkPool { wq := &WorkPool{ workers: workers, ctx: ctx, }
for _, opt := range opts { opt(wq) }
if wq.workers <= 0 { wq.workers = 1 }
wq.sem = semaphore.NewWeighted(int64(wq.workers))
return wq}
// Run starts the workers and waits for them to finish.func (wq *WorkPool) Run() { wq.work.Range(func(key, value any) bool { id := key.(string) fn := value.(func()) if err := wq.sem.Acquire(wq.ctx, 1); err != nil { wq.logf("workpool: %v", err) return false }
go func(id string, fn func()) { defer wq.sem.Release(1) fn() wq.work.Delete(id) }(id, fn)
return true })
if err := wq.sem.Acquire(wq.ctx, int64(wq.workers)); err != nil { wq.logf("workpool: %v", err) }}
// Add adds a new job to the pool.// If the job already exists, it is a no-op.func (wq *WorkPool) Add(id string, fn func()) { if _, ok := wq.work.Load(id); ok { return } wq.work.Store(id, fn)}
// Status checks if a job is in the queue.func (wq *WorkPool) Status(id string) bool { _, ok := wq.work.Load(id) return ok}
func (wq *WorkPool) logf(format string, args ...interface{}) { if wq.logger != nil { wq.logger(format, args...) }}