package workerpool import ( "sync" "github.com/panjf2000/ants/v2" ) const defaultWorkerPoolSize = 2000 type WorkerPool interface { Submit(task func()) error SubmitAndWaitAll(tasks ...func() error) (taskErr chan error, submitErr error) } type workerPool struct { p *ants.Pool } func NewWorkerPool(size int) WorkerPool { if size <= 0 { size = defaultWorkerPoolSize } p, err := ants.NewPool( size, ants.WithDisablePurge(true), ) if err != nil { return &workerPool{p: nil} } return &workerPool{p: p} } func (p *workerPool) Submit(task func()) error { if p.p == nil { return ants.Submit(task) } return p.p.Submit(task) } func (p *workerPool) SubmitAndWaitAll(tasks ...func() error) (chan error, error) { taskErrCh := make(chan error, len(tasks)) submitErrCh := make(chan error, len(tasks)) wg := sync.WaitGroup{} wg.Add(len(tasks)) for i := range tasks { task := tasks[i] err := p.Submit(func() { defer wg.Done() if err := task(); err != nil { taskErrCh <- err } }) if err != nil { submitErrCh <- err wg.Done() } } wg.Wait() if len(submitErrCh) != 0 { return nil, <-submitErrCh } return taskErrCh, nil }