Golang源码分析之golang/sync之singleflight

func main() {
   var flight singleflight.Group
   var errGroup errgroup.Group

   // 模拟并发获取数据缓存
   for i := 0; i < 10; i++ {
      i := i
      errGroup.Go(func() error {
         fmt.Printf("协程%v准备获取缓存\n", i)
         v, err, shared := flight.Do("getCache", func() (interface{}, error) {
            // 模拟获取缓存操作
            fmt.Printf("协程%v正在读数据库获取缓存\n", i)
            time.Sleep(100 * time.Millisecond)
            fmt.Printf("协程%v读取数据库生成缓存成功\n", i)
            return "mockCache", nil
         })
         if err != nil {
            fmt.Printf("err = %v", err)
            return err
         }
         fmt.Printf("协程%v获取缓存成功, v = %v, shared = %v\n", i, v, shared)
         return nil
      })
   }
   if err := errGroup.Wait(); err != nil {
      fmt.Printf("errGroup wait err = %v", err)
   }
}
// 输出：只有0号协程实际生成了缓存，其余协程读取生成的结果
协程0准备获取缓存
协程4准备获取缓存
协程3准备获取缓存
协程2准备获取缓存
协程6准备获取缓存
协程5准备获取缓存
协程7准备获取缓存
协程1准备获取缓存
协程8准备获取缓存
协程9准备获取缓存
协程0正在读数据库获取缓存
协程0读取数据库生成缓存成功
协程0获取缓存成功, v = mockCache, shared = true
协程8获取缓存成功, v = mockCache, shared = true
协程2获取缓存成功, v = mockCache, shared = true
协程6获取缓存成功, v = mockCache, shared = true
协程5获取缓存成功, v = mockCache, shared = true
协程7获取缓存成功, v = mockCache, shared = true
协程9获取缓存成功, v = mockCache, shared = true
协程1获取缓存成功, v = mockCache, shared = true
协程4获取缓存成功, v = mockCache, shared = true
协程3获取缓存成功, v = mockCache, shared = true

func testDoChan() {
   var flight singleflight.Group
   var errGroup errgroup.Group

   // 模拟并发获取数据缓存
   for i := 0; i < 10; i++ {
      i := i
      errGroup.Go(func() error {
         fmt.Printf("协程%v准备获取缓存\n", i)
         ch := flight.DoChan("getCache", func() (interface{}, error) {
            // 模拟获取缓存操作
            fmt.Printf("协程%v正在读数据库获取缓存\n", i)
            time.Sleep(100 * time.Millisecond)
            fmt.Printf("协程%v读取数据库获取缓存成功\n", i)
            return "mockCache", nil
         })
         res := <-ch
         if res.Err != nil {
            fmt.Printf("err = %v", res.Err)
            return res.Err
         }
         fmt.Printf("协程%v获取缓存成功, v = %v, shared = %v\n", i, res.Val, res.Shared)
         return nil
      })
   }
   if err := errGroup.Wait(); err != nil {
      fmt.Printf("errGroup wait err = %v", err)
   }
}
// 输出结果
协程9准备获取缓存
协程0准备获取缓存
协程1准备获取缓存
协程6准备获取缓存
协程5准备获取缓存
协程2准备获取缓存
协程7准备获取缓存
协程8准备获取缓存
协程4准备获取缓存
协程9正在读数据库获取缓存
协程9读取数据库获取缓存成功
协程3准备获取缓存
协程3获取缓存成功, v = mockCache, shared = true
协程8获取缓存成功, v = mockCache, shared = true
协程0获取缓存成功, v = mockCache, shared = true
协程1获取缓存成功, v = mockCache, shared = true
协程6获取缓存成功, v = mockCache, shared = true
协程5获取缓存成功, v = mockCache, shared = true
协程2获取缓存成功, v = mockCache, shared = true
协程7获取缓存成功, v = mockCache, shared = true
协程4获取缓存成功, v = mockCache, shared = true
协程9获取缓存成功, v = mockCache, shared = true

// singleflight.Group
type Group struct {
   mu sync.Mutex       // map的锁
   m  map[string]*call // 保存每个key的调用
}

// 一次Do对应的响应结果
type Result struct {
   Val    interface{}
   Err    error
   Shared bool
}

// 一个key会对应一个call
type call struct {
   wg sync.WaitGroup
   val interface{} // 保存调用的结果
   err error       // 调用出现的err
   // 该call被调用的次数
   dups  int
   // 每次DoChan时都会追加一个chan在该列表
   chans []chan<- Result
}

func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, err error, shared bool) {
   g.mu.Lock()
   if g.m == nil {
      // 第一次执行Do的时候创建map
      g.m = make(map[string]*call)
   }
   // 已经存在该key，对应后续的并发调用
   if c, ok := g.m[key]; ok {
      // 执行次数自增
      c.dups++
      g.mu.Unlock()
      // 等待执行fn的协程完成
      c.wg.Wait()
      // ...
      // 返回执行结果
      return c.val, c.err, true
   }
   
   // 不存在该key，说明第一次调用，初始化一个call
   c := new(call)
   // wg添加1，后续其他协程在该wg上阻塞
   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
}

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

   // 第三步、最后的设置和清理工作
   defer func() {
      // ...
      g.mu.Lock()
      defer g.mu.Unlock()
      // 执行完成，调用wg.Done，其他协程此时不再阻塞，读到fn执行结果
      c.wg.Done()
      // 二次校验map中key的值是否为当前call，并删除该key
      if g.m[key] == c {
         delete(g.m, key)
      }
      // ...
      // 如果c.chans存在，则遍历并写入执行结果
      for _, ch := range c.chans {
          ch <- Result{c.val, c.err, c.dups > 0}
        }
      }
   }()

   // 第一步、执行fn获取结果
   func() {
      // 3、如果fn执行过程中panic，将c.err设置为PanicError
      defer func() {
         if !normalReturn {
            if r := recover(); r != nil {
               c.err = newPanicError(r)
            }
         }
      }()
      // 1、执行fn，获取到执行结果
      c.val, c.err = fn()
      // 2、设置正常返回结果标识
      normalReturn = true
   }()

   // 第二步、fn执行出错，将recovered标识设置为true
   if !normalReturn {
      recovered = true
   }
}

func (g *Group) DoChan(key string, fn func() (interface{}, error)) <-chan Result {
   // 一次调用对应一个chan
   ch := make(chan Result, 1)
   g.mu.Lock()
   if g.m == nil {
      // 第一次调用，初始化map
      g.m = make(map[string]*call)
   }
   // 后续调用，已存在key
   if c, ok := g.m[key]; ok {
      // 调用次数自增
      c.dups++
      // 将chan添加到chans列表
      c.chans = append(c.chans, ch)
      g.mu.Unlock()
      // 直接返回chan，不等待fn执行完成
      return ch
   }

   // 第一次调用，初始化call及chans列表
   c := &call{chans: []chan<- Result{ch}}
   // wg加一
   c.wg.Add(1)
   // 保存key及call的关系
   g.m[key] = c
   g.mu.Unlock()

   // 异步执行fn函数
   go g.doCall(c, key, fn)

   // 直接返回该chan
   return ch
}