go sync.WaitGroup源码分析

 

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
		}
	}
}

提示

• Add()操作必须早于Wait(), 否则会panic
• Add()设置的值必须与实际等待的goroutine个数一致，否则会panic
• WaitGroup只可保持一份，不可拷贝给其他变量，否则会造成意想不到的BUG