refactor(scheduler): run ticks on one daemon thread instead of chained Timers

codecarbon/external/scheduler.py

@@ -1,52 +1,96 @@
-from threading import Lock, Timer
+import time
+from threading import Event, Thread, current_thread
+
+from codecarbon.external.logger import logger
 
 
 class PeriodicScheduler:
     """
-    A periodic task running in threading.Timers
-    From https://stackoverflow.com/a/18906292/14541668
+    Run ``function`` every ``interval`` seconds on a single daemon thread.
+
+    The deadline is absolute, so the cadence does not drift with the time the
+    function itself takes. A tick that overruns its slot is skipped rather than
+    queued, so the function is never re-entered.
     """
 
     def __init__(self, interval, function, *args, **kwargs):
         """
-        Init the scheduler. You have to call start() after initialization.
-        ::interval:: interval in seconds to run the function.
-        ::function:: function to run.
+        ::interval:: in seconds, the delay between two calls to function.
+        ::function:: the function to call.
         ::args:: args to pass to the function.
         ::kwargs:: kwargs to pass to the function.
         """
-        self._lock = Lock()
-        self._timer = None
-        self.function = function
         self.interval = interval
+        self.function = function
         self.args = args
         self.kwargs = kwargs
-        self._stopped = True
+        self._thread = None
+        self._stop_event = None
 
-    def start(self, from_run=False):
+    @property
+    def _stopped(self):
+        return self._stop_event is None or self._stop_event.is_set()
+
+    def _join_timeout(self, timeout):
+        # ponytail: 5s cap is a guess at "a measurement should never take
+        # longer than this"; make it configurable if a slow hardware
+        # backend ever needs more.
+        return min(self.interval, 5.0) if timeout is None else timeout
+
+    def start(self):
         """
-        Start the scheduler.
-        ::from_run:: For internal purposes to allow re-scheduling
-        Please do not use from_run=True until you know what you do !
+        Start the scheduler. Calling it on a running scheduler is a no-op.
+        If a previous run is still finishing its callback (a timed-out
+        stop()), wait for it; if it is still busy, refuse to start a second
+        loop so the function is never run concurrently.
         """
-        self._lock.acquire()
-        if from_run or self._stopped:
-            self._stopped = False
-            self._timer = Timer(self.interval, self._run)
-            self._timer.daemon = True
-            self._timer.start()
-        self._lock.release()
-
-    def _run(self):
-        self.start(from_run=True)
-        self.function(*self.args, **self.kwargs)
-
-    def stop(self):
+        if self._thread is not None:
+            if not self._stopped:
+                return
+            self._thread.join(self._join_timeout(None))
+            if self._thread.is_alive():
+                logger.warning(
+                    "Previous scheduled measurement is still running;"
+                    " not starting a second loop"
+                )
+                return
+        self._stop_event = Event()
+        self._thread = Thread(
+            target=self._loop,
+            args=(self._stop_event,),
+            daemon=True,
+            name=f"codecarbon-{getattr(self.function, '__name__', 'scheduler')}",
+        )
+        self._thread.start()
+
+    def _loop(self, stop_event):
+        next_call = time.monotonic() + self.interval
+        while not stop_event.wait(max(0.0, next_call - time.monotonic())):
+            try:
+                self.function(*self.args, **self.kwargs)
+            except Exception:  # noqa: BLE001 - must not kill the only thread
+                logger.error("Scheduled measurement failed", exc_info=True)
+            # Absolute deadline, but never a burst of catch-up ticks.
+            next_call = max(next_call + self.interval, time.monotonic())
+
+    def stop(self, timeout=None):
         """
-        Stop the scheduler.
+        Stop the scheduler and wait for the in-flight call to return.
+        ::timeout:: seconds to wait for the running function, bounded by
+        default so a wedged measurement cannot hang the caller for long.
         """
-        if not self._stopped:
-            self._lock.acquire()
-            self._stopped = True
-            self._timer.cancel()
-            self._lock.release()
+        if self._stop_event is not None:
+            self._stop_event.set()
+        thread = self._thread
+        if thread is None or thread is current_thread():
+            self._thread = None
+            return
+        thread.join(self._join_timeout(timeout))
+        if thread.is_alive():
+            # Keep the reference so start() will not launch a second loop.
+            logger.warning(
+                "Scheduled measurement still running after stop(); it will"
+                " finish in the background and may overlap the next one"
+            )
+        else:
+            self._thread = None

tests/test_scheduler.py

