使用场景

当缓存失效时，对于同一个key的查询，可以使用singleflight控制只有一个请求去访问数据库，防止缓存击穿。

var g singleflight.Group

func singlefilghtDemo() {
   var wg sync.WaitGroup
   wg.Add(10)

   for i := 0; i < 10; i++ {
      go func() {
         defer wg.Done()
         query("key")
      }()
   }
}

func query(key string) string {
   val, err := queryFromRedis("key")
   if err != nil {
      val = queryFromDB("key")
   }
   return val
}

func queryFromRedis(key string) (string, error) {
   return "", errors.New("not found")
}

func queryFromDB(key string) string {
   val, _, _ := g.Do(key, func() (interface{}, error) {
      return "val", nil
   })
   return val.(string)
}

源码分析

首先看看定义的三个结构体。

Group

Group代表了一个工作类，并形成一个工作空间，在这个工作空间中的工作单元可以被抑制重复执行。每个Group由一个互斥锁和一个映射表组成，映射表保存了当前正在执行或者已经执行完成的调用信息。

type Group struct {
   mu sync.Mutex         // 控制并发
   m  map[string]*call   // 保存了正在执行或者已经执行完成的调用
}

call

val和err保存了singleflight.Do调用方法返回的结果。这个结果只会被赋值一次，在WaitGroup结束前被赋值，在WaitGroup结束后可以被同一个WaitGroup阻塞的其他请求读取。

type call struct {
   wg sync.WaitGroup

   val interface{}  // 记录调用结果
   err error

   forgotten bool

   dups  int       // 重复请求的数量
   chans []chan<- Result  // 用于结果同步
}

Result

Result保存了Do返回的结果。

// Result holds the results of Do, so they can be passed
// on a channel.
type Result struct {
   Val    interface{}
   Err    error
   Shared bool   // 表示返回数据是调用 fn 得到的还是其他相同 key 调用返回的
}

Do

当同一时间有大量重复请求打进来时，只有一个goroutine调用doCall执行fn方法，其他goroutine则会被WaitGroup阻塞，直到doCall方法执行完成。

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

doCall

doCall处理了key的单次调用。当调用完成后，会把key对应的调用结果从map中删除，避免后续请求无法获取到最新的结果。

func (g *Group) doCall(c *call, key string, fn func() (interface{}, error)) {
   c.val, c.err = fn()
   c.wg.Done()

   g.mu.Lock()
   if !c.forgotten {
      delete(g.m, key)
   }
   for _, ch := range c.chans {
      ch <- Result{c.val, c.err, c.dups > 0}
   }
   g.mu.Unlock()
}

Forget

Forget用来通知singleflight移除某个Key，让后续请求调用doCall的逻辑，而不是直接沿用上次调用返回的结果。当有多个请求同时打进来时，如果doCall执行的方法刚好超时或者执行失败，可以调用Forget主动移除map中的key，重新调用doCall执行请求。

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

DoChan

DoChan() 通过 channel 返回结果。因此可以使用 select 语句实现超时控制。

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

参考资料

www.cyningsun.com/01-11-2021/golang-concurrency-singleflight.html

www.lixueduan.com/posts/go/si…

draveness.me/golang/docs…

go并发编程-singleflight