From efc1113406814d04550c6cd05b92f1f2a80e6e4a Mon Sep 17 00:00:00 2001
From: Sean Chittenden <sean@chittenden.org>
Date: Sat, 26 Mar 2016 23:41:01 -0700
Subject: [PATCH] Expose ServerManager.ResetRebalanceTimer

Move the rebalance timer from ServerManager.Start's stack to struct ServerManager.  This makes it possible to shuffle during tests without actually waiting >120s.
---
 consul/server_manager/server_manager.go | 25 +++++++++++++++++--------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/consul/server_manager/server_manager.go b/consul/server_manager/server_manager.go
index eab5ebbd46..7f11098242 100644
--- a/consul/server_manager/server_manager.go
+++ b/consul/server_manager/server_manager.go
@@ -76,6 +76,9 @@ type ServerManager struct {
 	serverConfigValue atomic.Value
 	serverConfigLock  sync.Mutex
 
+	// rebalanceTimer controls the duration of the rebalance interval
+	rebalanceTimer *time.Timer
+
 	// shutdownCh is a copy of the channel in consul.Client
 	shutdownCh chan struct{}
 
@@ -405,10 +408,8 @@ func (sm *ServerManager) RemoveServer(server *server_details.ServerDetails) {
 	}
 }
 
-// refreshServerRebalanceTimer is only called once the rebalanceTimer
-// expires.  Historically this was an expensive routine and is intended to be
-// run in isolation in a dedicated, non-concurrent task.
-func (sm *ServerManager) refreshServerRebalanceTimer(timer *time.Timer) time.Duration {
+// refreshServerRebalanceTimer is only called once sm.rebalanceTimer expires.
+func (sm *ServerManager) refreshServerRebalanceTimer() time.Duration {
 	serverCfg := sm.getServerConfig()
 	numConsulServers := len(serverCfg.servers)
 	// Limit this connection's life based on the size (and health) of the
@@ -420,10 +421,18 @@ func (sm *ServerManager) refreshServerRebalanceTimer(timer *time.Timer) time.Dur
 	numLANMembers := sm.clusterInfo.NumNodes()
 	connRebalanceTimeout := lib.RateScaledInterval(clusterWideRebalanceConnsPerSec, connReuseLowWatermarkDuration, numLANMembers)
 
-	timer.Reset(connRebalanceTimeout)
+	sm.rebalanceTimer.Reset(connRebalanceTimeout)
 	return connRebalanceTimeout
 }
 
+// ResetRebalanceTimer resets the rebalance timer.  This method primarily
+// exists for testing and should not be used directly.
+func (sm *ServerManager) ResetRebalanceTimer() {
+	sm.serverConfigLock.Lock()
+	defer sm.serverConfigLock.Unlock()
+	sm.rebalanceTimer.Reset(clientRPCMinReuseDuration)
+}
+
 // Start is used to start and manage the task of automatically shuffling and
 // rebalancing the list of consul servers.  This maintenance only happens
 // periodically based on the expiration of the timer.  Failed servers are
@@ -431,14 +440,14 @@ func (sm *ServerManager) refreshServerRebalanceTimer(timer *time.Timer) time.Dur
 // the list.  The order of the server list must be shuffled periodically to
 // distribute load across all known and available consul servers.
 func (sm *ServerManager) Start() {
-	var rebalanceTimer *time.Timer = time.NewTimer(clientRPCMinReuseDuration)
+	sm.rebalanceTimer = time.NewTimer(clientRPCMinReuseDuration)
 
 	for {
 		select {
-		case <-rebalanceTimer.C:
+		case <-sm.rebalanceTimer.C:
 			sm.logger.Printf("[INFO] server manager: Rebalancing server connections")
 			sm.RebalanceServers()
-			sm.refreshServerRebalanceTimer(rebalanceTimer)
+			sm.refreshServerRebalanceTimer()
 
 		case <-sm.shutdownCh:
 			sm.logger.Printf("[INFO] server manager: shutting down")