Go并发聚合请求利器——singleflight的源码详解

// 普通调用方法
func callFunc(i int) (int,error) {
	time.Sleep(500 * time.Millisecond)
	return i, nil
}

// 使用singleflight

// 1. 定义全局变量
var sf singleflight.Group

func callFuncBySF(key string, i int) (int, error) {
	// 2. 调用sf.Do方法
	value, err, shared := sf.Do(key, func() (interface{}, error) {
		return callFunc(i)
	})
	res, _ := value.(int)
	return res, err
}

// 使用DoChan进行超时控制 
func CtrTimeout(ctx context.Context, req interface{}){
	ch := g.DoChan(key, func() (interface{}, error) {
	    return call(ctx, req)
	})

	select {
	case <-time.After(500 * time.Millisecond): 
		    return
	case <-ctx.Done()
				return
	case ret := <-ch: 
	    go handle(ret)
	}
}

// 另外启用协程定时删除key，提高请求下游次数，提高成功率
func CtrRate(ctx context.Context, req interface{}){
	res, _, shared := g.Do(key, func() (interface{}, error) {
			// 另外其一个goroutine，等待一段时间后，删除key
			// 删除key后的调用，会重新执行Do
	    go func() {
	        time.Sleep(10 * time.Millisecond)
	        g.Forget(key)
	    }()

	    return call(ctx, req)
		})

	handle(res)
}

// Group 对外的核心结构体
type Group struct {
	mu sync.Mutex       // 保护 m
	m  map[string]*call // lazily initialized
}

// Do 执行函数, 对同一个 key 多次调用的时候，在第一次调用没有执行完的时候
// 只会执行一次 fn，其他的调用会**阻塞**住等待这次调用返回
// v, err 是传入的 fn 的返回值
// shared 表示fn的结果是否被共享
func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, err error, shared bool)

// DoChan 和 Do 类似，只是 DoChan 返回一个 channel，也就是同步与异步的区别
func (g *Group) DoChan(key string, fn func() (interface{}, error)) <-chan Result

// Forget 用于通知 Group 删除某个 key 这样后面继续这个 key 的调用的时候就不会在阻塞等待了
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()
}

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

	// val和err是用户调用返回的字段，（err包括panic err）
  // 其在wg.done之前只被写入一次
	// 在wg.done之后，只能被读取
	val interface{}
	err error

	// 如果执行过Forget，则会被设置为true，避免重复delete
	forgotten bool

	dups  int  // 此call被额外调用的次数，也是结果被额外共享的次数（不算首次）
	chans []chan<- Result //DoChan的结果数组，用于一次执行后，给所有结果传值
}

// Do 内部没有额外的goroutine执行，故panic可以被捕获
func (g *Group) Do(key string, fn func() (interface{}, error)) (v interface{}, err error, shared bool) {
	// 1. 加锁并懒加载内部变量
	g.mu.Lock()
	if g.m == nil {
		g.m = make(map[string]*call)
	}

	// 2. 如果callMap里存在这个key，说明直接在此之前此函数**正在被调用**
	// c为singleflight封装的call结构体
	if c, ok := g.m[key]; ok {
		
		c.dups++      // 2.1 记录此call被执行的次数+1
		g.mu.Unlock() // 2.2 释放锁，让之后被调用的函数也进来
		c.wg.Wait()   // 2.3 阻塞在这里，等待函数调用完成
		
		// 2.4 函数调用完成，进行错误处理，区分系统错误or用户错误
		if e, ok := c.err.(*panicError); ok {
			panic(e)
		} else if c.err == errGoexit {
			runtime.Goexit()
		}
		// 2.5 返回函数执行结果，其结果必定是共享的结果
		return c.val, c.err, true
	}

	// 3. 首次调用，新建call
	c := new(call)
	c.wg.Add(1)
	g.m[key] = c
	g.mu.Unlock() // 完成对map操作
	
	// 4. 执行fn函数，并将执行结果返回，并调用c.wg.Done(), 也会删除callMap里对应的key
	g.doCall(c, key, fn)
	return c.val, c.err, c.dups > 0 // c.dups代表此call被额外调用次数，>0说明结果共享
}

// 给DoChan的返回封装的结果，
type Result struct {
	Val    interface{}
	Err    error
	Shared bool
}

func (g *Group) DoChan(key string, fn func() (interface{}, error)) <-chan Result {
	// 本次执行的结果，一次调用只有1个
	ch := make(chan Result, 1)

	// 1. 加锁并懒加载内部变量
	g.mu.Lock()
	if g.m == nil {
		g.m = make(map[string]*call)
	}
	
	// 2. 如果此key对应的函数已在执行 
	if c, ok := g.m[key]; ok {
		c.dups++
		// 将此结果加入call（此处的c）的结果数组
		c.chans = append(c.chans, ch)
		g.mu.Unlock()
		return ch
	}

	// 3. 如果是首次执行，创建call，加入到Group的CallMap里
	// 之所以是结果数组，是为了用于doCall一次执行后，给所有结果传值
	c := &call{chans: []chan<- Result{ch}}	// doCall内部只准写入值
	c.wg.Add(1)
	g.m[key] = c
	g.mu.Unlock()
	
	// 4. 开goroutine通过doCall执行fn
	// 执行完成后会调用c.wg.Done(), 也会删除callMap里对应的key
	go g.doCall(c, key, fn)

	return ch
}

