golang sync.WaitGroup

go语言刚入门碰到第一个程序人懵了

package main

import (
	"fmt"
	"sync"
)

var wg sync.WaitGroup

func printer(ch chan int) {
	for i := range ch {
		fmt.Printf("Received %d ", i)
	}
	wg.Done()
}

// main is the entry point for the program.
func main() {
	c := make(chan int)
	go printer(c)
	wg.Add(1)

	// Send 10 integers on the channel.
	for i := 1; i <= 10; i++ {
		c <- i
	}
	print("\nhello world!")
	close(c)
	wg.Wait()
}

对于sync.WaitGroup,小白表示完全不认识，只得查询资料，索性查完有了一点点理解。

WaitGroup 用法

官方文档对 WaitGroup 的描述是：一个 WaitGroup 对象可以等待一组协程结束。使用方法是：

1. main协程通过调用 wg.Add(delta int) 设置worker协程的个数，然后创建worker协程；
2. worker协程执行结束以后，都要调用 wg.Done()；
3. main协程调用 wg.Wait() 且被block，直到所有worker协程全部执行结束后返回。

这里先看一个典型的例子：

func  main() {
  // 省略部分代码 ...
  var wg sync.WaitGroup
  for _, task := range tasks {
    task  := task
    wg.Add(1)
    go func() {
      task()
      wg.Done()
    }()
  }
  wg.Wait()
  // 省略部分代码...
}

然后在这又学到了很多新知识。

这个例子具备了 WaitGroup 正确使用的大部分要素，包括：

1. wg.Done 必须在所有 wg.Add 之后执行，所以要保证两个函数都在main协程中调用；
2. wg.Done 在 worker协程里调用，尤其要保证调用一次，不能因为 panic 或任何原因导致没有执行（建议使用 defer wg.Done()）；
3. wg.Done 和 wg.Wait 在时序上是没有先后。
4. Go 对 array/slice 进行遍历时，runtime 会把 task[i] 拷贝到 task 的内存地址，下标 i 会变，而 task 的内存地址不会变。如果不进行这次赋值操作，所有 goroutine 可能读到的都是最后一个task。

为了让大家有一个直观的感觉，我们用下面这段代码做实验：

package main

import (
  "fmt"
  "unsafe"
)

func main() {
  tasks := []func(){
    func() { fmt.Printf("1. ") },
    func() { fmt.Printf("2. ") },
  }

  for idx, task := range tasks {
    task()
    fmt.Printf("遍历 = %v, ", unsafe.Pointer(&task))
    fmt.Printf("下标 = %v, ", unsafe.Pointer(&tasks[idx]))
    task  := task
    fmt.Printf("局部变量 = %v\\n", unsafe.Pointer(&task))
  }
}

这段代码的打印结果是：

1. 遍历 = 0x40c140, 下标 = 0x40c138, 局部变量 = 0x40c150
2. 遍历 = 0x40c140, 下标 = 0x40c13c, 局部变量 = 0x40c158

不同机器上执行打印结果有所不同，但共同点是：

1. 遍历时，数据的内存地址不变
2. 通过下标取数时，内存地址不同
3. for-loop 内创建的局部变量，即便名字相同，内存地址也不会复用

WaitGroup 实现

在 Go 源码里，WaitGroup 在逻辑上包含：

1. worker 计数器：main协程调用 wg.Add(delta int) 时增加 delta，调用 wg.Done时减一。
2. waiter 计数器：调用 wg.Wait 时，计数器加一; worker计数器降低到0时，重置waiter计数器。
3. 信号量：用于阻塞 main协程。调用 wg.Wait 时，通过 runtime_Semacquire 获取信号量；降低 waiter 计数器时，通过 runtime_Semrelease 释放信号量。

例子：

package main

import (
  "fmt"
  "sync"
  "time"
)

func main() {
  tasks  := []func(){
    func() { time.Sleep(time.Second); fmt.Println("1 sec later") },
    func() { time.Sleep(time.Second *  2); fmt.Println("2 sec later") },
}

  var wg sync.WaitGroup // 1-1
  wg.Add(len(tasks))    // 1-2
  for _, task := range tasks {
    task  := task
    go func() {       // 1-3-1
      defer wg.Done() // 1-3-2
      task()          // 1-3-3
    }()               // 1-3-1
  }
  wg.Wait()           // 1-4
  fmt.Println("exit")
}

上面这段代码中，

1. 1-1 创建一个 WaitGroup 对象，worker计数器和waiter计数器默认值均为0。
2. 1-2 设置 worker计数器为 len(tasks)。
3. 1-3-1 创建 worker协程，并启动任务。
4. 1-4 设置 waiter计数器，获取信号量，main协程被阻塞。

5. 1-3-3 中执行结束后，1-3-2 降低worker计数器。当worker计数器降低到0时，

   • 重置 waiter计数器
   • 释放信号量，main 协程被激活，1-4 wg.Wait 返回

尽管 Add(delta int) 里 delta 可以是正数、0、负数。我们在使用时，delta 总是正数。

wg.Done 等价于 wg.Add(-1)。在本文中，我们提到 wg.Add时，默认 delta > 0。