是什么

singleflight 是一个 golang 的 package, 主要用于限制函数的并发执行。我们可以将 singleflight 理解为单一飞行, 即对于同一个 key 指定的函数, 同一时刻只能有一个函数在执行; 如果函数执行期间有相同 key 的请求过来, 则等待第一个请求执行, 并复用其结果。

下面将通过一个简单的例子测试它的功能：

package main

import (
    "fmt"
    "sync"
    "sync/atomic"
    "time"

    "golang.org/x/sync/singleflight"
)

func main() {
    g := singleflight.Group{}
    key := "key"                // 对相同的 key, 同一时刻只会执行一次
    calls := int32(0)           // fn 函数执行的次数
    wg1 := sync.WaitGroup{}     // 控制代码执行位置
    wg2 := sync.WaitGroup{}     // 控制所有 goroutine 执行完毕
    c := make(chan struct{}, 1) // 控制 fn 继续执行

    fn := func() (interface{}, error) {
        if atomic.AddInt32(&calls, 1) == 1 { // 第一次执行
            wg1.Done()
        }
        <-c
        c <- struct{}{}
        time.Sleep(10 * time.Millisecond) // 让更多的 goroutine 进入 g.Do(key, fn) 函数

        return "fn", nil
    }

    num := 10
    wg1.Add(1)
    for i := 0; i < num; i++ {
        wg1.Add(1)
        wg2.Add(1)
        go func() { // 启动 10 个 goroutine 调用 singleflight.Do() 函数
            defer wg2.Done()
            wg1.Done()
            g.Do(key, fn) // fn 函数只会执行一次
        }()
    }

    wg1.Wait()                            // 全部的 goroutine 均已经启动且至少执行到了 g.Do(key, fn)这一行, 已经有 goroutine 进入 fn 函数
    c <- struct{}{}                       // fn 函数继续执行
    wg2.Wait()                            // 所有 goroutine 均已执行完
    fmt.Println(atomic.LoadInt32(&calls)) // fn 共执行一次
    g.Do(key, fn)
    fmt.Println(atomic.LoadInt32(&calls)) // fn 共执行两次
}

上面的例子中验证了 singleflight 用途。如下图所示, 同一时刻只能有一个 goroutine 执行 fn 函数。此时, 如果 g2、g3 调用 singleflight.Do(key, fn), 则等待 g1 执行完毕, 并复用其结果; 对于 groutine4, 此时无 goroutine 占用函数, 则直接执行。

应用场景

从上面的例子可以看出, singleflight 的作用是保证相同 key 的函数在同一时刻只能有一个 goroutine 再执行。因此, 最常见的应用场景就是限流。在高并发的场景下, 为了防止将下游打垮, 往往对下游进行限流访问。业务上, 大家习惯于先读 redis, 缓存不命中就访问 db。当某个热 key 不命中时, 大量的请求访问 db 从而导致缓存击穿。使用 singleflight 可以有效的解决缓存击穿问题, 它保证了相同 key 在同一时刻只有一个 goroutine 会执行回源函数。

singleflight 的优点是可以进行并发控制, 且其他请求复用同一个结果。但是, 我们需要知道 singleflight 毕竟是单实例的, 如果线上有很多实例「上万」, 则下游的并发请求就是实例的数量。如果有大量不重复 key 无法命中缓存, 那请求数据就是 distinct(key) * count(pods)。

动手实现

我们已经对 singleflight 的功能有了细致的了解。如果让你来设计一个 pkg 来实现这个需求, 你会如何做呢？在这个需求中, 核心能力是对 key 的执行进行并发控制、并发请求的结果复用。

定义一个结构体, 用于初始化 singleflight 对象

type Group struct {
    mutext sync.Mutex       // 对 map 并发访问控制
    m      map[string]*call // 保存正在执行的函数
}

type call struct {
    wg    sync.WaitGroup // 等待函数执行结束，复用其结果
    dumps int            // 调用次数
    val   interface{}
    err   error
}

实现 do 函数, 实现对函数访问的并发控制

type fn func() (interface{}, error) // 需要执行的函数

