poll() when available to surpass 1024 file descriptor limit with select() (mjs#377)

Author: jhatch28

imapclient/imapclient.py

@@ -28,6 +28,13 @@
 from .util import to_bytes, to_unicode, assert_imap_protocol, chunk
 xrange = moves.xrange
 
+try:
+    from select import poll
+    POLL_SUPPORT = True
+except:
+    # Fallback to select() on systems that don't support poll()
+    POLL_SUPPORT = False
+
 if PY3:
     long = int  # long is just int in python3
 
@@ -760,6 +767,26 @@ def idle(self):
         if resp is not None:
             raise exceptions.IMAPClientError('Unexpected IDLE response: %s' % resp)
 
+    def _poll_socket(self, sock, timeout=None):
+        """
+        Polls the socket for events telling us it's available to read.
+        This implementation is more scalable because it ALLOWS your process
+        to have more than 1024 file descriptors.
+        """
+        poller = select.poll()
+        poller.register(sock.fileno(), select.POLLIN)
+        timeout = timeout * 1000 if timeout is not None else None
+        return poller.poll(timeout)
+
+    def _select_poll_socket(self, sock, timeout=None):
+        """
+        Polls the socket for events telling us it's available to read.
+        This implementation is a fallback because it FAILS if your process
+        has more than 1024 file descriptors.
+        We still need this for Windows and some other niche systems.
+        """
+        return select.select([sock], [], [], timeout)[0]
+
     @require_capability('IDLE')
     def idle_check(self, timeout=None):
         """Check for any IDLE responses sent by the server.
@@ -785,10 +812,16 @@ def idle_check(self, timeout=None):
         # implemented for this call
         sock.settimeout(None)
         sock.setblocking(0)
+
+        if POLL_SUPPORT:
+            poll_func = self._poll_socket
+        else:
+            poll_func = self._select_poll_socket
+
         try:
             resps = []
-            rs, _, _ = select.select([sock], [], [], timeout)
-            if rs:
+            events = poll_func(sock, timeout)
+            if events:
                 while True:
                     try:
                         line = self._imap._get_line()

tests/test_imapclient.py

@@ -11,6 +11,7 @@
 import logging
 
 import six
+from select import POLLIN
 
 from imapclient.exceptions import (
     CapabilityError, IMAPClientError, ProtocolError
@@ -379,14 +380,23 @@ def setUp(self):
         super(TestIdleAndNoop, self).setUp()
         self.client._cached_capabilities = [b'IDLE']
 
-    def assert_sock_calls(self, sock):
+    def assert_sock_select_calls(self, sock):
         self.assertListEqual(sock.method_calls, [
             ('settimeout', (None,), {}),
             ('setblocking', (0,), {}),
             ('setblocking', (1,), {}),
             ('settimeout', (None,), {}),
         ])
 
+    def assert_sock_poll_calls(self, sock):
+        self.assertListEqual(sock.method_calls, [
+            ('settimeout', (None,), {}),
+            ('setblocking', (0,), {}),
+            ('fileno', (), {}),
+            ('setblocking', (1,), {}),
+            ('settimeout', (None,), {}),
+        ])
+
     def test_idle(self):
         self.client._imap._command.return_value = sentinel.tag
         self.client._imap._get_response.return_value = None
@@ -396,6 +406,7 @@ def test_idle(self):
         self.client._imap._command.assert_called_with('IDLE')
         self.assertEqual(self.client._idle_tag, sentinel.tag)
 
+    @patch('imapclient.imapclient.POLL_SUPPORT', False)
     @patch('imapclient.imapclient.select.select')
     def test_idle_check_blocking(self, mock_select):
         mock_sock = Mock()
@@ -416,9 +427,10 @@ def fake_get_line():
         responses = self.client.idle_check()
 
         mock_select.assert_called_once_with([mock_sock], [], [], None)
-        self.assert_sock_calls(mock_sock)
+        self.assert_sock_select_calls(mock_sock)
         self.assertListEqual([(1, b'EXISTS'), (0, b'EXPUNGE')], responses)
 
+    @patch('imapclient.imapclient.POLL_SUPPORT', False)
     @patch('imapclient.imapclient.select.select')
     def test_idle_check_timeout(self, mock_select):
         mock_sock = Mock()
@@ -428,9 +440,10 @@ def test_idle_check_timeout(self, mock_select):
         responses = self.client.idle_check(timeout=0.5)
 
         mock_select.assert_called_once_with([mock_sock], [], [], 0.5)
-        self.assert_sock_calls(mock_sock)
+        self.assert_sock_select_calls(mock_sock)
         self.assertListEqual([], responses)
 
+    @patch('imapclient.imapclient.POLL_SUPPORT', False)
     @patch('imapclient.imapclient.select.select')
     def test_idle_check_with_data(self, mock_select):
         mock_sock = Mock()
@@ -449,7 +462,80 @@ def fake_get_line():
         responses = self.client.idle_check()
 
         mock_select.assert_called_once_with([mock_sock], [], [], None)
-        self.assert_sock_calls(mock_sock)
+        self.assert_sock_select_calls(mock_sock)
+        self.assertListEqual([(99, b'EXISTS')], responses)
+
+    @patch('imapclient.imapclient.POLL_SUPPORT', True)
+    @patch('imapclient.imapclient.select.poll')
+    def test_idle_check_blocking(self, mock_poll_module):
+        mock_sock = Mock(fileno=Mock(return_value=1))
+        self.client._imap.sock = self.client._imap.sslobj = mock_sock
+
+        mock_poller = Mock(poll=Mock(return_value=[(1, POLLIN)]))
+        mock_poll_module.return_value = mock_poller
+        counter = itertools.count()
+
+        def fake_get_line():
+            count = six.next(counter)
+            if count == 0:
+                return b'* 1 EXISTS'
+            elif count == 1:
+                return b'* 0 EXPUNGE'
+            else:
+                raise socket.timeout
+
+        self.client._imap._get_line = fake_get_line
+
+        responses = self.client.idle_check()
+
+        assert mock_poll_module.call_count == 1
+        mock_poller.register.assert_called_once_with(1, POLLIN)
+        mock_poller.poll.assert_called_once_with(None)
+        self.assert_sock_poll_calls(mock_sock)
+        self.assertListEqual([(1, b'EXISTS'), (0, b'EXPUNGE')], responses)
+
+    @patch('imapclient.imapclient.POLL_SUPPORT', True)
+    @patch('imapclient.imapclient.select.poll')
+    def test_idle_check_timeout(self, mock_poll_module):
+        mock_sock = Mock(fileno=Mock(return_value=1))
+        self.client._imap.sock = self.client._imap.sslobj = mock_sock
+
+        mock_poller = Mock(poll=Mock(return_value=[]))
+        mock_poll_module.return_value = mock_poller
+
+        responses = self.client.idle_check(timeout=0.5)
+
+        assert mock_poll_module.call_count == 1
+        mock_poller.register.assert_called_once_with(1, POLLIN)
+        mock_poller.poll.assert_called_once_with(500)
+        self.assert_sock_poll_calls(mock_sock)
+        self.assertListEqual([], responses)
+
+    @patch('imapclient.imapclient.POLL_SUPPORT', True)
+    @patch('imapclient.imapclient.select.poll')
+    def test_idle_check_with_data(self, mock_poll_module):
+        mock_sock = Mock(fileno=Mock(return_value=1))
+        self.client._imap.sock = self.client._imap.sslobj = mock_sock
+
+        mock_poller = Mock(poll=Mock(return_value=[(1, POLLIN)]))
+        mock_poll_module.return_value = mock_poller
+        counter = itertools.count()
+
+        def fake_get_line():
+            count = six.next(counter)
+            if count == 0:
+                return b'* 99 EXISTS'
+            else:
+                raise socket.timeout
+
+        self.client._imap._get_line = fake_get_line
+
+        responses = self.client.idle_check()
+
+        assert mock_poll_module.call_count == 1
+        mock_poller.register.assert_called_once_with(1, POLLIN)
+        mock_poller.poll.assert_called_once_with(None)
+        self.assert_sock_poll_calls(mock_sock)
         self.assertListEqual([(99, b'EXISTS')], responses)
 
     def test_idle_done(self):
@@ -835,4 +921,4 @@ def test_tagged_response_with_parse_error(self):
         client._imap._get_response = lambda: b'NOT-A-STAR 99 EXISTS'
 
         with self.assertRaises(ProtocolError):
-            client._consume_until_tagged_response(sentinel.tag, b'IDLE')
+            client._consume_until_tagged_response(sentinel.tag, b'IDLE')