Go 官方扩展库 Errgroup 详解

389 阅读2分钟

1. WaitGroup

WaitGroup 是 Go 的 sync 包下的一个并发原语,主要用于主 goroutines 等待一组 goroutines 执行完毕后再继续执行的业务场景。

使用

其使用方式可以参照如下:

func main() {
	var wg sync.WaitGroup
	
	wg.Add(2)
        
	go func() {
		defer wg.Done()
		// TODO
	}()
	go func() {
		defer wg.Done()
		// TODO
	}()
	wg.Wait()
	
	fmt.Print("finish")
}
  • wg.add():主 goroutines 传入一个计数,计数减到0的时候才继续执行wg.Wait()之后的逻辑。
  • wg.Done():子 goroutines 调用,使wg的计数-1.
  • wg.Wait():wg的计数减为0才继续往后执行。

2. Errgroup

地址

go get golang.org/x/sync/errgroup

与 WaitGroup 相比的优势

  • 不需要自己手动调用 Add() 和 Done() 方法
  • 可以抛出错误给调用者
  • 可以和 context 集成

使用

一般使用如下所示:

func main() {
   var eg errgroup.Group

   eg.Go(func() error {
      // TODO
      return nil
   })
   
   // 如果需要传参的话,可以参考如下形式
    eg.Go(
       func(i int) func() error {
          fmt.Printf("task %d", i)
          return func() error {
             return nil
          }
       }(1))


   if err := eg.Wait(); err != nil {
      // TODO
   }
}

WithContext()

源码如下:

func WithContext(ctx context.Context) (*Group, context.Context) {
   ctx, cancel := context.WithCancel(ctx)
   return &Group{cancel: cancel}, ctx
}

ctx对应的cancel()方法会在以下两种情况下触发:

  • eg.Wait()执行
  • eg.Done()第一次报错(eg.Go()底层会执行eg.Done())

数据结构

type Group struct {
   cancel func()

   wg sync.WaitGroup

   sem chan token // 用来计数

   errOnce sync.Once // eg.Done()第一次报错的时候做记录
   err     error // 记录的err
}

sem 与计数操作

Errgroup 底层用一个chan token来计数,token是空结构体。

SetLimit() 方法会初始化一个长度为n的chan,当调用TryGo()时会向chan中传入一个空结构体token起到计数的作用。

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)
}
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
}

本文正在参加技术专题18期-聊聊Go语言框架