From 1a30313ace7e8975b47da5e046097e18222357f8 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sat, 6 Aug 2016 14:41:47 +0700
Subject: [PATCH] only iterate over open Streams in StreamsMap

ref #256
---
 session.go     | 10 ----------
 streams_map.go | 12 +++++++++++-
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/session.go b/session.go
index 516543dc..d839f824 100644
--- a/session.go
+++ b/session.go
@@ -450,9 +450,6 @@ func (s *Session) closeStreamsWithError(err error) {
 	defer s.streamsMutex.Unlock()
 
 	fn := func(str *stream) (bool, error) {
-		if str == nil {
-			return true, nil
-		}
 		s.closeStreamWithError(str, err)
 		return true, nil
 	}
@@ -644,9 +641,6 @@ func (s *Session) newStreamImpl(id protocol.StreamID) (*stream, error) {
 // from the streams map.
 func (s *Session) garbageCollectStreams() {
 	fn := func(str *stream) (bool, error) {
-		if str == nil {
-			return true, nil
-		}
 		id := str.StreamID()
 		if str.finished() {
 			atomic.AddUint32(&s.openStreamsCount, ^uint32(0)) // decrement
@@ -697,10 +691,6 @@ func (s *Session) getWindowUpdateFrames() ([]*frames.WindowUpdateFrame, error) {
 	var res []*frames.WindowUpdateFrame
 
 	fn := func(str *stream) (bool, error) {
-		if str == nil {
-			return true, nil
-		}
-
 		id := str.StreamID()
 		doUpdate, offset, err := s.flowControlManager.MaybeTriggerStreamWindowUpdate(id)
 		if err != nil {
diff --git a/streams_map.go b/streams_map.go
index d9ea8d63..d51d412e 100644
--- a/streams_map.go
+++ b/streams_map.go
@@ -44,7 +44,14 @@ func (m *streamsMap) Iterate(fn streamLambda) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	for _, str := range m.streams {
+	for _, streamID := range m.openStreams {
+		str, ok := m.streams[streamID]
+		if !ok {
+			return errMapAccess
+		}
+		if str == nil {
+			return fmt.Errorf("BUG: Stream %d is closed, but still in openStreams map", streamID)
+		}
 		cont, err := fn(str)
 		if err != nil {
 			return err
@@ -69,6 +76,9 @@ func (m *streamsMap) RoundRobinIterate(fn streamLambda) error {
 		if !ok {
 			return errMapAccess
 		}
+		if str == nil {
+			return fmt.Errorf("BUG: Stream %d is closed, but still in openStreams map", streamID)
+		}
 		cont, err := fn(str)
 		m.roundRobinIndex = (m.roundRobinIndex + 1) % numStreams
 		if err != nil {