nettools/data_structures/worker_pool_func.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()
		}
	}
}