Make rate_limited use cvs.

[python_utils.git] / state_tracker.py

diff --git a/state_tracker.py b/state_tracker.py
index 16d2f595cf12ee34e3c28cdb9f1c522f8ac0b978..4f77ff42241474d214f222c428ae0f60e1da1e51 100644 (file)
--- a/state_tracker.py
+++ b/state_tracker.py
@@ -4,11 +4,11 @@ from abc import ABC, abstractmethod
 import datetime
 import logging
 import time
+import threading
 from typing import Dict, Optional
 
 import pytz
 
-import math_utils
 from thread_utils import background_thread
 
 logger = logging.getLogger(__name__)
@@ -113,7 +113,7 @@ class AutomaticStateTracker(StateTracker):
                 logger.debug('pace_maker noticed event; shutting down')
                 return
             self.heartbeat()
-            logger.debug(f'page_maker is sleeping for {self.sleep_delay}s')
+            logger.debug(f'pace_maker is sleeping for {self.sleep_delay}s')
             time.sleep(self.sleep_delay)
 
     def __init__(
@@ -122,6 +122,7 @@ class AutomaticStateTracker(StateTracker):
         *,
         override_sleep_delay: Optional[float] = None,
     ) -> None:
+        import math_utils
         super().__init__(update_ids_to_update_secs)
         if override_sleep_delay is not None:
             logger.debug(f'Overriding sleep delay to {override_sleep_delay}')
@@ -144,3 +145,32 @@ class AutomaticStateTracker(StateTracker):
         self.should_terminate.set()
         self.updater_thread.join()
         logger.debug('Background thread terminated.')
+
+
+class WaitableAutomaticStateTracker(AutomaticStateTracker):
+
+    def __init__(
+            self,
+            update_ids_to_update_secs: Dict[str, float],
+            *,
+            override_sleep_delay: Optional[float] = None,
+    ) -> None:
+        self._something_changed = threading.Event()
+        super().__init__(update_ids_to_update_secs,
+                         override_sleep_delay=override_sleep_delay)
+
+    def something_changed(self):
+        self._something_changed.set()
+
+    def did_something_change(self) -> bool:
+        return self._something_changed.is_set()
+
+    def reset(self):
+        self._something_changed.clear()
+
+    def wait(self,
+             *,
+             timeout=None):
+        return self._something_changed.wait(
+            timeout=timeout
+        )