From ffbd68770e34e8c84e7a78e9550ebe55091e3266 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Fri, 27 May 2016 19:08:05 +0700
Subject: [PATCH] remove closed streams from WindowUpdateManager

fixes #99
---
 session.go                    |  3 +++
 session_test.go               | 11 +++++++++++
 window_update_manager.go      |  8 ++++++++
 window_update_manager_test.go |  9 +++++++++
 4 files changed, 31 insertions(+)

diff --git a/session.go b/session.go
index f72200a4..4ddf78ac 100644
--- a/session.go
+++ b/session.go
@@ -624,6 +624,9 @@ func (s *Session) garbageCollectStreams() {
 		if v.finishedWriting() {
 			s.blockedManager.RemoveBlockedStream(k)
 		}
+		if v.finishedReading() {
+			s.windowUpdateManager.RemoveStream(k)
+		}
 		if v.finished() {
 			s.streams[k] = nil
 		}
diff --git a/session_test.go b/session_test.go
index 8157750d..c86b783e 100644
--- a/session_test.go
+++ b/session_test.go
@@ -201,6 +201,17 @@ var _ = Describe("Session", func() {
 			Expect(session.blockedManager.blockedStreams).ToNot(HaveKey(protocol.StreamID(5)))
 		})
 
+		It("removes closed streams from WindowUpdateManager", func() {
+			session.handleStreamFrame(&frames.StreamFrame{
+				StreamID: 5,
+				Data:     []byte{0xde, 0xca, 0xfb, 0xad},
+			})
+			session.updateReceiveFlowControlWindow(5, 0x1337)
+			session.streams[5].eof = 1
+			session.garbageCollectStreams()
+			Expect(session.windowUpdateManager.streamOffsets).ToNot(HaveKey(protocol.StreamID(5)))
+		})
+
 		It("closes empty streams with error", func() {
 			testErr := errors.New("test")
 			session.handleStreamFrame(&frames.StreamFrame{
diff --git a/window_update_manager.go b/window_update_manager.go
index e66e0c9b..74ed779b 100644
--- a/window_update_manager.go
+++ b/window_update_manager.go
@@ -64,3 +64,11 @@ func (m *windowUpdateManager) GetWindowUpdateFrames() []*frames.WindowUpdateFram
 
 	return wuf
 }
+
+// RemoveStream should be called when a stream is closed for receiving
+func (m *windowUpdateManager) RemoveStream(streamID protocol.StreamID) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	delete(m.streamOffsets, streamID)
+}
diff --git a/window_update_manager_test.go b/window_update_manager_test.go
index 2759763a..5a2b805b 100644
--- a/window_update_manager_test.go
+++ b/window_update_manager_test.go
@@ -74,4 +74,13 @@ var _ = Describe("WindowUpdateManager", func() {
 			Expect(frames).To(BeEmpty())
 		})
 	})
+
+	Context("removing streams", func() {
+		It("deletes the map entry", func() {
+			wum.SetStreamOffset(7, 0x1000)
+			Expect(wum.streamOffsets).To(HaveKey(protocol.StreamID(7)))
+			wum.RemoveStream(7)
+			Expect(wum.streamOffsets).ToNot(HaveKey(protocol.StreamID(7)))
+		})
+	})
 })