From f00dfdd34d6903f5aba191432d70d493099b8c1a Mon Sep 17 00:00:00 2001
From: Dreamacro <305009791@qq.com>
Date: Thu, 18 Jul 2019 00:12:01 +0800
Subject: [PATCH] Fix(picker): add WithTimeout for some situation

---
 adapters/outbound/base.go    |  4 ++++
 adapters/outbound/urltest.go | 11 ++++++-----
 common/picker/picker.go      |  9 +++++++++
 common/picker/picker_test.go |  4 ++--
 hub/route/proxies.go         |  3 +--
 5 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/adapters/outbound/base.go b/adapters/outbound/base.go
index 7c22549..6fb1806 100644
--- a/adapters/outbound/base.go
+++ b/adapters/outbound/base.go
@@ -12,6 +12,10 @@ import (
 	C "github.com/Dreamacro/clash/constant"
 )
 
+var (
+	defaultURLTestTimeout = time.Second * 5
+)
+
 type Base struct {
 	name string
 	tp   C.AdapterType
diff --git a/adapters/outbound/urltest.go b/adapters/outbound/urltest.go
index 9c6bd7d..351c2e4 100644
--- a/adapters/outbound/urltest.go
+++ b/adapters/outbound/urltest.go
@@ -102,15 +102,14 @@ func (u *URLTest) speedTest() {
 	}
 	defer atomic.StoreInt32(&u.once, 0)
 
-	ctx, cancel := context.WithTimeout(context.Background(), u.interval)
+	picker, ctx, cancel := picker.WithTimeout(context.Background(), defaultURLTestTimeout)
 	defer cancel()
-	picker, ctx := picker.WithContext(ctx)
 	for _, p := range u.proxies {
 		proxy := p
 		picker.Go(func() (interface{}, error) {
-			t, err := proxy.URLTest(ctx, u.rawURL)
-			if err != nil || t == 0 {
-				return nil, errors.New("speed test error")
+			_, err := proxy.URLTest(ctx, u.rawURL)
+			if err != nil {
+				return nil, err
 			}
 			return proxy, nil
 		})
@@ -120,6 +119,8 @@ func (u *URLTest) speedTest() {
 	if fast != nil {
 		u.fast = fast.(C.Proxy)
 	}
+
+	<-ctx.Done()
 }
 
 func NewURLTest(option URLTestOption, proxies []C.Proxy) (*URLTest, error) {
diff --git a/common/picker/picker.go b/common/picker/picker.go
index f420679..a66b9bc 100644
--- a/common/picker/picker.go
+++ b/common/picker/picker.go
@@ -3,6 +3,7 @@ package picker
 import (
 	"context"
 	"sync"
+	"time"
 )
 
 // Picker provides synchronization, and Context cancelation
@@ -18,11 +19,19 @@ type Picker struct {
 }
 
 // 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
 }
 
+// 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) {
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	return &Picker{}, ctx, cancel
+}
+
 // 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{} {
diff --git a/common/picker/picker_test.go b/common/picker/picker_test.go
index 7b225d3..7ff3712 100644
--- a/common/picker/picker_test.go
+++ b/common/picker/picker_test.go
@@ -30,9 +30,9 @@ func TestPicker_Basic(t *testing.T) {
 }
 
 func TestPicker_Timeout(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*5)
+	picker, ctx, cancel := WithTimeout(context.Background(), time.Millisecond*5)
 	defer cancel()
-	picker, ctx := WithContext(ctx)
+
 	picker.Go(sleepAndSend(ctx, 20, 1))
 
 	number := picker.Wait()
diff --git a/hub/route/proxies.go b/hub/route/proxies.go
index e4c77dc..7044f05 100644
--- a/hub/route/proxies.go
+++ b/hub/route/proxies.go
@@ -111,9 +111,8 @@ func getProxyDelay(w http.ResponseWriter, r *http.Request) {
 
 	proxy := r.Context().Value(CtxKeyProxy).(C.Proxy)
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*time.Duration(timeout))
+	picker, ctx, cancel := picker.WithTimeout(context.Background(), time.Millisecond*time.Duration(timeout))
 	defer cancel()
-	picker, ctx := picker.WithContext(ctx)
 	picker.Go(func() (interface{}, error) {
 		return proxy.URLTest(ctx, url)
 	})