go sync.WaitGroup源码分析

原创

wx63993a9e4baf6 2022-12-14 11:02:07 博主文章分类：Go语言 ©著作权

文章标签 WaitGroup源码 waitgroup源码 go waitgroup源码分析信号量 Group 文章分类 虚拟化云计算

©著作权归作者所有：来自51CTO博客作者wx63993a9e4baf6的原创作品，请联系作者获取转载授权，否则将追究法律责任

go版本：1.10.3

原理实现：信号量

信号量是Unix系统提供的一种保护共享资源的机制，用于防止多个线程同时访问某个资源。

可简单理解为信号量为一个数值：

当信号量>0时，表示资源可用，获取信号量时系统自动将信号量减1；
当信号量==0时，表示资源暂不可用，获取信号量时，当前线程会进入睡眠，当信号量为正时被唤醒

WaitGroup的定义

type WaitGroup struct {
  noCopy noCopy // noCopy用来标记不可复制,只能用指针传递,保证全局唯一.其实即使复制了,编译,运行都没问题,只有用go vet检测时才会显示出错误

  // 只需要64位,即8个字节,其中高32位是counter值,低32位值是waiter值
  // 不直接使用uint64,是因为uint64的原子操作需要64位系统,而32位系统下,可能会出现崩溃
  // 所以这里用byte数组来实现,32位系统下4字节对齐,64位系统下8字节对齐,所以申请12个字节,其中必定有8个字节是符合8字节对齐的,下面的state()函数中有进行判断
  state1 [12]byte

  sema   uint32 // 信号量
}

// 得到counter值（uint64的高32位），waiter值（uint64的低32位）
func (wg *WaitGroup) state() *uint64 {
  // 根据state1的起始地址分析,若是8字节对齐的,则直接用前8个字节作为*uint64类型
  // 若不是,说明是4字节对齐,则后移4个字节后,这样必为8字节对齐,然后取后面8个字节作为*uint64类型
  if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
    return (*uint64)(unsafe.Pointer(&wg.state1))
  } else {
    return (*uint64)(unsafe.Pointer(&wg.state1[4]))
  }
}

增加或减少counter的值

func (wg *WaitGroup) Add(delta int) {
  //当前的 counter值  waiter值
  statep := wg.state()

  // 把delta值加到counter上
  state := atomic.AddUint64(statep, uint64(delta)<<32)
  v := int32(state >> 32) // counter值
  w := uint32(state)      // waiter值

  // counter为负,则panic
  if v < 0 {
    panic("sync: negative WaitGroup counter")
  }

  // waiter值不为0,累加后的counter值和delta相等,说明Add()和Wait()同时调用了,panic,因为正确的做法是先Add()后Wait()
  if w != 0 && delta > 0 && v == int32(delta) {
    panic("sync: WaitGroup misuse: Add called concurrently with Wait")
  }

  // 正常Add()后的情况
  // 1. counter > 0,说明还不需要释放信号量,返回
  // 2. waiter = 0,说明没有等待的goruntine,也不需要释放信号量,返回
  if v > 0 || w == 0 {
    return
  }

  // 下面是counter==0,且w>0的情况
  // 现在若原state和新state不等,有以下两种可能
  // 1. add和wait同时调用
  // 2. counter已经为0,但waiter还为正值,这种情况永远不可能触发信号量了
  // 都是出错了
  if *statep != state {
    panic("sync: WaitGroup misuse: Add called concurrently with Wait")
  }

  // 把counter,waiter都置为0，因为已经触发信号，通知所有等待的goroutine即可，此时不可以再Add()或者Wait()了
  *statep = 0

  // 原子地递增信号量,并通知等待的goroutine
  for ; w != 0; w-- {
    runtime_Semrelease(&wg.sema, false)
  }
}

// Done函数即简单将counter值减1
func (wg *WaitGroup) Done() {
  wg.Add(-1)
}

增加waiter值

func (wg *WaitGroup) Wait() {
  //当前的 counter值  waiter值
  statep := wg.state()

  // 一直等待,直到无需等待或信号量触发时,才返回
  for {
    state := atomic.LoadUint64(statep)
    v := int32(state >> 32) // counter值
    w := uint32(state)      // waiter值
    
    // 若counter值为0,说明所有goroutine都退出了,无需等待,直接返回即可
    if v == 0 {
      return
    }

    // 原子地增加waiter的值,CAS方法,外面有for循环会一直尝试,保证多个goroutine同时调用Wait()也能正确累加waiter
    if atomic.CompareAndSwapUint64(statep, state, state+1) {

      // 一直等待信号量sema,直到>0,信号量触发,然后以原子的方式递减它
      runtime_Semacquire(&wg.sema)
      // 看上面的Add()函数，触发信号量前会先将counter和waiter置0，所以此时必定为0
            // 若不为0，说明WaitGroup此时又被执行Add()或者Wait()操作了，应panic
      if *statep != 0 {
        panic("sync: WaitGroup is reused before previous Wait has returned")
      }

      // 可以返回了
      return
    }
  }
}