Go singleflight 源码剖析

原创

陈明勇

修改于 2024-04-16 11:11:46

1680

修改于 2024-04-16 11:11:46

文章被收录于专栏：Go 技术 Go技术干货Go 技术

前言

前面的一篇文章 Go singleflight：防缓存击穿利器详细介绍 singleflight 包的使用，展示如何利用它来避免缓存击穿。而本篇文章，我们来剖析 singleflight 包的源码实现和工作原理，探索单飞的奥秘。

准备好了吗？准备一杯你最喜欢的咖啡或茶，随着本文一探究竟吧。

singleflight 版本：golang.org/x/sync v0.6.0

结构体解析

Group

Group 是 singleflight 包的一个核心结构体，它管理着所有的请求，确保同一时刻，对同一资源的请求只会被执行一次。该结构体的源码如下所示：

// Group represents a class of work and forms a namespace in  
// which units of work can be executed with duplicate suppression.
type Group struct {
    mu sync.Mutex       // protects m
    m  map[string]*call // lazily initialized
}

Group 结构体有两个字段：

mu sync.Mutex：一个互斥锁，用于保护下面的 m 映射。在并发环境下，多个 goroutine 可能会同时对 m 进行读写操作，所以需要通过互斥锁来确保对 m 的操作是安全的。
m map[string]*call：一个 map 映射，其键是字符串，值是指向 call 结构体的指针。Do 和 DoCHan 方法的参数里，其中一个是 key，m 的键保存的就是这个 key。m 是惰性初始化的，意味着它在第一次使用时才会被创建。

Group 通过维护 m 字段来跟踪每个 key 的调用状态，从而实现将多个请求合并成一个请求，多个请求共享同一个结果。

call

call 结构体表示一个针对特定 key 的正在进行中或者已完成的请求，它确保所有同时对该key调用 Do 或 DoChan 方法的 goroutine 共享同一个执行结果。该结构体的源码如下所示：

// call is an in-flight or completed singleflight.Do call
type call struct {
    wg sync.WaitGroup

    // These fields are written once before the WaitGroup is done
    // and are only read after the WaitGroup is done.
    val interface{}
    err error

    // These fields are read and written with the singleflight
    // mutex held before the WaitGroup is done, and are read but
    // not written after the WaitGroup is done.
    dups  int
    chans []chan<- Result
}

call 结构体有五个字段：

wg sync.WaitGroup：一个等待组，用于等待当前 call 的完成。当调用 Do 或 DoChan 方法后，内部会增加 WaitGroup 的计数器，当调用完成后，会减少计数器。在调用完成之前，其他想要获取当前 call 的结果的 goroutine 会等待 WaitGroup 的完成。
val interface{}：调用 Do 或 DoChan 方法的结果值之一，对应着 fn 函数（Do 或 DoChan 方法的参数）的第一个返回值 val。这个字段在 WaitGroup 完成之前被写入一次，只有在 WaitGroup 完成后才会被读取。
err error：这是在调用 Do 或者 DoChan 方法时可能发生的错误。和 val 类似，这个字段在 WaitGroup 完成之前被写入一次，只有在 WaitGroup 完成后才会被读取。
dups int：用于计数当前 call 的重复调用数量。这个计数是在 singleflight 的互斥锁保护下进行的，在 WaitGroup 完成之前可以读写，在 WaitGroup 完成后只能读取。目前该字段的作用是判断 call 的结果是否被共享。
chans []chan<- Result：一个通道切片，用于存储所有等待当前 call 结果的通道。这些通道在 call 完成时接收到结果。这个字段同样是在 singleflight 的互斥锁保护下进行的，在 WaitGroup 完成之前可以读写，在 WaitGroup 完成后只能读取。

一句话概括就是：call 结构体用于跟踪 Do 或 DoChan 方法的调用状态，包括等待其完成的 goroutine、调用的结果、发生的错误以及跟踪重复的调用次数，对于 singleflight 在共享调用结果中起到关键作用。

Result

Result 是一个封装了请求调用结果的结构体，在DoChan 方法返回结果时使用。该结构体的源码如下所示：

// Result holds the results of Do, so they can be passed
// on a channel.
type Result struct {
    Val    interface{}
    Err    error
    Shared bool
}

