From fbfaa6fd57f8dc8a3da808acb8422370fad2f27b Mon Sep 17 00:00:00 2001
From: Giampaolo Rodola <g.rodola@gmail.com>
Date: Fri, 9 Jun 2017 22:20:41 +0200
Subject: bpo-30014: make poll-like selector's modify() method faster (#1030)

* #30014: make selectors.DefaultSelector.modify() faster by relying on selector's modify() method instead of un/register()ing the fd

* #30014: add unit test

* speedup poll/epoll/devpoll modify() method by using internal modify() call

* update doc

* address PR comments

* update NEWS entries

* use != instead of 'is not'
---
 Doc/whatsnew/3.7.rst       |  3 +++
 Lib/selectors.py           | 30 +++++++++++++++++++++++++++++-
 Lib/test/test_selectors.py | 27 +++++++++++++++++++++++++++
 Misc/NEWS                  |  3 +++
 4 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/Doc/whatsnew/3.7.rst b/Doc/whatsnew/3.7.rst
index 602db8c..3df3093 100644
--- a/Doc/whatsnew/3.7.rst
+++ b/Doc/whatsnew/3.7.rst
@@ -234,6 +234,9 @@ Optimizations
   expressions <re>`.  Searching some patterns can now be up to 20 times faster.
   (Contributed by Serhiy Storchaka in :issue:`30285`.)
 
+* :meth:`selectors.EpollSelector.modify`, :meth:`selectors.PollSelector.modify`
+  and :meth:`selectors.DevpollSelector.modify` may be around 10% faster under
+  heavy loads. (Contributed by Giampaolo Rodola' in :issue:`30014`)
 
 Build and C API Changes
 =======================
diff --git a/Lib/selectors.py b/Lib/selectors.py
index f29d11f..edde22c 100644
--- a/Lib/selectors.py
+++ b/Lib/selectors.py
@@ -252,7 +252,6 @@ class _BaseSelectorImpl(BaseSelector):
         return key
 
     def modify(self, fileobj, events, data=None):
-        # TODO: Subclasses can probably optimize this even further.
         try:
             key = self._fd_to_key[self._fileobj_lookup(fileobj)]
         except KeyError:
@@ -342,6 +341,8 @@ class SelectSelector(_BaseSelectorImpl):
 class _PollLikeSelector(_BaseSelectorImpl):
     """Base class shared between poll, epoll and devpoll selectors."""
     _selector_cls = None
+    _EVENT_READ = None
+    _EVENT_WRITE = None
 
     def __init__(self):
         super().__init__()
@@ -371,6 +372,33 @@ class _PollLikeSelector(_BaseSelectorImpl):
             pass
         return key
 
+    def modify(self, fileobj, events, data=None):
+        try:
+            key = self._fd_to_key[self._fileobj_lookup(fileobj)]
+        except KeyError:
+            raise KeyError(f"{fileobj!r} is not registered") from None
+
+        changed = False
+        if events != key.events:
+            selector_events = 0
+            if events & EVENT_READ:
+                selector_events |= self._EVENT_READ
+            if events & EVENT_WRITE:
+                selector_events |= self._EVENT_WRITE
+            try:
+                self._selector.modify(key.fd, selector_events)
+            except Exception:
+                super().unregister(fileobj)
+                raise
+            changed = True
+        if data != key.data:
+            changed = True
+
+        if changed:
+            key = key._replace(events=events, data=data)
+            self._fd_to_key[key.fd] = key
+        return key
+
     def select(self, timeout=None):
         # This is shared between poll() and epoll().
         # epoll() has a different signature and handling of timeout parameter.
diff --git a/Lib/test/test_selectors.py b/Lib/test/test_selectors.py
index 852b2fe..f2594a6 100644
--- a/Lib/test/test_selectors.py
+++ b/Lib/test/test_selectors.py
@@ -175,6 +175,33 @@ class BaseSelectorTestCase(unittest.TestCase):
         self.assertFalse(s.register.called)
         self.assertFalse(s.unregister.called)
 
+    def test_modify_unregister(self):
+        # Make sure the fd is unregister()ed in case of error on
+        # modify(): http://bugs.python.org/issue30014
+        if self.SELECTOR.__name__ == 'EpollSelector':
+            patch = unittest.mock.patch(
+                'selectors.EpollSelector._selector_cls')
+        elif self.SELECTOR.__name__ == 'PollSelector':
+            patch = unittest.mock.patch(
+                'selectors.PollSelector._selector_cls')
+        elif self.SELECTOR.__name__ == 'DevpollSelector':
+            patch = unittest.mock.patch(
+                'selectors.DevpollSelector._selector_cls')
+        else:
+            raise self.skipTest("")
+
+        with patch as m:
+            m.return_value.modify = unittest.mock.Mock(
+                side_effect=ZeroDivisionError)
+            s = self.SELECTOR()
+            self.addCleanup(s.close)
+            rd, wr = self.make_socketpair()
+            s.register(rd, selectors.EVENT_READ)
+            self.assertEqual(len(s._map), 1)
+            with self.assertRaises(ZeroDivisionError):
+                s.modify(rd, selectors.EVENT_WRITE)
+            self.assertEqual(len(s._map), 0)
+
     def test_close(self):
         s = self.SELECTOR()
         self.addCleanup(s.close)
diff --git a/Misc/NEWS b/Misc/NEWS
index f126dd4..a8aff86 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -350,6 +350,9 @@ Extension Modules
 Library
 -------
 
+- bpo-30014: modify() method of poll(), epoll() and devpoll() based classes of
+  selectors module is around 10% faster.  Patch by Giampaolo Rodola'.
+
 - bpo-30418: On Windows, subprocess.Popen.communicate() now also ignore EINVAL
   on stdin.write() if the child process is still running but closed the pipe.
 
-- 
cgit v0.12