github.com/kelleygo/clashcore@v1.0.2/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[T any] struct {
    13  	ctx    context.Context
    14  	cancel func()
    15  
    16  	wg sync.WaitGroup
    17  
    18  	once    sync.Once
    19  	errOnce sync.Once
    20  	result  T
    21  	err     error
    22  }
    23  
    24  func newPicker[T any](ctx context.Context, cancel func()) *Picker[T] {
    25  	return &Picker[T]{
    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[T any](ctx context.Context) (*Picker[T], context.Context) {
    34  	ctx, cancel := context.WithCancel(ctx)
    35  	return newPicker[T](ctx, cancel), ctx
    36  }
    37  
    38  // WithTimeout returns a new Picker and an associated Context derived from ctx with timeout.
    39  func WithTimeout[T any](ctx context.Context, timeout time.Duration) (*Picker[T], context.Context) {
    40  	ctx, cancel := context.WithTimeout(ctx, timeout)
    41  	return newPicker[T](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[T]) Wait() T {
    47  	p.wg.Wait()
    48  	if p.cancel != nil {
    49  		p.cancel()
    50  		p.cancel = nil
    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  					p.cancel = nil
    74  				}
    75  			})
    76  		} else {
    77  			p.errOnce.Do(func() {
    78  				p.err = err
    79  			})
    80  		}
    81  	}()
    82  }
    83  
    84  // Close cancels the picker context and releases resources associated with it.
    85  // If Wait has been called, then there is no need to call Close.
    86  func (p *Picker[T]) Close() error {
    87  	if p.cancel != nil {
    88  		p.cancel()
    89  		p.cancel = nil
    90  	}
    91  	return nil
    92  }