From 5e2c02c1529b43af59e9ef4cab967dff63e12860 Mon Sep 17 00:00:00 2001
From: gansidui <lijiegm@gmail.com>
Date: Sun, 1 Nov 2015 12:35:44 +0800
Subject: [PATCH] change the way of use waitgroup

---
 conn.go   | 19 +++++++++++--------
 server.go |  5 ++++-
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/conn.go b/conn.go
index df00c28..f3778af 100644
--- a/conn.go
+++ b/conn.go
@@ -121,22 +121,19 @@ func (c *Conn) AsyncWritePacket(p Packet, timeout time.Duration) (err error) {
 
 // Do it
 func (c *Conn) Do() {
-	defer c.srv.waitGroup.Done()
 	if !c.srv.callback.OnConnect(c) {
 		return
 	}
 
-	c.srv.waitGroup.Add(3)
-	go c.handleLoop()
-	go c.readLoop()
-	go c.writeLoop()
+	asyncDo(c.handleLoop, c.srv.waitGroup)
+	asyncDo(c.readLoop, c.srv.waitGroup)
+	asyncDo(c.writeLoop, c.srv.waitGroup)
 }
 
 func (c *Conn) readLoop() {
 	defer func() {
 		recover()
 		c.Close()
-		c.srv.waitGroup.Done()
 	}()
 
 	for {
@@ -163,7 +160,6 @@ func (c *Conn) writeLoop() {
 	defer func() {
 		recover()
 		c.Close()
-		c.srv.waitGroup.Done()
 	}()
 
 	for {
@@ -189,7 +185,6 @@ func (c *Conn) handleLoop() {
 	defer func() {
 		recover()
 		c.Close()
-		c.srv.waitGroup.Done()
 	}()
 
 	for {
@@ -210,3 +205,11 @@ func (c *Conn) handleLoop() {
 		}
 	}
 }
+
+func asyncDo(fn func(), wg *sync.WaitGroup) {
+	wg.Add(1)
+	go func() {
+		fn()
+		wg.Done()
+	}()
+}
diff --git a/server.go b/server.go
index 4bd2a57..4bc6e74 100644
--- a/server.go
+++ b/server.go
@@ -54,7 +54,10 @@ func (s *Server) Start(listener *net.TCPListener, acceptTimeout time.Duration) {
 		}
 
 		s.waitGroup.Add(1)
-		go newConn(conn, s).Do()
+		go func() {
+			newConn(conn, s).Do()
+			s.waitGroup.Done()
+		}()
 	}
 }