executor.go

package executor

import (
	"context"
	"sync"
)

const (
	defaultNumOfGoroutines = 10
)

type WorkerService struct {
	ctx        context.Context
	pool       chan *worker
	workersNum int
	start      *sync.Once
	stop       *sync.Once
	exit       chan struct{}
}

func NewService(ctx context.Context, workerNums int) WorkerPool {
	if workerNums <= 0 {
		workerNums = defaultNumOfGoroutines
	}

	return &WorkerService{
		ctx:        ctx,
		pool:       make(chan *worker, workerNums),
		workersNum: workerNums,
		start:      &sync.Once{},
		stop:       &sync.Once{},
		exit:       make(chan struct{}),
	}
}

func (s *WorkerService) Start() {
	s.start.Do(func() {
		for i := 0; i < s.workersNum; i++ {
			s.pool <- newWorker(i + 1)
		}
	})
}

func (s *WorkerService) Stop() {
	s.stop.Do(func() {
		close(s.exit)
	})
}

func (s *WorkerService) AddTasks(tasks ...Runnable) {
	for i := 0; i < len(tasks); i++ {
		select {
		case <-s.ctx.Done():
			return
		case <-s.exit:
			return
		case w := <-s.pool:
			go s.process(w, tasks[i])
		}
	}
}

func (s *WorkerService) process(w *worker, task Runnable) {
	select {
	case <-s.ctx.Done():
		return
	case <-s.exit:
		return
	default:
		defer s.backWorkerToPool(w)

		if err := task.Execute(); err != nil {
			task.OnError(err)
			return
		}

		if err := task.OnComplete(); err != nil {
			task.OnError(err)
			return
		}
	}
}

func (s *WorkerService) backWorkerToPool(w *worker) {
	s.pool <- w
}