golang实现简单的并发任务消费

1,600 阅读2分钟

又是一个golang func的简单运用,不得不说这个语法糖真的很棒。

并发任务消费的本质其实就是多个自旋的goroutine来监听多个func 类型的chan,当收到从chan传递过来的func后,就执行之。

核心代码

type Work func()// 重新定义func的类型

var workChan []chan Work// 接收工作的chan
var wg = &sync.WaitGroup{}

func Start(size int) {// size规定了自旋的goroutine的个数
	for i := 0; i < size; i++ {
		workChan[i] = make(chan Work)
		go process(i)
	}
}
func process(i int) {
	wg.Add(1)
	defer wg.Done()
	if works, ok := workChan[i]; ok {
		for {
			work, ok := <-works// 收到工作
			if ok {
				work()// 执行工作
			}
		}
	}
}

通过定义一个slice类型的chan Work,使用多个goroutine来接收chan的工作,收到工作便执行。同时使用sync.WaitGroup来保证goroutine的无限自旋。

添加工作

外部发送工作其实就是往chan内发送工作的逻辑,考虑到为外部提供更为良好的使用,于是简单封装一下,如下

type Work func()

type workerGroup struct {
	workChan map[int]chan Work
	*sync.WaitGroup
	isStop bool
}

func NewWorkerGroup() *workerGroup {
	wg := &workerGroup{WaitGroup: &sync.WaitGroup{}}
	wg.workChan = make(map[int]chan Work)
	return wg
}
// 发送任务
func (w *workerGroup) SendWork(work Work, i int) error {
	if !w.isStop {
		if works, ok := w.workChan[i]; ok {
			works <- work
			return nil
		} else {
			return errors.New("error i")
		}
	}else{
		return errors.New("the worker group has been closed")
	}
}
func (w *workerGroup) process(i int) {
	w.Add(1)
	defer w.Done()
	if works, ok := w.workChan[i]; ok {
		for {
			work, ok := <-works
			if ok {
				work()
			}
		}
	}
}
func (w *workerGroup) Start(size int) {
	for i := 0; i < size; i++ {
		w.workChan[i] = make(chan Work)
		go w.process(i)
	}
}
// 监听退出
func (w *workerGroup) OnStop() {
	go func() {
		for {
			sig := make(chan os.Signal, 1)
			signal.Notify(sig, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT)
			select {
			case <-sig:
				w.isStop = true
				func() {
					for _, works := range w.workChan {
						if len(works) <=0 {
							close(works)
						}
					}
				}()
			}
			break
		}
		log.Println("worker group succeed to close")
		os.Exit(1)
	}()
	w.Wait()
}

外部代码需要使用到这个模块的时候,只需要创建工作组,开启一下,调用sendwork来发送工作即可,如下

func main() {
	wg := NewWorkerGroup()
	wg.Start(6)
	tick := time.Tick(time.Second)
	wg.OnStop()
	wg.SendWork(func() {
		fmt.Println("work")
	}, 5)
}

总结

以上代码只是简单的使用,并没有考虑很多的异常情况,比如在SendWork函数中,若是chan里面已经存在值了,那么sendwork就会被阻塞;其次,若是系统退出时,goroutine中存在未完成的工作,也会一并退出;此外,若是工作具有返回值,需要通知结果,是直接逐层返回结果还是用异步通知回调函数的方式,都有待考量。还有更多的问题,欢迎拍砖指正。