go sync Waitgroup数据结构实现基本操作详解

func TestWaitgroup(t *testing.T) {
   var wg sync.WaitGroup
   // 计数器 +2
   wg.Add(2)
   go func() {
      sendHttpRequest("https://baidu.com")
      // 计数器 -1
      wg.Done()
   }()
   go func() {
      sendHttpRequest("https://baidu.com")
      // 计数器 -1
      wg.Done()
   }()
   // 阻塞。计数器为 0 的时候，Wait 返回
   wg.Wait()
}
// 发起 HTTP GET 请求
func sendHttpRequest(url string) (string, error) {
   method := "GET"
   client := &http.Client{}
   req, err := http.NewRequest(method, url, nil)
   if err != nil {
      return "", err
   }
   res, err := client.Do(req)
   if err != nil {
      return "", err
   }
   defer res.Body.Close()
   body, err := io.ReadAll(res.Body)
   if err != nil {
      return "", err
   }
   return string(body), err
}

wait(S) {
    while (S <= 0)
        ; // 忙等待
    S--;
}

signal(S) {
    S++;
}

type WaitGroup struct {
   noCopy noCopy
   // 高 32 位为计数器，低 32 位为等待者数量
   state atomic.Uint64
   sema  uint32
}

func test(wg sync.WaitGroup) {
   wg.Done()
}
func TestWaitGroup(t *testing.T) {
   var wg sync.WaitGroup
   wg.Add(1)
   test(wg)
   wg.Wait()
}

func test(wg *sync.WaitGroup) {
   wg.Done()
}

func (wg *WaitGroup) Add(delta int) {
   // wg.state 的计数器加上 delta
   //（加到 state 的高 32 上）
   state := wg.state.Add(uint64(delta) << 32) // 高 32 位加上 delta
   v := int32(state >> 32)                    // 高 32 位（counter）
   w := uint32(state)                         // 低 32 位（waiter）
   // 计数器不能为负数（加上 delta 之后不能为负数，最小只能到 0）
   if v < 0 {
      panic("sync: negative WaitGroup counter")
   }
   // 正常使用情况下，是先调用 Add 再调用 Wait 的，这种情况下，w 是 0，v > 0
   if w != 0 && delta > 0 && v == int32(delta) {
      panic("sync: WaitGroup misuse: Add called concurrently with Wait")
   }
   // v > 0，计数器大于 0
   // w == 0，没有在 Wait 的协程
   // 说明还没有到唤醒 waiter 的时候
   if v > 0 || w == 0 {
      return
   }
   // Add 负数的时候，v 会减去对应的数值，减到最后 v 是 0。
   // 计数器是 0，并且有等待的协程，现在要唤醒这些协程。
   // 存在等待的协程时，goroutine 已将计数器设置为0。
   // 现在不可能同时出现状态突变：
   // - Add 不能与 Wait 同时发生，
   // - 如果看到计数器==0，则 Wait 不会增加等待的协程。
   // 仍然要做一个廉价的健康检查，以检测 WaitGroup 的误用。
   if wg.state.Load() != state { // 不能在 Add 的同时调用 Wait
      panic("sync: WaitGroup misuse: Add called concurrently with Wait")
   }
   // 将等待的协程数量设置为 0。
   wg.state.Store(0)
   for ; w != 0; w-- {
      // signal，调用 Wait 的地方会解除阻塞
      runtime_Semrelease(&wg.sema, false, 0) // goyield
   }
}

func (wg *WaitGroup) Wait() {
   for {
      // 获取当前计数器
      state := wg.state.Load()
      // 计数器
      v := int32(state >> 32)
      // waiter 数量
      w := uint32(state)
      // v 为 0，不需要等待，直接返回
      if v == 0 {
         // 计数器是 0，不需要等待
         return
      }
      // 增加 waiter 数量。
      // 调用一次 Wait，waiter 数量会加 1。
      if wg.state.CompareAndSwap(state, state+1) {
         // 这会阻塞，直到 sema (信号量)大于 0
         runtime_Semacquire(&wg.sema) // goparkunlock
         // state 不等 0
         // wait 还没有返回又继续使用了 WaitGroup
         if wg.state.Load() != 0 {
            panic("sync: WaitGroup is reused before previous Wait has returned")
         }
         // 解除阻塞状态了，可以返回了
         return
      }
      // 状态没有修改成功（state 没有成功 +1），开始下一次尝试。
   }
}