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

codecarbon/external/scheduler.py

@@ -1,10 +1,16 @@
-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):
@@ -15,38 +21,65 @@ def __init__(self, interval, function, *args, **kwargs):
         ::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._stop_event = Event()
+        self._thread = None
+
+    @property
+    def _stopped(self):
+        return self._thread is None or not self._thread.is_alive()
 
-    def start(self, from_run=False):
+    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.
         """
-        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 not self._stopped:
+            return
+        self._stop_event.clear()
+        self._thread = Thread(
+            target=self._loop,
+            daemon=True,
+            name=f"codecarbon-{getattr(self.function, '__name__', 'scheduler')}",
+        )
+        self._thread.start()
+
+    def _loop(self):
+        next_call = time.monotonic() + self.interval
+        while not self._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 so the cadence does not drift, but if we
+            # overran (or the process was suspended) skip ahead instead of
+            # firing a burst of catch-up ticks.
+            next_call += self.interval
+            now = time.monotonic()
+            if next_call <= now:
+                next_call = now + self.interval
+
+    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()
+        self._stop_event.set()
+        thread = self._thread
+        if thread is not None and thread is not current_thread():
+            # 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.
+            thread.join(min(self.interval, 5.0) if timeout is None else timeout)
+            if thread.is_alive():
+                logger.warning(
+                    "Scheduled measurement did not return in time; the "
+                    "scheduler thread is still running and will exit on its "
+                    "own. A new start() is ignored until then."
+                )
+        # `_thread` is deliberately kept: `_stopped` tests `is_alive()`, so a
+        # still-running thread keeps `start()` a no-op instead of letting it
+        # clear `_stop_event` and release the old loop alongside a new one.

tests/test_scheduler.py

@@ -0,0 +1,124 @@
+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_does_not_run_two_loops(self):
+        """A wedged callback must not be released back into its loop.
+
+        `stop()` gives up after a bounded join. If it forgot the thread anyway,
+        the next `start()` would clear `_stop_event` and the still-blocked
+        thread would resume ticking alongside the freshly started one.
+        """
+        release = threading.Event()
+        idents = set()
+        blocked_once = threading.Event()
+
+        def wedged():
+            idents.add(threading.get_ident())
+            if not blocked_once.is_set():
+                blocked_once.set()
+                release.wait(5)
+
+        scheduler = PeriodicScheduler(INTERVAL, wedged)
+        scheduler.start()
+        first = scheduler._thread
+        self.assertTrue(blocked_once.wait(2), "callback never ran")
+
+        scheduler.stop()  # join times out, callback still blocked
+        self.assertTrue(first.is_alive())
+
+        scheduler.start()  # must not arm a second loop
+        release.set()
+        time.sleep(INTERVAL * 6)
+        scheduler.stop()
+
+        self.assertEqual(len(idents), 1, f"two scheduler loops ran: {idents}")
+        self.assertFalse(first.is_alive())
+
+    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)