第五章 go并发原语之WaitGroup

package main

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

func main() {
    var wg sync.WaitGroup

    // 假设要并发启动 3 个任务
    taskCount := 3
    wg.Add(taskCount) // 将内部计数器 +3

    for i := 1; i <= taskCount; i++ {
        go func(id int) {
            defer wg.Done() // 在 goroutine 退出时将计数器 -1
            // 模拟实际工作
            time.Sleep(time.Duration(id) * 500 * time.Millisecond)
            fmt.Printf("任务 %d 执行完毕\n", id)
        }(i)
    }

    // 在 main goroutine 中等待所有子 goroutine 完成
    wg.Wait() // 阻塞，直到内部计数器归零
    fmt.Println("所有任务已完成")
}

type WaitGroup struct {
	noCopy noCopy
	// state: 64 位无符号整数
    // 高 32 位 (state1) 用来记录等待的 goroutine（调用 Wait 并阻塞的 goroutine）数目
    // 低 32 位 (state0) 用来记录计数器（Add/Done 相加减的结果）
	state atomic.Uint64 // high 32 bits are counter, low 32 bits are waiter count.  内部字段，具体布局随版本而变化
	sema  uint32 // futex 信号量，供被阻塞的 Wait() & Done() 使用
}

// Add 将 WaitGroup 的计数器增加或减少 delta。
// 当 delta 为正时，表示有新的任务加入；当 delta 为负时，表示有任务完成。
// 如果在有等待者时将计数器从正数减到 0，会释放所有阻塞在 Wait() 上的 goroutine。
func (wg *WaitGroup) Add(delta int) {
    // 下面是针对 Go race detector 的额外逻辑，用于捕捉竞态情况：
    if race.Enabled {
        if delta < 0 {
            // 当 delta < 0（即 Done 操作）时，需要与 Wait() 的释放进行同步。
            // race.ReleaseMerge 用于告诉检测器，“我在此处与 Wait 进行合并同步”。
            race.ReleaseMerge(unsafe.Pointer(wg))
        }
        // 在执行后续对 wg.state 的原子操作时，先禁用 race 检测
        race.Disable()
        defer race.Enable() // 在函数结束时重新启用 race 检测
    }

    // =============================== 核心原子操作 ===============================
    // state 是一个 64 位整数，由高 32 位和低 32 位组成。
    // - 高 32 位 (v) 表示 WaitGroup 的当前计数，即还有多少任务未完成。
    // - 低 32 位 (w) 表示当前有多少个 goroutine 正在 Wait() 中被阻塞，等待计数归零。
    //
    // 这里执行 Add(delta) 时，将 delta << 32，左移到“高 32 位”上，然后原子累加到 wg.state 上。
    // Add 返回加完之后的完整 64 位值。
    state := wg.state.Add(uint64(delta) << 32)

    // 从 state 中解析出：
    // v：高 32 位，表示“未完成任务”的数量
    // w：低 32 位，表示“等待者”的数量
    v := int32(state >> 32) // 高 32 位转换为 int32，即当前的计数值
    w := uint32(state)      // 低 32 位，即当前等待者的数量

    // race 检测：当 delta>0，并且 v==delta（即原来计数是 0，现在是 delta）
    // 说明这是第一次真正将计数从 0 推到正数。为了与可能已有的 Wait() 竞态，这里做一次“读”操作触发同步
    if race.Enabled && delta > 0 && v == int32(delta) {
        // 标记一下，表示“我读了一下 wg.sema”，让 race detector 能知道 Add 与 Wait 之间的依赖
        race.Read(unsafe.Pointer(&wg.sema))
    }

    // ============ 检查负数计数的 panic 情况 ============
    if v < 0 {
        // 如果 v<0，说明调用者做了多余的 Done（Add(-1) 超过了 Add() 的次数），导致计数器变为负数。
        panic("sync: negative WaitGroup counter")
    }

    // ============ 检查 Add 与 Wait 并发使用的 panic 情况 ============
    // 当 w != 0，说明当前已经有至少 1 个 goroutine 在 Wait() 阻塞（等待者 != 0）。
    // 如果又发生了 delta > 0 且 v == delta（即原来的计数为 0，现在正好被设为 delta），
    // 也就意味着“在有等待者挂起时，又从 0 重新开始 Add”，这属于 WaitGroup 不允许的误用，
    // 会导致竞态：因为 Wait() 假定一旦看见计数是0，就不再把自己计入等待者。
    // 既然此刻已有等待者，却又“首次 Add”，这种用法被认为是不安全的，故 panic。
    if w != 0 && delta > 0 && v == int32(delta) {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }

    // ============ 如果当前计数仍然 > 0，或者当前没有等待者，直接返回 ============
    // v > 0 表示计数器不为 0，说明还有子任务尚未完成，此时无需触发任何等待者唤醒逻辑。
    // w == 0 则表示没有任何 goroutine 在 Wait() 上阻塞，也无需做后续的释放操作。
    if v > 0 || w == 0 {
        return
    }

    // ============ 下面的逻辑只在 “v == 0 && w > 0” 时执行 ============
    // 这种场景：刚好把计数从 1 减到了 0（v == 0），并且此时有多个 Wait() 正在阻塞（w > 0）。
    // 根据 WaitGroup 的设计，需要一次性唤醒所有的等待者 goroutine。

    // 先做一个小的 sanity check（健全性检查）：
    // 确保在我们读取到 state 之后，没有其它 concurrent 的 Add/Done 对 wg.state 产生变化。
    // 如果 wg.state.Load() 不等于刚刚计算得到的 state，就说明有人在此时刻同时修改了状态，
    // 那就说明 WaitGroup 又被误用（Add 与 Wait 并发），panic 提示用户。
    if wg.state.Load() != state {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }

    // 因为 state 的高 32 位（计数器）已经是 0，而低 32 位（等待者计数）等于 w > 0。
    // 既然要唤醒完所有等待者，就要重置整个状态为 0（因为新的计数器 = 0，且等待者将被清空）。
    wg.state.Store(0)

    // 释放 semaphore，让所有在 Semacquire 阻塞的 Wait() goroutine 依次被唤醒
    // 注意：需要循环 w 次释放，因为低 32 位记录了有多少个 goroutine 阻塞在信号量上
    for ; w != 0; w-- {
        runtime_Semrelease(&wg.sema, false, 0)
    }
}

