Fix(picker): add WithTimeout for some situation

This commit is contained in:
gVisor bot 2019-07-18 00:12:01 +08:00
parent bd0b6a2e40
commit eba263a731
5 changed files with 22 additions and 9 deletions

View File

@ -12,6 +12,10 @@ import (
C "github.com/Dreamacro/clash/constant" C "github.com/Dreamacro/clash/constant"
) )
var (
defaultURLTestTimeout = time.Second * 5
)
type Base struct { type Base struct {
name string name string
tp C.AdapterType tp C.AdapterType

View File

@ -102,15 +102,14 @@ func (u *URLTest) speedTest() {
} }
defer atomic.StoreInt32(&u.once, 0) defer atomic.StoreInt32(&u.once, 0)
ctx, cancel := context.WithTimeout(context.Background(), u.interval) picker, ctx, cancel := picker.WithTimeout(context.Background(), defaultURLTestTimeout)
defer cancel() defer cancel()
picker, ctx := picker.WithContext(ctx)
for _, p := range u.proxies { for _, p := range u.proxies {
proxy := p proxy := p
picker.Go(func() (interface{}, error) { picker.Go(func() (interface{}, error) {
t, err := proxy.URLTest(ctx, u.rawURL) _, err := proxy.URLTest(ctx, u.rawURL)
if err != nil || t == 0 { if err != nil {
return nil, errors.New("speed test error") return nil, err
} }
return proxy, nil return proxy, nil
}) })
@ -120,6 +119,8 @@ func (u *URLTest) speedTest() {
if fast != nil { if fast != nil {
u.fast = fast.(C.Proxy) u.fast = fast.(C.Proxy)
} }
<-ctx.Done()
} }
func NewURLTest(option URLTestOption, proxies []C.Proxy) (*URLTest, error) { func NewURLTest(option URLTestOption, proxies []C.Proxy) (*URLTest, error) {

View File

@ -3,6 +3,7 @@ package picker
import ( import (
"context" "context"
"sync" "sync"
"time"
) )
// Picker provides synchronization, and Context cancelation // 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. // 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) { func WithContext(ctx context.Context) (*Picker, context.Context) {
ctx, cancel := context.WithCancel(ctx) ctx, cancel := context.WithCancel(ctx)
return &Picker{cancel: cancel}, 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, // Wait blocks until all function calls from the Go method have returned,
// then returns the first nil error result (if any) from them. // then returns the first nil error result (if any) from them.
func (p *Picker) Wait() interface{} { func (p *Picker) Wait() interface{} {

View File

@ -30,9 +30,9 @@ func TestPicker_Basic(t *testing.T) {
} }
func TestPicker_Timeout(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() defer cancel()
picker, ctx := WithContext(ctx)
picker.Go(sleepAndSend(ctx, 20, 1)) picker.Go(sleepAndSend(ctx, 20, 1))
number := picker.Wait() number := picker.Wait()

View File

@ -111,9 +111,8 @@ func getProxyDelay(w http.ResponseWriter, r *http.Request) {
proxy := r.Context().Value(CtxKeyProxy).(C.Proxy) 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() defer cancel()
picker, ctx := picker.WithContext(ctx)
picker.Go(func() (interface{}, error) { picker.Go(func() (interface{}, error) {
return proxy.URLTest(ctx, url) return proxy.URLTest(ctx, url)
}) })