@@ -0,0 +1,130 @@
+import threading
+import time
+import unittest
+
+from codecarbon.external.scheduler import PeriodicScheduler
+
+INTERVAL = 0.05
+
+
+class TestPeriodicScheduler(unittest.TestCase):
+    def test_ticks_run_on_a_single_thread(self):
+        """Regression guard: one long-lived thread, not one thread per tick."""
+        names = set()
+        scheduler = PeriodicScheduler(
+            INTERVAL, lambda: names.add(threading.current_thread().name)
+        )
+        scheduler.start()
+        time.sleep(INTERVAL * 10)
+        scheduler.stop()
+        # A one-element set also proves at least one tick happened.
+        self.assertEqual(len(names), 1, f"expected one worker thread, got {names}")
+
+    def test_slow_function_is_never_re_entered(self):
+        """A function slower than the interval must not overlap with itself."""
+        state = {"in_flight": 0, "overlaps": 0, "calls": 0}
+
+        def slow():
+            state["calls"] += 1
+            state["in_flight"] += 1
+            if state["in_flight"] > 1:
+                state["overlaps"] += 1
+            time.sleep(INTERVAL * 2.4)
+            state["in_flight"] -= 1
+
+        scheduler = PeriodicScheduler(INTERVAL, slow)
+        scheduler.start()
+        time.sleep(INTERVAL * 20)
+        scheduler.stop()
+        self.assertGreater(state["calls"], 1)
+        self.assertEqual(state["overlaps"], 0)
+
+    def test_stop_is_prompt_and_final(self):
+        calls = []
+        scheduler = PeriodicScheduler(10.0, lambda: calls.append(1))
+        scheduler.start()
+        before = time.monotonic()
+        scheduler.stop()
+        self.assertLess(time.monotonic() - before, 1.0)
+        self.assertTrue(scheduler._stopped)
+        time.sleep(0.1)
+        self.assertEqual(calls, [])
+
+    def test_stop_before_start_and_double_stop_do_not_raise(self):
+        scheduler = PeriodicScheduler(INTERVAL, lambda: None)
+        scheduler.stop()
+        scheduler.start()
+        scheduler.stop()
+        scheduler.stop()
+        self.assertTrue(scheduler._stopped)
+
+    def test_start_is_idempotent_and_restartable(self):
+        calls = []
+        scheduler = PeriodicScheduler(INTERVAL, lambda: calls.append(1))
+        scheduler.start()
+        first = scheduler._thread
+        scheduler.start()
+        self.assertIs(scheduler._thread, first, "start() armed a second thread")
+        scheduler.stop()
+        self.assertFalse(first.is_alive())
+        stopped_at = len(calls)
+
+        scheduler.start()
+        second = scheduler._thread
+        time.sleep(INTERVAL * 4)
+        scheduler.stop()
+        self.assertGreater(len(calls), stopped_at, "restart did not resume ticking")
+        self.assertFalse(second.is_alive())
+
+    def test_start_after_a_timed_out_stop_never_runs_two_loops(self):
+        """A wedged callback blocks a restart until it returns."""
+        release = threading.Event()
+        in_flight = {"now": 0, "max": 0}
+        blocked_once = threading.Event()
+
+        def wedged():
+            in_flight["now"] += 1
+            in_flight["max"] = max(in_flight["max"], in_flight["now"])
+            if not blocked_once.is_set():
+                blocked_once.set()
+                release.wait(5)
+            in_flight["now"] -= 1
+
+        scheduler = PeriodicScheduler(INTERVAL, wedged)
+        scheduler.start()
+        first = scheduler._thread
+        self.assertTrue(blocked_once.wait(2), "callback never ran")
+
+        with self.assertLogs("codecarbon", level="WARNING"):
+            scheduler.stop()  # join times out, callback still blocked
+        self.assertTrue(first.is_alive())
+        self.assertIs(scheduler._thread, first, "stop() dropped the wedged thread")
+
+        with self.assertLogs("codecarbon", level="WARNING"):
+            scheduler.start()  # refused: the old callback is still running
+        self.assertIs(scheduler._thread, first, "start() launched a second loop")
+
+        release.set()
+        first.join(2)
+        self.assertFalse(first.is_alive(), "the wedged thread never exited")
+
+        scheduler.start()  # now the restart goes through
+        second = scheduler._thread
+        self.assertIsNot(second, first)
+        time.sleep(INTERVAL * 6)
+        scheduler.stop()
+        self.assertEqual(in_flight["max"], 1, "two loops ran the callback at once")
+
+    def test_exception_does_not_kill_the_loop(self):
+        calls = []
+
+        def flaky():
+            calls.append(1)
+            if len(calls) == 1:
+                raise ValueError("boom")
+
+        scheduler = PeriodicScheduler(INTERVAL, flaky)
+        scheduler.start()
+        time.sleep(INTERVAL * 6)
+        scheduler.stop()
+        self.assertGreater(len(calls), 1)