Golang 之 WaitGroup 源码解析

Posted Go开发大全

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了Golang 之 WaitGroup 源码解析相关的知识,希望对你有一定的参考价值。

【导读】写并发代码使用waitgroup时考虑过它的实现吗?本文对waitgroup进行了详细解读。

如果我们有一个大的任务要做,我们会尝试将这个任务分解,分解完成之后并发交由 goroutine 去做,并且我需要当全部的任务完成之后再进行下面的步骤,在 sync 包下,就有这样一个东西适合上述情况,WaitGroup,今天我们来看看具体它是怎么实现的。

PS:在下面我统一用 wg 来简称 WaitGroup

使用

它的使用非常简单,如下:

func main () {
    wg := sync.WaitGroup {}
    for i := 0; i < 10; i++ {
        wg.Add (1)
        go func (job int) {
            defer wg.Done ()
            //do something
            fmt.Printf ("job % d done\n", job)
        }(i)
    }
    wg.Wait ()
    fmt.Println ("all done")
}

输出:

job 9 done
job 1 done
job 0 done
job 8 done
job 7 done
job 3 done
job 6 done
job 2 done
job 4 done
job 5 done
all done

我们可以看到,使用非常简单,每次有一个任务就使用 Add 方法加一个,每次做完任务就使用 Done 方法告诉它已经完成了,而 Wait 就是等着所有的任务完成。

思考问题

在看 wg 的实现之前,首先来问几个问题,来考考自己。

  1. Wait 方法能否被多次调用,比如再开一个 goroutine 去 wait
  2. Wait 方法调用后是否还能再继续调用 Add 添加任务
  3. 每次只能 Done 一个任务,能否一次性 Done 多个任务呢
  4. wg 能否被拷贝或作为参数传递
  5. 如果让你自己实现一个,你会如何实现

前几个问题,如果你都能很清楚的回答,那么你对 wg 的了解可以说已经非常熟悉了。首选我来说一下对于最后的一个问题的回答,因为在看源码之前我都会想想如果是我,我会如何去实现,那么我想的也很简单。

  • 使用一个变量进行计数
  • 每次任务数量变更时使用 atom 原子操作 + 1 或者 - 1
  • -1 时判断任务数量是否已经为 0
  • 如果为 0 向一个 channel 里面发送消息
  • 所有 wait 的地方监听 channel 的消息,收到消息则证明任务全部完成

源码分析

结构