// 原始写法，存在panic未处理和死锁问题
func (g *Group) doCall(c *call, key string, fn func() (any, error)) {
	// 1. 执行
	c.val, c.err = fn()
	// 2. 执行完成后，立刻通知其他等待的goroutine获取结果
	c.wg.Done()
	
	// 3. 删除key，之后的函数重新执行调用
	g.mu.Lock()
	delete(g.m, key)
	// 4. 将得到的结果写入结果数组里的每个chan
	for _, ch := range c.chans {
		ch <- Result{c.val, c.err, c.dups > 0}
	}
	g.mu.Unlock()
}

// doCall handles the single call for a key.
func (g *Group) doCall(c *call, key string, fn func() (interface{}, error)) {
	normalReturn := false
	recovered := false

	// 3. 执行第二个defer，对panic或Goexit统一处理
	defer func() {
		// 既不是普通返回，也不是recover，排除法得到Goexit
		if !normalReturn && !recovered {
			c.err = errGoexit
		}
	
		// 3.1 执行完成，返阻塞在调用Group.Do调用
		// Group.Do的部分到这里就返回结果了
		c.wg.Done()

		// 3.2 加锁删除key
		g.mu.Lock()
		defer g.mu.Unlock()
		if !c.forgotten {
			delete(g.m, key)
		}
		
		// 3.3 处理panic
		if e, ok := c.err.(*panicError); ok {
			// DoChan调用
			//    由于ch发生了panic，ch不会被写入，goroutine会一直阻塞，导致死锁
			//    这里通过go panic的方式保证必定panic，
			//    预防调用Group.DoChan后，外部recover导致死锁的问题
			if len(c.chans) > 0 {
				go panic(e)
				select {} // 保留这个goroutine到核心存储
			
			} else { // Do 调用，直接panic 
				panic(e)
			}

		// 3.4 处理Goexit
		} else if c.err == errGoexit {
			// 已经准备退出了，没有啥要处理的，资源在之前就已经释放完成了
		
		// 3.5 处理正常情况下的DoChan结果
		} else {
			for _, ch := range c.chans {
				ch <- Result{c.val, c.err, c.dups > 0}
			}
		}
	}()

	// 1. 首次执行此函数
	func() {
		defer func() {
			// 1.3 判断是否为正常返回
			if !normalReturn {
				// 非正常返回则recover，保留堆栈信息
				// 在下一个defer里统一处理 panic or Goexit
				if r := recover(); r != nil { 
					c.err = newPanicError(r)
				}
			}
		}()

		// 1.1 执行fn
		c.val, c.err = fn()
		
		// 1.2 执行成功则设置正常返回
		normalReturn = true // 如果fn() panic 或者 Goexit则不会执行此步骤
	}()

	// 2. 由于panic被第一个defer recover了，可以执行到此
	// Goexit会直接执行下一个defer
	if !normalReturn {
		recovered = true
	}
}

// 模拟并发请求
go func(wg *sync.WaitGroup) {
	// 3. 某个并发调用，释放锁
  defer wg.Done()
	
	// 4. 处理panic
	// Do内部没有另外执行的goroutine，所以可以被捕获
  defer func() {
    if r := recover(); r != nil {
		// 此处可以捕获到panic. 由 doCall 方法中捕获后再次抛出的异常
      fmt.Println("[[[", r, "]]]") 
    }
  }()

	// 1. 执行Do,返回结果
  val, err, _ = g.Do("getdata", func() (interface{}, error) {
    panic("panic")
    return "hello", nil
  })
	
	// 2. 处理结果
	handle(val)
	.....
}(wg)

// 模拟并发请求
go func(wg *sync.WaitGroup) {
	
	// 4. 处理panic
	// doCall内部对于DoChan使用了go panic，此时无法被recover
  defer func() {
    if r := recover(); r != nil {
      fmt.Println("[[[", r, "]]]") 
    }
  }()

	// 1. 执行Do,返回结果
  ch := g.DoChan("getdata", func() (interface{}, error) {
    panic("panic")
    return "hello", nil
  })
	
	// 2. 处理结果
	go func(){
		// 3. 此结果完成，执行done
	  defer wg.Done()

		str <- ch
	}
	.....
}(wg)

    > 当业务函数fn发生panic或者goexit时，可能**影响框架对chan传值，
    从而导致有其他协程阻塞等待chan，进而死锁**

    > 

func example(fn func()) {
	normalReturn := false
	recovered := false
	defer func() {
		if !normalReturn && !recovered {
			log.Info("get go exit")
		}
	}()

	func() {
		defer func() {
			if !normalReturn { // 减少recover调用次数
				if r := recover(); r != nil {
					recovered = true
					log.Info("get panic")
				}
			}
		}()

		fn()
		normalReturn = true
	}()

}