clash/common/picker/picker.go

90 lines
2.1 KiB
Go
Raw Normal View History

2018-12-05 13:13:29 +00:00
package picker
2019-07-02 11:18:03 +00:00
import (
"context"
"sync"
"time"
2019-07-02 11:18:03 +00:00
)
// Picker provides synchronization, and Context cancelation
// for groups of goroutines working on subtasks of a common task.
// Inspired by errGroup
type Picker struct {
ctx context.Context
2019-07-02 11:18:03 +00:00
cancel func()
wg sync.WaitGroup
once sync.Once
result interface{}
firstDone chan struct{}
}
func newPicker(ctx context.Context, cancel func()) *Picker {
return &Picker{
ctx: ctx,
cancel: cancel,
firstDone: make(chan struct{}, 1),
}
2019-07-02 11:18:03 +00:00
}
// WithContext returns a new Picker and an associated Context derived from ctx.
// and cancel when first element return.
2019-07-02 11:18:03 +00:00
func WithContext(ctx context.Context) (*Picker, context.Context) {
ctx, cancel := context.WithCancel(ctx)
return newPicker(ctx, cancel), ctx
2019-07-02 11:18:03 +00:00
}
// WithTimeout returns a new Picker and an associated Context derived from ctx with timeout.
func WithTimeout(ctx context.Context, timeout time.Duration) (*Picker, context.Context) {
ctx, cancel := context.WithTimeout(ctx, timeout)
return newPicker(ctx, cancel), ctx
}
// WithoutAutoCancel returns a new Picker and an associated Context derived from ctx,
// but it wouldn't cancel context when the first element return.
func WithoutAutoCancel(ctx context.Context) *Picker {
return newPicker(ctx, nil)
}
2019-07-02 11:18:03 +00:00
// Wait blocks until all function calls from the Go method have returned,
// then returns the first nil error result (if any) from them.
func (p *Picker) Wait() interface{} {
p.wg.Wait()
if p.cancel != nil {
p.cancel()
}
return p.result
}
// WaitWithoutCancel blocks until the first result return, if timeout will return nil.
func (p *Picker) WaitWithoutCancel() interface{} {
select {
case <-p.firstDone:
return p.result
case <-p.ctx.Done():
return p.result
}
}
2019-07-02 11:18:03 +00:00
// Go calls the given function in a new goroutine.
// The first call to return a nil error cancels the group; its result will be returned by Wait.
func (p *Picker) Go(f func() (interface{}, error)) {
p.wg.Add(1)
2018-12-05 13:13:29 +00:00
go func() {
2019-07-02 11:18:03 +00:00
defer p.wg.Done()
2018-12-05 13:13:29 +00:00
2019-07-02 11:18:03 +00:00
if ret, err := f(); err == nil {
p.once.Do(func() {
p.result = ret
p.firstDone <- struct{}{}
2019-07-02 11:18:03 +00:00
if p.cancel != nil {
p.cancel()
}
})
2018-12-05 13:13:29 +00:00
}
}()
}