github.com/rudderlabs/rudder-go-kit@v0.30.0/sync/group.go (about)

     1  package sync
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  // A EagerGroup is a collection of goroutines working on subtasks that are part of
     9  // the same overall task.
    10  //
    11  // Use NewEagerGroup to create a new group.
    12  type EagerGroup struct {
    13  	ctx     context.Context
    14  	cancel  context.CancelCauseFunc
    15  	wg      sync.WaitGroup
    16  	sem     chan struct{}
    17  	errOnce sync.Once
    18  	err     error
    19  }
    20  
    21  // NewEagerGroup returns a new eager group and an associated Context derived from ctx.
    22  //
    23  // The derived Context is canceled the first time a function passed to Go
    24  // returns a non-nil error or the first time Wait returns, whichever occurs
    25  // first.
    26  //
    27  // limit < 1 means no limit on the number of active goroutines.
    28  func NewEagerGroup(ctx context.Context, limit int) (*EagerGroup, context.Context) {
    29  	ctx, cancel := context.WithCancelCause(ctx)
    30  	g := &EagerGroup{
    31  		ctx:    ctx,
    32  		cancel: cancel,
    33  	}
    34  	if limit > 0 {
    35  		g.sem = make(chan struct{}, limit)
    36  	}
    37  	return g, ctx
    38  }
    39  
    40  // Go calls the given function in a new goroutine.
    41  // It blocks until the new goroutine can be added without the number of
    42  // active goroutines in the group exceeding the configured limit.
    43  //
    44  // The first call to return a non-nil error cancels the group's context.
    45  // The error will be returned by Wait.
    46  //
    47  // If the group was created by calling NewEagerGroup with limit < 1, there is no
    48  // limit on the number of active goroutines.
    49  //
    50  // If the group's context is canceled, routines that have not executed yet due to the limit won't be executed.
    51  // Additionally, there is a best effort not to execute `f()` once the context is canceled
    52  // and that happens whether or not a limit has been specified.
    53  func (g *EagerGroup) Go(f func() error) {
    54  	if err := g.ctx.Err(); err != nil {
    55  		g.errOnce.Do(func() {
    56  			g.err = g.ctx.Err()
    57  			g.cancel(g.err)
    58  		})
    59  		return
    60  	}
    61  
    62  	if g.sem != nil {
    63  		select {
    64  		case <-g.ctx.Done():
    65  			g.errOnce.Do(func() {
    66  				g.err = g.ctx.Err()
    67  				g.cancel(g.err)
    68  			})
    69  			return
    70  		case g.sem <- struct{}{}:
    71  		}
    72  	}
    73  
    74  	g.wg.Add(1)
    75  	go func() {
    76  		err := g.ctx.Err()
    77  		if err == nil {
    78  			err = f()
    79  		}
    80  		if err != nil {
    81  			g.errOnce.Do(func() {
    82  				g.err = err
    83  				g.cancel(g.err)
    84  			})
    85  		}
    86  		if g.sem != nil {
    87  			<-g.sem
    88  		}
    89  		g.wg.Done()
    90  	}()
    91  }
    92  
    93  // Wait blocks until all function calls from the Go method have returned, then
    94  // returns the first non-nil error (if any) from them.
    95  func (g *EagerGroup) Wait() error {
    96  	g.wg.Wait()
    97  	g.cancel(g.err)
    98  	return g.err
    99  }