singleflight中入分析

singleflight

面试被问过两次：有没有看过singleflight源码？在同一个坑里我栽了两次，下一次一定要站起来。

源码分析

Group

当我们调用 singleflight.Group{}是，返回的是一个Group类型的对象。

type Group struct {
	mu sync.Mutex       // 避免m发生并发读写问题
	m  map[string]*call // 懒加载
}

// call is an in-flight or completed singleflight.Do call
type call struct {
	wg sync.WaitGroup

    // 存储call中函数执行的结果和可能发生的错误。只在wg.Wait()返回之后才被读取。
	val interface{}
	err error

	// 记录与当前call关联的重复请求的数量
	dups  int
    
    // 只写通道的切片，用于存储那些等待当前call结果的goroutine的通道
	chans []chan<- Result
}

Do

当key已存在时，调用Do方法时会陷入阻塞，直至方法结束，返回结果。接下来看一下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)
	}
    if c, ok := g.m[key]; ok {	// key存在，等待c.wg.Wait()结束,返回val,err
		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
	}
	c := new(call)	// key不存在，新建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
}

当key已存在时，通过c.wg.Wait()阻塞协程，直至协程恢复返回 c.val，c.err；当key不存在时，新建call，调用doCall方法执行fn。

doCall

fn在g.doCall中执行，当执行结束时执行c.wg.Done()方法，阻塞中的goroutine恢复。

doCall方法内部定义了两个布尔类型的参数：

• normalReturn：用于标记函数fn是否通过正常路径返回。
• recovered：用于标记是否从panic中恢复。

执行fn，并捕获fn中可能存在的panic。如果没有发生panic，此时已经对c.val、c.err赋值。fn中可能存在runtime.Goexit()，会立即终止当前协程,并且runtime.Goexit不属于panic。

func() {
		defer func() {
			if !normalReturn {
                // 捕获panic，设置c.err
				if r := recover(); r != nil {
					c.err = newPanicError(r)
				}
			}
		}()

		c.val, c.err = fn()
		// fn 正常执行
		normalReturn = true
	}()

判断fn是否发生panic。

	if !normalReturn {
		recovered = true
	}

根据normalReturn、recovered结果处理。

defer func() {
  
  // 判断是否为 rutime.Goexit
	if !normalReturn && !recovered {
		c.err = errGoexit
	}  
  
   
   g.mu.Lock()
   defer g.mu.Unlock()
   c.wg.Done()	// 此时，其他具有相同key阻塞的协程恢复
   if g.m[key] == c {	
      delete(g.m, key)
   }

   if e, ok := c.err.(*panicError); ok {
      if len(c.chans) > 0 { // 如果调用了DoChan方法，为了避免channel死锁，需要抛出panic
         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 {
      // runtime错误不需要做任何处理，已经退出了
   } else {
      // 正常返回
      for _, ch := range c.chans {
         ch <- Result{c.val, c.err, c.dups > 0}
      }
   }
}()

整体代码如下：

// doCall handles the single call for a key.
func (g *Group) doCall(c *call, key string, fn func() (interface{}, error)) {
	normalReturn := false	// 是否正常返回
	recovered := false	// 是否从panic中恢复


	defer func() {
		// 判断错误是否为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 {
			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 {
			// runtime错误不需要做任何处理，已经退出了
		} else {
			// 正常返回
			for _, ch := range c.chans {
				ch <- Result{c.val, c.err, c.dups > 0}
			}
		}
	}()

	func() {
		defer func() {
			if !normalReturn {
                // 如果发生panic，捕获并处理
				if r := recover(); r != nil {
					c.err = newPanicError(r)
				}
			}
		}()
        // 执行fn函数
		c.val, c.err = fn()
        // fn正常执行
		normalReturn = true
	}()

	if !normalReturn {
		recovered = true
	}
}

DoChan

调用DoChan方法会返回一个只读的chan，后续结果会通过这个chan返回。call中会维持需要通知的chan切片，当fn执行结束后遍历返回。

func (g *Group) DoChan(key string, fn func() (interface{}, error)) <-chan Result {
	ch := make(chan Result, 1)	// 缓存区大小为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
}

Forget

删就完了。

func (g *Group) Forget(key string) {
	g.mu.Lock()
	delete(g.m, key)
	g.mu.Unlock()
}

小结

singleflight基于Group实现，Group 通过map[key]*call，快速判断key是否存在，通过sync.Metu保证map并发安全。

基于call结构体中的sync.WaitGroup实现了Do方法（相同key阻塞等待），基于 []chan <- Result 实现了 DoChan方法。

ps：感觉小结的不怎么滴，但是不小结又好像少点什么