Result 结构体有三个字段：

Val interface{}：存储 DoChan 方法调用的结果值之一，对应着 fn 函数（DoChan 方法的参数）的第一个返回值 val。
Err error：存储 DoChan 方法调用过程中可能发生的错误。
Shared bool：表示调用结果是否被多个请求（goroutine）共享。该字段的值，取决于 call 结构体的 dups 字段值，如果 dups 大于 0，Shared 的值则为 true，否则为 false。

panicError

panicError 用于封装从 panic 中恢复的任意值和在给定函数执行期间产生的堆栈跟踪信息。该结构体的源码如下所示：

// A panicError is an arbitrary value recovered from a panic
// with the stack trace during the execution of given function.
type panicError struct {
    value interface{}
    stack []byte
}

// Error implements error interface.
func (p *panicError) Error() string {
    return fmt.Sprintf("%v\n\n%s", p.value, p.stack)
}

func (p *panicError) Unwrap() error {
    err, ok := p.value.(error)
    if !ok {
        return nil
    }

    return err
}

func newPanicError(v interface{}) error {
    stack := debug.Stack()

    // The first line of the stack trace is of the form "goroutine N [status]:"
    // but by the time the panic reaches Do the goroutine may no longer exist
    // and its status will have changed. Trim out the misleading line.
    if line := bytes.IndexByte(stack[:], '\n'); line >= 0 {
            stack = stack[line+1:]
    }
    return &panicError{value: v, stack: stack}
}

字段

panicError 结构体有两个字段：

value interface{}：存储从 panic 中恢复的值，这个值是任意类型的，可能是 error 类型，也可能是其它类型。
stack []byte：存储堆栈跟踪信息的字节切片，这个堆栈跟踪提供了 panic 发生时的函数调用层次结构和顺序，有助于调试和诊断问题。

方法

panicError 结构体有两个方法：

Error() string：实现了 error 接口的 Error 方法。它将 panicError 结构体的 value 和 stack 字段的格式化成一个字符串。
Unwrap() error：实现了 Wrapper 接口的 Unwrap 接包方法，尝试将 value 字段断言为 error 类型并返回。如果 value 不是一个 error 类型，它将返回 nil。这个方法使得 panicError 能够与 Go 的错误处理机制（如 errors.Is 和 errors.As）更好地集成。初始化函数newPanicError(v interface{}) error：这个函数用于创建一个新的 panicError 实例。它接受从 panic 中恢复的值作为参数，然后通过 debug.Stack 获取堆栈信息，并移除堆栈信息的第一行（如 goroutine 的编号和状态），因为这一行包含的信息可能会因为 panic 的恢复而变得不准确。最后，返回指向 panicError 实例的指针。

核心方法解析

Do

func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, err error, shared bool) {
    // 获取锁
    g.mu.Lock()
    // 懒初始化 map
    if g.m == nil {
        g.m = make(map[string]*call)
    }
    // 判断特定 key 的 call 是否正在进行调用
    if c, ok := g.m[key]; ok {
        // 重复调用次数加 1
        c.dups++
        // 解锁
        g.mu.Unlock()
        // 挂起，等待调用的完成
        c.wg.Wait()
        // 判断是否发生了 panic
        if e, ok := c.err.(*panicError); ok {
                // panic
                panic(e)
        } else if c.err == errGoexit { // 判断是否发生了 runtime.Goexit
                // 执行 runtime.Goexit，停止当前 goroutine 的执行，并确保所有 defer 语句的执行
                runtime.Goexit()
        }
        // 返回结果
        return c.val, c.err, true
    }
    // 创建一个新的调用
    c := new(call)
    // 等待组加 1
    c.wg.Add(1)
    // key 和 call 映射
    g.m[key] = c
    // 释放锁
    g.mu.Unlock()
    // 调用开始，执行所接受的函数 fn
    g.doCall(c, key, fn)
    // 返回结果
    return c.val, c.err, c.dups > 0
}

Do 方法的执行流程如下所示：

