zq.go

package zq

import (
	"git.barsukov.pro/barsukov/zgo/zatomic"
	"git.barsukov.pro/barsukov/zgo/zwg"
)

type ZQ struct {
	WorkerCounter *zatomic.Int
	WorkerFn      func(any)
	JobsWg        *zwg.Zwg

	complete chan bool
	jobs     chan any
}

func New(qtyWorkers int, w func(any)) *ZQ {
	q := &ZQ{
		WorkerCounter: zatomic.New(),
		WorkerFn:      w,
		JobsWg:        zwg.New(),
		jobs:          make(chan any),
		complete:      make(chan bool),
	}

	for i := 0; i < qtyWorkers; i++ {
		q.AddWorker(1)
	}

	return q
}

func (q *ZQ) AddWorker(delta int) {
	if delta < 1 {
		for i := 0; i < delta*-1; i++ {
			q.complete <- true
		}

		return
	}

	for i := 0; i < delta; i++ {
		q.WorkerCounter.Add(1)

		go func() {
			defer q.WorkerCounter.Add(-1)

			for {
				select {
				case <-q.complete:
					return
				case j := <-q.jobs:
					q.WorkerFn(j)
					q.JobsWg.Done()
				}
			}
		}()
	}
}

func (q *ZQ) RemoveWorker() {
	q.AddWorker(-1)
}

func (q *ZQ) AddJob(a any) *ZQ {
	q.JobsWg.Inc()

	go func() {
		q.jobs <- a
	}()

	return q
}

func (q *ZQ) AddJobSync(a any) *ZQ {
	q.JobsWg.Inc()
	q.jobs <- a

	return q
}

func (q *ZQ) Wait() {
	q.JobsWg.Wait()
}

func (q *ZQ) Quit() {
	q.Wait()
	q.AddWorker(q.WorkerCounter.LoadInverse())
}