Golang sync.WaitGroup源码详细分析

迷南。 2023-02-22 08:28 10阅读 0赞

### 一、介绍 ###

WaitGroup是多个goroutine之间协作的一种实现方式，主要功能就是阻塞等待一组goroutine执行完成。

常用的使用场景：主goroutine调用Add函数设置需要等待的goroutine的数量,当每个goroutine执行完成后调用Done函数(将counter减1)，Wait函数用于阻塞等待直到该组中的所有goroutine都执行完成。

源码中主要设计到的三个概念：counter、waiter和semaphore

counter： 当前还未执行结束的goroutine计数器

waiter : 等待goroutine-group结束的goroutine数量，即有多少个等候者

semaphore: 信号量

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

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

当信号量>0时，表示资源可用，获取信号量时系统自动将信号量减1；

当信号量=0时，表示资源暂不可用，获取信号量时，当前线程会进入睡眠，当信号量为正时被唤醒。

### 二、源码分析 ###

Golang源码版本 ：1.10.3

**1.结构体**

type WaitGroup struct {
    	noCopy noCopy  //该WaitGroup对象不允许拷贝使用,只能用指针传递
    
    	// 64-bit value: high 32 bits are counter, low 32 bits are waiter count.
    	// 64-bit atomic operations require 64-bit alignment, but 32-bit
    	// compilers do not ensure it. So we allocate 12 bytes and then use
    	// the aligned 8 bytes in them as state.
    	//用于存储计数器(counter)和waiter的值
    	// 只需要64位,即8个字节,其中高32位是counter值,低32位值是waiter值
    	// 不直接使用uint64,是因为uint64的原子操作需要64位系统,而32位系统下,可能会出现崩溃
    	// 所以这里用byte数组来实现,32位系统下4字节对齐,64位系统下8字节对齐,所以申请12个字节,其中必定有8个字节是符合8字节对齐的,下面的state()函数中有进行判断
    	state1 [12]byte
    	sema   uint32  //信号量
    }

从结构体中我们看到

state1是一个12位长度的byte数组，用于存储counter和waiter的值

sema就是传说中的信号量

**2.state函数**

state是一个内部函数，用于获取counter和 waiter的值

//获取counter  、 waiter的值  (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]))
    	}
    }

**3.Add方法**

//用于增加或减少计数器(counter)的值
    //如果计数器为0，则释放调用Wait方法时的阻塞，如果计数器为负，则panic
    //Add()方法应该在Wait()方法调用之前
    func (wg *WaitGroup) Add(delta int) {
    	//获取当前counter和 waiter的值
    	statep := wg.state()
    	if race.Enabled {
    		_ = *statep // trigger nil deref early
    		if delta < 0 {
    			// Synchronize decrements with Wait.
    			race.ReleaseMerge(unsafe.Pointer(wg))
    		}
    		race.Disable()
    		defer race.Enable()
    	}
    
    	//将delta的值添加到counter上
    	state := atomic.AddUint64(statep, uint64(delta)<<32)
    	v := int32(state >> 32)  //counter值
    	w := uint32(state)  //waiter值
    	if race.Enabled && delta > 0 && v == int32(delta) {
    		// The first increment must be synchronized with Wait.
    		// Need to model this as a read, because there can be
    		// several concurrent wg.counter transitions from 0.
    		race.Read(unsafe.Pointer(&wg.sema))
    	}
    
    	//counter为负数，则触发panic
    	if v < 0 {
    		panic("sync: negative WaitGroup counter")
    	}
    
    	// waiter值不为0,累加后的counter值和delta相等,说明Wait()方法没有在Add()方法之后调用,触发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 ,说明没有等待的goroutine，也不需要释放信号量，可以直接返回
    	if v > 0 || w == 0 {
    		return
    	}
    	// This goroutine has set counter to 0 when waiters > 0.
    	// Now there can't be concurrent mutations of state:
    	// - Adds must not happen concurrently with Wait,
    	// - Wait does not increment waiters if it sees counter == 0.
    	// Still do a cheap sanity check to detect WaitGroup misuse.
    	//下面是 counter == 0 并且 waiter > 0的情况
    	//现在若原state和新的state不等，则有以下两种可能
    	//1. Add 和 Wait方法同时调用
    	//2. counter已经为0，但waiter值有增加，这种情况永远不会触发信号量了
    	// 以上两种情况都是错误的，所以触发异常
    	//注：state := atomic.AddUint64(statep, uint64(delta)<<32)  这一步调用之后，state和*statep的值应该是相等的，除非有以上两种情况发生
    	if *statep != state {
    		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    	}
    	// Reset waiters count to 0.
    	//将waiter 和 counter都置为0
    	*statep = 0
    	//原子递减信号量，并通知等待的goroutine
    	for ; w != 0; w-- {
    		runtime_Semrelease(&wg.sema, false)
    	}
    }

