From c7536f88bc69bb5dd0a4811d8dc96ab80f9cce65 Mon Sep 17 00:00:00 2001 From: Niv Yehezkel Date: Sun, 5 Apr 2020 22:35:16 +0300 Subject: [PATCH] connection: replace select by poll to support fds higher than 1024 --- aiopg/connection.py | 5 ++++- tests/test_connection.py | 22 ++++++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/aiopg/connection.py b/aiopg/connection.py index 4c5feef5..4deee53e 100755 --- a/aiopg/connection.py +++ b/aiopg/connection.py @@ -110,7 +110,10 @@ def _ready(weak_self): except (psycopg2.Warning, psycopg2.Error) as exc: if self._fileno is not None: try: - select.select([self._fileno], [], [], 0) + poll = select.poll() + poll.register(self._fileno) + poll.poll(0) + poll.unregister(self._fileno) except OSError as os_exc: if _is_bad_descriptor_error(os_exc): with contextlib.suppress(OSError): diff --git a/tests/test_connection.py b/tests/test_connection.py index c6a189e4..dbb958e7 100755 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -642,3 +642,25 @@ async def test_connection_on_server_restart(connect, pg_server, docker): delay *= 2 else: pytest.fail("Cannot connect to the restarted server") + + +async def test_no_poll_error_on_high_fd(connect): + # The connection file descriptor when higher than 1024 should not raise any + # further exception when OperationalError is raised + conn = await connect() + high_fd = 1025 + + impl = mock.Mock() + exc = psycopg2.OperationalError('Test') + impl.poll.side_effect = exc + conn._conn = impl + conn._fileno = high_fd + + m_remove_reader = mock.Mock() + conn._loop.remove_reader = m_remove_reader + + conn._ready(conn._weakref) + assert not m_remove_reader.called + + conn.close() + assert m_remove_reader.called_with(high_fd)