From 1c6414e80692abcd2a8b57ead299e348db782cb4 Mon Sep 17 00:00:00 2001
From: James Phillips <james@hashicorp.com>
Date: Thu, 13 Aug 2015 10:01:05 -0700
Subject: [PATCH] Cleans up locking and factors markForUse into a Conn method.

---
 consul/pool.go | 30 ++++++++++++------------------
 1 file changed, 12 insertions(+), 18 deletions(-)

diff --git a/consul/pool.go b/consul/pool.go
index e6c7220eb..56fa80e0a 100644
--- a/consul/pool.go
+++ b/consul/pool.go
@@ -114,6 +114,12 @@ func (c *Conn) returnClient(client *StreamClient) {
 	}
 }
 
+// markForUse does all the bookkeeping required to ready a connection for use.
+func (c *Conn) markForUse() {
+	c.lastUsed = time.Now()
+	atomic.AddInt32(&c.refCount, 1)
+}
+
 // ConnPool is used to maintain a connection pool to other
 // Consul servers. This is used to reduce the latency of
 // RPC requests between servers. It is only used to pool
@@ -192,20 +198,13 @@ func (p *ConnPool) Shutdown() error {
 // and will return that one if it succeeds. If all else fails, it will return a
 // newly-created connection and add it to the pool.
 func (p *ConnPool) acquire(dc string, addr net.Addr, version int) (*Conn, error) {
-	// markforUse does all the bookkeeping for an existing connection we wish
-	// to use from the pool.
-	markForUse := func(c *Conn) {
-		c.lastUsed = time.Now()
-		atomic.AddInt32(&c.refCount, 1)
-	}
-
 	// Check to see if there's a pooled connection available. This is up
 	// here since it should the the vastly more common case than the rest
 	// of the code here.
 	p.Lock()
 	c := p.pool[addr.String()]
 	if c != nil {
-		markForUse(c)
+		c.markForUse()
 		p.Unlock()
 		return c, nil
 	}
@@ -225,20 +224,15 @@ func (p *ConnPool) acquire(dc string, addr net.Addr, version int) (*Conn, error)
 	// If we are the lead thread, make the new connection and then wake
 	// everybody else up to see if we got it.
 	if isLeadThread {
-		defer func() {
-			p.Lock()
-			delete(p.limiter, addr.String())
-			p.Unlock()
-
-			close(wait)
-		}()
-
 		c, err := p.getNewConn(dc, addr, version)
+		p.Lock()
+		delete(p.limiter, addr.String())
+		close(wait)
 		if err != nil {
+			p.Unlock()
 			return nil, err
 		}
 
-		p.Lock()
 		p.pool[addr.String()] = c
 		p.Unlock()
 		return c, nil
@@ -255,7 +249,7 @@ func (p *ConnPool) acquire(dc string, addr net.Addr, version int) (*Conn, error)
 	// See if the lead thread was able to get us a connection.
 	p.Lock()
 	if c := p.pool[addr.String()]; c != nil {
-		markForUse(c)
+		c.markForUse()
 		p.Unlock()
 		return c, nil
 	}