From 8e7215aaeb931eb0c5c18761290cf70ad31264a0 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sat, 6 Aug 2016 17:00:33 +0700
Subject: [PATCH] only increase round-robin index when lambda returns true

---
 streams_map.go      |  2 +-
 streams_map_test.go | 18 +++++++++++++-----
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/streams_map.go b/streams_map.go
index d51d412e..fc271fb8 100644
--- a/streams_map.go
+++ b/streams_map.go
@@ -80,13 +80,13 @@ func (m *streamsMap) RoundRobinIterate(fn streamLambda) error {
 			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 {
 			return err
 		}
 		if !cont {
 			break
 		}
+		m.roundRobinIndex = (m.roundRobinIndex + 1) % numStreams
 	}
 	return nil
 }
diff --git a/streams_map_test.go b/streams_map_test.go
index 31d60dc3..6cc1ff99 100644
--- a/streams_map_test.go
+++ b/streams_map_test.go
@@ -212,7 +212,7 @@ var _ = Describe("Streams Map", func() {
 			fn := func(str *stream) (bool, error) {
 				lambdaCalledForStream = append(lambdaCalledForStream, str.StreamID())
 				numIterations++
-				if str.StreamID() == 2 || str.StreamID() == 4 {
+				if str.StreamID() == 2 {
 					return false, nil
 				}
 				return true, nil
@@ -221,13 +221,21 @@ var _ = Describe("Streams Map", func() {
 			Expect(err).ToNot(HaveOccurred())
 			Expect(numIterations).To(Equal(2))
 			Expect(lambdaCalledForStream).To(Equal([]protocol.StreamID{1, 2}))
-			Expect(m.roundRobinIndex).To(Equal(2))
+			Expect(m.roundRobinIndex).To(Equal(1))
 			numIterations = 0
 			lambdaCalledForStream = lambdaCalledForStream[:0]
-			err = m.RoundRobinIterate(fn)
+			fn2 := func(str *stream) (bool, error) {
+				lambdaCalledForStream = append(lambdaCalledForStream, str.StreamID())
+				numIterations++
+				if str.StreamID() == 4 {
+					return false, nil
+				}
+				return true, nil
+			}
+			err = m.RoundRobinIterate(fn2)
 			Expect(err).ToNot(HaveOccurred())
-			Expect(numIterations).To(Equal(2))
-			Expect(lambdaCalledForStream).To(Equal([]protocol.StreamID{3, 4}))
+			Expect(numIterations).To(Equal(3))
+			Expect(lambdaCalledForStream).To(Equal([]protocol.StreamID{2, 3, 4}))
 		})
 
 		It("adjust the RoundRobinIndex when deleting an element in front", func() {