优雅の使用sync.WaitGroup

Gracefully Use sync.WaitGroup

Posted by Hawkingrei on August 26, 2017

“keep it simple, stupid”

Background

自从上次参加2017 GopherChina被安利了NSQ后,阅读了NSQ的源代码,从中学到了不少代码技巧。于是乎,我就把这些代码技巧运用到了veda上,提高了代码质量。

WaitGroup介绍

WaitGroup用来等待一群goroutine结束,主Goroutine调用Add来设置有多少goroutines需要去等待。然后各个goroutines开始运行,当结束时,调用Done。同时,可以使用Wait来堵塞程序,直到全部Goroutines已经结束。下面是一个WaitGroup的小例子。

var wg sync.WaitGroup
var urls = []string{
        "http://www.golang.org/",
        "http://www.google.com/",
        "http://www.somestupidname.com/",
}
for _, url := range urls {
        // Increment the WaitGroup counter.
        wg.Add(1)
        // Launch a goroutine to fetch the URL.
        go func(url string) {
                // Decrement the counter when the goroutine completes.
                defer wg.Done()
                // Fetch the URL.
                http.Get(url)
        }(url)
}
// Wait for all HTTP fetches to complete.
wg.Wait()

好写法

一个子长度最佳长度在50~150行。而在使用中,按照上面官方的例子中的写法,开发者需要自己在启动前加上Add,结束后加上Done,无形中增加代码行数。其次在启动子goroutines时,如果启动的goroutines逻辑比较复杂,可以单独写一个子函数,提高代码的可读性。

在NSQ中有一个很好的写法,首先要实现一个代理类,把WaitGroup的行为封装进去:

package util

import (
	"sync"
)

type WaitGroupWrapper struct {
	sync.WaitGroup
}

func (w *WaitGroupWrapper) Wrap(cb func()) {
	w.Add(1)
	go func() {
		cb()
		w.Done()
	}()
}

使用的时候(完整代码):

type NSQD struct {
	waitGroup            util.WaitGroupWrapper
}

func New(opts *Options) *NSQD {
	n := &NSQD{
		startTime:            time.Now(),
		topicMap:             make(map[string]*Topic),
		exitChan:             make(chan int),
		notifyChan:           make(chan interface{}),
		optsNotificationChan: make(chan struct{}, 1),
		dl:                   dirlock.New(dataPath),
	}
}

func (n *NSQD) Main() {
	n.waitGroup.Wrap(func() {
		protocol.TCPServer(n.tcpListener, tcpServer,	n.logf)
	})
	n.waitGroup.Wrap(func() { n.queueScanLoop() })
	n.waitGroup.Wrap(func() { n.lookupLoop() })
}

func (n *NSQD) Exit() {
	n.waitGroup.Wait()
}

于是乎,golang的WaitGroup现在只要Wrap和Wait就能完成,是不是干净简单很多了呢?