Jeremy

LanguageENZH

Common Routine pool

package main

import (
	"fmt"
	"sync"
	"time"
)

type RoutinePool struct {
	PoolNum  int
	QueueNum int
	tasks    chan func()
	mtx      sync.Mutex
	wg       sync.WaitGroup
}

func NewRoutinePool(poolNum int, queueNum int) *RoutinePool {
	pool := RoutinePool{
		PoolNum:  poolNum,
		QueueNum: queueNum,
		tasks:    make(chan func(), queueNum),
		mtx:      sync.Mutex{},
		wg:       sync.WaitGroup{},
	}

	return &pool
}

func (r *RoutinePool) Submit(task func()) bool {
	if r.tasks == nil {
		return false
	}

	r.tasks <- task

	return true
}

func (r *RoutinePool) Start() {
	if r.tasks == nil {
		r.tasks = make(chan func(), r.QueueNum)
	}

	for i := 0; i < r.PoolNum; i++ {
		r.wg.Add(1)
		go func() {
			defer r.wg.Done()
			for task := range r.tasks {
				task()
			}
		}()
	}
}

func (r *RoutinePool) Close() {
	close(r.tasks)
	r.wg.Wait()
}

func main() {
	pool := NewRoutinePool(2, 10)

	pool.Start()

	for i := 0; i < 20; i++ {
		pool.Submit(func() {
			fmt.Printf("task %d\n", i)
			time.Sleep(500 * time.Millisecond)
		})
	}

	pool.Close()

}

/**
Possible Output
task 0
task 1
task 2
task 3
task 4
task 5
task 7
task 6
task 8
task 9
task 10
task 11
task 12
task 13
task 14
task 15
task 16
task 17
task 18
task 19
*/