go学习笔记 sync/once源码 和 sync/WaitGroup源码

sync.Once.Do(f func())是一个挺有趣的东西,能保证once只执行一次，无论你是否更换once.Do(xx)这里的方法,这个sync.Once块只会执行一次。

package sync
 
import (
    "sync/atomic"
)
 
// Once is an object that will perform exactly one action.
type Once struct {
    // done indicates whether the action has been performed.
    // It is first in the struct because it is used in the hot path.
    // The hot path is inlined at every call site.
    // Placing done first allows more compact instructions on some architectures (amd64/x86),
    // and fewer instructions (to calculate offset) on other architectures.
    done uint32 // 初始值为0表示还未执行过，1表示已经执行过
    m    Mutex
}
 
// Do calls the function f if and only if Do is being called for the
// first time for this instance of Once. In other words, given
//     var once Once
// if once.Do(f) is called multiple times, only the first call will invoke f,
// even if f has a different value in each invocation. A new instance of
// Once is required for each function to execute.
//
// Do is intended for initialization that must be run exactly once. Since f
// is niladic, it may be necessary to use a function literal to capture the
// arguments to a function to be invoked by Do:
//     config.once.Do(func() { config.init(filename) })
//
// Because no call to Do returns until the one call to f returns, if f causes
// Do to be called, it will deadlock.
//
// If f panics, Do considers it to have returned; future calls of Do return
// without calling f.
//
func (o *Once) Do(f func()) {
    // Note: Here is an incorrect implementation of Do:
    //
    //    if atomic.CompareAndSwapUint32(&o.done, 0, 1) {
    //        f()
    //    }
    //
    // Do guarantees that when it returns, f has finished.
    // This implementation would not implement that guarantee:
    // given two simultaneous calls, the winner of the cas would
    // call f, and the second would return immediately, without
    // waiting for the first's call to f to complete.
    // This is why the slow path falls back to a mutex, and why
    // the atomic.StoreUint32 must be delayed until after f returns.
    // 每次一进来先读标识位 0 标识没有被执行过，1 标识已经被执行过
    if atomic.LoadUint32(&o.done) == 0 {
        // Outlined slow-path to allow inlining of the fast-path.
        o.doSlow(f)
    }
}
 
func (o *Once) doSlow(f func()) {
    o.m.Lock() // 施加互斥锁
    defer o.m.Unlock()
    if o.done == 0 {
        defer atomic.StoreUint32(&o.done, 1)
        f()
    }
}

从上面我们可以看出，once只有一个 Do 方法；once的结构体中只定义了两个字段：一个mutex的m，一个代表标识位的done。

下面我们来看看Do方法的流程：

WaitGroup用于等待一组线程的结束。父线程调用Add 方法来设定应等待的线程数量。每个被等待的线程在结束时应调用Done方法。同时，主线程里可以调用wait方法阻塞至所有线程结束。 注意：Add和创建协程的数量一定要匹配，否则会产出panic
主要函数：
func (wg *WaitGroup) Add(delta int)：等待协程的数量。
func (wg *WaitGroup) Done()： 减少waitgroup线程等待线程数量的值，一般在协程完成之后执行。
func (wg *WaitGroup) Wait()：wait方法一般在主线程调用，阻塞直到group计数减少为0。

 
package sync
 
import (
    "internal/race"
    "sync/atomic"
    "unsafe"
)
 
// A WaitGroup waits for a collection of goroutines to finish.
// The main goroutine calls Add to set the number of
// goroutines to wait for. Then each of the goroutines
// runs and calls Done when finished. At the same time,
// Wait can be used to block until all goroutines have finished.
//
// A WaitGroup must not be copied after first use.
type WaitGroup struct {
    noCopy noCopy // noCopy可以嵌入到结构中，在第一次使用后不可复制
 
    // 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, and the other 4 as storage
    // for the sema.
    // 64 bit：高32 bit是计数器，低32位是 阻塞的goroutine计数。
    // 64位的原子操作需要64位的对齐，但是32位。
    // 编译器不能确保它,所以分配了12个byte对齐的8个byte作为状态。其他4个作为信号量
    state1 [3]uint32
}
 