type WaitGroup struct {
 noCopy noCopy
 // 64-bit value: high 32 bits are counter, low 32 bits are waiter count.
 // 64-bit atomic operations require 64-bit alignment, but 32-bit
 //compilers do not ensure it. So we allocate 12 bytes and then use
 //the aligned 8 bytes in them as state, and the other 4 as storage
 //for the sema.
 state1 [3uint32
}

结构非常简单,就只有两个熟悉,一个 noCopy 还有一个 state1(我也很好奇为什么要用 1 来结尾命名,大佬的想法总是很奇妙)

noCopy:sync 包下的一个特殊标记吧,vet 检查,如果有拷贝的变量则会报错

func main () {
    wg := sync.WaitGroup {}
    w := wg
    fmt.Println (w, wg)
}

你 run 肯定没问题的,但是如果你使用 go vet 做个检查就有警告了

➜  go vet main.go
# command-line-arguments
./main.go:10:10: assignment copies lock value to w: sync.WaitGroup contains sync.noCopy
./main.go:11:17: call of fmt.Println copies lock value: sync.WaitGroup contains sync.noCopy
./main.go:11:20: call of fmt.Println copies lock value: sync.WaitGroup contains sync.noCopy

state1:是用来存放任务计数器和等待者计数器的(我一看到这个结构就明白肯定后面又是位操作这样的高端操作了)

state [0] state [1] state [2]
-64 位- - waiter- - counter-
-32 位- - sema - -waiter-

其中 waiter 是等待者计数,counter 是任务计数,sema 是信号量

奇怪的是在 64 位还 32 位操作系统上是不一样的,具体原因以及对于它操作请继续看下去

state
//state returns pointers to the state and sema fields stored within wg.state1.
func (wg *WaitGroup) state () (statep *uint64, semap *uint32) {
 if uintptr (unsafe.Pointer (&wg.state1))%8 == 0 {
  return (*uint64)(unsafe.Pointer (&wg.state1)), &wg.state1 [2]
 } else {
  return (*uint64)(unsafe.Pointer (&wg.state1 [1])), &wg.state1 [0]
 }
}

这个方法是一个内部方法,就是将 state1 中存储的状态取出来,返回值 statep 就是计数器的状态,semap 是信号量

Done

func (wg *WaitGroup) Done () { wg.Add (-1) } 没想到吧~居然 Done 就是调用 Add 并传递一个 - 1

所以其实我们完全可以再外部调用 Add 传递一个 - 3 一次性结束 3 个任务

Add

func (wg *WaitGroup) Add (delta int) {
 // 首先获取状态值
 statep, semap := wg.state ()
 // 对于 statep 中 counter + delta
 state := atomic.AddUint64 (statep, uint64 (delta)<<32)
 // 获取任务计数器的值
 v := int32 (state >> 32)
 // 获取等待者计数器的值
 w := uint32 (state)
 
 // 任务计数器不能为负数
 if v < 0 {
  panic ("sync: negative WaitGroup counter")
 }
 // 已经有人在等待,但是还在添加任务
 if w != 0 && delta > 0 && v == int32 (delta) {
  panic ("sync: WaitGroup misuse: Add called concurrently with Wait")
 }
 // 没有等待者或者任务还有没做完的
 if v > 0 || w == 0 {
  return
 }
 // 有等待者,但是在这个过程中数据还在变动
 if *statep != state {
  panic ("sync: WaitGroup misuse: Add called concurrently with Wait")
 }

 // Reset waiters count to 0.
 // 重置状态,并用发出等同于等待者数量的信号量,告诉所有等待者任务已经完成
 *statep = 0
 for ; w != 0; w-- {
  runtime_Semrelease (semap, false0)
 }
}

这里有几个要点我们其实已经看到了:

Wait 的 过程中 是不能 Add 的,不然就会 panic,要注意 虽然我们可以借助 Add 一个负数来一次性结束多个任务,但是如果任务数量控制的不好,变成负数也会 panic,Done 次数多了也一样 wg 是通过信号量来通知的,当然可以有很多人在等,wg 它都会一一通知到位的

Wait

func (wg *WaitGroup) Wait () {
 // 先获取状态
 statep, semap := wg.state ()
  
 for {
  // 这里注意要用 atomic 的 Load 来保证一下写操作已经完成
  state := atomic.LoadUint64 (statep)
  // 同样的,这里是任务计数
  v := int32 (state >> 32)
  // 这里是等待者计数
  w := uint32 (state)
  // 如果没有任务,那么直接结束,不用等待了
  if v == 0 {
   return
  }
  // 使用 cas 操作,如果不相等,证明中间已经被其他人修改了状态,重新走 for 循环
  // 注意这里 if 进去之后等待者的数量就 +1 了
  if atomic.CompareAndSwapUint64 (statep, state, state+1) {
   // 等待信号量
   runtime_Semacquire (semap)
   // 如果信号量来了,但是状态还不是 0,则证明 wait 之后还是在人在 add,证明有人想充分利用 wg 但是时机不对
   if *statep != 0 {
    panic ("sync: WaitGroup is reused before previous Wait has returned")
   }
   return
  }
 }
}

其实 wait 虽然简单,也有要点

通过 load 和 cas 操作 + 循环来避免了锁,其实这个操作可以学一下 其实这里也说明明白了,wg 可以重用,但是你必须等到 wait 全部完成之后再说 其他注意点

func main () {
    wg := sync.WaitGroup {}
    for i := 0; i < 10; i++ {
        wg.Add (1)
        go func (job int) {
            doJob (job, wg)
        }(i)
    }
    wg.Wait ()
    fmt.Println ("all done")
}

func doJob (job int, wg sync.WaitGroup) {
    fmt.Printf ("job % d done\n", job)
    wg.Done ()
}

上面的代码有问题吗?问题在哪呢?

其实很简单,wg 作为一个参数传递的时候,wg 还是一个普通的结构体,我们在函数中操作的时候还是操作的一个拷贝的变量而已,对于原来的 wg 是不会改变的,所以这里需要传递指针才是正确的

func main () {
    wg := &sync.WaitGroup {}
        for i := 0; i < 10; i++ {
            wg.Add (1)
            go func (job int) {
                doJob (job, wg)
            }(i)
    }
    wg.Wait ()
    fmt.Println ("all done")
}

func doJob (job int, wg *sync.WaitGroup) {
    fmt.Printf ("job % d done\n", job)
    wg.Done ()
}

但是其实并不推荐这样去传递 wg,因为这样很容易出现问题,一个不好就出问题了,个人还是建议直接在使用 goroutine 之后马上接一个 defer wg.Done () 来的更加靠谱一些

总结

回过头来看看,之前的问题也都有了答案:

  1. Wait 可以被调用多次,并且每个都会收到完成的通知
  2. Wait 之后,如果再 Wait 的过程中不能在 Add,否则会 panic,但是 Wait 结束之后可以继续使用 Add 进行重用
  3. 可以使用 Add 传递负数的方式一次性结束多个任务,但是需要保证任务计数器非负,否则会 panic
  4. wg作为参数传递的时候需要注意传递指针,或者尽量避免传递
  5. 官方利用位操作节约了空间,存在在同一个地方;利用信号量来实现任务结束的通知….

总的来说 wg 的实现还是非常简单的,需要注意的就是几个使用上的点不要出现意外即可。

linkinstar.wiki/2020/03/15/golang/source-code/sync-waitgroup-source-code

 - EOF -

推荐阅读(点击标题可打开)

1、

2、

3、


Go 开发大全

参与维护一个非常全面的Go开源技术资源库。日常分享 Go, 云原生、k8s、Docker和微服务方面的技术文章和行业动态。

关注后获取

回复 Go 获取6万star的Go资源库



分享、点赞和在看

支持我们分享更多好文章,谢谢!

以上是关于Golang 之 WaitGroup 源码解析的主要内容,如果未能解决你的问题,请参考以下文章

深入源码分析golang之WaitGroup

Golang WaitGroup源码分析

Go语言自学系列 | golang并发编程之WaitGroup实现同步

Go并发控制之sync.WaitGroup

golang中sync.WaitGroup的使用

golang 中 sync包的 WaitGroup