library-go/worker_pool/worker_pool.go

72 lines
1.2 KiB
Go
Raw Permalink Normal View History

2024-08-20 15:33:20 +00:00
package workerpool
2024-08-20 15:11:32 +00:00
import (
"sync"
2024-08-20 15:33:20 +00:00
"github.com/panjf2000/ants/v2"
2024-08-20 15:11:32 +00:00
)
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
}