github.com/laof/lite-speed-test@v0.0.0-20230930011949-1f39b7037845/common/picker/picker.go (about)

     1  package picker
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  // Picker provides synchronization, and Context cancelation
    10  // for groups of goroutines working on subtasks of a common task.
    11  // Inspired by errGroup
    12  type Picker struct {
    13  	ctx    context.Context
    14  	cancel func()
    15  
    16  	wg sync.WaitGroup
    17  
    18  	once    sync.Once
    19  	errOnce sync.Once
    20  	result  interface{}
    21  	err     error
    22  }
    23  
    24  func newPicker(ctx context.Context, cancel func()) *Picker {
    25  	return &Picker{
    26  		ctx:    ctx,
    27  		cancel: cancel,
    28  	}
    29  }
    30  
    31  // WithContext returns a new Picker and an associated Context derived from ctx.
    32  // and cancel when first element return.
    33  func WithContext(ctx context.Context) (*Picker, context.Context) {
    34  	ctx, cancel := context.WithCancel(ctx)
    35  	return newPicker(ctx, cancel), ctx
    36  }
    37  
    38  // WithTimeout returns a new Picker and an associated Context derived from ctx with timeout.
    39  func WithTimeout(ctx context.Context, timeout time.Duration) (*Picker, context.Context) {
    40  	ctx, cancel := context.WithTimeout(ctx, timeout)
    41  	return newPicker(ctx, cancel), ctx
    42  }
    43  
    44  // Wait blocks until all function calls from the Go method have returned,
    45  // then returns the first nil error result (if any) from them.
    46  func (p *Picker) Wait() interface{} {
    47  	p.wg.Wait()
    48  	if p.cancel != nil {
    49  		p.cancel()
    50  	}
    51  	return p.result
    52  }
    53  
    54  // Error return the first error (if all success return nil)
    55  func (p *Picker) Error() error {
    56  	return p.err
    57  }
    58  
    59  // Go calls the given function in a new goroutine.
    60  // The first call to return a nil error cancels the group; its result will be returned by Wait.
    61  func (p *Picker) Go(f func() (interface{}, error)) {
    62  	p.wg.Add(1)
    63  
    64  	go func() {
    65  		defer p.wg.Done()
    66  
    67  		if ret, err := f(); err == nil {
    68  			p.once.Do(func() {
    69  				p.result = ret
    70  				if p.cancel != nil {
    71  					p.cancel()
    72  				}
    73  			})
    74  		} else {
    75  			p.errOnce.Do(func() {
    76  				p.err = err
    77  			})
    78  		}
    79  	}()
    80  }