Go 防缓存击穿 —— singleflight

缓存击穿：缓存在某个时间点过期时，突然在这个时间点出现对这个key的大量并发请求，此时缓存已过期，请求会直接落在DB上，使得DB瞬间请求量增大，压力骤增。singleflight能够在同一时间有大量针对同一key的请求这种情况，只让一个请求执行去获取数据，而其他协程阻塞等待结果的返回 

一、数据结构

type Group struct {
	mu sync.Mutex       // 互斥锁
	m  map[string]*call // 对于每一个要获取的key有一个对应的call
}

type call struct {
	wg sync.WaitGroup // 用于阻塞这个call的其他调用请求

	// 调用函数后返回的值和error，只会写一次，并且在wg的Done方法执行后才能被读取
	val interface{}
	err error

	// 用来标识fn方法执行完成之后结果是否立马删除还是保留在singleflight中
	forgotten bool

	// 同时执行一个key的对应call的协程数量
	dups  int
	chans []chan<- Result
}

type Result struct {
	Val    interface{}
	Err    error
	Shared bool
}

二、使用方法

1. Do

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

    // 查看是否是首次执行key的call方法，是的话则阻塞当前协程，等待其他执行call方法的协程唤醒返回数据
	if c, ok := g.m[key]; ok {
		c.dups++
		g.mu.Unlock()
		c.wg.Wait()

		if e, ok := c.err.(*panicError); ok {
			panic(e)
		} else if c.err == errGoexit {
			runtime.Goexit()
		}
		return c.val, c.err, true
	}
    // 首次尝试获取key，则将key放入map中
	c := new(call)
	c.wg.Add(1)
	g.m[key] = c
	g.mu.Unlock()

    // 延迟执行fn方法
	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 func() {
		// 既没有正常执行fn，也没有recover
		if !normalReturn && !recovered {
			c.err = errGoexit
		}

		c.wg.Done()
		g.mu.Lock()
		defer g.mu.Unlock()
        // 如果已经被forget则不用再删除key
		if !c.forgotten {
			delete(g.m, key)
		}

		if e, ok := c.err.(*panicError); ok {
			// 为了避免channel死锁，要保证panic不能被恢复
			if len(c.chans) > 0 {
				go panic(e)
				select {}
			} else {
				panic(e)
			}
		} else if c.err == errGoexit {
			// 已经准备退出就不用做其他操作了
		} else {
			// 正常情况下向channel传递数据
			for _, ch := range c.chans {
				ch <- Result{c.val, c.err, c.dups > 0}
			}
		}
	}()

	func() {
		defer func() {
			if !normalReturn {
				// 若fn panic了则进行recover，并new一个panic型的error
				if r := recover(); r != nil {
					c.err = newPanicError(r)
				}
			}
		}()

        // 执行fn方法
		c.val, c.err = fn()
        // 若fn没有panic则会执行到这
		normalReturn = true
	}()

    // fn panic了，设置recovered
	if !normalReturn {
		recovered = true
	}
}

2. DoChan

// 与Do类似，但是Do是同步返回结果，DoChan返回一个channel，异步将结果返回
func (g *Group) DoChan(key string, fn func() (interface{}, error)) <-chan Result {
	ch := make(chan Result, 1)
	g.mu.Lock()
	if g.m == nil {
		g.m = make(map[string]*call)
	}
	if c, ok := g.m[key]; ok {
		c.dups++
        // 每个等待的协程都有一个对应的channel
		c.chans = append(c.chans, ch)
		g.mu.Unlock()
		return ch
	}
	c := &call{chans: []chan<- Result{ch}}
	c.wg.Add(1)
	g.m[key] = c
	g.mu.Unlock()

	go g.doCall(c, key, fn)

	return ch
}

3. Forget

// 将key从map中删除
func (g *Group) Forget(key string) {
	g.mu.Lock()
	if c, ok := g.m[key]; ok {
		c.forgotten = true
	}
	delete(g.m, key)
	g.mu.Unlock()
}

三、例子

package main

import (
	"fmt"
	"sync"

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

var count int
var mu sync.Mutex

func main() {
	group := new(singleflight.Group)
	var wg sync.WaitGroup
	wg.Add(10)

    // 模拟多个协程同时调用GetData方法获取数据
	for i := 0; i < 10; i++ {
		tmp := i
		go func(tmp int) {
			_, err, _ := group.Do("test", GetData)
			if err != nil {
				return
			}
			fmt.Printf("[%d]get data successfully\n", tmp)
			wg.Done()
		}(tmp)
	}
	wg.Wait()
    fmt.Println("count: ", count)
}

func GetData() (interface{}, error) {

    // 计算GetData的调用次数
	mu.Lock()
	count++
	mu.Unlock()
	fmt.Println("trying to get data...")
	return 1, nil
}



// 输出，可见GetData只调用了一次
trying to get data...
[0]get data successfully
[1]get data successfully
[6]get data successfully
[9]get data successfully
[8]get data successfully
[3]get data successfully
[7]get data successfully
[2]get data successfully
[5]get data successfully
[4]get data successfully
count:  1