package zq import ( "fmt" "git.barsukov.pro/barsukov/zgo/zatomic" "git.barsukov.pro/barsukov/zgo/zquit" "runtime" ) type ZQ struct { WorkerCounter *zatomic.Int WorkerFn func(int, any) ZQuit *zquit.ZQuit Jobs chan any workerID *zatomic.Int complete chan bool } func New(qtyWorkers int, buff int, fn func(int, any)) *ZQ { q := &ZQ{ WorkerCounter: zatomic.New(), WorkerFn: fn, ZQuit: zquit.Default(), Jobs: make(chan any, buff), complete: make(chan bool), workerID: zatomic.New(), } q.AddWorker(qtyWorkers) return q } func Default(fn func(int, any)) *ZQ { return DefaultN(runtime.NumCPU(), fn) } func DefaultN(qtyWorkers int, fn func(int, any)) *ZQ { return New(qtyWorkers, qtyWorkers*2, fn) } func Default2(fn func(int, any)) *ZQ { return DefaultN(2, fn) } func Default4(fn func(int, any)) *ZQ { return DefaultN(4, fn) } func Default6(fn func(int, any)) *ZQ { return DefaultN(6, fn) } func Default8(fn func(int, any)) *ZQ { return DefaultN(8, fn) } func Default12(fn func(int, any)) *ZQ { return DefaultN(12, fn) } func Default16(fn func(int, any)) *ZQ { return DefaultN(16, fn) } func Default24(fn func(int, any)) *ZQ { return DefaultN(24, fn) } func Default32(fn func(int, any)) *ZQ { return DefaultN(32, fn) } func Default48(fn func(int, any)) *ZQ { return DefaultN(48, fn) } func Default64(fn func(int, any)) *ZQ { return DefaultN(64, fn) } func (q *ZQ) WithZQuit(quit *zquit.ZQuit) *ZQ { q.ZQuit = quit 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.Inc() go func() { defer q.WorkerCounter.Dec() workerID := q.workerID.Inc() for { select { case <-q.complete: return case j := <-q.Jobs: q.WorkerFn(workerID, j) q.ZQuit.Done() } } }() } } func (q *ZQ) RemoveWorker() { q.AddWorker(-1) } func (q *ZQ) AddJob(a any) error { if q.ZQuit.IsQuit() { return fmt.Errorf("quit") } q.ZQuit.Inc() q.Jobs <- a return nil } func (q *ZQ) AddJobForce(a any) error { q.ZQuit.Inc() q.Jobs <- a return nil } func (q *ZQ) Wait() { q.ZQuit.Wait() } func (q *ZQ) Quit() { q.Wait() q.AddWorker(q.WorkerCounter.LoadInverse()) }