**  4.Done方法**

// Done decrements the WaitGroup counter by one.
    //将计数器(counter)的值减1
    func (wg *WaitGroup) Done() {
    	wg.Add(-1)
    }

**5.Wait方法**

// Wait blocks until the WaitGroup counter is zero.
    //调用Wait方法会阻塞当前调用的goroutine直到 counter的值为0
    //也会增加waiter的值
    func (wg *WaitGroup) Wait() {
    	//获取当前counter和 waiter的值
    	statep := wg.state()
    	if race.Enabled {
    		_ = *statep // trigger nil deref early
    		race.Disable()
    	}
    	//一直等待，直到无需等待或信号量触发，才返回
    	for {
    		state := atomic.LoadUint64(statep)
    		v := int32(state >> 32)  //counter值
    		w := uint32(state)  //waiter值
    		//如果counter值为0，则说明所有goroutine都退出了，无需等待，直接退出
    		if v == 0 {
    			// Counter is 0, no need to wait.
    			if race.Enabled {
    				race.Enable()
    				race.Acquire(unsafe.Pointer(wg))
    			}
    			return
    		}
    
    		// Increment waiters count.
    		//原子增加waiter的值，CAS方法，外面for循环会一直尝试，保证多个goroutine同时调用Wait()也能正常累加waiter
    		if atomic.CompareAndSwapUint64(statep, state, state+1) {
    			if race.Enabled && w == 0 {
    				// Wait must be synchronized with the first Add.
    				// Need to model this is as a write to race with the read in Add.
    				// As a consequence, can do the write only for the first waiter,
    				// otherwise concurrent Waits will race with each other.
    				race.Write(unsafe.Pointer(&wg.sema))
    			}
    			//一直等待信号量sema，直到信号量触发，
    			runtime_Semacquire(&wg.sema)
    			//从上面的Add()方法看到，触发信号量之前会将seatep置为0(即counter和waiter都置为0)，所以此时应该也为0
    			//如果不为0，说明WaitGroup此时又执行了Add()或者Wait()操作，所以会触发panic
    			if *statep != 0 {
    				panic("sync: WaitGroup is reused before previous Wait has returned")
    			}
    			if race.Enabled {
    				race.Enable()
    				race.Acquire(unsafe.Pointer(wg))
    			}
    			return
    		}
    	}
    }

## ##

### 三、注意点 ###

1.Add()必须在Wait()前调用

2.Add()设置的值必须与实际等待的goroutine个数一致，如果设置的值大于实际的goroutine数量，可能会一直阻塞。如果小于会触发panic

3. WaitGroup不可拷贝，可以通过指针传递，否则很容易造成BUG

以下为值拷贝引起的Bug示例

**demo1：因为值拷贝引起的死锁**

func main() {
    	var wg sync.WaitGroup
    	wg.Add(5)
    	for i := 0 ; i < 5 ; i++ {
    	 	test(wg)
    	}
    	wg.Wait()
    }
    
    func test(wg sync.WaitGroup) {
    	go func() {
    		fmt.Println("hello")
    		wg.Done()
    	}()
    }

**demo2:因为值拷贝引起的不会阻塞等待现象**

func main() {
    	var wg sync.WaitGroup
    	for i := 0 ; i < 5 ; i++ {
    	 	test(wg)
    	}
    	wg.Wait()
    }
    
    func test(wg sync.WaitGroup) {
    	go func() {
    		wg.Add(1)
    		fmt.Println("hello")
    		time.Sleep(time.Second*5)
    		wg.Done()
    	}()
    }

**demo3:因为值拷贝引发的panic**

type person struct {
    	wg sync.WaitGroup
    }
    
    func (t *person) say()  {
    	go func() {
    		fmt.Println("say Hello!")
    		time.Sleep(time.Second*5)
    		t.wg.Done()
    	}()
    }
    
    
    func main() {
    	var wg sync.WaitGroup
    	t := person{wg:wg}
    	wg.Add(5)
    	for  i := 0 ; i< 5 ;i++ {
    		t.say()
    	}
    	wg.Wait()
    }

感谢：https://blog.csdn.net/yzf279533105/article/details/97302666