// uintptr和unsafe.Pointer的区别就是：unsafe.Pointer只是单纯的通用指针类型，用于转换不同类型指针，它不可以参与指针运算；
// 而uintptr是用于指针运算的，GC 不把 uintptr 当指针，也就是说 uintptr 无法持有对象，uintptr类型的目标会被回收。
// state()函数可以获取到wg.state1数组中元素组成的二进制对应的十进制的值 和信号量
// 根据编译器位数，获得标志位和等待次数的数据域
// state returns pointers to the state and sema fields stored within wg.state1.
func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
    if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
       // 是否是 64位机器：因为64位机器站高8位 信号量在后面
        return (*uint64)(unsafe.Pointer(&wg.state1)), &wg.state1[2]
    } else {
        // 如果是 32位机器，型号量在最前面
        return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
    }
}
 
// Add adds delta, which may be negative, to the WaitGroup counter.
// If the counter becomes zero, all goroutines blocked on Wait are released.
// If the counter goes negative, Add panics.
//
// Note that calls with a positive delta that occur when the counter is zero
// must happen before a Wait. Calls with a negative delta, or calls with a
// positive delta that start when the counter is greater than zero, may happen
// at any time.
// Typically this means the calls to Add should execute before the statement
// creating the goroutine or other event to be waited for.
// If a WaitGroup is reused to wait for several independent sets of events,
// new Add calls must happen after all previous Wait calls have returned.
// See the WaitGroup example.
func (wg *WaitGroup) Add(delta int) {
    // 获取到wg.state1数组中元素组成的二进制对应的十进制的值的指针 和信号量
    statep, semap := 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 因为高32位是计数器 所以把 delta的值左移32位，并从数组的首元素处开始赋值
    state := atomic.AddUint64(statep, uint64(delta)<<32)
    v := int32(state >> 32) // 获取计数器的值：转int32
    //获得调用 wait（）等待次数：转uint32
    w := uint32(state)
    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(semap))
    }
    // 计数器为负数，报panic
    //标记位不能小于0（done过多或者Add（）负值太多）
    if v < 0 {
        panic("sync: negative WaitGroup counter")
    }
    // 不能Add 与Wait 同时调用
    if w != 0 && delta > 0 && v == int32(delta) {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }
    // Add 完毕
    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.
    // 当等待计数器> 0时，而goroutine将设置为0。
    // 此时不可能有同时发生的状态突变:
    // - Add()不能与 Wait() 同时发生，
    // - 如果计数器counter == 0，不再增加等待计数器
    // 不能Add 与Wait 同时调用
    if *statep != state {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }
    // Reset waiters count to 0.
    *statep = 0 // 所有状态位清零
    for ; w != 0; w-- {
        // 目的是作为一个简单的wakeup原语，以供同步使用。true为唤醒排在等待队列的第一个goroutine
        runtime_Semrelease(semap, false, 0)
    }
}
 
// Done decrements the WaitGroup counter by one.
// Done方法其实就是Add（-1）
func (wg *WaitGroup) Done() {
    wg.Add(-1)
}
 
// Wait blocks until the WaitGroup counter is zero.
// Wait 会一直阻塞到 计数器值为0为止
func (wg *WaitGroup) Wait() {
    statep, semap := wg.state()
    if race.Enabled {
        _ = *statep // trigger nil deref early
        race.Disable()
    }
    //循环检查计数器V啥时候等于0
    for {
        state := atomic.LoadUint64(statep)
        v := int32(state >> 32)
        w := uint32(state)
        if v == 0 {
            // Counter is 0, no need to wait.
            if race.Enabled {
                race.Enable()
                race.Acquire(unsafe.Pointer(wg))
            }
            return
        }
        // Increment waiters count.
        // 尚有未执行完的go程，等待标志位+1（直接在低位处理，无需移位）
        // 增加等待goroution计数，对低32位加1，不需要移位
        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(semap))
            }
            // 目的是作为一个简单的sleep原语，以供同步使用
            runtime_Semacquire(semap)
            // 在上一次Wait返回之前重新使用WaitGroup，即在之前的Done 中没有清空 计数量就会有问题
            if *statep != 0 {
                panic("sync: WaitGroup is reused before previous Wait has returned")
            }
            if race.Enabled {
                race.Enable()
                race.Acquire(unsafe.Pointer(wg))
            }
            return
        }
    }
}

Add：

Wait: