From 4175cf423554597c504256daa5fa88fa8d80f931 Mon Sep 17 00:00:00 2001
From: xtaci <daniel820313@gmail.com>
Date: Mon, 13 May 2019 21:37:45 +0800
Subject: [PATCH] move heavy ops out and fix once

---
 session.go | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/session.go b/session.go
index c41bed5..24f469f 100644
--- a/session.go
+++ b/session.go
@@ -168,19 +168,18 @@ func (s *Session) AcceptStream() (*Stream, error) {
 // Close is used to close the session and all streams.
 func (s *Session) Close() error {
 	var once bool
-	var err error
 	s.dieOnce.Do(func() {
-		err = s.conn.Close()
+		close(s.die)
+		once = true
+	})
+
+	if once {
 		s.streamLock.Lock()
 		for k := range s.streams {
 			s.streams[k].sessionClose()
 		}
 		s.streamLock.Unlock()
-		close(s.die)
-	})
-
-	if once {
-		return err
+		return s.conn.Close()
 	} else {
 		return errors.WithStack(io.ErrClosedPipe)
 	}