fix: stop PeriodicScheduler from re-arming itself after stop()

codecarbon/external/scheduler.py

@@ -29,24 +29,32 @@ def start(self, from_run=False):
         ::from_run:: For internal purposes to allow re-scheduling
         Please do not use from_run=True until you know what you do !
         """
-        self._lock.acquire()
-        if from_run or self._stopped:
+        with self._lock:
+            if from_run:
+                if self._stopped:
+                    # stop() ran while the timer was firing: do not re-arm.
+                    return
+            elif not self._stopped:
+                # Already running, do not arm a second timer.
+                return
             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)
+        if self._stopped:
+            return
         self.function(*self.args, **self.kwargs)
 
     def stop(self):
         """
         Stop the scheduler.
         """
-        if not self._stopped:
-            self._lock.acquire()
+        with self._lock:
+            if self._stopped:
+                return
             self._stopped = True
-            self._timer.cancel()
-            self._lock.release()
+            if self._timer is not None:
+                self._timer.cancel()

tests/test_scheduler.py

@@ -0,0 +1,64 @@
+import threading
+import time
+import unittest
+
+from codecarbon.external.scheduler import PeriodicScheduler
+
+
+class TestPeriodicScheduler(unittest.TestCase):
+    def test_stop_during_timer_firing_does_not_rearm(self):
+        """stop() landing between the timer firing and _run rescheduling must win."""
+        calls = []
+        scheduler = PeriodicScheduler(0.05, calls.append, 1)
+        gate = threading.Event()
+        original_start = scheduler.start
+
+        def gated_start(from_run=False):
+            if from_run:
+                gate.wait(2)
+            return original_start(from_run=from_run)
+
+        scheduler.start = gated_start
+        original_start(from_run=False)
+        time.sleep(0.2)  # the timer fired, _run is blocked before start()
+
+        scheduler.stop()
+        calls_at_stop = len(calls)
+        gate.set()
+        time.sleep(0.3)
+
+        self.assertTrue(scheduler._stopped)
+        self.assertEqual(calls_at_stop, len(calls))
+
+    def test_stop_while_payload_is_running(self):
+        """No further invocation once stop() returned, even mid-payload."""
+        calls = []
+        in_payload = threading.Event()
+
+        def slow_payload():
+            calls.append(1)
+            in_payload.set()
+            time.sleep(0.2)
+
+        scheduler = PeriodicScheduler(0.01, slow_payload)
+        scheduler.start()
+        self.assertTrue(in_payload.wait(2))
+        scheduler.stop()
+        time.sleep(0.3)
+
+        self.assertEqual(1, len(calls))
+
+    def test_start_twice_arms_a_single_timer(self):
+        scheduler = PeriodicScheduler(10, lambda: None)
+        scheduler.start()
+        timer = scheduler._timer
+        scheduler.start()
+        self.assertIs(timer, scheduler._timer)
+        scheduler.stop()
+
+    def test_stop_is_safe_before_start_and_twice(self):
+        scheduler = PeriodicScheduler(10, lambda: None)
+        scheduler.stop()
+        scheduler.start()
+        scheduler.stop()
+        scheduler.stop()