From 58e1343da1d3190c2930aae2560887317d39f8bb Mon Sep 17 00:00:00 2001
From: fangdingjun <fangdingjun@szchuzhong.com>
Date: Thu, 12 Jul 2018 14:10:06 +0800
Subject: [PATCH] close stream and delete from conn.streams

---
 callbacks.go |  4 ++--
 conn.go      | 17 ++++++++++-------
 stream.go    | 16 +++++++++++-----
 3 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/callbacks.go b/callbacks.go
index 9d90dcb..38b3df1 100644
--- a/callbacks.go
+++ b/callbacks.go
@@ -198,7 +198,7 @@ func onServerStreamClose(ptr unsafe.Pointer, streamID C.int) C.int {
 	//conn.lock.Lock()
 	delete(conn.streams, int(streamID))
 	//conn.lock.Unlock()
-	s.Close()
+	go s.Close()
 	return NGHTTP2_NO_ERROR
 }
 
@@ -406,7 +406,7 @@ func onClientStreamClose(ptr unsafe.Pointer, streamID C.int) C.int {
 
 	stream, ok := conn.streams[int(streamID)]
 	if ok {
-		stream.Close()
+		go stream.Close()
 		//conn.lock.Lock()
 		delete(conn.streams, int(streamID))
 		//go stream.Close()
diff --git a/conn.go b/conn.go
index eb7e754..5fa8e12 100644
--- a/conn.go
+++ b/conn.go
@@ -68,16 +68,18 @@ func (c *ClientConn) Error() error {
 
 // Close close the http2 connection
 func (c *ClientConn) Close() error {
-	c.lock.Lock()
-	defer c.lock.Unlock()
 	if c.closed {
 		return nil
 	}
-	//log.Println("close client connection")
 	c.closed = true
+
 	for _, s := range c.streams {
 		s.Close()
 	}
+
+	c.lock.Lock()
+	defer c.lock.Unlock()
+	//log.Println("close client connection")
 	C.nghttp2_session_terminate_session(c.session, 0)
 	C.nghttp2_session_del(c.session)
 	close(c.exitch)
@@ -428,20 +430,21 @@ func (c *ServerConn) serve(s *ServerStream) {
 
 // Close close the server connection
 func (c *ServerConn) Close() error {
-	c.lock.Lock()
-	defer c.lock.Unlock()
-
 	if c.closed {
 		return nil
 	}
+	c.closed = true
 	for _, s := range c.streams {
 		s.Close()
 	}
+
+	c.lock.Lock()
+	defer c.lock.Unlock()
+
 	C.nghttp2_session_terminate_session(c.session, 0)
 	C.nghttp2_session_del(c.session)
 	close(c.exitch)
 	c.conn.Close()
-	c.closed = true
 	return nil
 }
 
diff --git a/stream.go b/stream.go
index 235316a..2bb25e2 100644
--- a/stream.go
+++ b/stream.go
@@ -89,11 +89,11 @@ func (s *ClientStream) Close() error {
 		s.cdp = nil
 	}
 
-	//s.conn.lock.Lock()
-	//defer s.conn.lock.Unlock()
-	//if _, ok := s.conn.streams[s.streamID]; ok {
-	//delete(s.conn.streams, s.streamID)
-	//}
+	s.conn.lock.Lock()
+	defer s.conn.lock.Unlock()
+	if _, ok := s.conn.streams[s.streamID]; ok {
+		delete(s.conn.streams, s.streamID)
+	}
 	return nil
 }
 
@@ -209,6 +209,12 @@ func (s *ServerStream) Close() error {
 		C.free(unsafe.Pointer(s.cdp))
 		s.cdp = nil
 	}
+
+	s.conn.lock.Lock()
+	s.conn.lock.Unlock()
+	if _, ok := s.conn.streams[s.streamID]; ok {
+		delete(s.conn.streams, s.streamID)
+	}
 	//log.Printf("stream %d closed", s.streamID)
 	return nil
 }