Fix: urltest get fastest node ehavior (#326)

This commit is contained in:
comwrg 2019-10-12 23:29:00 +08:00 committed by Dreamacro
parent 4cd8b6f24f
commit 0cdc40beb3
5 changed files with 75 additions and 29 deletions

View file

@ -110,8 +110,9 @@ func (u *URLTest) speedTest() {
} }
defer atomic.StoreInt32(&u.once, 0) defer atomic.StoreInt32(&u.once, 0)
picker, ctx, cancel := picker.WithTimeout(context.Background(), defaultURLTestTimeout) ctx, cancel := context.WithTimeout(context.Background(), defaultURLTestTimeout)
defer cancel() defer cancel()
picker := picker.WithoutAutoCancel(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) {
@ -123,10 +124,12 @@ func (u *URLTest) speedTest() {
}) })
} }
fast := picker.Wait() fast := picker.WaitWithoutCancel()
if fast != nil { if fast != nil {
u.fast = fast.(C.Proxy) u.fast = fast.(C.Proxy)
} }
picker.Wait()
} }
func NewURLTest(option URLTestOption, proxies []C.Proxy) (*URLTest, error) { func NewURLTest(option URLTestOption, proxies []C.Proxy) (*URLTest, error) {

View file

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

View file

@ -4,6 +4,8 @@ import (
"context" "context"
"testing" "testing"
"time" "time"
"github.com/stretchr/testify/assert"
) )
func sleepAndSend(ctx context.Context, delay int, input interface{}) func() (interface{}, error) { func sleepAndSend(ctx context.Context, delay int, input interface{}) func() (interface{}, error) {
@ -24,19 +26,41 @@ func TestPicker_Basic(t *testing.T) {
picker.Go(sleepAndSend(ctx, 20, 1)) picker.Go(sleepAndSend(ctx, 20, 1))
number := picker.Wait() number := picker.Wait()
if number != nil && number.(int) != 1 { assert.NotNil(t, number)
t.Error("should recv 1", number) assert.Equal(t, number.(int), 1)
}
} }
func TestPicker_Timeout(t *testing.T) { func TestPicker_Timeout(t *testing.T) {
picker, ctx, cancel := WithTimeout(context.Background(), time.Millisecond*5) picker, ctx := WithTimeout(context.Background(), time.Millisecond*5)
defer cancel()
picker.Go(sleepAndSend(ctx, 20, 1)) picker.Go(sleepAndSend(ctx, 20, 1))
number := picker.Wait() number := picker.Wait()
if number != nil { assert.Nil(t, number)
t.Error("should recv nil")
} }
func TestPicker_WaitWithoutAutoCancel(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*60)
defer cancel()
picker := WithoutAutoCancel(ctx)
trigger := false
picker.Go(sleepAndSend(ctx, 10, 1))
picker.Go(func() (interface{}, error) {
timer := time.NewTimer(time.Millisecond * time.Duration(30))
select {
case <-timer.C:
trigger = true
return 2, nil
case <-ctx.Done():
return nil, ctx.Err()
}
})
elm := picker.WaitWithoutCancel()
assert.NotNil(t, elm)
assert.Equal(t, elm.(int), 1)
elm = picker.Wait()
assert.True(t, trigger)
assert.Equal(t, elm.(int), 1)
} }

View file

@ -171,10 +171,7 @@ func (r *Resolver) IsFakeIP() bool {
} }
func (r *Resolver) batchExchange(clients []resolver, m *D.Msg) (msg *D.Msg, err error) { func (r *Resolver) batchExchange(clients []resolver, m *D.Msg) (msg *D.Msg, err error) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second) fast, ctx := picker.WithTimeout(context.Background(), time.Second)
defer cancel()
fast, ctx := picker.WithContext(ctx)
for _, client := range clients { for _, client := range clients {
r := client r := client
fast.Go(func() (interface{}, error) { fast.Go(func() (interface{}, error) {

View file

@ -9,7 +9,6 @@ import (
"time" "time"
A "github.com/Dreamacro/clash/adapters/outbound" A "github.com/Dreamacro/clash/adapters/outbound"
"github.com/Dreamacro/clash/common/picker"
C "github.com/Dreamacro/clash/constant" C "github.com/Dreamacro/clash/constant"
T "github.com/Dreamacro/clash/tunnel" T "github.com/Dreamacro/clash/tunnel"
@ -111,21 +110,17 @@ func getProxyDelay(w http.ResponseWriter, r *http.Request) {
proxy := r.Context().Value(CtxKeyProxy).(C.Proxy) proxy := r.Context().Value(CtxKeyProxy).(C.Proxy)
picker, ctx, cancel := picker.WithTimeout(context.Background(), time.Millisecond*time.Duration(timeout)) ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*time.Duration(timeout))
defer cancel() defer cancel()
picker.Go(func() (interface{}, error) {
return proxy.URLTest(ctx, url)
})
elm := picker.Wait() delay, err := proxy.URLTest(ctx, url)
if elm == nil { if ctx.Err() != nil {
render.Status(r, http.StatusGatewayTimeout) render.Status(r, http.StatusGatewayTimeout)
render.JSON(w, r, ErrRequestTimeout) render.JSON(w, r, ErrRequestTimeout)
return return
} }
delay := elm.(uint16) if err != nil || delay == 0 {
if delay == 0 {
render.Status(r, http.StatusServiceUnavailable) render.Status(r, http.StatusServiceUnavailable)
render.JSON(w, r, newError("An error occurred in the delay test")) render.JSON(w, r, newError("An error occurred in the delay test"))
return return