check only once after idle

RainZhai · May 26, 2016 · 374cd38 · 374cd38
1 parent 9ba90c9
commit 374cd38
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 12 deletions.
diff --git a/src/github.com/getlantern/balancer/balancer.go b/src/github.com/getlantern/balancer/balancer.go
@@ -32,7 +32,7 @@ type Balancer struct {
 	mu           sync.RWMutex
 	dialers      dialerHeap
 	trusted      dialerHeap
-	lastDialTime atomic.Value // time.Time
+	lastDialTime int64 // Time.Unix()
 }
 
 // New creates a new Balancer using the supplied Strategy and Dialers.
@@ -53,8 +53,6 @@ func New(st Strategy, dialers ...*Dialer) *Balancer {
 	bal := &Balancer{dialers: st(dls), trusted: st(tdls)}
 	heap.Init(&bal.dialers)
 	heap.Init(&bal.trusted)
-	// Force checking all dialers on first Dial()
-	bal.lastDialTime.Store(time.Time{})
 	return bal
 }
 
@@ -73,13 +71,13 @@ func (b *Balancer) OnRequest(req *http.Request) {
 // either manages to connect, or runs out of dialers in which case it returns an
 // error.
 func (b *Balancer) Dial(network, addr string) (net.Conn, error) {
-	lastDialTime := b.lastDialTime.Load().(time.Time)
-	idled := time.Since(lastDialTime)
-	if idled > recheckAfterIdleFor {
-		log.Debugf("Balancer idled for %s, start checking all dialers", idled)
+	now := time.Now()
+	lastDialTime := time.Unix(atomic.SwapInt64(&b.lastDialTime, now.Unix()), 0)
+	idlePeriod := now.Sub(lastDialTime)
+	if idlePeriod > recheckAfterIdleFor {
+		log.Debugf("Balancer idle for %s, start checking all dialers", idlePeriod)
 		b.checkDialers()
 	}
-	defer b.lastDialTime.Store(time.Now())
 	var dialers dialerHeap
 
 	_, port, _ := net.SplitHostPort(addr)

diff --git a/src/github.com/getlantern/balancer/balancer_test.go b/src/github.com/getlantern/balancer/balancer_test.go
@@ -158,6 +158,7 @@ func TestCheck(t *testing.T) {
 
 	var wg sync.WaitGroup
 	var failToDial uint32
+	var checkCount uint32
 	d := &Dialer{
 		DialFN: func(network, addr string) (net.Conn, error) {
 			if atomic.LoadUint32(&failToDial) == 1 {
@@ -166,6 +167,7 @@ func TestCheck(t *testing.T) {
 			return nil, nil
 		},
 		Check: func() bool {
+			atomic.AddUint32(&checkCount, 1)
 			wg.Done()
 			return atomic.LoadUint32(&failToDial) == 0
 		},
@@ -191,13 +193,15 @@ func TestCheck(t *testing.T) {
 	assert.NoError(t, err)
 
 	// recheck failed dialer
+	// The check count is time sensitive, can't rely on WaitGroup
+	before := atomic.LoadUint32(&checkCount)
+	wg.Add(5)
 	atomic.StoreUint32(&failToDial, 1)
-	wg.Add(1)
-	_, err = bal.Dial("tcp", "does-not-exist.com:80")
-	assert.Error(t, err)
 	_, err = bal.Dial("tcp", "does-not-exist.com:80")
 	assert.Error(t, err)
-	wg.Wait()
+	time.Sleep(100 * time.Millisecond)
+	after := atomic.LoadUint32(&checkCount)
+	assert.True(t, after > before, "should recheck failed dialer")
 }
 
 func newDialer(id int) *Dialer {