Fix: urltest get fastest node ehavior (#326)

common/picker/picker.go

@@ -10,26 +10,42 @@ import (
 // for groups of goroutines working on subtasks of a common task.
 // Inspired by errGroup
 type Picker struct {
+	ctx    context.Context
 	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),
+	}
 }
 
 // WithContext returns a new Picker and an associated Context derived from ctx.
 // and cancel when first element return.
 func WithContext(ctx context.Context) (*Picker, context.Context) {
 	ctx, cancel := context.WithCancel(ctx)
-	return &Picker{cancel: cancel}, ctx
+	return newPicker(ctx, cancel), ctx
 }
 
-// WithTimeout returns a new Picker and an associated Context derived from ctx with timeout,
-// but it doesn't cancel when first element return.
-func WithTimeout(ctx context.Context, timeout time.Duration) (*Picker, context.Context, context.CancelFunc) {
+// 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 &Picker{}, ctx, cancel
+	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)
 }
 
 // Wait blocks until all function calls from the Go method have returned,
@@ -42,6 +58,16 @@ func (p *Picker) Wait() interface{} {
 	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
+	}
+}
+
 // 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)) {
@@ -53,6 +79,7 @@ func (p *Picker) Go(f func() (interface{}, error)) {
 		if ret, err := f(); err == nil {
 			p.once.Do(func() {
 				p.result = ret
+				p.firstDone <- struct{}{}
 				if p.cancel != nil {
 					p.cancel()
 				}