Go singleflight 源码剖析

前言

前面的一篇文章 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 方法发起时，每个请求会立即获得一个专属的通道。这些请求可以继续执行其他操作，直到它们准备好从各自的通道接收结果。在接收结果时，如果结果尚未发送过来，也会暂时处于阻塞状态。

除了以上两个关键点，还需要考虑错误的处理，singleflight 通过使用双重 defer 的机制，用于辨别 函数正常执行完成 、函数里发生了 panic 以及 函数里调用了 runtime.Goexit() 函数 三种情况，每种情况采取不同的处理机制。

作者：陈明勇

个人网站：chenmingyong.cn

文章持续更新，如果本文能让您有所收获，欢迎点赞收藏加关注本掘金号。 微信阅读可搜《程序员陈明勇》。 这篇文章已被收录于 GitHub github.com/chenmingyon...，欢迎大家 Star 催更并持续关注。