1
2
3
4
5
6

func TestSingleFlight(t *testing.T){
	var g singleflight.Group
	g.Do("key",func() (interface{}, error){
		 return "test",nil
	})
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

package singleflight

import "sync"


type call struct {
	wg  sync.WaitGroup  // 阻塞并发函数请求
	val interface{}			// 函数返回值
	err error						// 函数返回值
}

// 一类工作
type Group struct {
	mu sync.Mutex       
	m  map[string]*call // 延迟初始化,用于临时存储函数
}

// 判断 key 是否第一次调用
// 是则调用函数获取结果
// 不是, 则阻塞在 c.wg.Wait() 等待函数返回值
func (g *Group) Do(key string, fn func() (interface{}, error)) (interface{}, error) {
  
	g.mu.Lock()
  // 创建 map
	if g.m == nil {
		g.m = make(map[string]*call)
	}
  
  // 如果已经有这个函数了, 则进入并解锁, 让其它协程全部进入
  // 等待函数返回值，注意 map 的值是指针
	if c, ok := g.m[key]; ok {
		g.mu.Unlock()
    // 等待函数执行结果
		c.wg.Wait()
		return c.val, c.err
	}
  
  // 初始化 call,写入 map 后解锁
	c := new(call)
	c.wg.Add(1)
	g.m[key] = c
	g.mu.Unlock()

  // 执行函数
	c.val, c.err = fn()
  // 已经获取结果, 后进来的协程都可以返回了
	c.wg.Done()
	
  // group 仅仅是执行函数时的临时存储空间
  // 操作结束后, 删除内容
	g.mu.Lock()
	delete(g.m, key)
	g.mu.Unlock()
	
	return c.val, c.err
}