1、获取锁：通过 g.mu.Lock() 加锁，确保对内部的 g.m（一个 map，用于跟踪 key 的调用状态）和 c.dups（对于该 key 的重复调用次数）的访问是并发安全的。
2、初始化 map：如果 g.m == nil，意味着是第一次调用 Do 方法且没有调用过 DoChan 方法，所以初始化 g.m。
3、检查是否有正在进行的调用：通过 c, ok := g.m[key]; ok 检查是否有一个对于该 key 的调用正在进行，如果 ok 为 true，则说明有一个对于该 key 的调用正在进行：
- 增加重复调用次数 c.dups，表示来了一个新的 goroutine 在等待这个调用结果。
- 释放锁 g.mu.Unlock()，因为不再需要修改共享资源。
- 等待 c.wg.Wait()，直到当前的调用完成。
- 检查错误类型，并按需处理（如果是 panicError 或 errGoexit，则分别触发 panic 或 Goexit）。
- 返回当前进行的调用的结果。
4、初始化并执行新的调用：如果没有一个对于该 key 的调用正在进行，则：
- 创建一个新的 call 实例。
- c.wg 等待组计数加 1，标记新操作的开始，后续有相同调用的请求将会等待该操作的完成并共享结果。
- 在 g.m 中注册 key 和新创建的 call 实例的映射 g.m[key] = c。
- 释放锁。
- 调用 g.doCall(c, key, fn) 执行实际的函数调用。
- 返回调用结果。

Do 方法的关键在于综合使用等待组（sync.WaitGroup）、互斥锁（sync.Mutex）以及一个映射（map），以确保：

对于相同的 key，fn 函数只会被执行一次。这是通过 map 检查当前 key 是否存在对应的 call 实例来实现的。如果已存在，意味着函数调用正在执行或已完成，不需要再次执行。
同一时刻，所有请求同一 key 的调用都能得到同一个结果。通过 map 追踪特定 key 对应的调用结果，确保所有的 goroutine 对同一 key 发起 Do 方法调用都能共享相同的结果。
正确地处理并发和同步。通过 sync.Mutex 保护并发环境下 map 的读写操作，避免并发冲突；通过 sync.WaitGroup 等待异步操作完成，保证所有请求都在函数执行完成后才返回结果。

doCall

doCall 方法负责执行给定 key 的函数 fn，并处理可能的错误和同步执行结果，确保所有请求该key的goroutine 得到统一的结果。该方法的源码如下所示：

func (g *Group) doCall(c *call, key string, fn func() (interface{}, error)) {
    // 定义正常返回标志
    normalReturn := false
    // 定义 panic 标志
    recovered := false

    // 使用双重 defer 来区分 panic 和 runtime.Goexit
    defer func() {
        // fn 函数里面调用了 runtime.Goexit 函数
        if !normalReturn && !recovered {
            // 将 errGoexit 的值赋给 c.err
            c.err = errGoexit
        }
        // 加锁
        g.mu.Lock()
        // 函数执行结束时释放锁
        defer g.mu.Unlock()
        // 标记 call 的完成
        c.wg.Done()
        // 保险起见，判断当前 key 对应的 call 是否被覆盖，没有被覆盖就从 map 中移除这个 key 
        if g.m[key] == c {
            delete(g.m, key)
        }
        // 判断执行 fn 的时候是否发生 panic
        if e, ok := c.err.(*panicError); ok {
            // 避免等待中的通道永久阻塞，如果发生了 panic，需要确保这个 panic 不能被捕获
            if len(c.chans) > 0 {
                // 开一个新的协程去 panic，这个 panic 就不会被捕获了
                go panic(e)
                // 保持当前 goroutine 的存活，这样等到 panic 之后，关于当前 goroutine 的信息就会出现在堆栈中
                select {}
            } else {
                // 直接 panic
                panic(e)
            }
        } else if c.err == errGoexit {
            // 如果是 errGoexit，什么都不用做，因为之前已经执行了 runtime.Goexit
        } else {
            // 向等待中的通道发送结果
            for _, ch := range c.chans {
                ch <- Result{c.val, c.err, c.dups > 0}
            }
        }
    }()

    func() {
        defer func() {
            // 如果 fn 没有正常执行完
            if !normalReturn {
                // 获取从 panic 中恢复的值
                if r := recover(); r != nil {
                    // 创建一个 `panicError` 实例并赋值给 c.err
                    c.err = newPanicError(r)
                }
            }
        }()
        // 执行函数调用
        c.val, c.err = fn()
        // 设置正常返回标志为 true
        normalReturn = true
    }()
    // 如果 fn 没有正常执行完，则发生了 panic
    if !normalReturn {
        // 设置 panic 标志为 true
        recovered = true
    }
}

