github.com/songzhibin97/gkit@v1.2.13/egroup/errgroup.go (about)

     1  package egroup
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/songzhibin97/gkit/goroutine"
     8  )
     9  
    10  type Group struct {
    11  	ctx    context.Context
    12  	cancel func()
    13  	wg     sync.WaitGroup
    14  	sync.Once
    15  	goroutine goroutine.GGroup
    16  	err       error
    17  }
    18  
    19  func WithContextGroup(ctx context.Context, group goroutine.GGroup) *Group {
    20  	g := &Group{}
    21  	g.ctx, g.cancel = context.WithCancel(ctx)
    22  	g.goroutine = group
    23  	return g
    24  }
    25  
    26  // WithContext 实例化方法
    27  func WithContext(ctx context.Context) *Group {
    28  	g := &Group{}
    29  	g.ctx, g.cancel = context.WithCancel(ctx)
    30  	g.goroutine = goroutine.NewGoroutine(ctx)
    31  	return g
    32  }
    33  
    34  // Wait 等待
    35  func (g *Group) Wait() error {
    36  	g.wg.Wait()
    37  	if g.cancel != nil {
    38  		g.cancel()
    39  	}
    40  	return g.err
    41  }
    42  
    43  // Go 异步调用
    44  func (g *Group) Go(f func() error) {
    45  	g.wg.Add(1)
    46  	g.goroutine.AddTask(func() {
    47  		defer g.wg.Done()
    48  		if err := f(); err != nil {
    49  			g.Do(func() {
    50  				// 级联取消
    51  				g.err = err
    52  				if g.cancel != nil {
    53  					g.cancel()
    54  				}
    55  			})
    56  		}
    57  	})
    58  }