library-go/worker_pool/worker_pool.go

72 lines
1.2 KiB
Go

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
}