// Done 是 Add(-1) 的简写，用于表示当前一个任务完成，计数器减 1。
// 通常在每个 goroutine 中使用 defer wg.Done() 来确保在 goroutine 结束时自动调用 Done。
func (wg *WaitGroup) Done() {
    wg.Add(-1)
}

// Wait 会阻塞，直到 WaitGroup 的计数器变为 0。
// 即：所有调用过 Add(+1) 的任务都执行 Done()（Add(-1)）后，Wait 才会返回。
func (wg *WaitGroup) Wait() {
	if race.Enabled {
		// 如果启用了竞态检测器，暂时禁用它（避免在原子操作期间误判）
		race.Disable()
	}

	for {
		// 读取当前的 state：64位，包含“任务计数器”（高32位）和“等待者数量”（低32位）
		state := wg.state.Load()
		v := int32(state >> 32) // 高 32 位：表示当前还有多少个子任务未完成
		w := uint32(state)      // 低 32 位：表示当前已经有多少个 goroutine 在等待
		// ============ 如果 v == 0，说明已经没有未完成任务，无需阻塞 ============
		if v == 0 {
			// 如果计数器已经为 0，说明不需要等待，立即返回
			if race.Enabled {
				// 如果启用了 race 检测器，恢复，并声明 Acquire（获取了 wg 的同步锁）
				race.Enable()
				race.Acquire(unsafe.Pointer(wg))
			}
			return
		}

		// 尝试将等待者数量加 1（低 32 位加 1）
		// 比如原始 state = (v<<32 | w)，现在希望变成 (v<<32 | w+1)
		if wg.state.CompareAndSwap(state, state+1) {
			// 加 1 成功（这段是 CAS 原子操作，防止并发 Wait 出错）

			if race.Enabled && w == 0 {
				// 如果这是第一个等待者（w == 0），且开启了竞态检测
				// 需要与 Add 中的 race.Read 配对，模拟一次“写入”
				// 否则 Wait() 与后续 Add() 可能存在未捕捉的并发竞态

				// 注意：只有第一个 Wait() 会做 race 写操作，多个并发 Wait() 会互相竞态
				race.Write(unsafe.Pointer(&wg.sema))
			}

			// 真正阻塞当前 goroutine，等待被 runtime_Semrelease 唤醒（在 Add 中）
			runtime_SemacquireWaitGroup(&wg.sema)

			// 被唤醒后，如果发现状态不为 0，说明 WaitGroup 被重复使用了
			// 例如：一个 Wait() 返回后，还没等下一轮 Add() 开始，用户就又调用 Wait()，属于误用
			if wg.state.Load() != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
			}

			// 恢复 race 检测器状态
			if race.Enabled {
				race.Enable()
				race.Acquire(unsafe.Pointer(wg))
			}
			return
		}

		// 如果 CompareAndSwap 失败（说明有其它 goroutine 同时修改了 state），重新尝试
	}
}