Fix: urltest get fastest node ehavior (#326)

This commit is contained in:
gVisor bot 2019-10-12 23:29:00 +08:00
parent 2adb586a78
commit 06e1f31a89
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)
picker, ctx, cancel := picker.WithTimeout(context.Background(), defaultURLTestTimeout)
ctx, cancel := context.WithTimeout(context.Background(), defaultURLTestTimeout)
defer cancel()
picker := picker.WithoutAutoCancel(ctx)
for _, p := range u.proxies {
proxy := p
picker.Go(func() (interface{}, error) {
@ -123,10 +124,12 @@ func (u *URLTest) speedTest() {
})
}
fast := picker.Wait()
fast := picker.WaitWithoutCancel()
if fast != nil {
u.fast = fast.(C.Proxy)
}
picker.Wait()
}
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.
// 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()
}

View file

@ -4,6 +4,8 @@ import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
)
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))
number := picker.Wait()
if number != nil && number.(int) != 1 {
t.Error("should recv 1", number)
}
assert.NotNil(t, number)
assert.Equal(t, number.(int), 1)
}
func TestPicker_Timeout(t *testing.T) {
picker, ctx, cancel := WithTimeout(context.Background(), time.Millisecond*5)
defer cancel()
picker, ctx := WithTimeout(context.Background(), time.Millisecond*5)
picker.Go(sleepAndSend(ctx, 20, 1))
number := picker.Wait()
if number != nil {
t.Error("should recv nil")
}
assert.Nil(t, number)
}
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) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
fast, ctx := picker.WithContext(ctx)
fast, ctx := picker.WithTimeout(context.Background(), time.Second)
for _, client := range clients {
r := client
fast.Go(func() (interface{}, error) {

View file

@ -9,7 +9,6 @@ import (
"time"
A "github.com/Dreamacro/clash/adapters/outbound"
"github.com/Dreamacro/clash/common/picker"
C "github.com/Dreamacro/clash/constant"
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)
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()
picker.Go(func() (interface{}, error) {
return proxy.URLTest(ctx, url)
})
elm := picker.Wait()
if elm == nil {
delay, err := proxy.URLTest(ctx, url)
if ctx.Err() != nil {
render.Status(r, http.StatusGatewayTimeout)
render.JSON(w, r, ErrRequestTimeout)
return
}
delay := elm.(uint16)
if delay == 0 {
if err != nil || delay == 0 {
render.Status(r, http.StatusServiceUnavailable)
render.JSON(w, r, newError("An error occurred in the delay test"))
return