Optimize graceful shutdown

peer.go

@@ -241,9 +241,11 @@ func (p *Peer) Close() (err error) {
 	)
 	p.sessHub.Range(func(sess *session) bool {
 		count++
-		Go(func() {
+		if !Go(func() {
 			errCh <- sess.Close()
-		})
+		}) {
+			errCh <- sess.Close()
+		}
 		return true
 	})
 	for i := 0; i < count; i++ {

session.go

@@ -450,10 +450,6 @@ var (
 
 // Close closes the session.
 func (s *session) Close() (err error) {
-	if atomic.LoadInt32(&s.closed) == 1 {
-		return nil
-	}
-
 	s.closeLock.Lock()
 	if atomic.LoadInt32(&s.closed) == 1 {
 		s.closeLock.Unlock()
@@ -472,11 +468,10 @@ func (s *session) Close() (err error) {
 
 		err = s.socket.Close()
 
+		s.peer.sessHub.Delete(s.Id())
 		s.closeLock.Unlock()
 	}()
 
-	s.peer.sessHub.Delete(s.Id())
-
 	if !s.isDisconnected() {
 		// make sure do not disconnect because of reading timeout.
 		// wait for the subsequent write to complete.