From 20b2ad3729e6a5a3373abf79fc44ce6ccfdcac7f Mon Sep 17 00:00:00 2001 From: Alex Dadgar Date: Thu, 26 Apr 2018 13:22:09 -0700 Subject: [PATCH 1/2] Track all client connections --- nomad/client_rpc.go | 49 ++++++++++++++++++++++++++++++++++++--------- nomad/server.go | 4 ++-- 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/nomad/client_rpc.go b/nomad/client_rpc.go index af59679b0937..e1bbb6317daf 100644 --- a/nomad/client_rpc.go +++ b/nomad/client_rpc.go @@ -30,7 +30,16 @@ type nodeConnState struct { func (s *Server) getNodeConn(nodeID string) (*nodeConnState, bool) { s.nodeConnsLock.RLock() defer s.nodeConnsLock.RUnlock() - state, ok := s.nodeConns[nodeID] + conns, ok := s.nodeConns[nodeID] + + // Return the latest conn + var state *nodeConnState + for _, conn := range conns { + if state == nil || state.Established.Before(conn.Established) { + state = conn + } + } + return state, ok } @@ -39,8 +48,12 @@ func (s *Server) connectedNodes() map[string]time.Time { s.nodeConnsLock.RLock() defer s.nodeConnsLock.RUnlock() nodes := make(map[string]time.Time, len(s.nodeConns)) - for nodeID, state := range s.nodeConns { - nodes[nodeID] = state.Established + for nodeID, conns := range s.nodeConns { + for _, conn := range conns { + if nodes[nodeID].Before(conn.Established) { + nodes[nodeID] = conn.Established + } + } } return nodes } @@ -54,11 +67,26 @@ func (s *Server) addNodeConn(ctx *RPCContext) { s.nodeConnsLock.Lock() defer s.nodeConnsLock.Unlock() - s.nodeConns[ctx.NodeID] = &nodeConnState{ + + // Capture the tracked connections so far + currentConns := s.nodeConns[ctx.NodeID] + + // Check if we already have the connection. If we do, just update the + // establish time. + for _, c := range currentConns { + if c.Ctx.Conn.LocalAddr().String() == ctx.Conn.LocalAddr().String() && + c.Ctx.Conn.RemoteAddr().String() == ctx.Conn.RemoteAddr().String() { + c.Established = time.Now() + return + } + } + + // Add the new conn + s.nodeConns[ctx.NodeID] = append(s.nodeConns[ctx.NodeID], &nodeConnState{ Session: ctx.Session, Established: time.Now(), Ctx: ctx, - } + }) } // removeNodeConn removes the mapping between a node and its session. @@ -70,7 +98,7 @@ func (s *Server) removeNodeConn(ctx *RPCContext) { s.nodeConnsLock.Lock() defer s.nodeConnsLock.Unlock() - state, ok := s.nodeConns[ctx.NodeID] + conns, ok := s.nodeConns[ctx.NodeID] if !ok { return } @@ -80,9 +108,12 @@ 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. - if state.Ctx.Conn.LocalAddr().String() == ctx.Conn.LocalAddr().String() && - state.Ctx.Conn.RemoteAddr().String() == ctx.Conn.RemoteAddr().String() { - delete(s.nodeConns, ctx.NodeID) + 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 + } } } diff --git a/nomad/server.go b/nomad/server.go index 5a214d71ea49..235e2988a3df 100644 --- a/nomad/server.go +++ b/nomad/server.go @@ -143,7 +143,7 @@ type Server struct { // nodeConns is the set of multiplexed node connections we have keyed by // NodeID - nodeConns map[string]*nodeConnState + nodeConns map[string][]*nodeConnState nodeConnsLock sync.RWMutex // peers is used to track the known Nomad servers. This is @@ -294,7 +294,7 @@ func NewServer(config *Config, consulCatalog consul.CatalogAPI, logger *log.Logg tlsWrap: tlsWrap, rpcServer: rpc.NewServer(), streamingRpcs: structs.NewStreamingRpcRegistry(), - nodeConns: make(map[string]*nodeConnState), + nodeConns: make(map[string][]*nodeConnState), peers: make(map[string][]*serverParts), localPeers: make(map[raft.ServerAddress]*serverParts), reconcileCh: make(chan serf.Member, 32), From 16272ad5459e039e7da1e8788f0c3ae39feb4d9f Mon Sep 17 00:00:00 2001 From: Alex Dadgar Date: Thu, 26 Apr 2018 13:28:24 -0700 Subject: [PATCH 2/2] Add test --- nomad/client_rpc_test.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/nomad/client_rpc_test.go b/nomad/client_rpc_test.go index c64eecec0293..d7edad6bc51d 100644 --- a/nomad/client_rpc_test.go +++ b/nomad/client_rpc_test.go @@ -57,6 +57,7 @@ func TestServer_removeNodeConn_differentAddrs(t *testing.T) { s1.addNodeConn(ctx1) s1.addNodeConn(ctx2) require.Len(s1.connectedNodes(), 1) + require.Len(s1.nodeConns[nodeID], 2) // Check that the value is the second conn. state, ok := s1.getNodeConn(nodeID) @@ -66,6 +67,7 @@ func TestServer_removeNodeConn_differentAddrs(t *testing.T) { // Delete the first s1.removeNodeConn(ctx1) require.Len(s1.connectedNodes(), 1) + require.Len(s1.nodeConns[nodeID], 1) // Check that the value is the second conn. state, ok = s1.getNodeConn(nodeID)