frostfs-node/pkg/util/worker_pool.go

package util

import (
	"context"
	"sync"
	"sync/atomic"

	"github.com/panjf2000/ants/v2"
)

// WorkerPool represents a tool to control
// the execution of go-routine pool.
type WorkerPool interface {
	// Submit queues a function for execution
	// in a separate routine.
	//
	// Implementation must return any error encountered
	// that prevented the function from being queued.
	Submit(func()) error

	// Release releases worker pool resources. All `Submit` calls will
	// finish with ErrPoolClosed. It doesn't wait until all submitted
	// functions have returned so synchronization must be achieved
	// via other means (e.g. sync.WaitGroup).
	Release()
}

// pseudoWorkerPool represents a pseudo worker pool which executes the submitted job immediately in the caller's routine.
type pseudoWorkerPool struct {
	closed atomic.Bool
}

// ErrPoolClosed is returned when submitting task to a closed pool.
var ErrPoolClosed = ants.ErrPoolClosed

// NewPseudoWorkerPool returns a new instance of a synchronous worker pool.
func NewPseudoWorkerPool() WorkerPool {
	return &pseudoWorkerPool{}
}

// Submit executes the passed function immediately.
//
// Always returns nil.
func (p *pseudoWorkerPool) Submit(fn func()) error {
	if p.closed.Load() {
		return ErrPoolClosed
	}

	fn()

	return nil
}

// Release implements the WorkerPool interface.
func (p *pseudoWorkerPool) Release() {
	p.closed.Store(true)
}

type WorkerTask func(ctx context.Context) error

type WorkerPoolSubmitError struct {
	err error
}

func (e *WorkerPoolSubmitError) Error() string {
	return e.err.Error()
}

func (e *WorkerPoolSubmitError) Unwrap() error {
	return e.err
}

// ExecuteWithWorkerPool runs tasks in parallel using a pool and waits for all
// tasks to be complete.
//
// Returns [WorkerPoolSubmitError] when it couldn't submit a task.
func ExecuteWithWorkerPool(ctx context.Context, pool WorkerPool, tasks []WorkerTask) error {
	taskCtx, taskCancel := context.WithCancelCause(ctx)
	defer taskCancel(nil)

	var wg sync.WaitGroup

loop:
	for _, task := range tasks {
		select {
		case <-ctx.Done():
			taskCancel(context.Cause(ctx))
			break loop
		default:
		}

		wg.Add(1)
		if err := pool.Submit(func() {
			defer wg.Done()

			if err := task(taskCtx); err != nil {
				taskCancel(err)
			}
		}); err != nil {
			wg.Done()
			taskCancel(err)
			wg.Wait()

			return &WorkerPoolSubmitError{err}
		}
	}
	wg.Wait()

	return context.Cause(taskCtx)
}