python · serhiy-storchaka · Oct 19, 2020 · Feb 27, 2024
diff --git a/Lib/sched.py b/Lib/sched.py
@@ -32,7 +32,8 @@
 
 __all__ = ["scheduler"]
 
-Event = namedtuple('Event', 'time, priority, sequence, action, argument, kwargs')
+class Event(namedtuple('Event', ('time', 'priority', 'sequence', 'action', 'argument', 'kwargs'))):
+    _cancelled = False
 Event.time.__doc__ = ('''Numeric type compatible with the return value of the
 timefunc function passed to the constructor.''')
 Event.priority.__doc__ = ('''Events scheduled for the same time will be executed
@@ -55,10 +56,16 @@ def __init__(self, timefunc=_time, delayfunc=time.sleep):
         functions"""
         self._queue = []
         self._lock = threading.RLock()
+        self._cancellations = 0
         self.timefunc = timefunc
         self.delayfunc = delayfunc
         self._sequence_generator = count()
 
+    def _reheapify(self):
+        self._cancellations = 0
+        self._queue = [x for x in self._queue if not x._cancelled]
+        heapq.heapify(self._queue)
+
     def enterabs(self, time, priority, action, argument=(), kwargs=_sentinel):
         """Enter a new event in the queue at an absolute time.
 
@@ -92,13 +99,16 @@ def cancel(self, event):
 
         """
         with self._lock:
-            self._queue.remove(event)
-            heapq.heapify(self._queue)
+            assert isinstance(event, Event), event
+            if event._cancelled:
+                raise ValueError("this event was already cancelled")
+            event._cancelled = True
+            self._cancellations += 1
 
     def empty(self):
         """Check whether the queue is empty."""
         with self._lock:
-            return not self._queue
+            return len(self._queue) <= self._cancellations
 
     def run(self, blocking=True):
         """Execute events until the queue is empty.
@@ -135,14 +145,22 @@ def run(self, blocking=True):
             with lock:
                 if not q:
                     break
-                (time, priority, sequence, action,
-                 argument, kwargs) = q[0]
+                time, _, _, action, argument, kwargs = event = q[0]
+                if event._cancelled:
+                    self._cancellations -= 1
+                    pop(q)
+                    continue
                 now = timefunc()
                 if time > now:
+                    if (self._cancellations > 50 and
+                        self._cancellations > (len(self._queue) >> 1)):
+                        self._reheapify()
+                        continue
                     delay = True
                 else:
                     delay = False
                     pop(q)
+
             if delay:
                 if not blocking:
                     return time - now
@@ -163,5 +181,7 @@ def queue(self):
         # With heapq, two events scheduled at the same time will show in
         # the actual order they would be retrieved.
         with self._lock:
+            if self._cancellations:
+                self._reheapify()
             events = self._queue[:]
         return list(map(heapq.heappop, [events]*len(events)))
diff --git a/Lib/test/test_sched.py b/Lib/test/test_sched.py
@@ -38,7 +38,7 @@ def advance(self, t):
             self._cond.notify_all()
 
 
-class TestCase(unittest.TestCase):
+class TestScheduler(unittest.TestCase):
 
     def test_enter(self):
         l = []
@@ -222,5 +222,26 @@ def test_run_non_blocking(self):
         self.assertEqual(l, [])
 
 
+class TestCancellations(unittest.TestCase):
+
+    def test_cancel_twice(self):
+        scheduler = sched.scheduler()
+        ev = scheduler.enter(0.01, 1, lambda: 0)
+        scheduler.cancel(ev)
+        self.assertRaises(ValueError, scheduler.cancel, ev)
+
+    def test_queue(self):
+        scheduler = sched.scheduler()
+        ev = scheduler.enter(0.01, 1, lambda: 0)
+        scheduler.cancel(ev)
+        self.assertEqual(list(scheduler.queue), [])
+
+    def test_empty(self):
+        scheduler = sched.scheduler()
+        ev = scheduler.enter(0.01, 1, lambda: 0)
+        scheduler.cancel(ev)
+        self.assertTrue(scheduler.empty())
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Library/2020-10-19-11-05-37.bpo-13451.zT8Jgb.rst b/Misc/NEWS.d/next/Library/2020-10-19-11-05-37.bpo-13451.zT8Jgb.rst
@@ -0,0 +1 @@
+Optimized :meth:`sched.scheduler.cancel`.