update connect method

hexl21 · Oct 19, 2015 · f12f0eb · f12f0eb
1 parent 31a8c01
commit f12f0eb
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/network/tcp_client.go b/network/tcp_client.go
@@ -86,17 +86,18 @@ func (client *TCPClient) connect() {
 		return
 	}
 
+	client.wg.Add(1)
+	defer client.wg.Done()
+
 	client.Lock()
 	if client.closeFlag {
-		conn.Close()
 		client.Unlock()
+		conn.Close()
 		return
 	}
 	client.conns[conn] = struct{}{}
 	client.Unlock()
 
-	client.wg.Add(1)
-
 	tcpConn := newTCPConn(conn, client.PendingWriteNum, client.msgParser)
 	agent := client.NewAgent(tcpConn)
 	agent.Run()
@@ -107,8 +108,6 @@ func (client *TCPClient) connect() {
 	delete(client.conns, conn)
 	client.Unlock()
 	agent.OnClose()
-
-	client.wg.Done()
 }
 
 func (client *TCPClient) Close() {