From 0c0fa40a3b1b2c47f388ef7f2e2d6d4edea96038 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Fri, 27 Apr 2018 10:16:03 -0700
Subject: [PATCH 1/2] Fix issue where node connection map wasn't being pruned

---
 nomad/client_rpc.go      | 15 ++++++++++++++-
 nomad/client_rpc_test.go |  3 +++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/nomad/client_rpc.go b/nomad/client_rpc.go
index e1bbb6317daf..fb47ed325f8b 100644
--- a/nomad/client_rpc.go
+++ b/nomad/client_rpc.go
@@ -40,6 +40,11 @@ func (s *Server) getNodeConn(nodeID string) (*nodeConnState, bool) {
 		}
 	}
 
+	// Shouldn't happen but rather be safe
+	if state == nil {
+		return nil, false
+	}
+
 	return state, ok
 }
 
@@ -108,13 +113,21 @@ func (s *Server) removeNodeConn(ctx *RPCContext) {
 	// dial various addresses that all route to the same server. The most common
 	// case for this is the original address the client uses to connect to the
 	// server differs from the advertised address sent by the heartbeat.
+	numConns := len(conns)
+	found := false
 	for i, conn := range conns {
 		if conn.Ctx.Conn.LocalAddr().String() == ctx.Conn.LocalAddr().String() &&
 			conn.Ctx.Conn.RemoteAddr().String() == ctx.Conn.RemoteAddr().String() {
 			s.nodeConns[ctx.NodeID] = append(s.nodeConns[ctx.NodeID][:i], s.nodeConns[ctx.NodeID][i+1:]...)
-			return
+			found = true
+			break
 		}
 	}
+
+	// If we just deleted the last conn, remove it from the map
+	if found && numConns == 1 {
+		delete(s.nodeConns, ctx.NodeID)
+	}
 }
 
 // serverWithNodeConn is used to determine which remote server has the most
diff --git a/nomad/client_rpc_test.go b/nomad/client_rpc_test.go
index d7edad6bc51d..e1a75c59693d 100644
--- a/nomad/client_rpc_test.go
+++ b/nomad/client_rpc_test.go
@@ -77,6 +77,9 @@ func TestServer_removeNodeConn_differentAddrs(t *testing.T) {
 	// Delete the second
 	s1.removeNodeConn(ctx2)
 	require.Len(s1.connectedNodes(), 0)
+
+	_, ok = s1.getNodeConn(nodeID)
+	require.False(ok)
 }
 
 func TestServerWithNodeConn_NoPath(t *testing.T) {

From 5cd1ff45018919284b72defbf333fd8029abff09 Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Fri, 27 Apr 2018 10:36:28 -0700
Subject: [PATCH 2/2] small cleanup and logging

---
 nomad/client_rpc.go | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/nomad/client_rpc.go b/nomad/client_rpc.go
index fb47ed325f8b..5a3498de7bff 100644
--- a/nomad/client_rpc.go
+++ b/nomad/client_rpc.go
@@ -42,6 +42,7 @@ func (s *Server) getNodeConn(nodeID string) (*nodeConnState, bool) {
 
 	// Shouldn't happen but rather be safe
 	if state == nil {
+		s.logger.Printf("[WARN] nomad.client_rpc: node %q exists in node connection map without any connection", nodeID)
 		return nil, false
 	}
 
@@ -113,20 +114,20 @@ func (s *Server) removeNodeConn(ctx *RPCContext) {
 	// dial various addresses that all route to the same server. The most common
 	// case for this is the original address the client uses to connect to the
 	// server differs from the advertised address sent by the heartbeat.
-	numConns := len(conns)
-	found := false
 	for i, conn := range conns {
 		if conn.Ctx.Conn.LocalAddr().String() == ctx.Conn.LocalAddr().String() &&
 			conn.Ctx.Conn.RemoteAddr().String() == ctx.Conn.RemoteAddr().String() {
-			s.nodeConns[ctx.NodeID] = append(s.nodeConns[ctx.NodeID][:i], s.nodeConns[ctx.NodeID][i+1:]...)
-			found = true
-			break
-		}
-	}
 
-	// If we just deleted the last conn, remove it from the map
-	if found && numConns == 1 {
-		delete(s.nodeConns, ctx.NodeID)
+			if len(conns) == 1 {
+				// We are deleting the last conn, remove it from the map
+				delete(s.nodeConns, ctx.NodeID)
+			} else {
+				// Slice out the connection we are deleting
+				s.nodeConns[ctx.NodeID] = append(s.nodeConns[ctx.NodeID][:i], s.nodeConns[ctx.NodeID][i+1:]...)
+			}
+
+			return
+		}
 	}
 }