Go中sync.WaitGroup处理协程同步

// 初始化 WaitGroup
var wg sync.WaitGroup

// 告诉 WaitGroup 有 2 个 goroutine 需要等待
wg.Add(2)

// 启动第一个 goroutine
go func() {
    defer wg.Done()

    // do something
    time.Sleep(time.Second * 1)
    fmt.Println("goroutine 1 done")
}()

// 启动第二个 goroutine
go func() {  
    defer wg.Done()

    // do something
    time.Sleep(time.Second * 2)
    fmt.Println("goroutine 2 done")
}()

// 等待所有注册过的 goroutine 都执行完
wg.Wait()

// 主 goroutine 等待 wg.Wait() 完成
// go next

package sync

import (
	"internal/race"
	"sync/atomic"
	"unsafe"
)

// WaitGroup等待一组协程完成。
// 主协程调用Add来设置
// 等待的协程。然后是每个协程
// 运行并在完成时调用Done。同时，
// Wait可以用来阻塞，直到所有的协程都完成。

// WaitGroup首次使用后不能复制。
type WaitGroup struct {
	noCopy noCopy

	// 64位值:高32位为计数器，低32位为等待计数。
    // 64位原子操作需要64位对齐，但是32位编译器只保证64位字段是32位对齐的。
	// 出于这个原因，在32位体系结构上，我们需要检查state()中state1是否对齐，并在需要时动态地“交换”字段顺序。
    state1 uint64
	state2 uint32
}

// State返回指向存储在wg.state*中的State和sema字段的指针。
func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
	if unsafe.Alignof(wg.state1) == 8 || uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
		// State1是64位对齐的:不做任何事情。
		return &wg.state1, &wg.state2
	} else {
		// State1是32位对齐，但不是64位对齐:这意味着(&state1)+4是64位对齐的。
		state := (*[3]uint32)(unsafe.Pointer(&wg.state1))
		return (*uint64)(unsafe.Pointer(&state[1])), &state[0]
	}
}

// Add向WaitGroup计数器添加增量，增量可能为负。
// 如果计数器变为零，则释放被Wait阻塞的所有协程。
// 如果计数器为负，则添加panics。

// 请注意，当计数器为零时，具有正增量的调用必须在Wait之前发生。
// 具有负增量的调用，或者在计数器大于零时开始的具有正增量的调用，可能在任何时候发生。
// 通常，这意味着对Add的调用应该在语句创建要等待的程序或其他事件之前执行。
// 如果重用WaitGroup来等待几个独立的事件集，则必须在所有先前的wait调用返回之后发生新的Add调用。
// 参见WaitGroup示例。
func (wg *WaitGroup) Add(delta int) {
	statep, semap := wg.state()
	if race.Enabled {
		_ = *statep // 提前触发nil延迟
		if delta < 0 {
			// 与Wait同步减量。
			race.ReleaseMerge(unsafe.Pointer(wg))
		}
		race.Disable()
		defer race.Enable()
	}
	state := atomic.AddUint64(statep, uint64(delta)<<32)
	v := int32(state >> 32)
	w := uint32(state)
	if race.Enabled && delta > 0 && v == int32(delta) {
		// 第一个增量必须与Wait同步。
        // 需要将其建模为读取，因为可能有多个并发的wg。计数器从0转换。
		race.Read(unsafe.Pointer(semap))
	}
	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
	}
	// 当 waiters > 0时，这个协程将counter设置为0。
    // 状态不可能同时发生突变
    // -添加不能与等待同时发生，
    // - Wait如果看到counter == 0，则不会增加waiters。
    // 仍然要做一个便宜的完整性检查来检测WaitGroup的误用。
	if *statep != state {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	// 将waiters计数重置为0。 
	*statep = 0
	for ; w != 0; w-- {
		runtime_Semrelease(semap, false, 0)
	}
}

// Done将WaitGroup counter减1。
func (wg *WaitGroup) Done() {
	wg.Add(-1)
}

// 等待阻塞直到WaitGroup counter为0。
func (wg *WaitGroup) Wait() {
	statep, semap := wg.state()
	if race.Enabled {
		_ = *statep // 提前触发nil延迟
		race.Disable()
	}
	for {
		state := atomic.LoadUint64(statep)
		v := int32(state >> 32)
		w := uint32(state)
		if v == 0 {
			// Counter is 0, no need to wait.
			if race.Enabled {
				race.Enable()
				race.Acquire(unsafe.Pointer(wg))
			}
			return
		}
		// Increment waiters count.
		if atomic.CompareAndSwapUint64(statep, state, state+1) {
			if race.Enabled && w == 0 {
				// Wait必须与第一个Add同步。
                // 需要将其建模为写操作与Add中的读操作竞争。
                // 因此，只能给第一个 waiter 写入，
                // 否则并发等待将相互竞争。
				race.Write(unsafe.Pointer(semap))
			}
			runtime_Semacquire(semap)
			if *statep != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
			}
			if race.Enabled {
				race.Enable()
				race.Acquire(unsafe.Pointer(wg))
			}
			return
		}
	}
}