func (g *Group) Do(key string, fn fn) (val interface{}, err error, shared bool) {
    g.mutext.Lock()
    if g.m == nil { // 初始化
        g.m = make(map[string]*call)
    }
    if call, ok := g.m[key]; ok { // 已经有 goroutine 在执行 fn
        call.dumps++
        g.mutext.Unlock()
        call.wg.Wait() // 等待执行 fn 函数的 goroutine 返回结果
        return call.val, call.err, true
    }

    call := &call{wg: sync.WaitGroup{}} // 初始化一次新的调用
    call.wg.Add(1)
    g.m[key] = call
    g.mutext.Unlock()

    func() { // 执行 fn 函数
        defer func() { // fn 执行完毕, 从 map 中删除
            g.mutext.Lock()
            defer g.mutext.Unlock()
            call.wg.Done()
            if g.m[key] == call {
                delete(g.m, key)
            }
        }()
        call.val, call.err = fn()
    }()

    return call.val, call.err, call.dumps > 0
}

源码解读

上面的示例中, 简单的实现了一个简单版本 singleflight。在官方的 pkg 中, 除了支持 Do(key, fn) 之外, 还支持 DoChan(key, fn), 将结果返回到 channel 中。此外, 官方的 pkg 还对 panic 和 goroutine 退出进行了处理。

在实现 Do 函数时, 源码中抽出了一个子函数 doCall, 同时在 Do 函数中处理 panic 和 goexit。

func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, err error, shared bool) {
    g.mu.Lock()
    if g.m == nil {
        g.m = make(map[string]*call)
    }
    if c, ok := g.m[key]; ok {
        c.dups++
        g.mu.Unlock()
        c.wg.Wait()

        if e, ok := c.err.(*panicError); ok { // fn 函数发生 panic 则所有的调用方也要 panic
            panic(e)
        } else if c.err == errGoexit { // goroutine 发生退出, 直接退出
            runtime.Goexit()
        }
        return c.val, c.err, true
    }
    c := new(call)
    c.wg.Add(1)
    g.m[key] = c
    g.mu.Unlock()

    g.doCall(c, key, fn)
    return c.val, c.err, c.dups > 0
}

func (g *Group) doCall(c *call, key string, fn func() (interface{}, error)) {
    normalReturn := false
    recovered := false

    // 使用两个 defer 来区分 panic 和 runtime.Goexit
    defer func() { // 第二个 defer 用于判断是否发生了 runtime.Goexit
        // 执行了 runtime.Goexit
        if !normalReturn && !recovered {
            c.err = errGoexit
        }

        g.mu.Lock()
        defer g.mu.Unlock()
        c.wg.Done()
        if g.m[key] == c {
            delete(g.m, key)
        }

        if e, ok := c.err.(*panicError); ok {
            // 为了防止等待中的 channels 永远被阻塞，需要确保 panic 不能被 recovered.
            if len(c.chans) > 0 {
                go panic(e)
                select {} // Keep this goroutine around so that it will appear in the crash dump.
            } else {
                panic(e)
            }
        } else if c.err == errGoexit {
            // Already in the process of goexit, no need to call again
        } else {
            // Normal return
            for _, ch := range c.chans {
                ch <- Result{c.val, c.err, c.dups > 0}
            }
        }
    }()

    func() {
        defer func() {  // 第一个 defer 判断是否发生了 panic. 如果发生了 panic 将设置 recovered 的值
            if !normalReturn {
                if r := recover(); r != nil {
                    c.err = newPanicError(r)
                }
            }
        }()

        c.val, c.err = fn()
        normalReturn = true
    }()

    if !normalReturn {
        recovered = true
    }
}

总结

singleflight 是一个很简单的 pkg, 代码量比较少, 很值得大家去学习源码。在真实的业务开发中, 通常将其用于对下游服务的限流访问, 可以保证同一时刻只有一个请求到下游。这里需要注意的是, 如果对于某些请求, 相同请求返回的结果不同, 那是无法使用 singleflight 的。此外, 通过源码发现, singleflight 使用的是一个全局的 mutext, 对于 key 比较多的请求, 会发生较多的抢锁操作, 产生性能问题。