singleflight 包原理解析

singleflight 包主要是用來做併發控制，常見的比如防止 快取擊穿 ，我們來模擬一下這種場景：

快取擊穿：快取在某個時間點過期的時候，恰好在這個時間點對這個 Key 有大量的併發請求過來，這些請求發現快取過期一般都會從後端 DB 載入資料並回設到快取，這個時候大併發的請求可能會瞬間把後端 DB 壓垮。

package main

import (
    "errors"
    "log"
    "sync"

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

var errorNotExist = errors.New("not exist")
func main() {
    var wg sync.WaitGroup
    wg.Add(10)

    //模擬10個併發
    for i := 0; i < 10; i++ {
        go func() {
            defer wg.Done()
            data, err := getData("key")
            if err != nil {
                log.Print(err)
                return
            }
            log.Println(data)
        }()
    }
    wg.Wait()
}

//獲取資料
func getData(key string) (string, error) {
    data, err := getDataFromCache(key)
    if err == errorNotExist {
        //模擬從db中獲取資料
        data, err = getDataFromDB(key)
        if err != nil {
            log.Println(err)
            return "", err
        }

        //TOOD: set cache
    } else if err != nil {
        return "", err
    }
    return data, nil
}

//模擬從cache中獲取值，cache中無該值
func getDataFromCache(key string) (string, error) {
    return "", errorNotExist
}

//模擬從資料庫中獲取值
func getDataFromDB(key string) (string, error) {
    log.Printf("get %s from database", key)
    return "data", nil
}

其中通過 getData(key)方法獲取資料，邏輯是：

1. 先嚐試從 cache 中獲取
2. 如果 cache 中不存在就從 db 中獲取

我們模擬了 10 個併發請求，來同時呼叫 getData 函式，執行結果如下：

2020/03/08 17:13:11 get key from database
2020/03/08 17:13:11 data
2020/03/08 17:13:11 get key from database
2020/03/08 17:13:11 data
2020/03/08 17:13:11 get key from database
2020/03/08 17:13:11 data
2020/03/08 17:13:11 get key from database
2020/03/08 17:13:11 data
2020/03/08 17:13:11 get key from database
2020/03/08 17:13:11 data
2020/03/08 17:13:11 get key from database
2020/03/08 17:13:11 data
2020/03/08 17:13:11 get key from database
2020/03/08 17:13:11 data
2020/03/08 17:13:11 get key from database
2020/03/08 17:13:11 data
2020/03/08 17:13:11 get key from database
2020/03/08 17:13:11 data
2020/03/08 17:13:11 get key from database
2020/03/08 17:13:11 data

可以看得到 10 個請求都是走的 db，因為 cache 中不存在該值，當我們利用上 singlefligth 包， getData 改動一下：

import "golang.org/x/sync/singleflight"
var g singleflight.Group
//獲取資料
func getData(key string) (string, error) {
    data, err := getDataFromCache(key)
    if err == errorNotExist {
        //模擬從db中獲取資料
        v, err, _ := g.Do(key, func() (interface{}, error) {
            return getDataFromDB(key)
            //set cache
        })
        if err != nil {
            log.Println(err)
            return "", err
        }

        //TOOD: set cache
        data = v.(string)
    } else if err != nil {
        return "", err
    }
    return data, nil
}

執行結果如下，可以看得到只有一個請求進入的 db，其他的請求也正常返回了值，從而保護了後端 DB。

2020/03/08 17:18:16 get key from database
2020/03/08 17:18:16 data
2020/03/08 17:18:16 data
2020/03/08 17:18:16 data
2020/03/08 17:18:16 data
2020/03/08 17:18:16 data
2020/03/08 17:18:16 data
2020/03/08 17:18:16 data
2020/03/08 17:18:16 data
2020/03/08 17:18:16 data
2020/03/08 17:18:16 data

原理解析

singleflight 在 golang.org/x/sync/singleflight 專案下，對外提供了以下幾個方法

//Do方法，傳入key，以及回撥函式，如果key相同，fn方法只會執行一次，同步等待
//返回值v:表示fn執行結果
//返回值err:表示fn的返回的err
//第三個返回值shared：表示是否是真實fn返回的還是從儲存的map[key]返回的，也就是共享的
func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, err error, shared bool) {
//DoChan方法類似Do方法，只是返回的是一個chan
func (g *Group) DoChan(key string, fn func() (interface{}, error)) <-chan Result {
//暫時未用到：設計Forget 控制key關聯的值是否失效，預設以上兩個方法只要fn方法執行完成後，內部維護的fn的值也刪除（即併發結束後就失效了）
func (g *Group) Forget(key string) {

程式碼也很簡單，我們開啟上看下，擷取了部分： singleflight/singleflight.go

package singleflight // import "golang.org/x/sync/singleflight"

import "sync"

// 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和err用來記錄fn發放執行的返回值
    val interface{}
    err error

    // forgotten indicates whether Forget was called with this call's key
    // while the call was still in flight.
    // 用來標識fn方法執行完成之後結果是否立馬刪除還是保留在singleflight中
    forgotten bool

    // 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 用來記錄fn方法執行的次數
    dups  int
    //用來記錄DoChan中呼叫次數以及需要返回的資料
    chans []chan<- Result
}

// 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
}


// Do executes and returns the results of the given function, making
// sure that only one execution is in-flight for a given key at a
// time. If a duplicate comes in, the duplicate caller waits for the
// original to complete and receives the same results.
// The return value shared indicates whether v was given to multiple callers.
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)
    }
    //check map是否已經存在值
    if c, ok := g.m[key]; ok {
        c.dups++
        g.mu.Unlock()
        c.wg.Wait()
        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
}
//執行fn方法，並且wg.Done
// doCall handles the single call for a key.
func (g *Group) doCall(c *call, key string, fn func() (interface{}, error)) {
    c.val, c.err = fn()
    c.wg.Done()

    ...
}

在 Do 方法中主要是通過 waitgroup 來控制的，主要流程如下：

1. 在 Group 中設定了一個 map，如果 key 不存在，則例項化 call(用來儲存值資訊)，並將 key=>call 的對應關係存入 map 中（通過 mutex 保證了併發安全）
2. 如果已經在呼叫中則 key 已經存在 map，則 wg.Wait
3. 在 fn 執行結束之後（在 doCall 方法中執行）執行 wg.Done
4. 卡在第 2 步的方法得到執行，返回結果

其他的 DoChan 方法也是類似的邏輯，只是返回的是一個 chan。

文中用到的例項程式碼放在 github 上：

https://github.com/go-demo/singleflight-demo

個人公眾號：