improve performance of _get_free_channel_id, fix channel max bug (#385)

* improve performance of _get_free_channel_id, fix channel max bug

* add integration tests for _get_free_channel_id performance improvement

3 files changed, 13 insertions, 4 deletions

diff --git a/amqp/connection.py b/amqp/connection.py
index 1dd52c1..24bfb66 100644
--- a/amqp/connection.py
+++ b/amqp/connection.py
@@ -482,8 +482,11 @@ class Connection(AbstractChannel):
         self._transport = self.connection = self.channels = None
 
     def _get_free_channel_id(self):
-        for channel_id in range(1, self.channel_max):
-            if channel_id not in self._used_channel_ids:
+        # Cast to a set for fast lookups, and keep stored as an array for lower memory usage.
+        used_channel_ids = set(self._used_channel_ids)
+
+        for channel_id in range(1, self.channel_max + 1):
+            if channel_id not in used_channel_ids:
                 self._used_channel_ids.append(channel_id)
                 return channel_id
 
diff --git a/t/integration/test_integration.py b/t/integration/test_integration.py
index 1ce3851..d441488 100644
--- a/t/integration/test_integration.py
+++ b/t/integration/test_integration.py
@@ -1,4 +1,5 @@
 import socket
+from array import array
 from struct import pack
 from unittest.mock import ANY, Mock, call, patch
 
@@ -534,9 +535,11 @@ class test_channel:
             frame_writer_mock.reset_mock()
 
             on_open_mock = Mock()
+            assert conn._used_channel_ids == array('H')
             ch = conn.channel(channel_id=channel_id, callback=on_open_mock)
             on_open_mock.assert_called_once_with(ch)
             assert ch.is_open is True
+            assert conn._used_channel_ids == array('H', (1,))
 
             ch.close()
             frame_writer_mock.assert_has_calls(
@@ -552,6 +555,7 @@ class test_channel:
                 ]
             )
             assert ch.is_open is False
+            assert conn._used_channel_ids == array('H')
 
     def test_received_channel_Close_during_connection_close(self):
         # This test verifies that library handles correctly closing channel
diff --git a/t/unit/test_connection.py b/t/unit/test_connection.py
index dfa9a74..13eb0b6 100644
--- a/t/unit/test_connection.py
+++ b/t/unit/test_connection.py
@@ -352,8 +352,10 @@ class test_Connection:
         assert self.conn._get_free_channel_id() == 1
         assert self.conn._get_free_channel_id() == 2
 
-    def test_get_free_channel_id__raises_IndexError(self):
-        self.conn._used_channel_ids = array('H', range(1, self.conn.channel_max))
+    def test_get_free_channel_id__raises_ResourceError(self):
+        self.conn.channel_max = 2
+        self.conn._get_free_channel_id()
+        self.conn._get_free_channel_id()
         with pytest.raises(ResourceError):
             self.conn._get_free_channel_id()