github.com/aaabigfish/gopkg@v1.1.0/syncx/errgroup/errgroup.go (about)

     1  // Copyright 2016 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package errgroup provides synchronization, error propagation, and Context
     6  // cancelation for groups of goroutines working on subtasks of a common task.
     7  package errgroup
     8  
     9  import (
    10  	"context"
    11  	"fmt"
    12  	"runtime"
    13  	"sync"
    14  )
    15  
    16  // A Group is a collection of goroutines working on subtasks that are part of
    17  // the same overall task.
    18  //
    19  // A zero Group is valid and does not cancel on error.
    20  type Group struct {
    21  	err     error
    22  	wg      sync.WaitGroup
    23  	errOnce sync.Once
    24  
    25  	workerOnce sync.Once
    26  	ch         chan func(ctx context.Context) error
    27  	chs        []func(ctx context.Context) error
    28  
    29  	ctx    context.Context
    30  	cancel func()
    31  }
    32  
    33  // WithContext create a Group.
    34  // given function from Go will receive this context,
    35  func WithContext(ctx context.Context) *Group {
    36  	return &Group{ctx: ctx}
    37  }
    38  
    39  // WithCancel create a new Group and an associated Context derived from ctx.
    40  //
    41  // given function from Go will receive context derived from this ctx,
    42  // The derived Context is canceled the first time a function passed to Go
    43  // returns a non-nil error or the first time Wait returns, whichever occurs
    44  // first.
    45  func WithCancel(ctx context.Context) *Group {
    46  	ctx, cancel := context.WithCancel(ctx)
    47  	return &Group{ctx: ctx, cancel: cancel}
    48  }
    49  
    50  func (g *Group) do(f func(ctx context.Context) error) {
    51  	ctx := g.ctx
    52  	if ctx == nil {
    53  		ctx = context.Background()
    54  	}
    55  	var err error
    56  	defer func() {
    57  		if r := recover(); r != nil {
    58  			buf := make([]byte, 64<<10)
    59  			buf = buf[:runtime.Stack(buf, false)]
    60  			err = fmt.Errorf("errgroup: panic recovered: %s\n%s", r, buf)
    61  		}
    62  		if err != nil {
    63  			g.errOnce.Do(func() {
    64  				g.err = err
    65  				if g.cancel != nil {
    66  					g.cancel()
    67  				}
    68  			})
    69  		}
    70  		g.wg.Done()
    71  	}()
    72  	err = f(ctx)
    73  }
    74  
    75  // GOMAXPROCS set max goroutine to work.
    76  func (g *Group) GOMAXPROCS(n int) {
    77  	if n <= 0 {
    78  		panic("errgroup: GOMAXPROCS must great than 0")
    79  	}
    80  	g.workerOnce.Do(func() {
    81  		g.ch = make(chan func(context.Context) error, n)
    82  		for i := 0; i < n; i++ {
    83  			go func() {
    84  				for f := range g.ch {
    85  					g.do(f)
    86  				}
    87  			}()
    88  		}
    89  	})
    90  }
    91  
    92  // Go calls the given function in a new goroutine.
    93  //
    94  // The first call to return a non-nil error cancels the group; its error will be
    95  // returned by Wait.
    96  func (g *Group) Go(f func(ctx context.Context) error) {
    97  	g.wg.Add(1)
    98  	if g.ch != nil {
    99  		select {
   100  		case g.ch <- f:
   101  		default:
   102  			g.chs = append(g.chs, f)
   103  		}
   104  		return
   105  	}
   106  	go g.do(f)
   107  }
   108  
   109  // Wait blocks until all function calls from the Go method have returned, then
   110  // returns the first non-nil error (if any) from them.
   111  func (g *Group) Wait() error {
   112  	if g.ch != nil {
   113  		for _, f := range g.chs {
   114  			g.ch <- f
   115  		}
   116  	}
   117  	g.wg.Wait()
   118  	if g.ch != nil {
   119  		close(g.ch) // let all receiver exit
   120  	}
   121  	if g.cancel != nil {
   122  		g.cancel()
   123  	}
   124  	return g.err
   125  }