Fix: urltest get fastest node ehavior (#326)
This commit is contained in:
parent
2adb586a78
commit
06e1f31a89
5 changed files with 75 additions and 29 deletions
|
@ -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) {
|
||||
|
|
|
@ -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()
|
||||
}
|
||||
|
|
|
@ -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)
|
||||
}
|
||||
|
|
|
@ -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) {
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in a new issue