1.介绍
在 Go 语言中,我们可以使用 errgroup 库处理 goroutine 中的错误。
errgroup 库最近更新了,新增支持限制并发数量的功能。
本文我们介绍 errgroup 库的使用方式和实现原理。
2.使用方式
errgroup 库使用非常简单,我们通过三个简单示例代码,分别介绍三种使用方式。
基础使用
func main() { eg := errgroup.Group{} eg.Go(func() error { fmt.Println("go1") return nil }) eg.Go(func() error { fmt.Println("go2") err := errors.New("go2 err") return err }) err := eg.Wait() if err != nil { fmt.Println("err =", err) } }
阅读上面这段代码,我们使用 errgroup 库的 Go() 方法启动两个 goroutine,分别模拟错误 goroutine 和正常 goroutine。
然后,使用 errgroup 库的 Wait() 方法判断是否有 goroutine 返回错误信息。
附加 cancel 功能
func main() { eg, ctx := errgroup.WithContext(context.Background()) eg.Go(func() error { time.Sleep(1 * time.Second) select { case <-ctx.Done(): fmt.Println("go1 cancel, err = ", ctx.Err()) default: fmt.Println("go1 run") } return nil }) eg.Go(func() error { err := errors.New("go2 err") return err }) err := eg.Wait() if err != nil { fmt.Println("err =", err) } }
阅读上面这段代码,我们使用 errgroup 库的 WithContext() 函数,可以附加 cancel 功能。
我们在第一个使用 Go() 方法启动的协程函数中,使用 select ... case ... default 监听其他协程是否返回错误并做出相应的逻辑处理。
限制并发数量
func main() { eg := errgroup.Group{} eg.SetLimit(2) eg.TryGo(func() error { fmt.Println("go1 run") return nil }) eg.TryGo(func() error { err := errors.New("go2 err") return err }) eg.TryGo(func() error { fmt.Println("go3 run") return nil }) err := eg.Wait() if err != nil { fmt.Println("err =", err) } }
阅读上面这段代码,我们使用 errgroup 库新增的限制并发数量的功能。
首先,使用 SetLimit() 方法设置并发数量,然后使用 TryGo() 方法替换 Go() 方法。
3.实现原理
我们通过阅读 errgroup 库的源码,简单介绍 errgroup 的实现原理。
我们先阅读 Group 结构体的源码。
type Group struct { cancel func() wg sync.WaitGroup sem chan token errOnce sync.Once err error }
在源码中,我们可以发现 Group 结构体包含的 5 个字段,其中 sem 字段是最近为了实现限制并发数量功能而新增的。
通过 Group 结构体的字段,我们可以看出 errgroup 实际上是对 sync 和 context 的封装。
其中,cancel 是使用 context 的 cancel 方法;wg 是使用 sync.WairGroup 的相关方法;sem 是通过 channel 实现控制并发数量;errOnce 是使用 sync.Once 的特性,只保存第一个返回的 goroutine 错误;err 是 goroutine 返回的错误。
func (g *Group) Go(f func() error) { if g.sem != nil { g.sem <- token{} } g.wg.Add(1) go func() { defer g.done() if err := f(); err != nil { g.errOnce.Do(func() { g.err = err if g.cancel != nil { g.cancel() } }) } }() }
我们阅读 errgroup 库的 Go() 方法,首先,通过判断 g.sem 的值是否是 nil,如果 g.sem 的值不是 nil,说明已设置并发数量,就通过向 g.sem 中发送一个空结构体 token{},来抢占资源。
如果抢到资源,就启动一个 goroutine,否则,就阻塞,等待其他正在执行的 goroutine 释放一个资源。
细心的读者可能已经发现,Go() 方法除了开头新增判断 g.sem 的值是否为 nil 的逻辑代码之外,defer 也发生了变化,由之前的直接调用 sync.WaitGroup 的 Done() 方法,改为调用 errgroup 库新增的 done() 方法。
done() 方法源码:
func (g *Group) done() { if g.sem != nil { <-g.sem } g.wg.Done() }
通过阅读 done() 方法的源码,我们可以发现,在调用 sync.WaitGroup 的 Done() 方法之前,先判断 g.sem 的值是否是 nil,如果不是 nil,则释放资源。
我们再阅读 Wait() 方法的源码:
func (g *Group) Wait() error { g.wg.Wait() if g.cancel != nil { g.cancel() } return g.err }
通过阅读 Wait() 方法的源码,我们可以发现它实际上是封装 sync.WaitGroup 的 Wait() 方法,和 context 包的 cancel,并且返回所有运行的 goroutine 中第一个返回的错误。
最后,我们阅读新增控制并发数量的功能 TryGo() 方法和 SetLimit() 方法的源码:
func (g *Group) TryGo(f func() error) bool { if g.sem != nil { select { case g.sem <- token{}: // Note: this allows barging iff channels in general allow barging. default: return false } } g.wg.Add(1) go func() { defer g.done() if err := f(); err != nil { g.errOnce.Do(func() { g.err = err if g.cancel != nil { g.cancel() } }) } }() return true }
通过阅读 TryGo() 方法的源码,我们可以发现,它和 Go() 方法的区别就是在处理 g.sem 的值上,使用的逻辑不同。
TryGo() 方法在处理 g.sem 的值时,使用 select ... case ... default 语句,先尝试一次抢占资源,当无法抢到资源时,不再阻塞,而是直接返回 false,表示执行失败。
SetLimit() 方法的源码:
func (g *Group) SetLimit(n int) { if n < 0 { g.sem = nil return } if len(g.sem) != 0 { panic(fmt.Errorf("errgroup: modify limit while %v goroutines in the group are still active", len(g.sem))) } g.sem = make(chan token, n) }
通过阅读 SetLimit() 方法的源码,我们可以看出当入参 n 的值小于 0 时,直接给 g.sem 赋值为 nil,表示不限制并发数量。
在调用 SetLimit() 方法时,g.sem 必须是一个空通道,否则程序会 panic。
除去 SetLimit() 方法的判断逻辑代码,实际上 SetLimit() 方法就是创建一个大小为 n 的有缓冲 channel。
SetLimit() 和 TryGo() 通常一起使用。
4.总结
本文我们介绍 Go 方法提供的 errgroup 库,该库最近新增了控制并发数量的功能。
我们先介绍了三种使用方式,然后通过阅读源码,分析其实现原理。
原文地址:https://mp.weixin.qq.com/s/R8f_-cl98th7ZiSVxaeUAQ