Move timer func to internal, improve log statement and address review…

… comments

balancer/pickfirst/internal/internal.go

@@ -18,7 +18,15 @@
 // Package internal contains code internal to the pickfirst package.
 package internal
 
-import "math/rand"
+import (
+	"math/rand"
+	"time"
+)
 
-// RandShuffle pseudo-randomizes the order of addresses.
-var RandShuffle = rand.Shuffle
+var (
+	// RandShuffle pseudo-randomizes the order of addresses.
+	RandShuffle = rand.Shuffle
+	// TimeAfterFunc allows mocking the timer for testing connection delay
+	// related functionality.
+	TimeAfterFunc = time.AfterFunc
+)

balancer/pickfirst/pickfirstleaf/pickfirstleaf.go

@@ -58,9 +58,6 @@ var (
 	// It is changed to "pick_first" in init() if this balancer is to be
 	// registered as the default pickfirst.
 	Name = "pick_first_leaf"
-	// timerAfterFunc allows mocking the timer for testing connection delay
-	// related functionality.
-	timerAfterFunc = time.AfterFunc
 )
 
 const (
@@ -438,25 +435,26 @@ func (b *pickfirstBalancer) scheduleNextConnectionLocked() {
 	if !envconfig.PickFirstHappyEyeballsEnabled {
 		return
 	}
+	curAddr := b.addressList.currentAddress()
 	b.cancelScheduled()
 	ctx, cancel := context.WithCancel(context.Background())
-	closeFn := timerAfterFunc(connectionDelayInterval, func() {
+	closeFn := internal.TimeAfterFunc(connectionDelayInterval, func() {
 		b.mu.Lock()
 		defer b.mu.Unlock()
 		// If the scheduled task is cancelled while acquiring the mutex, return.
 		if ctx.Err() != nil {
 			return
 		}
 		if b.logger.V(2) {
-			b.logger.Infof("Happy Eyeballs timer expired.")
+			b.logger.Infof("Happy Eyeballs timer expired while waiting for connection to %q.", curAddr.Addr)
 		}
 		if b.addressList.increment() {
 			b.requestConnectionLocked()
 		}
 	}).Stop
 	b.cancelScheduled = sync.OnceFunc(func() {
-		closeFn()
 		cancel()
+		closeFn()
 	})
 }
 

balancer/pickfirst/pickfirstleaf/pickfirstleaf_test.go

@@ -27,6 +27,7 @@ import (
 
 	"google.golang.org/grpc/attributes"
 	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/pickfirst/internal"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/internal/envconfig"
 	"google.golang.org/grpc/internal/grpctest"
@@ -261,8 +262,8 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TriggerConnectionDelay(t *testing.T) {
 	}()
 
 	timerCh := make(chan struct{})
-	originalTimer := timerAfterFunc
-	timerAfterFunc = func(_ time.Duration, f func()) *time.Timer {
+	originalTimer := internal.TimeAfterFunc
+	internal.TimeAfterFunc = func(_ time.Duration, f func()) *time.Timer {
 		// Set a really long expiration to prevent it from triggering
 		// automatically.
 		ret := time.AfterFunc(time.Hour, f)
@@ -277,7 +278,7 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TriggerConnectionDelay(t *testing.T) {
 	}
 
 	defer func() {
-		timerAfterFunc = originalTimer
+		internal.TimeAfterFunc = originalTimer
 	}()
 
 	cc := testutils.NewBalancerClientConn(t)
@@ -358,8 +359,8 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TFAfterEndOfList(t *testing.T) {
 	}()
 
 	timerCh := make(chan struct{})
-	originalTimer := timerAfterFunc
-	timerAfterFunc = func(_ time.Duration, f func()) *time.Timer {
+	originalTimer := internal.TimeAfterFunc
+	internal.TimeAfterFunc = func(_ time.Duration, f func()) *time.Timer {
 		// Set a really long expiration to prevent it from triggering
 		// automatically.
 		ret := time.AfterFunc(time.Hour, f)
@@ -374,11 +375,11 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TFAfterEndOfList(t *testing.T) {
 	}
 
 	defer func() {
-		timerAfterFunc = originalTimer
+		internal.TimeAfterFunc = originalTimer
 	}()
 
 	defer func() {
-		timerAfterFunc = originalTimer
+		internal.TimeAfterFunc = originalTimer
 	}()
 
 	cc := testutils.NewBalancerClientConn(t)
@@ -505,8 +506,8 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TFThenTimerFires(t *testing.T) {
 
 	timerMu := sync.Mutex{}
 	timerCh := make(chan struct{})
-	originalTimer := timerAfterFunc
-	timerAfterFunc = func(_ time.Duration, f func()) *time.Timer {
+	originalTimer := internal.TimeAfterFunc
+	internal.TimeAfterFunc = func(_ time.Duration, f func()) *time.Timer {
 		// Set a really long expiration to prevent it from triggering
 		// automatically.
 		ret := time.AfterFunc(time.Hour, f)
@@ -524,7 +525,7 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TFThenTimerFires(t *testing.T) {
 	}
 
 	defer func() {
-		timerAfterFunc = originalTimer
+		internal.TimeAfterFunc = originalTimer
 	}()
 
 	cc := testutils.NewBalancerClientConn(t)