代码剖析：

标志位定义：定义 normalReturn 和 recovered 用来区分 fn 是否正常执行完成或者发生了 panic。
双重 defer 机制：目的是为了能够区分 fn 函数的正常执行完成、fn 函数里发生的 panic 以及 fn 函数里调用runtime.Goexit 终止协程的情况。
- 第一个 defer 用于清理资源和处理结果。
  - 如果非正常函数执行完成并且没有发生 panic，则 fn 里执行了 runtime.Goexit 函数。
  - 加锁，调用 c.wg.Done() 以标记 call 调用完成，然后从 g.m 映射中移除当前 key。
  - 错误处理。
    - 如果 fn 函数中发生了 panic，先判断是否有通道正在等待结果，有的话，新开一个协程去 panic，确保 panic 不能被恢复，这里还用到了 select{} 来阻塞当前线程，保证 panic 之后，当前 goroutine 的信息会出现在堆栈中。如果没有通道正在等待结果，则直接 panic。
    - 如果是 errGoexit 错误，说明 fn 函数中执行了 runtime.Goexit，这时什么都不用做。
  - 结果同步。如果没有发生 error，就向正在等待的通道发送结果。
- 第二个 defer 在一个匿名函数里，它的目的是执行 fn 函数和捕获 panic。如果 fn 函数正常执行完成，normalReturn 就会被设置为 true；在 defer 里，如果 normalReturn 为 false，则说明可能发生了 panic，通过 recover() 函数尝试恢复 panic 并新建一个 panicError 存储信息。
recovered 标志更新：如果 fn 函数非正常执行成功（normalReturn 为 false），则将 recovered 赋值为 true，表示发生了 panic。

call 方法的关键在于使用了双重 defer 机制，结合标志 normalReturn 和 recovered 来判断 fn 函数的状态。normalReturn 和 recovered 有三组值：

normalReturn 为 true，recovered 为 false：表明 fn 函数执行成功，后续执行第一个 defer 时，除了资源清理以外，还会向等待中的通道发送调用完成的结果。
normalReturn 为 false，recovered 为 true：表明在 fn 函数里发生了 panic，并且这个 panic 被成功捕获并恢复。后续执行第一个 defer 时，除了资源清理以外，会再次触发 panic。
normalReturn 为 false，recovered 为 false：这种情况说明在 fn 函数里，调用了 runtime.Goexit 函数终止当前协程，不再执行后续的代码。这意味着 normalReturn = true 和 recovered = true 代码都不可能被执行，因此 normalReturn 和 recovered 的值都为 false。后续执行第一个 defer 时不会向等待的通道发送任何结果，仅仅是进行资源清理。

DoChan

DoChan 方法与 Do 方法类似，但是它返回的是一个通道，通道在操作完成时接收到结果。返回值是通道，意味着我们能以非阻塞的方式等待结果。该方法的源码如下所示：

func (g *Group) DoChan(key string, fn func() (interface{}, error)) <-chan Result {
    // 创建一个通道，类型为 Result
    ch := make(chan Result, 1)
    // 加锁
    g.mu.Lock()
    // 懒初始化 map
    if g.m == nil {
        g.m = make(map[string]*call)
    }
    // 判定该 key 是否有正在进行的调用
    if c, ok := g.m[key]; ok {
        // 重复调用次数加 1
        c.dups++
        // 将新通道添加到通道切片里
        c.chans = append(c.chans, ch)
        // 释放锁
        g.mu.Unlock()
        // 返回通道
        return ch
    }
    // 创建一个 call 实例，并将 ch 通道作为参数传递
    c := &call{chans: []chan<- Result{ch}}
    // 等待组加 1
    c.wg.Add(1)
    // key 和 call 映射
    g.m[key] = c
    // 释放锁
    g.mu.Unlock()
    // 异步执行调用
    go g.doCall(c, key, fn)
    // 返回通道
    return ch
}

