nettools/data_structures/worker_pool_func.go

59 lines
1.4 KiB
Go

package ds_utils
import (
"context"
common "git.mic.pp.ua/anderson/nettools/common"
)
// RunWorkerPool is a single function worker pool implementation
//
// Create a channel where u will send tasks, send some tasks
// there and they will execute.
// If there are too many tasks -- they will be saved to the queue.
//
// Also this function may be extended to have resizable capacity,
// more callbacks, tasks priority, but those modifications would
// harm performance.
//
// If you need out-of-the-box more flexible approach -- use WorkerPool struct instead.
// Also this wp is worse in terms of performance. IDK why it exists, it's just sucks.
func RunWorkerPool(
ctx context.Context,
tasksChan chan common.Task,
capacity int,
errHanlder func(error) error,
) (err error) {
countBusy := 0
errChan := make(chan error, 3)
tasksQueue := []common.Task{}
for {
select {
case task := <-tasksChan:
if countBusy >= capacity {
tasksQueue = append(tasksQueue, task)
continue
}
go func() { errChan <- task(context.Background()) }()
countBusy++
case err := <-errChan:
if err != nil {
if err = errHanlder(err); err != nil {
return err
}
}
if len(tasksQueue) == 0 {
countBusy--
continue
}
task := tasksQueue[0]
tasksQueue = tasksQueue[1:]
go func() { errChan <- task(context.Background()) }()
case <-ctx.Done():
return ctx.Err()
}
}
}