better handling of token bucket notify

xtaci · Apr 17, 2019 · 6e0e0d2 · 6e0e0d2
1 parent 5eb9306
commit 6e0e0d2
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/session.go b/session.go
@@ -152,7 +152,6 @@ func (s *Session) Close() (err error) {
 			s.streams[k].sessionClose()
 		}
 		s.streamLock.Unlock()
-		s.notifyBucket()
 		return s.conn.Close()
 	}
 }
@@ -217,7 +216,11 @@ func (s *Session) recvLoop() {
 
 	for {
 		for atomic.LoadInt32(&s.bucket) <= 0 && !s.IsClosed() {
-			<-s.bucketNotify
+			select {
+			case <-s.bucketNotify:
+			case <-s.die:
+				return
+			}
 		}
 
 		// read header first