github.com/yaling888/clash@v1.53.0/common/picker/picker.go (about)

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