From 6d38158e6b81d23072dfaf2baefa52cce27043ef Mon Sep 17 00:00:00 2001
From: Max Hirschhorn <max.hirschhorn@mongodb.com>
Date: Tue, 6 Aug 2019 10:52:55 -0400
Subject: [PATCH] SERVER-42623 Remove events from scheduler using `is`
 operator.

---
 buildscripts/resmokelib/logging/flush.py      |  5 ++-
 buildscripts/resmokelib/utils/scheduler.py    | 30 +++++++++++++
 .../tests/resmokelib/utils/test_scheduler.py  | 43 +++++++++++++++++++
 3 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 buildscripts/resmokelib/utils/scheduler.py
 create mode 100644 buildscripts/tests/resmokelib/utils/test_scheduler.py

diff --git a/buildscripts/resmokelib/logging/flush.py b/buildscripts/resmokelib/logging/flush.py
index b812fc2b606..1880d866f0a 100644
--- a/buildscripts/resmokelib/logging/flush.py
+++ b/buildscripts/resmokelib/logging/flush.py
@@ -4,10 +4,11 @@ These instances are used to send logs to buildlogger.
 """
 
 import logging
-import sched
 import threading
 import time
 
+from ..utils import scheduler
+
 _FLUSH_THREAD_LOCK = threading.Lock()
 _FLUSH_THREAD = None
 
@@ -93,7 +94,7 @@ class _FlushThread(threading.Thread):
             self.__schedule_updated.wait(secs)
             self.__schedule_updated.clear()
 
-        self.__scheduler = sched.scheduler(time.time, interruptible_sleep)
+        self.__scheduler = scheduler.Scheduler(time.monotonic, interruptible_sleep)
         self.__schedule_updated = threading.Event()
         self.__should_stop = threading.Event()
         self.__terminated = threading.Event()
diff --git a/buildscripts/resmokelib/utils/scheduler.py b/buildscripts/resmokelib/utils/scheduler.py
new file mode 100644
index 00000000000..12ebd47c35d
--- /dev/null
+++ b/buildscripts/resmokelib/utils/scheduler.py
@@ -0,0 +1,30 @@
+"""Version of sched.scheduler with a fixed cancel() method."""
+
+import heapq
+import sched
+
+
+class Scheduler(sched.scheduler):
+    """A thread-safe, general purpose event scheduler."""
+
+    def cancel(self, event):
+        """Remove an event from the queue.
+
+        Raises a ValueError if the event is not in the queue.
+        """
+
+        # The changes from https://hg.python.org/cpython/rev/d8802b055474 made it so sched.Event
+        # instances returned by sched.scheduler.enter() and sched.scheduler.enterabs() are treated
+        # as equal if they have the same (time, priority). It is therefore possible to remove the
+        # wrong event from the list when sched.scheduler.cancel() is called. Note that this is still
+        # true even with time.monotonic being the default timefunc, as GetTickCount64() on Windows
+        # only has a resolution of ~15ms. We therefore use the `is` operator to remove the correct
+        # event from the list.
+        with self._lock:
+            for i in range(len(self._queue)):
+                if self._queue[i] is event:
+                    del self._queue[i]
+                    heapq.heapify(self._queue)
+                    return
+
+            raise ValueError("event not in list")
diff --git a/buildscripts/tests/resmokelib/utils/test_scheduler.py b/buildscripts/tests/resmokelib/utils/test_scheduler.py
new file mode 100644
index 00000000000..58ac32b3c54
--- /dev/null
+++ b/buildscripts/tests/resmokelib/utils/test_scheduler.py
@@ -0,0 +1,43 @@
+"""Unit tests for buildscripts/resmokelib/utils/scheduler.py."""
+
+import sched
+import unittest
+
+from buildscripts.resmokelib.utils import scheduler as _scheduler
+
+# pylint: disable=missing-docstring
+
+
+def noop():
+    pass
+
+
+class TestScheduler(unittest.TestCase):
+    """Unit tests for the Scheduler class."""
+    scheduler = _scheduler.Scheduler
+
+    def setUp(self):
+        self.__scheduler = self.scheduler()
+
+    def test_cancel_with_identical_time_and_priority(self):
+        event1 = self.__scheduler.enterabs(time=0, priority=0, action=noop)
+        event2 = self.__scheduler.enterabs(time=0, priority=0, action=noop)
+
+        self.__scheduler.cancel(event1)
+        self.assertIs(self.__scheduler.queue[0], event2)
+
+        # Attempting to cancel the same event should fail because it has already been removed.
+        with self.assertRaises(ValueError):
+            self.__scheduler.cancel(event1)
+
+        self.__scheduler.cancel(event2)
+        self.assertEqual(self.__scheduler.queue, [])
+
+
+class TestBuiltinScheduler(TestScheduler):
+    """Unit tests for the sched.scheduler class."""
+    scheduler = sched.scheduler
+
+    def test_cancel_with_identical_time_and_priority(self):
+        with self.assertRaises(AssertionError):
+            super().test_cancel_with_identical_time_and_priority()