reduce memory usage of Connection

2 files changed, 24 insertions, 8 deletions

diff --git a/librabbitmq/__init__.py b/librabbitmq/__init__.py
index 6a4189f..a1d7e3b 100644
--- a/librabbitmq/__init__.py
+++ b/librabbitmq/__init__.py
@@ -201,7 +201,7 @@ class Connection(_librabbitmq.Connection):
             client_properties=client_properties,
         )
         self.channels = {}
-        self._avail_channel_ids = array('H', xrange(self.channel_max, 0, -1))
+        self._used_channel_ids = array('H')
         if not lazy:
             self.connect()
 
@@ -247,15 +247,21 @@ class Connection(_librabbitmq.Connection):
             pass
         self.channels.pop(channel.channel_id, None)
         self.callbacks.pop(channel.channel_id, None)
-        self._avail_channel_ids.append(channel.channel_id)
+        try:
+            self._used_channel_ids.remove(channel.channel_id)
+        except ValueError:
+            # channel id already removed
+            pass
 
     def _get_free_channel_id(self):
-        try:
-            return self._avail_channel_ids.pop()
-        except IndexError:
-            raise ConnectionError(
-                'No free channel ids, current=%d, channel_max=%d' % (
-                    len(self.channels), self.channel_max))
+        for channel_id in range(1, self.channel_max):
+            if channel_id not in self._used_channel_ids:
+                self._used_channel_ids.append(channel_id)
+                return channel_id
+
+        raise ConnectionError(
+            'No free channel ids, current=%d, channel_max=%d' % (
+                len(self.channels), self.channel_max))
 
     def close(self):
         try:
diff --git a/tests/test_functional.py b/tests/test_functional.py
index e93dfa0..7b8b17f 100644
--- a/tests/test_functional.py
+++ b/tests/test_functional.py
@@ -5,6 +5,7 @@ from six.moves import xrange
 
 import socket
 import unittest
+from array import array
 
 from librabbitmq import Message, Connection, ConnectionError, ChannelError
 TEST_QUEUE = 'pyrabbit.testq'
@@ -122,6 +123,15 @@ class test_Channel(unittest.TestCase):
             self.connection.drain_events(timeout=0.1)
         self.assertEqual(len(messages), 1)
 
+    def test_get_free_channel_id(self):
+        self.connection._used_channel_ids = array('H')
+        assert self.connection._get_free_channel_id() == 1
+
+    def test_get_free_channel_id__channels_full(self):
+        self.connection._used_channel_ids = array('H', range(1, self.connection.channel_max))
+        with self.assertRaises(ConnectionError):
+            self.connection._get_free_channel_id()
+
     def tearDown(self):
         if self.channel and self.connection.connected:
             self.channel.queue_purge(TEST_QUEUE)