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

codecarbon/emissions_tracker.py

@@ -283,6 +283,9 @@ def _initialize_runtime_state(self) -> None:
         self._ram_used_history: List[float] = []
         # The monitor thread appends while the measurement thread clears/averages
         self._utilization_history_lock = threading.Lock()
+        # A scheduled measurement still running when stop() gave up joining it
+        # must not overlap the final one
+        self._measurement_lock = threading.RLock()
         self._total_cpu_energy: Energy = Energy.from_energy(kWh=0)
         self._total_gpu_energy: Energy = Energy.from_energy(kWh=0)
         self._total_ram_energy: Energy = Energy.from_energy(kWh=0)
@@ -1357,15 +1360,20 @@ def _do_measurements(self) -> None:
 
     def _measure_power_and_energy_if_stale(self, min_interval_s: float = 0.05) -> None:
         """Measure only if the last sample is older than ``min_interval_s``."""
-        if time.perf_counter() - self._last_measured_time >= min_interval_s:
-            self._measure_power_and_energy()
+        with self._measurement_lock:
+            if time.perf_counter() - self._last_measured_time >= min_interval_s:
+                self._measure_power_and_energy()
 
     def _measure_power_and_energy(self) -> None:
         """
         A function that is periodically run by the `BackgroundScheduler`
         every `self._measure_power_secs` seconds.
         :return: None
         """
+        with self._measurement_lock:
+            self._measure_power_and_energy_locked()
+
+    def _measure_power_and_energy_locked(self) -> None:
         try:
             last_duration = time.perf_counter() - self._last_measured_time
         except AttributeError as e:

codecarbon/external/scheduler.py

@@ -1,52 +1,111 @@
-from threading import Lock, Timer
+import math
+import time
+from threading import Event, Lock, Thread, current_thread
+
+from codecarbon.external.logger import logger
+
+# Longest stop() waits for the in-flight call, a guess at "a measurement should
+# never take longer than this", so a wedged one cannot hang the caller for long.
+STOP_JOIN_TIMEOUT_SECONDS = 5.0
+
+
+class MonotonicClock:
+    """The real clock. Tests swap in a fake one to drive the loop."""
+
+    monotonic = staticmethod(time.monotonic)
+
+    @staticmethod
+    def wait(event, timeout):
+        return event.wait(timeout)
 
 
 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
+        # Replace before start() to control time, e.g. in tests.
+        self.clock = MonotonicClock()
+        self._lock = Lock()
+        self._thread = None
+        self._stop_event = None
+
+    @property
+    def _stopped(self):
+        return self._stop_event is None or self._stop_event.is_set()
 
-    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.
+        If a previous run is still finishing its callback (a timed-out
+        stop()), the new loop waits for it before its first tick, so the
+        function is never run concurrently and the restart is not lost.
         """
-        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)
+        with self._lock:
+            if self._thread is not None and not self._stopped:
+                return
+            self._stop_event = Event()
+            self._thread = Thread(
+                target=self._loop,
+                args=(self._stop_event, self._thread),
+                daemon=True,
+                name=f"codecarbon-{getattr(self.function, '__name__', 'scheduler')}",
+            )
+            self._thread.start()
+
+    def _loop(self, stop_event, previous=None):
+        if previous is not None:
+            previous.join()
+        clock = self.clock
+        next_call = clock.monotonic() + self.interval
+        while not clock.wait(stop_event, max(0.0, next_call - clock.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)
+            next_call += self.interval
+            overrun = clock.monotonic() - next_call
+            if overrun > 0 and self.interval > 0:
+                # Skip the slots the call overran: no burst of catch-up ticks,
+                # and no back-to-back call either.
+                next_call += math.ceil(overrun / self.interval) * self.interval
 
     def stop(self):
         """
-        Stop the scheduler.
+        Stop the scheduler and wait for the in-flight call to return, at most
+        min(interval, STOP_JOIN_TIMEOUT_SECONDS) seconds.
         """
-        if not self._stopped:
-            self._lock.acquire()
-            self._stopped = True
-            self._timer.cancel()
-            self._lock.release()
+        with self._lock:
+            if self._stop_event is not None:
+                self._stop_event.set()
+            thread = self._thread
+        if thread is None or thread is current_thread():
+            # Called from the callback itself: it is still running, so keep
+            # the reference for the next start() to wait on.
+            return
+        thread.join(min(self.interval, STOP_JOIN_TIMEOUT_SECONDS))
+        if thread.is_alive():
+            # Keep the reference so the next start() waits for it.
+            logger.warning(
+                "Scheduled measurement still running after stop(); it will"
+                " finish in the background"
+            )
+        else:
+            with self._lock:
+                if self._thread is thread:
+                    self._thread = None

tests/test_emissions_tracker_constant.py

@@ -1,5 +1,6 @@
 import os
 import tempfile
+import threading
 import time
 import unittest
 from unittest import mock
@@ -12,6 +13,8 @@
     OfflineEmissionsTracker,
     track_emissions,
 )
+from codecarbon.external import scheduler as scheduler_module
+from tests.testutils import FakeClock
 
 
 def heavy_computation(run_time_secs: int = 3):
@@ -202,3 +205,55 @@ def test_carbon_tracker_offline_region_error(self):
 
         finally:
             tracker.stop()
+
+
+class TestFinalMeasurementWaitsForInFlightTick(unittest.TestCase):
+    @mock.patch.object(scheduler_module, "STOP_JOIN_TIMEOUT_SECONDS", 0.01)
+    def test_stop_does_not_overlap_a_slow_scheduled_measurement(self):
+        tracker = OfflineEmissionsTracker(
+            country_iso_code="USA",
+            measure_power_secs=1,
+            output_methods=[],
+            allow_multiple_runs=True,
+        )
+        clock = tracker._scheduler.clock = FakeClock()
+        tracker.start()
+        main = threading.current_thread()
+        do_measurements = tracker._do_measurements
+        in_flight = {"now": 0, "max": 0}
+        entered = threading.Event()
+        # Set once stop() reaches its final measurement, whether it waits for
+        # the lock (fixed) or measures right away (the overlap).
+        final_measurement_reached = threading.Event()
+
+        def slow_do_measurements():
+            in_flight["now"] += 1
+            in_flight["max"] = max(in_flight["max"], in_flight["now"])
+            if threading.current_thread() is main:
+                final_measurement_reached.set()
+            else:
+                entered.set()
+                final_measurement_reached.wait(5)
+            do_measurements()
+            in_flight["now"] -= 1
+
+        class LockSpy:
+            def __init__(self, lock):
+                self.lock = lock
+
+            def __enter__(self):
+                if threading.current_thread() is main:
+                    final_measurement_reached.set()
+                return self.lock.__enter__()
+
+            def __exit__(self, *exc):
+                return self.lock.__exit__(*exc)
+
+        tracker._do_measurements = slow_do_measurements
+        tracker._measurement_lock = LockSpy(tracker._measurement_lock)
+        clock.advance(1, settle=False)
+        self.assertTrue(entered.wait(5), "the scheduled measurement never ran")
+        # Make the last sample stale so stop() takes its final measurement.
+        tracker._last_measured_time -= 1
+        tracker.stop()  # the join times out while the tick is still running
+        self.assertEqual(in_flight["max"], 1, "final measurement overlapped a tick")