Optimization: use context in fallback speed test (#357)

This commit is contained in:
comwrg 2019-10-13 18:11:02 +08:00 committed by Dreamacro
parent 2c82a2bfc8
commit e22ff74e79
1 changed files with 23 additions and 11 deletions

View File

@ -5,9 +5,10 @@ import (
"encoding/json" "encoding/json"
"errors" "errors"
"net" "net"
"sync" "sync/atomic"
"time" "time"
"github.com/Dreamacro/clash/common/picker"
C "github.com/Dreamacro/clash/constant" C "github.com/Dreamacro/clash/constant"
) )
@ -17,6 +18,7 @@ type Fallback struct {
rawURL string rawURL string
interval time.Duration interval time.Duration
done chan struct{} done chan struct{}
once int32
} }
type FallbackOption struct { type FallbackOption struct {
@ -72,12 +74,15 @@ func (f *Fallback) Destroy() {
func (f *Fallback) loop() { func (f *Fallback) loop() {
tick := time.NewTicker(f.interval) tick := time.NewTicker(f.interval)
go f.validTest() ctx, cancel := context.WithCancel(context.Background())
defer cancel()
go f.validTest(ctx)
Loop: Loop:
for { for {
select { select {
case <-tick.C: case <-tick.C:
go f.validTest() go f.validTest(ctx)
case <-f.done: case <-f.done:
break Loop break Loop
} }
@ -93,18 +98,24 @@ func (f *Fallback) findAliveProxy() C.Proxy {
return f.proxies[0] return f.proxies[0]
} }
func (f *Fallback) validTest() { func (f *Fallback) validTest(ctx context.Context) {
wg := sync.WaitGroup{} if !atomic.CompareAndSwapInt32(&f.once, 0, 1) {
wg.Add(len(f.proxies)) return
}
defer atomic.StoreInt32(&f.once, 0)
ctx, cancel := context.WithTimeout(ctx, defaultURLTestTimeout)
defer cancel()
picker := picker.WithoutAutoCancel(ctx)
for _, p := range f.proxies { for _, p := range f.proxies {
go func(p C.Proxy) { proxy := p
p.URLTest(context.Background(), f.rawURL) picker.Go(func() (interface{}, error) {
wg.Done() return proxy.URLTest(ctx, f.rawURL)
}(p) })
} }
wg.Wait() picker.Wait()
} }
func NewFallback(option FallbackOption, proxies []C.Proxy) (*Fallback, error) { func NewFallback(option FallbackOption, proxies []C.Proxy) (*Fallback, error) {
@ -128,6 +139,7 @@ func NewFallback(option FallbackOption, proxies []C.Proxy) (*Fallback, error) {
rawURL: option.URL, rawURL: option.URL,
interval: interval, interval: interval,
done: make(chan struct{}), done: make(chan struct{}),
once: 0,
} }
go Fallback.loop() go Fallback.loop()
return Fallback, nil return Fallback, nil