From 68e51de0f8be9081ac1b3e711625aa214041d77a Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 10 Dec 2013 10:34:35 -0800
Subject: [PATCH] packer/rpc: MuxConn.NextId properly increments

---
 packer/rpc/muxconn.go      | 12 ++++++------
 packer/rpc/muxconn_test.go | 13 +++++++++++++
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/packer/rpc/muxconn.go b/packer/rpc/muxconn.go
index 24e57fd63..0e7c08ca0 100644
--- a/packer/rpc/muxconn.go
+++ b/packer/rpc/muxconn.go
@@ -72,7 +72,7 @@ func (m *MuxConn) Accept(id uint32) (io.ReadWriteCloser, error) {
 	stream.mu.Lock()
 	if stream.state != streamStateSynRecv && stream.state != streamStateClosed {
 		stream.mu.Unlock()
-		return nil, fmt.Errorf("Stream already open in bad state: %d", stream.state)
+		return nil, fmt.Errorf("Stream %d already open in bad state: %d", id, stream.state)
 	}
 
 	if stream.state == streamStateSynRecv {
@@ -124,7 +124,7 @@ func (m *MuxConn) Dial(id uint32) (io.ReadWriteCloser, error) {
 	stream.mu.Lock()
 	if stream.state != streamStateClosed {
 		stream.mu.Unlock()
-		return nil, fmt.Errorf("Stream already open in bad state: %d", stream.state)
+		return nil, fmt.Errorf("Stream %d already open in bad state: %d", id, stream.state)
 	}
 
 	// Open a connection
@@ -157,11 +157,11 @@ func (m *MuxConn) NextId() uint32 {
 	defer m.mu.Unlock()
 
 	for {
-		if _, ok := m.streams[m.curId]; !ok {
-			return m.curId
-		}
-
+		result := m.curId
 		m.curId++
+		if _, ok := m.streams[result]; !ok {
+			return result
+		}
 	}
 }
 
diff --git a/packer/rpc/muxconn_test.go b/packer/rpc/muxconn_test.go
index fce29b3af..e6c23f26f 100644
--- a/packer/rpc/muxconn_test.go
+++ b/packer/rpc/muxconn_test.go
@@ -159,3 +159,16 @@ func TestMuxConn_serverClosesStreams(t *testing.T) {
 		t.Fatalf("err: %s", err)
 	}
 }
+
+func TestMuxConnNextId(t *testing.T) {
+	client, server := testMux(t)
+	defer client.Close()
+	defer server.Close()
+
+	a := client.NextId()
+	b := client.NextId()
+
+	if a != 0 || b != 1 {
+		t.Fatalf("IDs should increment")
+	}
+}