Go协程安全用法

Go语言内置协程，这极大降低了并发编程的门槛。但这并不意味着并发编程的难度降低了。为此官方提供了 errgroup 包。但这个包无法处理协程 panic 的问题。

在 Go 语言中，我们无法在父协程里捕获子协程的异常(panic)。如果服务有异常没有处理，整个进程都会退出。这会产生很多不可预知的问题。

这里仿照errgroup分享一种跨协程的的异常处理方案，当出现error或panic的时候及时返回不再继续执行未执行方法。

package threading

import (
	"context"
	"fmt"
	"runtime/debug"
	"sync"
)

type token struct{}

// RoutineGroup 协程分组 一组内的协程发生错误或异常可以进行统一处理
type RoutineGroup struct {
	cancel func()         // 发生错误或异常进行取消通知
	wg     sync.WaitGroup // 控制并发协程
	sem    chan token     // 用于限制协程的并发数量
	once   Once           // 确保错误或异常仅执行一次
	err    error          // 第一个 error 信息
	panic  *Panic         // panic 信息
	skip   bool           // 发生错误是否跳过未执行代码 默认跳过
}

// Panic 子协程 panic 会被重新包装，添加调用栈信息
type Panic struct {
	R     interface{} // recover() 返回值
	Stack []byte      // 当时的调用栈
}

func (p Panic) String() string {
	return fmt.Sprintf("%v\n%s", p.R, p.Stack)
}

type RoutineGroupOptionFunc func(*RoutineGroup)

// WithRoutineGroupSkip 发生错误是否跳过未执行代码
func WithRoutineGroupSkip(notSkip bool) RoutineGroupOptionFunc {
	return func(g *RoutineGroup) {
		g.skip = notSkip
	}
}

// WithRoutineGroupLimit 用于限制协程的并发数量
func WithRoutineGroupLimit(n int) RoutineGroupOptionFunc {
	return func(g *RoutineGroup) {
		if n < 0 {
			g.sem = nil
			return
		}
		if len(g.sem) != 0 {
			panic(fmt.Errorf("routine-group: modify limit while %v goroutines in the group are still active", len(g.sem)))
		}
		g.sem = make(chan token, n)
	}
}

// NewRoutineGroup .
func NewRoutineGroup(ctx context.Context, opts ...RoutineGroupOptionFunc) (*RoutineGroup, context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	rg := &RoutineGroup{cancel: cancel, skip: true}
	for _, o := range opts {
		o(rg)
	}
	return rg, ctx
}

func (g *RoutineGroup) done() {
	if g.sem != nil {
		<-g.sem
	}
	if r := recover(); r != nil {
		g.doOnce(func() { g.panic = &Panic{R: r, Stack: debug.Stack()} })
	}
	g.wg.Done()
}

// 确保仅执行一次
// 同时会触发 *RoutineGroup.cancel
func (g *RoutineGroup) doOnce(f func()) {
	g.once.Do(func() {
		f()
		if g.cancel != nil {
			g.cancel()
		}
	})
}

// Wait 等待一组协程结束
func (g *RoutineGroup) Wait() error {
	g.wg.Wait()
	if g.cancel != nil {
		g.cancel()
	}
	if g.panic != nil {
		panic(*g.panic)
	}
	return g.err
}

// Go 新建协程并执行f()，需要跟 Wait 在同一协程使用
// 会根据限制进行携程的阻塞。第一次调用返回非 nil 错误会取消该组； Wait 会返回它的错误。
func (g *RoutineGroup) Go(f func() error) {
	if g.sem != nil {
		g.sem <- token{}
	}

	g.wg.Add(1)
	go func() {
		defer g.done()
		if g.skip && g.once.Did() {
			return
		}
		if err := f(); err != nil {
			g.doOnce(func() { g.err = err })
		}
	}()
}

// TryGo 只在一个新的协程中调用给定的函数
// RoutineGroup 中协程数量需要当前低于配置的限制。
// 返回值报告协程是否启动。
func (g *RoutineGroup) TryGo(f func() error) bool {
	if g.sem != nil {
		select {
		case g.sem <- token{}:
			// Note: this allows barging iff channels in general allow barging.
		default:
			return false
		}
	}

	g.wg.Add(1)
	go func() {
		defer g.done()
		if g.skip && g.once.Did() {
			return
		}
		if err := f(); err != nil {
			g.doOnce(func() { g.err = err })
		}
	}()
	return true
}

package threading

import (
	"sync"
	"sync/atomic"
)

type Once struct {
	done uint32
	m    sync.Mutex
}

// Do 仅执行一次代码
func (o *Once) Do(f func()) {
	if atomic.LoadUint32(&o.done) == 0 {
		o.doSlow(f)
	}
}

func (o *Once) doSlow(f func()) {
	o.m.Lock()
	defer o.m.Unlock()
	if o.done == 0 {
		defer atomic.StoreUint32(&o.done, 1)
		f()
	}
}

// Did 判断是否已经执行过 true则是已经执行过 false未执行过
func (o *Once) Did() bool {
	return atomic.LoadUint32(&o.done) != 0
}

参考：

• https://github.com/zeromicro/go-zero/blob/422f401153/core/threading/routinegroup.go
• golang.org/x/sync/errgroup
• https://taoshu.in/go/goroutine-panic.html