clash/common/picker/picker.go

81 lines
1.8 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
2022-04-23 18:07:57 +00:00
type Picker[T any] struct {
ctx context.Context
2019-07-02 11:18:03 +00:00
cancel func()
wg sync.WaitGroup
2020-04-16 10:31:40 +00:00
once sync.Once
errOnce sync.Once
2022-04-23 18:07:57 +00:00
result T
2020-04-16 10:31:40 +00:00
err error
}
2022-04-23 18:07:57 +00:00
func newPicker[T any](ctx context.Context, cancel func()) *Picker[T] {
return &Picker[T]{
ctx: ctx,
cancel: cancel,
}
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.
2022-04-23 18:07:57 +00:00
func WithContext[T any](ctx context.Context) (*Picker[T], context.Context) {
2019-07-02 11:18:03 +00:00
ctx, cancel := context.WithCancel(ctx)
2022-04-23 18:07:57 +00:00
return newPicker[T](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.
2022-04-23 18:07:57 +00:00
func WithTimeout[T any](ctx context.Context, timeout time.Duration) (*Picker[T], context.Context) {
ctx, cancel := context.WithTimeout(ctx, timeout)
2022-04-23 18:07:57 +00:00
return newPicker[T](ctx, cancel), ctx
}
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.
2022-04-23 18:07:57 +00:00
func (p *Picker[T]) Wait() T {
2019-07-02 11:18:03 +00:00
p.wg.Wait()
if p.cancel != nil {
p.cancel()
}
return p.result
}
2020-04-16 10:31:40 +00:00
// Error return the first error (if all success return nil)
2022-04-23 18:07:57 +00:00
func (p *Picker[T]) Error() error {
2020-04-16 10:31:40 +00:00
return p.err
}
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.
2022-04-23 18:07:57 +00:00
func (p *Picker[T]) Go(f func() (T, error)) {
2019-07-02 11:18:03 +00:00
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
if p.cancel != nil {
p.cancel()
}
})
2020-04-16 10:31:40 +00:00
} else {
p.errOnce.Do(func() {
p.err = err
})
2018-12-05 13:13:29 +00:00
}
}()
}