DoChan 方法的执行流程如下所示：

1、创建一个大小为 1 的缓冲通道。
2、获取锁：通过 g.mu.Lock() 加锁，确保对内部的 g.m（一个 map，用于跟踪 key 的调用状态）和 c.dups（对于该 key 的重复调用次数）以及 c.chans（通道切片）的访问是并发安全的。
3、初始化 map：如果 g.m == nil，意味着是第一次调用 Do 方法且没有调用过 DoChan 方法，所以初始化 g.m。
4、检查是否有正在进行的调用：通过 c, ok := g.m[key]; ok 检查是否有一个对于该 key 的调用正在进行，如果 ok 为 true，则说明有一个对于该 key 的调用正在进行：
- 增加重复调用次数 c.dups，表示来了一个新的 goroutine 在等待这个调用结果。
- 将新创建的通道追加到当前 call 的通道切片里。
- 释放锁 g.mu.Unlock()，因为不再需要修改共享资源。
- 返回新创建的通道。
5、初始化并异步执行新的调用：如果没有一个对于该 key 的调用正在进行，则：
- 创建一个新的 call 实例，并关联新创建的通道。
- c.wg 等待组计数加 1，标记新操作的开始，后续有相同调用的请求将会等待该操作的完成并共享结果。
- 在 g.m 中注册 key 和新创建的 call 实例的映射 g.m[key] = c。
- 释放锁。
- 异步调用 g.doCall(c, key, fn) 执行实际的函数调用。
- 返回新创建的通道。

DoChan 与 Do 方法的区别在于同步共享结果的方式：

Do 方法：
- 如果有其他请求正在进行（对同一个key），它会使用 sync.WaitGroup 等待这个请求完成以共享结果。
- 如果是针对给定 key 的新请求，它将直接启动 doCall 来执行函数调用，等待执行完成且 call 实例的更新，然后返回结果。
DoChan 方法：为每个调用创建一个新的通道，将其加入到对应 key 的 call 实例的通道切片里，然后返回一个通道。这样，等 g.doCall 正常异步调用完成后，会向各个通道发送结果。

Forget

Forget 方法用于从 g.m 移除特定 key 的调用。

func (g *Group) Forget(key string) {
    // 加锁
    g.mu.Lock()
    // 移除特定的 key
    delete(g.m, key)
    // 释放锁
    g.mu.Unlock()
}

该方法在删除特定 key 前执行加锁操作，保护并发环境下 map 的读写操作，避免并发冲突。

小结

本文对 Go singleflight 的源码进行剖析，该包的主要作用是用于防止重复的请求，它确保给定的 key，函数在同一时间内只执行一次，多个请求共享同一结果。singleflight 能实现这种效果，关键点在于：

将多个相同请求合并成一个请求，确保函数只执行一次：singleflight 为了解决这个问题，引入了互斥锁 sync.Mutex 和 map。

互斥锁用于保护在并发环境下 map 的读写操作，避免并发冲突。

map 则负责将每一个唯一的 key 映射到 call 实例上，该实例包含了fn 函数的返回值和可能的错误等。

遇到一个尚未在 map 中记录的 key 请求时，创建并执行一个新的 call 实例。
如果 map 中已存在该 key 对应的 call 实例，表明有一个相同的请求正在执行或已完成，此时仅需等待此 call 完成并直接其共享结果。
结果共享机制：singleflight 通过阻塞式和非阻塞式两种方式，实现了结果的共享。
- 阻塞式机制：当多个请求通过 Do 方法进行相同的调用时，它们处于等待状态（里面借助了 sync.WaitGroup 来实现阻塞的效果），直到首个请求的 fn 函数的执行完毕。此后，等待的请求会接收到已完成的请求结果。
- 非阻塞式机制：相比于阻塞等待，当请求通过 DoChan 方法发起时，每个请求会立即获得一个专属的通道。这些请求可以继续执行其他操作，直到它们准备好从各自的通道接收结果。在接收结果时，如果结果尚未发送过来，也会暂时处于阻塞状态。