71 lines
1.2 KiB
Go
71 lines
1.2 KiB
Go
package worker_pool
|
|
|
|
import (
|
|
"github.com/panjf2000/ants/v2"
|
|
"sync"
|
|
)
|
|
|
|
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
|
|
}
|