library-go/worker_pool/worker_pool.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
}