From 935a06b9506c7fa596863de3fc6f369ae1205bdc Mon Sep 17 00:00:00 2001 From: Paul Brown Date: Wed, 22 Dec 2021 11:17:03 +0000 Subject: 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 --- amqp/connection.py | 7 +++++-- t/integration/test_integration.py | 4 ++++ t/unit/test_connection.py | 6 ++++-- 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() -- cgit v1.2.1