126 lines
3.6 KiB
Go
126 lines
3.6 KiB
Go
// Package pool provides a protocol-agnostic recursive worker pool.
|
|
//
|
|
// It replaces the three near-identical semaphore+WaitGroup
|
|
// implementations that existed in the FTP, SFTP, and WebDAV packages
|
|
// with a single generic type. Each protocol supplies its own task
|
|
// type and handler function; the pool takes care of concurrency
|
|
// limiting, goroutine lifecycle, and error collection.
|
|
//
|
|
// Tasks may submit additional tasks from inside the handler (tree
|
|
// recursion). Submit never blocks the caller—each task gets its
|
|
// own goroutine that waits for a semaphore slot—so recursive
|
|
// submissions cannot deadlock even when all pool slots are occupied.
|
|
package pool
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
)
|
|
|
|
// ctxKey is an unexported type for context value keys in this package.
|
|
type ctxKey struct{}
|
|
|
|
// Pool manages a bounded set of goroutines that execute tasks of
|
|
// type T. Tasks are submitted via Submit and executed by the
|
|
// handler function provided to New. The pool is safe for
|
|
// concurrent use from multiple goroutines.
|
|
type Pool[T any] struct {
|
|
handler func(ctx context.Context, task T)
|
|
|
|
sem chan struct{}
|
|
wg sync.WaitGroup
|
|
errMu sync.Mutex
|
|
errs []error
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
}
|
|
|
|
// New creates a pool with max parallelism workers. handler is
|
|
// called for every submitted task, running in its own goroutine
|
|
// bounded by the semaphore. The returned pool is ready to use.
|
|
//
|
|
// ctx is the parent context for all handler goroutines. When ctx
|
|
// is cancelled, waiting and running handlers observe the
|
|
// cancellation via the context passed to them.
|
|
//
|
|
// The context passed to handler carries a reference to the pool
|
|
// itself, retrievable via Get. This allows handlers to submit
|
|
// child tasks without an explicit closure capture.
|
|
func New[T any](ctx context.Context, maxParallelism int, handler func(ctx context.Context, task T)) *Pool[T] {
|
|
if maxParallelism < 1 {
|
|
maxParallelism = 1
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(ctx)
|
|
|
|
return &Pool[T]{
|
|
handler: handler,
|
|
sem: make(chan struct{}, maxParallelism),
|
|
cancel: cancel,
|
|
ctx: ctx,
|
|
}
|
|
}
|
|
|
|
// Get retrieves the pool of type *Pool[T] from the context.
|
|
// It returns nil when the context does not carry a pool (for
|
|
// example, when called outside a handler goroutine).
|
|
func Get[T any](ctx context.Context) *Pool[T] {
|
|
v, _ := ctx.Value(ctxKey{}).(*Pool[T])
|
|
return v
|
|
}
|
|
|
|
// Submit enqueues a task for execution. Each submitted task gets
|
|
// its own goroutine that will wait for a semaphore slot before
|
|
// running the handler. Submit may be called from inside a handler
|
|
// (for recursive tasks).
|
|
//
|
|
// Submit is a no-op if the pool has already been waited on.
|
|
func (p *Pool[T]) Submit(task T) {
|
|
// After Wait returns the context is cancelled; skip work that
|
|
// can never run.
|
|
select {
|
|
case <-p.ctx.Done():
|
|
return
|
|
default:
|
|
}
|
|
|
|
p.wg.Add(1)
|
|
go func() {
|
|
defer p.wg.Done()
|
|
select {
|
|
case p.sem <- struct{}{}:
|
|
defer func() { <-p.sem }()
|
|
handlerCtx := context.WithValue(p.ctx, ctxKey{}, p)
|
|
p.handler(handlerCtx, task)
|
|
case <-p.ctx.Done():
|
|
return
|
|
}
|
|
}()
|
|
}
|
|
|
|
// Wait blocks until all submitted tasks (including tasks submitted
|
|
// by handlers) have completed. It returns all errors collected
|
|
// during execution. After Wait returns the pool must not be used.
|
|
func (p *Pool[T]) Wait() []error {
|
|
p.wg.Wait()
|
|
p.cancel()
|
|
|
|
p.errMu.Lock()
|
|
defer p.errMu.Unlock()
|
|
if len(p.errs) == 0 {
|
|
return nil
|
|
}
|
|
return p.errs
|
|
}
|
|
|
|
// AddError records a task-level error. It is safe to call from
|
|
// any goroutine, typically from inside a handler.
|
|
func (p *Pool[T]) AddError(err error) {
|
|
if err == nil {
|
|
return
|
|
}
|
|
p.errMu.Lock()
|
|
p.errs = append(p.errs, err)
|
|
p.errMu.Unlock()
|
|
}
|